refactor: simplified soma loading (#411)

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>

Author: anna-grim

src/deep_neurographs/fragments_graph.py

@@ -71,7 +71,7 @@ def __init__(
         remove_high_risk_merges=False,
         segmentation_path=None,
         smooth_bool=True,
-        somas_path=None,
+        soma_centroids=None,
         verbose=False,
     ):
         """
@@ -99,9 +99,8 @@ def __init__(
         smooth_bool : bool, optional
             Indication of whether to smooth xyz coordinates from SWC files.
             The default is True.
-        somas_path : str, optional
-            Path to a txt file containing xyz coordinates of detected somas.
-            The default is None.
+        soma_centroids : List[Tuple[float]] or None, optional
+            Physcial coordinates of soma centroids. The default is None.
         verbose : bool, optional
             Indication of whether to display a progress bar while building
             FragmentsGraph. The default is True.
@@ -123,7 +122,7 @@ def __init__(
             remove_high_risk_merges=remove_high_risk_merges,
             segmentation_path=segmentation_path,
             smooth_bool=smooth_bool,
-            somas_path=somas_path,
+            soma_centroids=soma_centroids,
             verbose=verbose,
         )
         self.swc_reader = swc_util.Reader(anisotropy, min_size)

src/deep_neurographs/inference.py

@@ -36,7 +36,6 @@
 from torch.nn.functional import sigmoid
 from tqdm import tqdm
 
-import ast
 import networkx as nx
 import numpy as np
 import os
@@ -68,7 +67,7 @@ def __init__(
         output_dir,
         config,
         segmentation_path=None,
-        somas_path=None,
+        soma_centroids=None,
         s3_dict=None,
     ):
         """
@@ -92,9 +91,8 @@ def __init__(
             for the inference pipeline.
         segmentation_path : str, optional
             Path to segmentation stored in GCS bucket. The default is None.
-        somas_path : str, optional
-            Path to a txt file containing xyz coordinates of detected somas.
-            The default is None.
+        soma_centroids : List[Tuple[float]] or None, optional
+            Physcial coordinates of soma centroids. The default is None.
         s3_dict : dict, optional
             ...
 
@@ -110,7 +108,7 @@ def __init__(
         self.brain_id = brain_id
         self.segmentation_id = segmentation_id
         self.segmentation_path = segmentation_path
-        self.somas_path = somas_path
+        self.soma_centroids = soma_centroids
         self.s3_dict = s3_dict
 
         # Extract config settings
@@ -148,14 +146,14 @@ def run(self, fragments_pointer):
 
         # Main
         self.build_graph(fragments_pointer)
-        self.connect_soma_fragments() if self.somas_path else None
+        self.connect_soma_fragments() if self.soma_centroids else None
         self.generate_proposals(self.graph_config.search_radius)
         self.classify_proposals(self.ml_config.threshold)
 
         # Finish
         t, unit = util.time_writer(time() - t0)
         self.report_graph(prefix="\nFinal")
-        self.report(f"Total Runtime: {round(t, 4)} {unit}\n")
+        self.report(f"Total Runtime: {t:.2f} {unit}\n")
         self.save_results()
 
     def run_schedule(
@@ -177,7 +175,7 @@ def run_schedule(
         # Finish
         t, unit = util.time_writer(time() - t0)
         self.report_graph(prefix="\nFinal")
-        self.report(f"Total Runtime: {round(t, 4)} {unit}\n")
+        self.report(f"Total Runtime: {t:.2f} {unit}\n")
         self.save_results()
 
     def build_graph(self, fragments_pointer):
@@ -207,7 +205,7 @@ def build_graph(self, fragments_pointer):
             remove_high_risk_merges=self.graph_config.remove_high_risk_merges,
             segmentation_path=self.segmentation_path,
             smooth_bool=self.graph_config.smooth_bool,
-            somas_path=self.somas_path,
+            soma_centroids=self.soma_centroids,
             verbose=True,
         )
         self.graph.load_fragments(fragments_pointer)
@@ -219,7 +217,7 @@ def build_graph(self, fragments_pointer):
 
         t, unit = util.time_writer(time() - t0)
         self.report_graph(prefix="\nInitial")
-        self.report(f"Module Runtime: {round(t, 4)} {unit}\n")
+        self.report(f"Module Runtime: {t:.2f} {unit}\n")
 
     def filter_fragments(self):
         self.graph = fragment_filtering.remove_curvy(self.graph, 200)
@@ -235,7 +233,7 @@ def connect_soma_fragments(self):
         # Parse locations
         nodes_list = list()
         merge_cnt, soma_cnt = 0, 0
-        for soma_xyz in util.load_soma_locations(self.somas_path):
+        for soma_xyz in self.soma_centroids:
             hits = self.graph.find_fragments_near_xyz(soma_xyz, 20)
             if len(hits) > 1:
                 # Determine new swc id
@@ -295,13 +293,13 @@ def generate_proposals(self, radius):
             proposals_per_leaf=self.graph_config.proposals_per_leaf,
             trim_endpoints_bool=self.graph_config.trim_endpoints_bool,
         )
-        n_proposals = util.reformat_number(self.graph.n_proposals())
+        n_proposals = format(self.graph.n_proposals(), ",")
 
         # Report results
         t, unit = util.time_writer(time() - t0)
         self.report(f"# Proposals: {n_proposals}")
         self.report(f"# Proposals Blocked: {self.graph.n_proposals_blocked}")
-        self.report(f"Module Runtime: {round(t, 4)} {unit}\n")
+        self.report(f"Module Runtime: {t:.2f} {unit}\n")
 
     def classify_proposals(self, accept_threshold):
         """
@@ -341,9 +339,9 @@ def classify_proposals(self, accept_threshold):
         # Report results
         t, unit = util.time_writer(time() - t0)
         self.report(f"# Merges Blocked: {self.graph.n_merges_blocked}")
-        self.report(f"# Accepted: {util.reformat_number(len(accepts))}")
-        self.report(f"% Accepted: {round(len(accepts) / n_proposals, 4)}")
-        self.report(f"Module Runtime: {round(t, 4)} {unit}\n")
+        self.report(f"# Accepted: {format(len(accepts), ',')}")
+        self.report(f"% Accepted: {len(accepts) / n_proposals:.4f}")
+        self.report(f"Module Runtime: {t:.4f} {unit}\n")
 
     def save_results(self):
         """
@@ -447,7 +445,7 @@ def write_metadata(self):
             "min_fragment_size": f"{self.graph_config.min_size}um",
             "node_spacing": self.graph_config.node_spacing,
             "remove_doubles": self.graph_config.remove_doubles,
-            "use_somas": self.segmentation_path and self.somas_path,
+            "use_somas": len(self.soma_centroids) > 0,
             "complex_proposals": self.graph_config.complex_bool,
             "long_range_bool": self.graph_config.long_range_bool,
             "proposals_per_leaf": self.graph_config.proposals_per_leaf,
@@ -486,17 +484,16 @@ def report_graph(self, prefix="\n"):
         """
         # Compute values
         n_components = nx.number_connected_components(self.graph)
-        n_components = util.reformat_number(n_components)
-        n_nodes = util.reformat_number(self.graph.number_of_nodes())
-        n_edges = util.reformat_number(self.graph.number_of_edges())
-        usage = round(util.get_memory_usage(), 2)
+        n_components = format(n_components, ",")
+        n_nodes = format(self.graph.number_of_nodes(), ",")
+        n_edges = format(self.graph.number_of_edges(), ",")
 
         # Report
         self.report(f"{prefix} Graph")
         self.report(f"# Connected Components: {n_components}")
         self.report(f"# Nodes: {n_nodes}")
         self.report(f"# Edges: {n_edges}")
-        self.report(f"Memory Consumption: {usage} GBs")
+        self.report(f"Memory Consumption: {util.get_memory_usage():.2f} GBs")
 
 
 class InferenceEngine:

src/deep_neurographs/machine_learning/groundtruth_generation.py

@@ -102,7 +102,7 @@ def find_alignments(gt_graph, pred_graph, kdtree):
         )
         if aligned:
             i = util.sample_once(nodes)
-            pred_id = pred_graph.nodes[i]["swc_id"]            
+            pred_id = pred_graph.nodes[i]["swc_id"]
             valid_ids.add(pred_id)
             pred_to_target[pred_id] = target_id
     return pred_to_target
@@ -150,7 +150,7 @@ def is_component_aligned(gt_graph, pred_graph, nodes, kdtree):
 
     intersects = True if percent_aligned > 0.5 else False
     aligned_score = np.mean(dists[dists < np.percentile(dists, 80)])
-    
+
     # Deterine whether aligned
     if (aligned_score < ALIGNED_THRESHOLD and hat_swc_id) and intersects:
         return True, hat_swc_id
@@ -196,7 +196,7 @@ def is_valid(gt_graph, pred_graph, kdtree, gt_id, proposal):
     if is_connected(hat_edge_i, hat_edge_j):
         # Orient ground truth edges
         hat_edge_xyz_i, hat_edge_xyz_j = orient_edges(
-            gt_graph.edges[hat_edge_i]["xyz"], 
+            gt_graph.edges[hat_edge_i]["xyz"],
             gt_graph.edges[hat_edge_j]["xyz"]
         )
 
@@ -303,5 +303,3 @@ def length_to_idx(xyz_list, idx):
     for i in range(0, idx):
         length += geometry_util.dist(xyz_list[i], xyz_list[i + 1])
     return length
-    
-        

src/deep_neurographs/utils/graph_util.py

@@ -34,8 +34,6 @@
 from scipy.spatial import KDTree
 from tqdm import tqdm
 
-import ast
-import multiprocessing
 import networkx as nx
 import numpy as np
 import os
@@ -61,7 +59,7 @@ def __init__(
         remove_high_risk_merges=False,
         segmentation_path=None,
         smooth_bool=True,
-        somas_path=None,
+        soma_centroids=None,
         verbose=False,
     ):
         """
@@ -90,9 +88,8 @@ def __init__(
         smooth_bool : bool, optional
             Indication of whether to smooth xyz coordinates from SWC files.
             The default is True.
-        somas_path : str, optional
-            Path to a txt file containing xyz coordinates of detected somas.
-            The default is None.
+        soma_centroids : List[Tuple[float]] or None, optional
+            Physcial coordinates of soma centroids. The default is None.
         verbose : bool, optional
             Indication of whether to display a progress bar while building
             FragmentsGraph. The default is True.
@@ -108,65 +105,54 @@ def __init__(
         self.node_spacing = node_spacing
         self.prune_depth = prune_depth
         self.smooth_bool = smooth_bool
-        self.soma_kdtree = None
+        self.soma_centroids = soma_centroids
         self.verbose = verbose
 
         # Set irreducibles extracter
-        if somas_path and remove_high_risk_merges:
+        if soma_centroids and remove_high_risk_merges:
             self.extracter = self.break_and_extract
         else:
             self.extracter = self.extract
 
         # Load somas
-        if segmentation_path and somas_path:
-            self.load_somas(segmentation_path, somas_path)
+        if segmentation_path and soma_centroids:
+            self.soma_kdtree = KDTree(self.soma_centroids)
+            self.ingest_somas(segmentation_path)
+        else:
+            self.soma_kdtree = None
 
-    def load_somas(self, segmentation_path, somas_path):
+    def ingest_somas(self, segmentation_path):
         """
-        Loads soma locations from a specified file and detects merges in a
-        segmentation.
+        Loads soma locations from a specified file and search for interestions
+        between soma locations and objects in segmentation mask.
 
         Parameters
         ----------
         segmentation_path : str
-            Path to segmentation stored in GCS bucket. The default is None.
-        somas_path : str
-            Path to a txt file containing xyz coordinates of detected somas.
+            Path to segmentation stored in GCS bucket.
 
         Returns
         -------
         None
 
         """
-        # Read soma locations txt file
-        if isinstance(somas_path, str):
-            xyz_list = util.read_txt(somas_path)
-        elif isinstance(somas_path, dict):
-            xyz_list = util.read_s3_txt_file(somas_path)
-        else:
-            raise Exception(f"Invalid format - somas_path={somas_path}")
-
-        # Process soma locations
         reader = img_util.TensorStoreReader(segmentation_path)
         with ThreadPoolExecutor() as executor:
             # Assign threads
             threads = list()
-            for xyz in  util.load_soma_locations(somas_path):
+            for xyz in self.soma_centroids:
                 voxel = img_util.to_voxels(xyz, (0.748, 0.748, 1.0))
                 threads.append(executor.submit(reader.read_voxel, voxel, xyz))
 
             # Store results
-            soma_xyz_list = list()
             for thread in as_completed(threads):
                 xyz, seg_id = thread.result()
                 if seg_id != 0:
                     self.id_to_soma[str(seg_id)].append(xyz)
-                    soma_xyz_list.append(xyz)
-        self.soma_kdtree = KDTree(soma_xyz_list)
 
         # Report results
         if self.verbose:
-            print("# Somas:", len(soma_xyz_list))
+            print("# Somas:", len(self.soma_centroids))
             print("# Soma-Fragment Intersections:", len(self.id_to_soma))
 
     # --- Irreducibles Extraction ---
@@ -198,7 +184,6 @@ def extract_irreducibles(self, swc_dicts):
         high_risk_cnt = 0
         desc = "Extract Graphs"
         pbar = tqdm(total=len(swc_dicts), desc=desc) if self.verbose else None
-        multiprocessing.set_start_method('spawn', force=True)
         with ProcessPoolExecutor() as executor:
             processes = list()
             while swc_dicts: