diff --git a/brainglobe_atlasapi/atlas_generation/atlas_scripts/praire_vole.py b/brainglobe_atlasapi/atlas_generation/atlas_scripts/praire_vole.py new file mode 100644 index 00000000..9d8cf59d --- /dev/null +++ b/brainglobe_atlasapi/atlas_generation/atlas_scripts/praire_vole.py @@ -0,0 +1,266 @@ +__version__ = "0" + +import random +import shutil +import time +import zipfile +from pathlib import Path + +import numpy as np +import pandas as pd +import pooch +import tifffile +from rich.progress import track + +from brainglobe_atlasapi import utils +from brainglobe_atlasapi.atlas_generation.mesh_utils import ( + Region, + create_region_mesh, +) +from brainglobe_atlasapi.atlas_generation.wrapup import wrapup_atlas_from_data +from brainglobe_atlasapi.structure_tree_util import get_structures_tree + + +def create_atlas(working_dir): + ATLAS_NAME = "prairie_vole" + SPECIES = "Microtus ochrogaster" + ATLAS_LINK = "https://doi.org/10.7554/eLife.87029" + CITATION = ( + "Morgan L Gustison, Rodrigo Muñoz-Castañeda, Pavel Osten, " + "Steven M Phelps (2024) Sexual coordination " + "in a whole-brain map of prairie vole pair bonding eLife 12:RP87029 " + "https://doi.org/10.7554/eLife.87029.3" + ) + ATLAS_BASE_URL = "https://ndownloader.figshare.com/files/" + ORIENTATION = "asr" + ROOT_ID = 997 + ATLAS_PACKAGER = "Sacha Hadaway-Andreae" + ADDITIONAL_METADATA = {} + + # setup folder for downloading + download_dir_path = working_dir / "downloads" + download_dir_path.mkdir(exist_ok=True) + atlas_path = download_dir_path / "atlas_files" + + files = { + "PrV_ReferenceBrain_bilateral.tif": "42917449", + "PrV_Annotation_bilateral.tif": "42917452", + "PrV_Atlas_Ontology.csv": "42917443", + "PrV_results_datasets.zip": "42501327", + } + + utils.check_internet_connection() + pooch.create( + path=atlas_path, + base_url=ATLAS_BASE_URL, + registry={}, + ) + + for filename, file_id in files.items(): + url = f"{ATLAS_BASE_URL}{file_id}" + pooch.retrieve( + url=url, + fname=filename, + path=atlas_path, + known_hash=None, + ) + + zip_file_path = atlas_path / "PrV_results_datasets.zip" + with zipfile.ZipFile(zip_file_path, "r") as zip_ref: + zip_ref.extract( + "results_datasets/neural_nomenclature_hierarchy.csv", atlas_path + ) + + if not (atlas_path / "neural_nomenclature_hierarchy.csv").exists(): + shutil.move( + ( + atlas_path + / "results_datasets" + / "neural_nomenclature_hierarchy.csv" + ), + atlas_path, + ) + else: + print( + f"File '{(atlas_path / 'neural_nomenclature_hierarchy.csv')}' " + f"already exists, skipping move operation." + ) + + shutil.rmtree(atlas_path / r"results_datasets") + + annotations_file = atlas_path / "PrV_Annotation_bilateral.tif" + reference_file = atlas_path / "PrV_ReferenceBrain_bilateral.tif" + meshes_dir_path = atlas_path / "meshes" + meshes_dir_path.mkdir(exist_ok=True) + + nomenclature_hierarchy = pd.read_csv( + atlas_path / "neural_nomenclature_hierarchy.csv" + ).values + reference_id = pd.read_csv(atlas_path / "PrV_Atlas_Ontology.csv") + + acronyms = reference_id["acronym"].values + names = reference_id["name"].values + ids = reference_id["id"].values + + acronym_to_id = dict(zip(acronyms, ids)) + name_to_id = dict(zip(names, ids)) + + ids = list(ids) + hierarchy = [ + { + "acronym": "root", + "id": 997, + "name": "root", + "structure_id_path": [997], + "rgb_triplet": [255, 255, 255], + } + ] + + for row in nomenclature_hierarchy: + structure_hierarchy = [] + unique_pairs = [] + pairs = [(row[i], row[i + 1]) for i in range(0, len(row), 2)] + pairs_seen = set() + + for pair in pairs: + if pair not in pairs_seen: + unique_pairs.append(pair) + pairs_seen.add(pair) + + for acronym, name in unique_pairs: + if acronym in acronym_to_id: + structure_hierarchy.append(acronym_to_id[acronym]) + elif name in name_to_id: + structure_hierarchy.append(name_to_id[name]) + + structure_hierarchy = [ROOT_ID] + structure_hierarchy + + result = [ + structure_hierarchy[: i + 1] + for i in range(len(structure_hierarchy)) + ] + + for a in result: + structure_template = { + "acronym": acronyms[ids.index(a[-1])], + "id": a[-1], + "name": names[ids.index(a[-1])], + "structure_id_path": a, + "rgb_triplet": [random.randint(0, 255) for _ in range(3)], + } + + hierarchy.append(structure_template) + + hierarchy = list( + {tuple(d["structure_id_path"]): d for d in hierarchy}.values() + ) + + # use tifffile to read annotated file + annotated_volume = tifffile.imread(annotations_file) + reference_volume = tifffile.imread(reference_file) + + print(f"Saving atlas data at {atlas_path}") + + tree = get_structures_tree(hierarchy) + print( + f"Number of brain regions: {tree.size()}, " + f"max tree depth: {tree.depth()}" + ) + labels = np.unique(annotated_volume).astype(np.int32) + for key, node in tree.nodes.items(): + if key in labels: + is_label = True + else: + is_label = False + node.data = Region(is_label) + + # Mesh creation + closing_n_iters = 2 + decimate_fraction = 0.2 + smooth = False # smooth meshes after creation + start = time.time() + + for node in track( + tree.nodes.values(), + total=tree.size(), + description="Creating meshes", + ): + create_region_mesh( + ( + meshes_dir_path, + node, + tree, + labels, + annotated_volume, + ROOT_ID, + closing_n_iters, + decimate_fraction, + smooth, + ) + ) + print( + "Finished mesh extraction in: ", + round((time.time() - start) / 60, 2), + " minutes", + ) + + # Create meshes dict & hierarchy to match + hierarchy_with_mesh = [] + meshes_dict = dict() + structs_with_mesh = [] + for s in hierarchy: + # Check if a mesh was created + 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 + else: + # Check that the mesh actually exists (i.e. not empty) + if mesh_path.stat().st_size < 512: + print(f"obj file for {s} is too small, ignoring it.") + continue + + structs_with_mesh.append(s) + meshes_dict[s["id"]] = mesh_path + + hierarchy_with_mesh.append(s) + + print( + f"In the end, {len(structs_with_mesh)} structures with mesh are kept" + ) + + for item in hierarchy_with_mesh: + item["id"] = int(item["id"]) + item["structure_id_path"] = [int(i) for i in item["structure_id_path"]] + + output_filename = wrapup_atlas_from_data( + atlas_name=ATLAS_NAME, + atlas_minor_version=__version__, + citation=CITATION, + atlas_link=ATLAS_LINK, + species=SPECIES, + orientation=ORIENTATION, + root_id=ROOT_ID, + reference_stack=reference_volume, + annotation_stack=annotated_volume, + structures_list=hierarchy_with_mesh, + working_dir=working_dir, + hemispheres_stack=None, + cleanup_files=False, + compress=True, + atlas_packager=ATLAS_PACKAGER, + additional_metadata=ADDITIONAL_METADATA, + resolution=(25, 25, 25), + meshes_dict=meshes_dict, + scale_meshes=True, + ) + + return output_filename + + +if __name__ == "__main__": + home = str(Path.home()) + bg_root_dir = Path.home() / "brainglobe_workingdir" + bg_root_dir.mkdir(exist_ok=True, parents=True) + + create_atlas(bg_root_dir)