Fix AZBA

brainglobe_atlasapi/atlas_generation/atlas_scripts/azba_zfish.py

@@ -1,34 +1,28 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 """
 Script to generate a Brainglobe compatible atlas object
 for the Adult Zebrafish Brain Atlas (AZBA)
 
 @author: Kailyn Fields, [email protected]
 """
 
-__version__ = "1"
+__version__ = "2"
 
 import csv
-import multiprocessing as mp
-import tarfile
 import time
 from pathlib import Path
 
 import numpy as np
+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
 
-PARALLEL = False  # Disable for debugging mesh creation
-
 
 def create_atlas(working_dir, resolution):
     # metadata
@@ -42,32 +36,26 @@
     ATLAS_PACKAGER = "Kailyn Fields, [email protected]"
     ADDITIONAL_METADATA = {}
 
-    # 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}"
-
-    # download atlas files
-    utils.check_internet_connection()
-    destination_path = download_dir_path / "atlas_download"
-    utils.retrieve_over_http(ATLAS_FILE_URL, destination_path)
-
-    # unpack the atlas download folder
-    tar = tarfile.open(destination_path)
-    tar.extractall(path=atlas_path)
-    tar.close()
-    destination_path.unlink()
+    download_path = working_dir / "downloads"
+    download_path.mkdir(exist_ok=True, parents=True)
+
+    atlas_path = pooch.retrieve(
+        url=ATLAS_FILE_URL,
+        known_hash="a14b09b88979bca3c06fa96d525e6c1ba8906fe08689239433eb72d8d3e2ba44",
+        path=download_path,
+        progressbar=True,
+        processor=pooch.Untar(extract_dir="."),
+    )
 
     print("Atlas files download completed")
 
     # paths
-    structures_file = atlas_path / "2021-08-22_AZBA_labels.csv"
-    annotations_file = atlas_path / "2021-08-22_AZBA_segmentation.tif"
-    reference_topro = atlas_path / "20180219_AZBA_topro_average_2020.tif"
-    reference_file = atlas_path / "20180628_AZBA_AF_average.tif"
-    meshes_dir_path = atlas_path / "meshes"
+    structures_file = download_path / "2021-08-22_AZBA_labels.csv"
+    annotations_file = download_path / "2021-08-22_AZBA_segmentation.tif"
+    reference_topro = download_path / "20180219_AZBA_topro_average_2020.tif"
+    reference_file = download_path / "20180628_AZBA_AF_average.tif"
+    meshes_dir_path = download_path / "meshes"
     meshes_dir_path.mkdir(exist_ok=True)
 
     # adding topro image as additional reference file,
@@ -91,24 +79,25 @@
     # 'id', 'structure_id_path', and 'rgb_triplet' key values
     for i in range(0, len(hierarchy)):
         hierarchy[i]["id"] = int(hierarchy[i]["id"])
-    for j in range(0, len(hierarchy)):
-        hierarchy[j]["structure_id_path"] = list(
-            map(int, hierarchy[j]["structure_id_path"].split("/"))
+        hierarchy[i]["structure_id_path"] = list(
+            map(int, hierarchy[i]["structure_id_path"].split("/"))
         )
-    for k in range(0, len(hierarchy)):
         try:
-            hierarchy[k]["rgb_triplet"] = list(
-                map(int, hierarchy[k]["rgb_triplet"].split("/"))
+            hierarchy[i]["rgb_triplet"] = list(
+                map(int, hierarchy[i]["rgb_triplet"].split("/"))
             )
         except ValueError:
-            hierarchy[k]["rgb_triplet"] = [255, 255, 255]
+            hierarchy[i]["rgb_triplet"] = [255, 255, 255]
 
     # remove clear label (id 0) from hierarchy.
     # ITK-Snap uses this to label unlabeled areas,
     # but this convention interferes with the root mask generation
     # and is unnecessary for this application
     hierarchy.remove(hierarchy[1])
 
+    # Set root mesh to white
+    hierarchy[0]["rgb_triplet"] = [255, 255, 255]
+
     # use tifffile to read annotated file
     annotated_volume = tifffile.imread(annotations_file)
 
@@ -136,51 +125,24 @@
     decimate_fraction = 0.3
     smooth = True
 
-    if PARALLEL:
-        print("Multiprocessing mesh creation...")
-        pool = mp.Pool(int(mp.cpu_count() / 2))
-
-        try:
-            pool.map(
-                create_region_mesh,
-                [
-                    (
-                        meshes_dir_path,
-                        node,
-                        tree,
-                        labels,
-                        annotated_volume,
-                        ROOT_ID,
-                        closing_n_iters,
-                    )
-                    for node in tree.nodes.values()
-                ],
-            )
-        except mp.pool.MaybeEncodingError:
-            pass
-
-    else:
-        print("Multiprocessing disabled")
-        # nodes = list(tree.nodes.values())
-        # nodes = choices(nodes, k=10)
-        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,
-                )
+    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 : ",
@@ -213,7 +175,7 @@
     # import reference file with tifffile so
     # it can be read in wrapup_atlas_from_data
     reference = tifffile.imread(reference_file)
-    # inspect_meshes_folder(meshes_dir_path)
+
     # wrap up atlas file
     print("Finalising atlas")
     output_filename = wrapup_atlas_from_data(
@@ -233,6 +195,7 @@
         atlas_packager=ATLAS_PACKAGER,
         additional_metadata=ADDITIONAL_METADATA,
         additional_references=ADDITIONAL_REFERENCES,
+        scale_meshes=True,
     )
 
     return output_filename
@@ -243,5 +206,4 @@
 
     # generated atlas path
     bg_root_dir = Path.home() / "brainglobe_workingdir"
-    bg_root_dir.mkdir(exist_ok=True, parents=True)
     create_atlas(bg_root_dir, resolution)