Allen bbp ccfv3a entire mouse brain, improve mesh utils, and add metadata filtering to wrapup

brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_bbp_ccfv3a_entire_mouse_brain.py

@@ -0,0 +1,213 @@
+import json
+import zipfile
+from pathlib import Path
+
+import nrrd
+import numpy as np
+import requests
+from brainglobe_utils.IO.image import load_any
+from scipy.ndimage import zoom
+from tqdm import tqdm
+
+from brainglobe_atlasapi.atlas_generation.mesh_utils import (
+    construct_meshes_from_annotation,
+)
+from brainglobe_atlasapi.atlas_generation.wrapup import wrapup_atlas_from_data
+
+VERSION = 0
+NAME = "allen_bbp_ccfv3a_entire_mouse_brain"
+CITATION = "https://doi.org/10.1101/2024.11.06.622212"
+ATLAS_LINK = "https://zenodo.org/api/records/14034334/files-archive"
+SPECIES = "Mus musculus"
+ORIENTATION = "spr"
+ROOT_ID = 997
+RESOLUTION = 25
+
+
+def download_resources(download_dir_path, atlas_file_url, atlas_name):
+    """
+    Download the necessary resources for the atlas.
+
+    Pooch for some reason refused to download the entire file.
+    """
+    if "download=1" not in atlas_file_url:
+        atlas_file_url = atlas_file_url + "?download=1"
+    file_path = download_dir_path / f"{atlas_name}-files-archive.zip"
+
+    if not file_path.exists():
+        # Pooch for some reason refused to download the entire file.
+        response = requests.get(atlas_file_url, stream=True)
+        response.raise_for_status()
+        # Get file size in bytes (21.5GB)
+        total_size = int(21.5 * 1024 * 1024 * 1024)
+        with tqdm(
+            total=total_size, unit="B", unit_scale=True, desc="Downloading"
+        ) as pbar:
+            with open(file_path, "wb") as f:
+                for chunk in response.iter_content(
+                    chunk_size=1024 * 1024
+                ):  # 1MB chunks
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+        with zipfile.ZipFile(file_path, "r") as zip_ref:
+            zip_ref.extractall(download_dir_path)
+    zipped_template_dir = download_dir_path / "average_nissl_template.zip"
+    if not (download_dir_path / "nissl_average_full.nii.gz").exists():
+        with zipfile.ZipFile(zipped_template_dir, "r") as zip_ref:
+            zip_ref.extractall(download_dir_path)
+
+
+def retrieve_reference_and_annotation(download_path, resolution):
+    """
+    Retrieve the desired reference and annotation as two numpy arrays.
+    When the volume is 10 microns the authors only provided a hemi
+    segmentation. I checked with the lead author and he confirmed this was
+    accidental. Here I correct this by mirroring if the resolution is 10um.
+    """
+    reference_path = download_path / f"arav3a_bbp_nisslCOR_{resolution}.nrrd"
+    reference, header = nrrd.read(reference_path)
+    annotation, header = nrrd.read(
+        download_path / f"annotv3a_bbp_{resolution}.nrrd"
+    )
+    if resolution == 10:
+        # mirror the volume
+        new_annot = np.zeros(reference.shape)
+        new_annot[:, :, 570:] = annotation
+        new_annot[:, :, :570] = annotation[:, :, ::-1]
+        annotation = new_annot
+    return reference, annotation
+
+
+def retrieve_additional_references(download_path, resolution):
+    """
+    This returns the population average nissl and converts it
+    to 16 bit unsigned int which bg requires. The developers
+    only provided a 10um template so I downsample here.
+    """
+    additional_reference_path = download_path / "nissl_average_full.nii.gz"
+    reference = load_any(additional_reference_path)
+    if resolution == 25:
+        reference = zoom(reference, 0.4, order=1)
+    reference = (reference * 65535).astype(np.uint16)
+    reference = reference.transpose(1, 2, 0)[::-1, ::-1]
+    return {"population_average_nissl": reference}
+
+
+def retrieve_hemisphere_map():
+    """
+    At least one of the templates is technically not
+    symmetrical as in pixels differ in left vs right
+    hemisphere. But it is registered to a symmetrical
+    template meaning there should not be morphological
+    differences.
+    """
+    return None
+
+
+def hex_to_rgb(hex_string):
+    """Convert hex color string to RGB list."""
+    hex_string = hex_string.lstrip("#")
+    return [int(hex_string[i : i + 2], 16) for i in (0, 2, 4)]
+
+
+def flatten_structure_tree(node, structure_id_path=None):
+    """Recursively flatten the tree structure."""
+    if structure_id_path is None:
+        structure_id_path = []
+
+    # Create current path
+    current_path = structure_id_path + [node["id"]]
+
+    # Create entry for current node
+    entry = {
+        "id": node["id"],
+        "name": node["name"],
+        "acronym": node["acronym"],
+        "structure_id_path": current_path,
+        "rgb_triplet": hex_to_rgb(node["color_hex_triplet"]),
+    }
+
+    # Start with current node's data
+    entries = [entry]
+
+    # Recursively process children
+    if "children" in node:
+        for child in node["children"]:
+            entries.extend(flatten_structure_tree(child, current_path))
+
+    return entries
+
+
+def retrieve_structure_information(download_path):
+    """
+    The authors provide a json which is nested so I unest it.
+    """
+    with open(download_path / "hierarchy_bbp_atlas_pipeline.json") as f:
+        metadata_json = json.load(f)
+
+    # Get the root of the hierarchy
+    root = metadata_json["msg"][0]
+
+    # Flatten the tree structure
+    flattened_structures = flatten_structure_tree(root)
+
+    return flattened_structures
+
+
+def retrieve_or_construct_meshes(download_path, annotated_volume, structures):
+    """
+    I use the construct_meshes_from_annotation helper function introduced in
+    this pull request
+    """
+    meshes_dict = construct_meshes_from_annotation(
+        download_path, annotated_volume, structures, ROOT_ID
+    )
+    return meshes_dict
+
+
+### If the code above this line has been filled correctly, nothing needs to be
+### edited below (unless variables need to be passed between the functions).
+if __name__ == "__main__":
+    bg_root_dir = Path.home() / "brainglobe_workingdir" / NAME
+    bg_root_dir.mkdir(exist_ok=True)
+    download_resources(
+        download_dir_path=bg_root_dir,
+        atlas_name=NAME,
+        atlas_file_url=ATLAS_LINK,
+    )
+    for resolution in [25, 10]:
+        reference_volume, annotated_volume = retrieve_reference_and_annotation(
+            bg_root_dir, resolution=resolution
+        )
+        additional_references = retrieve_additional_references(
+            bg_root_dir, resolution=resolution
+        )
+
+        hemispheres_stack = retrieve_hemisphere_map()
+        if resolution == 25:
+            structures = retrieve_structure_information(bg_root_dir)
+            meshes_dict = retrieve_or_construct_meshes(
+                bg_root_dir, annotated_volume, structures
+            )
+        reference_volume = (reference_volume * 65535).astype(np.uint16)
+        output_filename = wrapup_atlas_from_data(
+            atlas_name=NAME,
+            atlas_minor_version=VERSION,
+            citation=CITATION,
+            atlas_link=ATLAS_LINK,
+            species=SPECIES,
+            resolution=(resolution,) * 3,
+            orientation=ORIENTATION,
+            root_id=ROOT_ID,
+            reference_stack=reference_volume,
+            annotation_stack=annotated_volume,
+            structures_list=structures,
+            meshes_dict=meshes_dict,
+            working_dir=bg_root_dir,
+            hemispheres_stack=None,
+            cleanup_files=False,
+            compress=True,
+            scale_meshes=True,
+            additional_references=additional_references,
+        )

brainglobe_atlasapi/atlas_generation/mesh_utils.py

@@ -20,10 +20,12 @@
 import numpy as np
 import scipy
 from loguru import logger
+from rich.progress import track
 
 from brainglobe_atlasapi.atlas_generation.volume_utils import (
     create_masked_array,
 )
+from brainglobe_atlasapi.structure_tree_util import get_structures_tree
 
 # ----------------- #
 #   MESH CREATION   #
@@ -245,6 +247,96 @@
                 )
 
 
+def construct_meshes_from_annotation(
+    download_path,
+    volume,
+    structures_list,
+    root_id,
+    closing_n_iters=2,
+    decimate_fraction=0,
+    smooth=False,
+):
+    """
+    Retrieve or construct atlas region meshes for a given annotation volume.
+
+    If an atlas is packaged with mesh files, reuse those. Otherwise, construct
+    the meshes using the existing volume and structure tree. Returns a
+    dictionary mapping structure IDs to their corresponding .obj mesh files.
+
+    Parameters
+    ----------
+    download_path : Path
+        Path to the directory where new mesh files will be saved.
+    volume : np.ndarray
+        3D annotation volume.
+    structures_list : list
+        List of structure dictionaries containing id information.
+    root_id : int
+        Identifier for the root structure.
+    smooth: bool
+        if True the surface mesh is smoothed
+    closing_n_iters: int
+        number of iterations of closing morphological operation.
+        set to None to avoid applying morphological operations
+    decimate_fraction: float  in range [0, 1].
+        What fraction of the original number of vertices is to be kept.
+        EG .5 means that 50% of the vertices are kept,
+        the others are removed.
+    Returns
+    -------
+    dict
+        Dictionary of structure IDs and paths to their .obj mesh files.
+    """
+    meshes_dir_path = download_path / "meshes"
+    meshes_dir_path.mkdir(exist_ok=True)
+
+    tree = get_structures_tree(structures_list)
+    labels = np.unique(volume).astype(np.int32)
+
+    for key, node in tree.nodes.items():
+        node.data = Region(key in labels)
+
+    for node in track(
+        tree.nodes.values(),
+        total=tree.size(),
+        description="Creating meshes",
+    ):
+        create_region_mesh(
+            (
+                meshes_dir_path,
+                node,
+                tree,
+                labels,
+                volume,
+                root_id,
+                closing_n_iters,
+                decimate_fraction,
+                smooth,
+            )
+        )
+    meshes_dict = {}
+    structures_with_mesh = []
+    for s in structures_list:
+        mesh_path = meshes_dir_path / f'{s["id"]}.obj'
+        if not mesh_path.exists():
+            print(f"No mesh file exists for: {s}, ignoring it")
+            continue
+        if mesh_path.stat().st_size < 512:
+            print(f"obj file for {s} is too small, ignoring it.")
+            continue
+
+        structures_with_mesh.append(s)
+        meshes_dict[s["id"]] = mesh_path
+
+    print(
+        (
+            f"In the end, {len(structures_with_mesh)}",
+            "structures with mesh are kept",
+        )
+    )
+    return meshes_dict
+
+
 class Region(object):
     """
     Class used to add metadata to treelib.Tree during atlas creation.

brainglobe_atlasapi/atlas_generation/wrapup.py

@@ -5,6 +5,7 @@
 
 import brainglobe_space as bgs
 import meshio as mio
+import numpy as np
 import tifffile
 
 import brainglobe_atlasapi.atlas_generation
@@ -25,13 +26,49 @@
 from brainglobe_atlasapi.atlas_generation.validate_atlases import (
     get_all_validation_functions,
 )
+from brainglobe_atlasapi.structure_tree_util import get_structures_tree
 from brainglobe_atlasapi.utils import atlas_name_from_repr
 
 # This should be changed every time we make changes in the atlas
 # structure:
 ATLAS_VERSION = brainglobe_atlasapi.atlas_generation.__version__
 
 
+def filter_structures_not_present_in_annotation(structures, annotation):
+    """
+    Filter out structures that are not present in the annotation volume,
+    or whose children are not present. Also prints removed structures.
+
+    Args:
+        structures (list of dict): List containing structure information
+        annotation (np.ndarray): Annotation volume
+
+    Returns:
+        list of dict: Filtered list of structure dictionaries
+    """
+    present_ids = set(np.unique(annotation))
+    # Create a structure tree for easy parent-child relationship traversal
+    tree = get_structures_tree(structures)
+
+    # Function to check if a structure or any of its descendants are present
+    def is_present(structure):
+        node = tree.get_node(structure["id"])
+        if structure["id"] in present_ids:
+            return True
+
+        # Check if any children are present
+        for child in tree.children(node.identifier):
+            if child.identifier in present_ids:
+                return True
+        return False
+
+    removed = [s for s in structures if not is_present(s)]
+    for r in removed:
+        print("Removed structure:", r["name"], "(ID:", r["id"], ")")
+
+    return [s for s in structures if is_present(s)]
+
+
 def wrapup_atlas_from_data(
     atlas_name,
     atlas_minor_version,
@@ -119,6 +156,11 @@
 
     # If no hemisphere file is given, assume the atlas is symmetric:
     symmetric = hemispheres_stack is None
+    if isinstance(annotation_stack, str) or isinstance(annotation_stack, Path):
+        annotation_stack = tifffile.imread(annotation_stack)
+    structures_list = filter_structures_not_present_in_annotation(
+        structures_list, annotation_stack
+    )
 
     # Instantiate BGSpace obj, using original stack size in um as meshes
     # are un um: