Mesh cleaning improvements (#15)

* gitignore: added vscode

* gitignore: added vscode

* added rich fancy traceback for easier debugging

* replaced mesh cleaning with decimation

* pre-commit cleanup

* added TEST mode to atlas creation

* added smoothing options to all atlas gen scripts

* removed ratlas.py since outdated and theres a better atlas

* atlas gen update

* added smoothing and decimation to adult zebrafish atlas

Co-authored-by: Luigi Petrucco <[email protected]>

.gitignore

@@ -1,4 +1,8 @@
+.vscode
+
 # Byte-compiled / optimized / DLL files
+__pycache__
+*.__pycache__
 __pycache__/
 *.py[cod]
 *$py.class

bg_atlasgen/__init__.py

@@ -1 +1,5 @@
 __version__ = "1"  # will set major version of all atlases
+
+from rich.traceback import install
+
+install()

bg_atlasgen/atlas_scripts/allen_cord.py

@@ -1,29 +1,48 @@
-__version__ = "0"
+__version__ = "1"
 
 import json
 import time
 import tifffile
 import zipfile
 
+
 import pandas as pd
 import numpy as np
 import multiprocessing as mp
-
+from random import choices
+from loguru import logger
 from rich.progress import track
 from pathlib import Path
 
+# import sys
+
+# sys.path.append("./")
+
 from bg_atlasapi import utils
-from bg_atlasgen.mesh_utils import create_region_mesh, Region
+from bg_atlasgen.mesh_utils import (
+    create_region_mesh,
+    Region,
+    inspect_meshes_folder,
+)
 from bg_atlasgen.wrapup import wrapup_atlas_from_data
 from bg_atlasapi.structure_tree_util import get_structures_tree
 
 PARALLEL = True
+TEST = False
 
 
-def download_atlas_files(download_dir_path, atlas_file_url):
+def download_atlas_files(download_dir_path: Path, atlas_file_url: str) -> Path:
     utils.check_internet_connection()
 
     atlas_files_dir = download_dir_path / "atlas_files"
+
+    # only download data if they weren't already downloaded
+    if atlas_files_dir.exists():
+        print("Not downloading atlas since it was downloaded already already")
+        return atlas_files_dir / "SC_P56_Atlas_10x10x20_v5_2020"
+    else:
+        print("Downloading atlas data")
+
     destination_path = download_dir_path / "atlas_download"
     utils.retrieve_over_http(atlas_file_url, destination_path)
 
@@ -96,9 +115,22 @@ def create_meshes(download_dir_path, structures, annotated_volume, root_id):
 
     # Mesh creation
     closing_n_iters = 2
+    decimate_fraction = 0.2
+    smooth = False  # smooth meshes after creation
     start = time.time()
-    if PARALLEL:
 
+    # check how many regions to create the meshes for
+    nodes = list(tree.nodes.values())
+    if TEST:
+        logger.info(
+            f"Creating atlas in test mode: selecting 10 random regions for mesh creation"
+        )
+        nodes = choices(nodes, k=10)
+
+    if PARALLEL:
+        print(
+            f"Creating {tree.size()} meshes in parallel with {mp.cpu_count() - 2} CPU cores"
+        )
         pool = mp.Pool(mp.cpu_count() - 2)
 
         try:
@@ -113,17 +145,18 @@ def create_meshes(download_dir_path, structures, annotated_volume, root_id):
                         annotated_volume,
                         root_id,
                         closing_n_iters,
+                        decimate_fraction,
+                        smooth,
                     )
-                    for node in tree.nodes.values()
+                    for node in nodes
                 ],
             )
         except mp.pool.MaybeEncodingError:
             pass
     else:
+        print(f"Creating {len(nodes)} meshes")
         for node in track(
-            tree.nodes.values(),
-            total=tree.size(),
-            description="Creating meshes",
+            nodes, total=len(nodes), description="Creating meshes",
         ):
             create_region_mesh(
                 (
@@ -134,6 +167,8 @@ def create_meshes(download_dir_path, structures, annotated_volume, root_id):
                     annotated_volume,
                     root_id,
                     closing_n_iters,
+                    decimate_fraction,
+                    smooth,
                 )
             )
 
@@ -142,6 +177,11 @@ def create_meshes(download_dir_path, structures, annotated_volume, root_id):
         round((time.time() - start) / 60, 2),
         " minutes",
     )
+
+    if TEST:
+        # create visualization of the various meshes
+        inspect_meshes_folder(meshes_dir_path)
+
     return meshes_dir_path
 
 
@@ -191,7 +231,7 @@ def create_atlas(working_dir):
     # Download atlas files from Mendeley
     atlas_files_dir = download_atlas_files(download_dir_path, ATLAS_FILE_URL)
 
-    ## Load files
+    # Load files
     structures_file = atlas_files_dir / "Atlas_Regions.csv"
     reference_file = atlas_files_dir / "Template.tif"
     annotations_file = atlas_files_dir / "Annotation.tif"
@@ -203,7 +243,7 @@ def create_atlas(working_dir):
     atlas_segments = pd.read_csv(segments_file)
     atlas_segments = dict(atlas_segments=atlas_segments.to_dict("records"))
 
-    ## Parse structure metadata
+    # Parse structure metadata
     structures = parse_structures(structures_file, ROOT_ID)
 
     # save regions list json:
@@ -247,7 +287,11 @@ def create_atlas(working_dir):
 
 
 if __name__ == "__main__":
+
     # Generated atlas path:
-    bg_root_dir = Path.home() / "brainglobe_workingdir" / "allen_cord"
+    bg_root_dir = Path.home() / "brainglobe_workingdir" / "allen_cord_smooth"
     bg_root_dir.mkdir(exist_ok=True, parents=True)
+
+    # generate atlas
+    print(f'Creating atlas and saving it at "{bg_root_dir}"')
     create_atlas(bg_root_dir)

bg_atlasgen/atlas_scripts/azba_zfish.py

@@ -13,13 +13,14 @@
 import time
 import tifffile
 import tarfile
+from random import choices
 
 import numpy as np
 import multiprocessing as mp
 
 from rich.progress import track
 from pathlib import Path
-from bg_atlasgen.mesh_utils import Region, create_region_mesh
+from bg_atlasgen.mesh_utils import Region, create_region_mesh, inspect_meshes_folder
 from bg_atlasapi.structure_tree_util import get_structures_tree
 from bg_atlasgen.wrapup import wrapup_atlas_from_data
 from bg_atlasapi import utils
@@ -42,6 +43,7 @@ def create_atlas(working_dir, resolution):
 
     #setup folder for downloading
     working_dir = working_dir  / ATLAS_NAME
+    working_dir.mkdir(exist_ok=True)
     download_dir_path = working_dir / "downloads"
     download_dir_path.mkdir(exist_ok=True)
     atlas_path = download_dir_path / f"{ATLAS_NAME}"
@@ -118,6 +120,9 @@ def create_atlas(working_dir, resolution):
     #mesh creation
     closing_n_iters = 2
     start = time.time()
+
+    decimate_fraction = 0.3
+    smooth = True
 
     if PARALLEL:
 
@@ -146,9 +151,9 @@ def create_atlas(working_dir, resolution):
     else:
 
         print("Multiprocessing disabled")
-
-        for node in track(
-            tree.nodes.values(),
+        # nodes = list(tree.nodes.values())
+        # nodes = choices(nodes, k=10)
+        for node in track(tree.nodes.values(),
             total=tree.size(),
             description = "Creating meshes",
         ):
@@ -161,6 +166,8 @@ def create_atlas(working_dir, resolution):
                     annotated_volume,
                     ROOT_ID,
                     closing_n_iters,
+                    decimate_fraction,
+                    smooth
                 )
             )
 
@@ -186,8 +193,8 @@ def create_atlas(working_dir, resolution):
     print(f"In the end, {len(structures_with_mesh)} structures with mesh are kept")    
 
     #import reference file with tifffile so it can be read in wrapup_atlas_from_data
-    reference = tifffile.imread(reference_file)    
-
+    reference = tifffile.imread(reference_file)
+    inspect_meshes_folder(meshes_dir_path)
     # wrap up atlas file
     print("Finalising atlas")
     output_filename = wrapup_atlas_from_data(

bg_atlasgen/atlas_scripts/humanatlas.py

@@ -13,7 +13,11 @@
 # import sys
 
 # sys.path.append("./")
-from bg_atlasgen.mesh_utils import create_region_mesh, Region
+from bg_atlasgen.mesh_utils import (
+    create_region_mesh,
+    Region,
+    inspect_meshes_folder,
+)
 from bg_atlasgen.wrapup import wrapup_atlas_from_data
 from bg_atlasapi.structure_tree_util import get_structures_tree
 
@@ -51,6 +55,7 @@ def prune_tree(tree):
 
 if __name__ == "__main__":
     PARALLEL = False  # disable parallel mesh extraction for easier debugging
+    TEST = False
 
     # ---------------------------------------------------------------------------- #
     #                                 PREP METADATA                                #
@@ -177,6 +182,8 @@ def prune_tree(tree):
 
     # Mesh creation
     closing_n_iters = 2
+    decimate_fraction = 0.2
+    smooth = False  # smooth meshes after creation
     start = time.time()
     if PARALLEL:
         print("Starting mesh creation in parallel")
@@ -195,6 +202,8 @@ def prune_tree(tree):
                         annotation,
                         ROOT_ID,
                         closing_n_iters,
+                        decimate_fraction,
+                        smooth,
                     )
                     for node in tree.nodes.values()
                 ],
@@ -224,6 +233,8 @@ def prune_tree(tree):
                     volume,
                     ROOT_ID,
                     closing_n_iters,
+                    decimate_fraction,
+                    smooth,
                 )
             )
 
@@ -233,6 +244,10 @@ def prune_tree(tree):
         " minutes",
     )
 
+    if TEST:
+        # create visualization of the various meshes
+        inspect_meshes_folder(meshes_dir_path)
+
     # Create meshes dict
     meshes_dict = dict()
     structures_with_mesh = []

bg_atlasgen/atlas_scripts/kim_mouse.py

@@ -1,4 +1,4 @@
-__version__ = "0"
+__version__ = "1"
 
 import json
 import time
@@ -13,8 +13,12 @@
 from pathlib import Path
 from scipy.ndimage import zoom
 from allensdk.core.reference_space_cache import ReferenceSpaceCache
-from bg_atlasapi import utils
 
+# import sys
+
+# sys.path.append("./")
+
+from bg_atlasapi import utils
 from bg_atlasgen.mesh_utils import create_region_mesh, Region
 from bg_atlasgen.wrapup import wrapup_atlas_from_data
 from bg_atlasapi.structure_tree_util import get_structures_tree
@@ -37,7 +41,7 @@ def create_atlas(working_dir, resolution):
     download_dir_path.mkdir(exist_ok=True)
     atlas_files_dir = download_dir_path / "atlas_files"
 
-    ## Download atlas_file
+    # Download atlas_file
     utils.check_internet_connection()
 
     destination_path = download_dir_path / "atlas_download"
@@ -126,6 +130,9 @@ def create_atlas(working_dir, resolution):
 
     # Mesh creation
     closing_n_iters = 2
+    decimate_fraction = 0.2
+    smooth = False  # smooth meshes after creation
+
     start = time.time()
     if PARALLEL:
 
@@ -143,6 +150,8 @@ def create_atlas(working_dir, resolution):
                         rotated_annotations,
                         ROOT_ID,
                         closing_n_iters,
+                        decimate_fraction,
+                        smooth,
                     )
                     for node in tree.nodes.values()
                 ],
@@ -164,6 +173,8 @@ def create_atlas(working_dir, resolution):
                     rotated_annotations,
                     ROOT_ID,
                     closing_n_iters,
+                    decimate_fraction,
+                    smooth,
                 )
             )
 
@@ -225,7 +236,7 @@ def create_atlas(working_dir, resolution):
 
 
 if __name__ == "__main__":
-    resolution = 50  # some resolution, in microns
+    resolution = 10  # some resolution, in microns (10, 25, 50, 100)
 
     # Generated atlas path:
     bg_root_dir = Path.home() / "brainglobe_workingdir" / "kim_mouse"

bg_atlasgen/atlas_scripts/mpin_zfish.py

@@ -48,10 +48,7 @@ def add_path_inplace(parent):
 
 
 def collect_all_inplace(
-    node,
-    traversing_list,
-    download_path,
-    mesh_dict,
+    node, traversing_list, download_path, mesh_dict,
 ):
     """Recursively traverse a region hierarchy, download meshes, and append
     regions to a list inplace.