refactor: update (#96)

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

Author: anna-grim

demo/evaluation_results.xls

@@ -26,7 +26,6 @@
     swc_util,
     util
 )
-from segmentation_skeleton_metrics.utils.graph_util import to_array
 
 MERGE_DIST_THRESHOLD = 100
 MIN_CNT = 40
@@ -201,13 +200,15 @@ def set_node_labels(self, key):
             # Assign threads
             threads = []
             for i in self.graphs[key].nodes:
-                voxel = tuple(self.graphs[key].nodes[i]["xyz"])
+                voxel = tuple(self.graphs[key].nodes[i]["voxel"])
                 threads.append(executor.submit(self.get_label, i, voxel))
 
             # Store label
+            pbar = tqdm(total=len(threads), desc="threads finished")
             for thread in as_completed(threads):
                 i, label = thread.result()
                 self.graphs[key].nodes[i].update({"label": label})
+                pbar.update(1)
 
     def get_label(self, i, voxel):
         """
@@ -227,7 +228,7 @@ def get_label(self, i, voxel):
 
         """
         # Read label
-        if type(self.label_mask) is ts.TensorStore:
+        if isinstance(self.label_mask, ts.TensorStore):
             label = int(self.label_mask[voxel].read().result())
         else:
             label = self.label_mask[voxel]
@@ -578,16 +579,16 @@ def is_fragment_merge(self, key, label, kdtree):
         None
 
         """
-        for voxel in to_array(self.fragment_graphs[label])[::2]:
+        for voxel in gutil.to_array(self.fragment_graphs[label])[::2]:
             if kdtree.query(voxel)[0] > MERGE_DIST_THRESHOLD:
-                # Check whether to take inverse of label
+                # Check whether to get inverse of label
                 if self.inverse_label_map:
                     equivalent_label = self.label_map[label]
                 else:
                     equivalent_label = label
 
                 # Record merge mistake
-                xyz = util.to_world(voxel)
+                xyz = util.to_physical(voxel)
                 self.merge_cnt[key] += 1
                 self.merged_labels.add((key, equivalent_label, tuple(xyz)))
                 if self.save_projections:

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -153,8 +153,8 @@ def compute_run_length(graph):
     """
     path_length = 0
     for i, j in nx.dfs_edges(graph):
-        xyz_1 = util.to_physical(graph.nodes[i]["xyz"], ANISOTROPY)
-        xyz_2 = util.to_physical(graph.nodes[j]["xyz"], ANISOTROPY)
+        xyz_1 = util.to_physical(graph.nodes[i]["voxel"], ANISOTROPY)
+        xyz_2 = util.to_physical(graph.nodes[j]["voxel"], ANISOTROPY)
         path_length += get_dist(xyz_1, xyz_2)
     return path_length
 
@@ -167,16 +167,16 @@ def to_array(graph):
     Parameters
     ----------
     graph : networkx.Graph
-        Graph that contains nodes with "xyz" attributes.
+        Graph that contains nodes with "voxel" attributes.
 
     Returns
     -------
     numpy.ndarray
         Array where each row represents the 3D coordinates of a node.
 
     """
-    xyz_coords = nx.get_node_attributes(graph, "xyz")
-    return np.array([xyz_coords[i] for i in graph.nodes])
+    voxels = nx.get_node_attributes(graph, "voxel")
+    return np.array([voxels[i] for i in graph.nodes])
 
 
 def sample_leaf(graph):

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -60,9 +60,7 @@ def __init__(self, anisotropy=(1.0, 1.0, 1.0), min_size=0):
             (1.0, 1.0, 1.0).
         min_size : int, optional
             Threshold on the number of nodes in swc file. Only swc files with
-            more than "min_size" nodes are stored in "xyz_coords". The default
-            is 0.
-
+            more than "min_size" nodes are returned.
         Returns
         -------
         None
@@ -157,8 +155,10 @@ def load_from_local_paths(self, swc_paths):
 
             # Store results
             graph_dict = dict()
+            pbar = tqdm(total=len(filenames), desc="Label Graph")
             for thread in as_completed(threads):
                 graph_dict.update(thread.result())
+                pbar.update(1)
         return graph_dict
 
     def load_from_local_path(self, path):
@@ -336,7 +336,7 @@ def process_content(self, content, filename):
             return {name: graph}
         return dict()
 
-    def read_xyz(self, xyz_str, offset):
+    def read_voxel(self, xyz_str, offset):
         """
         Reads the coordinates from a string, then transforms them to image
         coordinates (if applicable).
@@ -378,13 +378,13 @@ def get_graph(self, content):
         for line in content:
             if line.startswith("# OFFSET"):
                 parts = line.split()
-                offset = self.read_xyz(parts[2:5])
+                offset = self.read_voxel(parts[2:5])
             if not line.startswith("#"):
                 parts = line.split()
                 child = int(parts[0])
                 parent = int(parts[-1])
-                xyz = self.read_xyz(parts[2:5], offset=offset)
-                graph.add_node(child, xyz=xyz)
+                voxel = self.read_voxel(parts[2:5], offset=offset)
+                graph.add_node(child, voxel=voxel)
                 if parent != -1:
                     graph.add_edge(parent, child)
 
@@ -432,30 +432,29 @@ def save(path, xyz_1, xyz_2, color=None):
         f.write(make_entry(2, 1, xyz_2))
 
 
-def make_entry(node_id, parent_id, xyz):
+def make_entry(node, parent, xyz):
     """
-    Makes an entry to be written in an swc file.
+    Makes an entry to be written in an SWC file.
 
     Parameters
     ----------
     graph : networkx.Graph
         Graph that "node_id" and "parent_id" belong to.
-    node_id : int
-        Node that entry corresponds to.
+    node : int
+        Node ID that entry corresponds to.
     parent_id : int
-         Parent of node "node_id".
-    xyz : ...
-        xyz coordinate of "node_id".
+         Parent ID of the given node.
+    voxel : ...
+        Voxel coordinate of the given node.
 
     Returns
     -------
     entry : str
         Entry to be written in an swc file.
 
     """
-    x, y, z = tuple(util.to_world(xyz))
-    entry = f"{node_id} 2 {x} {y} {z} 3 {parent_id}"
-    return entry
+    x, y, z = tuple(util.to_physical(voxel, (0.748, 0.748, 1.0)))
+    return f"{node} 2 {x} {y} {z} 3 {parent}"
 
 
 def to_zipped_swc(zip_writer, graph, color=None):
@@ -488,7 +487,7 @@ def to_zipped_swc(zip_writer, graph, color=None):
         r = 5 if color else 3
         for i, j in nx.dfs_edges(graph):
             # Special Case: Root
-            x, y, z = tuple(util.to_world(graph.nodes[i]["xyz"]))
+            x, y, z = tuple(util.to_world(graph.nodes[i]["voxel"]))
             if n_entries == 0:
                 parent = -1
                 node_to_idx[i] = 1

src/segmentation_skeleton_metrics/utils/swc_util.py

@@ -461,4 +461,7 @@ def load_valid_labels(path):
         reconstruction process.
 
     """
-    return set(map(int, read_txt(path)))
+    valid_labels = set()
+    for label_str in read_txt(path):
+        valid_labels.add(int(label_str.split(".")[0]))
+    return valid_labels