refactor: updated anisotropy handling (#95)

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

demo/demo.py

@@ -1,5 +1,3 @@
-import os
-
 import numpy as np
 from tifffile import imread
 from xlwt import Workbook

demo/evaluation_results.xls

@@ -0,0 +1 @@
+ Label   -   xyz

demo/merged_ids-segmentation.txt

@@ -7,22 +7,26 @@
 
 """
 
-import os
-from collections import defaultdict
+
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from copy import deepcopy
+from scipy.spatial import KDTree
 from time import time
+from tqdm import tqdm
 from zipfile import ZipFile
 
 import networkx as nx
 import numpy as np
+import os
 import tensorstore as ts
-from scipy.spatial import KDTree
-from tqdm import tqdm
 
-from segmentation_skeleton_metrics import graph_utils as gutils
-from segmentation_skeleton_metrics import split_detection, swc_utils, utils
-from segmentation_skeleton_metrics.graph_utils import to_xyz_array
+from segmentation_skeleton_metrics import split_detection
+from segmentation_skeleton_metrics.utils import (
+    graph_util as gutil,
+    swc_util,
+    util
+)
+from segmentation_skeleton_metrics.utils.graph_util import to_array
 
 MERGE_DIST_THRESHOLD = 100
 MIN_CNT = 40
@@ -62,7 +66,7 @@ def __init__(
         Parameters
         ----------
         gt_pointer : dict/str/list[str]
-            Pointer to ground truth swcs, see "swc_utils.Reader" for further
+            Pointer to ground truth swcs, see "swc_util.Reader" for further
             documentation. Note these swc files are assumed to be stored in
             image coordinates.
         pred_labels : numpy.ndarray or tensorstore.TensorStore
@@ -75,12 +79,12 @@ def __init__(
             that were merged into a single segment. The default is None.
         fragments_pointer : dict/str/list[str], optional
             Pointer to fragments (i.e. swcs) corresponding to "pred_labels",
-            see "swc_utils.Reader" for further documentation. Note these swc
+            see "swc_util.Reader" for further documentation. Note these swc
             files may be stored in either world or image coordinates. If the
             swcs are stored in world coordinates, then provide the world to
             image coordinates anisotropy factor. Note the filename of each swc
             is assumed to "segment_id.swc" where segment_id cooresponds to the
-            segment id from "pred_labels". The default is None. 
+            segment id from "pred_labels". The default is None.
         output_dir : str, optional
             Path to directory that mistake sites are written to. The default
             is None.
@@ -103,7 +107,7 @@ def __init__(
 
         """
         # Instance attributes
-        self.anisotropy = [1.0 / a for a in anisotropy]
+        self.anisotropy = anisotropy
         self.connections_path = connections_path
         self.output_dir = output_dir
         self.preexisting_merges = preexisting_merges
@@ -142,7 +146,7 @@ def init_label_map(self, path):
         """
         if path:
             assert self.valid_labels is not None, "Must provide valid labels!"
-            self.label_map, self.inverse_label_map = utils.init_label_map(
+            self.label_map, self.inverse_label_map = util.init_label_map(
                 path, self.valid_labels
             )
         else:
@@ -166,14 +170,14 @@ def init_graphs(self, paths):
 
         """
         # Read graphs
-        self.graphs = swc_utils.Reader().load(paths)
+        self.graphs = swc_util.Reader().load(paths)
         self.fragment_graphs = None
 
         # Label nodes
         self.key_to_label_to_nodes = dict()  # {id: {label: nodes}}
         for key in tqdm(self.graphs, desc="Labeling Graphs"):
             self.set_node_labels(key)
-            self.key_to_label_to_nodes[key] = gutils.init_label_to_nodes(
+            self.key_to_label_to_nodes[key] = gutil.init_label_to_nodes(
                 self.graphs[key]
             )
 
@@ -318,7 +322,7 @@ def load_fragments(self, fragments_pointer):
 
         """
         # Read fragments
-        reader = swc_utils.Reader(anisotropy=self.anisotropy, min_size=40)
+        reader = swc_util.Reader(anisotropy=self.anisotropy, min_size=40)
         fragment_graphs = reader.load(fragments_pointer)
         self.fragment_ids = set(fragment_graphs.keys())
 
@@ -349,13 +353,13 @@ def init_zip_writer(self):
         """
         # Initialize output directory
         output_dir = os.path.join(self.output_dir, "projections")
-        utils.mkdir(output_dir)
+        util.mkdir(output_dir)
 
         # Save intial graphs
         self.zip_writer = dict()
         for key in self.graphs.keys():
             self.zip_writer[key] = ZipFile(f"{output_dir}/{key}.zip", "w")
-            swc_utils.to_zipped_swc(
+            swc_util.to_zipped_swc(
                 self.zip_writer[key], self.graphs[key],
             )
 
@@ -419,7 +423,7 @@ def adjust_metrics(self, key):
 
                 # Adjust metrics
                 n_edges = subgraph.number_of_edges()
-                rls = gutils.compute_run_lengths(subgraph)
+                rls = gutil.compute_run_lengths(subgraph)
                 self.graphs[key].graph["run_length"] -= np.sum(rls)
                 self.graphs[key].graph["n_edges"] -= n_edges
 
@@ -448,13 +452,13 @@ def detect_splits(self):
             graph = split_detection.run(graph, self.graphs[key])
 
             # Update graph by removing omits (i.e. nodes labeled 0)
-            self.graphs[key] = gutils.delete_nodes(graph, 0)
-            self.key_to_label_to_nodes[key] = gutils.init_label_to_nodes(
+            self.graphs[key] = gutil.delete_nodes(graph, 0)
+            self.key_to_label_to_nodes[key] = gutil.init_label_to_nodes(
                 self.graphs[key]
             )
 
         # Report runtime
-        t, unit = utils.time_writer(time() - t0)
+        t, unit = util.time_writer(time() - t0)
         print(f"Runtime: {round(t, 2)} {unit}\n")
 
     def quantify_splits(self):
@@ -478,7 +482,7 @@ def quantify_splits(self):
             n_pred_edges = self.graphs[key].number_of_edges()
             n_target_edges = self.graphs[key].graph["n_edges"]
 
-            self.split_cnt[key] = gutils.count_splits(self.graphs[key])
+            self.split_cnt[key] = gutil.count_splits(self.graphs[key])
             self.omit_cnts[key] = n_target_edges - n_pred_edges
             self.omit_percent[key] = 1 - n_pred_edges / n_target_edges
 
@@ -507,7 +511,7 @@ def detect_merges(self):
         if self.fragment_graphs:
             for key, graph in self.graphs.items():
                 if graph.number_of_nodes() > 0:
-                    kdtree = KDTree(gutils.to_xyz_array(graph))
+                    kdtree = KDTree(gutil.to_array(graph))
                     self.count_merges(key, kdtree)
 
         # Process merges
@@ -574,20 +578,20 @@ def is_fragment_merge(self, key, label, kdtree):
         None
 
         """
-        for xyz in to_xyz_array(self.fragment_graphs[label])[::5]:
-            if kdtree.query(xyz)[0] > MERGE_DIST_THRESHOLD:
+        for voxel in to_array(self.fragment_graphs[label])[::2]:
+            if kdtree.query(voxel)[0] > MERGE_DIST_THRESHOLD:
                 # Check whether to take inverse of label
                 if self.inverse_label_map:
                     equivalent_label = self.label_map[label]
                 else:
                     equivalent_label = label
 
                 # Record merge mistake
-                xyz = utils.to_world(xyz)
+                xyz = util.to_world(voxel)
                 self.merge_cnt[key] += 1
                 self.merged_labels.add((key, equivalent_label, tuple(xyz)))
                 if self.save_projections:
-                    swc_utils.to_zipped_swc(
+                    swc_util.to_zipped_swc(
                         self.zip_writer[key], self.fragment_graphs[label]
                     )
                 return
@@ -871,7 +875,7 @@ def compute_erl(self):
         total_run_length = 0
         for key in self.graphs:
             run_length = self.get_run_length(key)
-            run_lengths = gutils.compute_run_lengths(self.graphs[key])
+            run_lengths = gutil.compute_run_lengths(self.graphs[key])
             total_run_length += run_length
             wgt = run_lengths / max(np.sum(run_lengths), 1)
 
@@ -924,7 +928,7 @@ def list_metrics(self):
         ]
         return metrics
 
-    # -- Utils --
+    # -- util --
     def init_counter(self):
         """
         Initializes a dictionary that is used to count some type of mistake
@@ -943,7 +947,7 @@ def init_counter(self):
         return {key: 0 for key in self.graphs}
 
 
-# -- utils --
+# -- util --
 def find_sites(graphs, get_labels):
     """
     Detects merges between ground truth graphs which are considered to be

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -11,8 +11,7 @@
 
 import networkx as nx
 
-from segmentation_skeleton_metrics import graph_utils as gutils
-from segmentation_skeleton_metrics import utils
+from segmentation_skeleton_metrics.utils import graph_util as gutil, util
 
 
 def run(target_graph, labeled_graph):
@@ -33,7 +32,7 @@ def run(target_graph, labeled_graph):
         Labeled graph with omit and split edges removed.
 
     """
-    r = gutils.sample_leaf(labeled_graph)
+    r = gutil.sample_leaf(labeled_graph)
     dfs_edges = list(nx.dfs_edges(labeled_graph, source=r))
     while len(dfs_edges) > 0:
         # Visit edge
@@ -94,7 +93,7 @@ def is_zero_misalignment(target_graph, labeled_graph, dfs_edges, nb, root):
             # Add nbs to queue
             nbs = target_graph.neighbors(j)
             for k in [k for k in nbs if k not in visited]:
-                if utils.check_edge(dfs_edges, (j, k)):
+                if util.check_edge(dfs_edges, (j, k)):
                     queue.append(k)
                     dfs_edges = remove_edge(dfs_edges, (j, k))
                 elif k == nb:
@@ -103,7 +102,7 @@ def is_zero_misalignment(target_graph, labeled_graph, dfs_edges, nb, root):
     # Upd zero nodes
     if len(collision_labels) == 1:
         label = collision_labels.pop()
-        labeled_graph = gutils.upd_labels(labeled_graph, visited, label)
+        labeled_graph = gutil.upd_labels(labeled_graph, visited, label)
 
     return dfs_edges, labeled_graph
 
@@ -150,7 +149,7 @@ def is_nonzero_misalignment(target_graph, labeled_graph, dfs_edges, nb, root):
         visited.add(j)
         if label_j == origin_label and len(queue) == 0:
             # misalignment
-            labeled_graph = gutils.upd_labels(
+            labeled_graph = gutil.upd_labels(
                 labeled_graph, visited, origin_label
             )
             return dfs_edges, labeled_graph
@@ -171,7 +170,7 @@ def is_nonzero_misalignment(target_graph, labeled_graph, dfs_edges, nb, root):
     return dfs_edges, labeled_graph
 
 
-# -- utils --
+# -- util --
 def remove_edge(dfs_edges, edge):
     """
     Checks whether "edge" is in "dfs_edges" and removes it.

src/segmentation_skeleton_metrics/split_detection.py

@@ -14,7 +14,9 @@
 import numpy as np
 from scipy.spatial.distance import euclidean as get_dist
 
-from segmentation_skeleton_metrics import utils
+from segmentation_skeleton_metrics.utils import util
+
+ANISOTROPY = (0.748, 0.748, 1.0)
 
 
 # --- Update graph ---
@@ -151,14 +153,14 @@ def compute_run_length(graph):
     """
     path_length = 0
     for i, j in nx.dfs_edges(graph):
-        xyz_1 = utils.to_world(graph.nodes[i]["xyz"])
-        xyz_2 = utils.to_world(graph.nodes[j]["xyz"])
+        xyz_1 = util.to_physical(graph.nodes[i]["xyz"], ANISOTROPY)
+        xyz_2 = util.to_physical(graph.nodes[j]["xyz"], ANISOTROPY)
         path_length += get_dist(xyz_1, xyz_2)
     return path_length
 
 
 # -- miscellaneous --
-def to_xyz_array(graph):
+def to_array(graph):
     """
     Converts node coordinates from a graph into a NumPy array.