-
Notifications
You must be signed in to change notification settings - Fork 31
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Script to generate Allen mouse brain atlas with barrel annotations #313
Changes from all commits
cf2a8fb
e0cc4a0
b617d3b
5fa35c4
a20fb46
831e9e2
31ca81a
83fac20
c8e1feb
1cbee61
8afd634
10441e1
308a975
44aecac
100f05a
bc3e472
533fd8f
d98b34a
9e1b1f7
84ce3b8
eb431ea
a2f6920
b951f20
ff1b1c1
be74f37
1d6f684
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,375 @@ | ||
__version__ = "0" | ||
import json | ||
import time | ||
from pathlib import Path | ||
Check warning on line 4 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
import meshio as mio | ||
import nrrd | ||
import numpy as np | ||
import pooch | ||
from allensdk.api.queries.ontologies_api import OntologiesApi | ||
from allensdk.api.queries.reference_space_api import ReferenceSpaceApi | ||
from allensdk.core.reference_space_cache import ReferenceSpaceCache | ||
from requests import exceptions | ||
from rich.progress import track | ||
from tqdm import tqdm | ||
Check warning on line 15 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
from brainglobe_atlasapi import descriptors, utils | ||
from brainglobe_atlasapi.atlas_generation.mesh_utils import ( | ||
Check warning on line 18 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
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 | ||
Check warning on line 23 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
|
||
def create_atlas(working_dir, resolution): | ||
# Specify information about the atlas: | ||
ATLAS_NAME = "allen_mouse_bluebrain_barrels" | ||
SPECIES = "Mus musculus" | ||
ATLAS_LINK = "http://www.brain-map.org" | ||
CITATION = "Bolaños-Puchet S., Teska A., et al. (2024). https://doi.org/10.1162/imag_a_00209" | ||
ATLAS_PACKAGER = "Axel Bisi" | ||
ORIENTATION = "asr" | ||
Check warning on line 33 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Parameters for mesh creation: | ||
ROOT_ID = 997 | ||
CLOSING_N_ITERS = 2 | ||
DECIMATE_FRACTION = 0.3 | ||
SMOOTH = True | ||
Check warning on line 39 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Temporary folder for nrrd files download: | ||
download_dir_path = working_dir / "downloads" | ||
download_dir_path.mkdir(exist_ok=True) | ||
Check warning on line 43 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Download original Allen template volume: | ||
######################################### | ||
spacecache = ReferenceSpaceCache( | ||
manifest=download_dir_path / "manifest.json", | ||
# downloaded files are stored relative to here | ||
resolution=resolution, | ||
reference_space_key="annotation/ccf_2017", | ||
# use the latest version of the CCF | ||
) | ||
|
||
template_volume, _ = spacecache.get_template_volume() | ||
|
||
# Download enhanced barrel-containing Allen annotation files by BlueBrain, | ||
# and hierarchy: | ||
######################################### | ||
annotation_dir_path = working_dir / "downloads/annotation_enhanced" | ||
annotation_dir_path.mkdir(exist_ok=True) | ||
Check warning on line 61 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
if resolution == 10: | ||
gin_url = "https://gin.g-node.org/BrainGlobe/bluebrain_barrel_materials/raw/master/annotation_barrels_10.nrrd" | ||
elif resolution == 25: | ||
gin_url = "https://gin.g-node.org/BrainGlobe/bluebrain_barrel_materials/raw/master/annotation_barrels_25.nrrd" | ||
Check warning on line 66 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
else: | ||
raise ValueError("Resolution {}um not supported.".format(resolution)) | ||
|
||
utils.check_internet_connection() | ||
annotation_file_path = pooch.retrieve( | ||
Check warning on line 71 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
gin_url, | ||
known_hash=None, | ||
path=annotation_dir_path, | ||
progressbar=True, | ||
) | ||
|
||
# Load enhanced annotation volume: | ||
annotated_volume = nrrd.read(annotation_file_path)[0] | ||
|
||
# Download structures tree and meshes: | ||
###################################### | ||
oapi = OntologiesApi() # ontologies | ||
struct_tree = spacecache.get_structure_tree() # structures tree | ||
Check warning on line 84 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Find id of set of regions with mesh: | ||
select_set = ( | ||
"Structures whose surfaces are represented by a precomputed mesh" | ||
) | ||
|
||
mesh_set_ids = [ | ||
s["id"] | ||
for s in oapi.get_structure_sets() | ||
if s["description"] == select_set | ||
] | ||
|
||
# Get structures with mesh for both versions | ||
structs_with_mesh = struct_tree.get_structures_by_set_id(mesh_set_ids) | ||
|
||
# Download hierarchy: | ||
gin_url = "https://gin.g-node.org/BrainGlobe/bluebrain_barrel_materials/raw/master/hierarchy.json" | ||
hierarchy_path = pooch.retrieve( | ||
Check warning on line 102 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
gin_url, | ||
known_hash=None, | ||
path=annotation_dir_path, | ||
fname="hierarchy.json", | ||
progressbar=True, | ||
) | ||
structs_with_barrels = json.load(open(hierarchy_path)) | ||
Check warning on line 109 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Add barrels structures to Allen structures | ||
def find_dicts_with_key_containing_substring(d, key, substring): | ||
Check warning on line 112 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
""" | ||
Recursively find all dictionaries within a nested dictionary that | ||
contain a specific substring in the value associated with a given key. | ||
|
||
Args: | ||
d (dict): The input dictionary. | ||
key (str): The key to search for. | ||
substring (str): The substring to search for in the value associated | ||
with the key. | ||
|
||
Returns: | ||
list: A list of dictionaries that contain the key with a value | ||
containing the substring. | ||
""" | ||
if not isinstance(d, dict): | ||
raise ValueError("Input should be a dictionary") | ||
Check warning on line 128 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
matching_dicts = [] | ||
Check warning on line 130 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
def recurse(sub_d): | ||
contains_substring = False | ||
Check warning on line 133 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
for k, v in sub_d.items(): | ||
if isinstance(v, dict): | ||
recurse(v) | ||
elif isinstance(v, list): | ||
for item in v: | ||
if isinstance(item, dict): | ||
recurse(item) | ||
if k == key and substring in str(v): | ||
contains_substring = True | ||
Check warning on line 143 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
if contains_substring: | ||
matching_dicts.append(sub_d) | ||
Check warning on line 146 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
recurse(d) | ||
Check warning on line 148 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
return matching_dicts | ||
Check warning on line 150 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
matching_dicts = find_dicts_with_key_containing_substring( | ||
Check warning on line 152 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
structs_with_barrels, key="name", substring="barrel" | ||
) | ||
matching_dicts = [ | ||
Check warning on line 155 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
d | ||
for d in matching_dicts | ||
if d["graph_order"] in [52, 53, 54, 55, 56, 57] | ||
] | ||
main_barrel_parent_struct = [ | ||
Check warning on line 160 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
s for s in structs_with_mesh if s["acronym"] == "SSp-bfd" | ||
][0] | ||
structures_present = [ | ||
Check warning on line 163 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"SSp-bfd1", | ||
"SSp-bfd2/3", | ||
"SSp-bfd4", | ||
"SSp-bfd5", | ||
"SSp-bfd6a", | ||
"SSp-bfd6b", | ||
] # keep laminar structures | ||
keys_to_keep = [ | ||
Check warning on line 171 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"acronym", | ||
"graph_id", | ||
"graph_order", | ||
"id", | ||
"name", | ||
"rgb_triplet", | ||
"structure_set_ids", | ||
"structure_id_path", | ||
] | ||
dict_to_add = [] | ||
for d in matching_dicts: | ||
Check warning on line 182 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Ignore parent-level SSp-bfd layers | ||
if d["acronym"] in structures_present: | ||
print("Skipping because already present:", d) | ||
continue | ||
Check warning on line 187 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
# Ignore sub-structures layer 2 and 3 to keep layer 2/3 structure | ||
if d["graph_order"] == 53 and d["acronym"] in ["SSp-bfd2", "SSp-bfd3"]: | ||
print("Excluding", d, "to keep layer 2/3 structure only.") | ||
continue | ||
elif d["graph_order"] == 54 and ( | ||
Check warning on line 192 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"layer 2" in d["name"] or "layer 3" in d["name"] | ||
): | ||
print("Excluding", d, "to keep layer 2/3 structure only.") | ||
continue | ||
Check warning on line 196 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Add desired barrel-related structures, with corresponding fields | ||
else: | ||
current_id = d["id"] | ||
Check warning on line 200 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
# Find corresponding parent structure | ||
if d["graph_order"] == 52: # barrel-level -> find SSp-bfd | ||
Check warning on line 202 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
# Create new structure_id_path for barrel structure | ||
d["structure_id_path"] = main_barrel_parent_struct[ | ||
Check warning on line 204 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"structure_id_path" | ||
] + [current_id] | ||
elif ( | ||
Check warning on line 207 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
d["graph_order"] == 53 | ||
): # barrel layer-level -> find SSp-bfd-barrel also | ||
parent_struct_id = d["parent_structure_id"] | ||
parent_struct = [ | ||
Check warning on line 211 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
s for s in matching_dicts if s["id"] == parent_struct_id | ||
][0] | ||
parent_struct["structure_id_path"] = main_barrel_parent_struct[ | ||
Check warning on line 214 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"structure_id_path" | ||
] + [parent_struct_id] | ||
# Create new structure_id_path for barrel-layer structure | ||
d["structure_id_path"] = ( | ||
Check warning on line 218 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
main_barrel_parent_struct["structure_id_path"] | ||
+ [d["parent_structure_id"]] | ||
+ [current_id] | ||
) | ||
|
||
# Complete with other keys | ||
d["rgb_triplet"] = main_barrel_parent_struct["rgb_triplet"] | ||
d["graph_id"] = 1 | ||
d["structure_set_ids"] = None | ||
dict_to_add.append({k: d[k] for k in keys_to_keep}) | ||
Check warning on line 228 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Add list of dicts to structs_with_mesh | ||
structs_with_mesh = structs_with_mesh + dict_to_add | ||
Check warning on line 231 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Directory for mesh saving: | ||
meshes_dir = ( | ||
Check warning on line 234 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
working_dir | ||
/ descriptors.MESHES_DIRNAME | ||
/ "meshes_{}".format(resolution) | ||
) | ||
# If directory exists, then skip | ||
if not meshes_dir.exists(): | ||
meshes_dir.mkdir(exist_ok=False, parents=True) | ||
Check warning on line 241 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Download existing Allen meshes: | ||
space = ReferenceSpaceApi() | ||
meshes_dict = dict() | ||
for s in tqdm(structs_with_mesh): | ||
name = s["id"] | ||
filename = meshes_dir / f"{name}.obj" | ||
Check warning on line 248 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
if filename.exists(): | ||
meshes_dict[name] = filename | ||
continue | ||
Check warning on line 252 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
try: | ||
space.download_structure_mesh( | ||
Check warning on line 255 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
structure_id=s["id"], | ||
ccf_version="annotation/ccf_2017", | ||
file_name=filename, | ||
) | ||
meshes_dict[name] = filename | ||
except (exceptions.HTTPError, ConnectionError): | ||
print(f"Failed to download mesh for {s['name']} ({s['id']})") | ||
Check warning on line 262 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Create missing meshes | ||
tree = get_structures_tree(structs_with_mesh) | ||
print( | ||
Check warning on line 266 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
f"Number of brain regions: {tree.size()}, " | ||
f"max tree depth: {tree.depth()}" | ||
) | ||
|
||
# generate binary mask for mesh creation | ||
labels = np.unique(annotated_volume).astype(np.int_) | ||
Check warning on line 272 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
for key, node in tree.nodes.items(): | ||
if key in labels: | ||
is_label = True | ||
Check warning on line 276 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
else: | ||
is_label = False | ||
Check warning on line 278 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
node.data = Region(is_label) | ||
Check warning on line 280 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
start = time.time() | ||
Check warning on line 282 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
for node in track( | ||
Check warning on line 284 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
tree.nodes.values(), | ||
total=tree.size(), | ||
description="Creating meshes", | ||
): | ||
|
||
# Check if mesh already exists | ||
file_name = meshes_dir / f"{node.identifier}.obj" | ||
if file_name.exists(): | ||
meshes_dict[node.identifier] = file_name | ||
continue | ||
Check warning on line 294 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
else: | ||
create_region_mesh( | ||
Check warning on line 297 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
( | ||
meshes_dir, | ||
node, | ||
tree, | ||
labels, | ||
annotated_volume, | ||
ROOT_ID, | ||
CLOSING_N_ITERS, | ||
DECIMATE_FRACTION, | ||
SMOOTH, | ||
) | ||
) | ||
meshes_dict[node.identifier] = file_name | ||
Check warning on line 310 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
print( | ||
Check warning on line 312 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
"Finished mesh extraction in : ", | ||
round((time.time() - start) / 60, 2), | ||
" minutes", | ||
) | ||
|
||
# Once mesh creation is over, rescale | ||
for mesh_id, meshfile in meshes_dict.items(): | ||
Check warning on line 319 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
# Check if mesh is barrel-related | ||
if mesh_id in [s["id"] for s in dict_to_add]: | ||
Check warning on line 321 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
try: | ||
mesh = mio.read(meshfile) | ||
mesh.points *= resolution | ||
mio.write(meshfile, mesh) | ||
except mio._exceptions.ReadError: | ||
print(f"Mesh file {meshfile} not found.") | ||
Check warning on line 328 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Loop over structures, remove entries not used: | ||
for struct in structs_with_mesh: | ||
[ | ||
Check warning on line 332 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
struct.pop(k) | ||
for k in ["graph_id", "structure_set_ids", "graph_order"] | ||
] | ||
|
||
# Remove problematic 545 structure | ||
if 545 in meshes_dict.keys(): | ||
meshes_dict.pop(545) | ||
Check warning on line 339 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
# Wrap up, compress, and remove file:0 | ||
print("Finalising atlas") | ||
Check warning on line 342 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
output_filename = wrapup_atlas_from_data( | ||
Check warning on line 344 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
atlas_name=ATLAS_NAME, | ||
atlas_minor_version=__version__, | ||
citation=CITATION, | ||
atlas_link=ATLAS_LINK, | ||
species=SPECIES, | ||
resolution=(resolution,) * 3, | ||
orientation=ORIENTATION, | ||
root_id=997, | ||
reference_stack=template_volume, | ||
annotation_stack=annotated_volume, | ||
structures_list=structs_with_mesh, | ||
meshes_dict=meshes_dict, | ||
working_dir=working_dir, | ||
hemispheres_stack=None, | ||
cleanup_files=False, | ||
compress=True, | ||
additional_metadata={"atlas_packager": ATLAS_PACKAGER}, | ||
) | ||
|
||
return output_filename | ||
Check warning on line 364 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
|
||
if __name__ == "__main__": | ||
RES_UM = 25 | ||
Check warning on line 368 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
# Generated atlas path: | ||
bg_root_dir = ( | ||
Check warning on line 370 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
Path.home() / "brainglobe_workingdir" / "allen_mouse_bluebrain_barrels" | ||
) | ||
bg_root_dir.mkdir(exist_ok=True) | ||
Check warning on line 373 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
||
|
||
create_atlas(bg_root_dir, RES_UM) | ||
Check warning on line 375 in brainglobe_atlasapi/atlas_generation/atlas_scripts/allen_mouse_bluebrain_barrels.py
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Needed to work on machines where
~/brainglobe_workingdir
doesn't exist yet.