refactor: valid labels (#51)

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

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -18,6 +18,7 @@
 
 from segmentation_skeleton_metrics import graph_utils as gutils
 from segmentation_skeleton_metrics import split_detection, utils
+from segmentation_skeleton_metrics import swc_utils
 from segmentation_skeleton_metrics.swc_utils import (
     get_xyz_coords,
     save,
@@ -49,65 +50,88 @@ def __init__(
         pred_swc_paths,
         target_swc_paths,
         anisotropy=[1.0, 1.0, 1.0],
-        ignore_boundary_mistakes=False,
         black_holes_xyz_id=None,
         black_hole_radius=24,
         equivalent_ids=None,
-        valid_ids=None,
-        write_to_swc=False,
+        ignore_boundary_mistakes=False,
         output_dir=None,
+        pred_on_cloud=False,
+        valid_size_threshold=40.0,
+        write_to_swc=False,
     ):
         """
         Constructs skeleton metric object that evaluates the quality of a
         predicted segmentation.
 
         Parameters
         ----------
+        pred_labels : numpy.ndarray or tensorstore.TensorStore
+            Predicted segmentation mask.
+        pred_swc_paths : list[str]
+            List of paths to swc files where each file corresponds to a
+            neuron in the prediction.
         target_swc_paths : list[str]
-            List of paths to swc files such that each file corresponds to a
+            List of paths to swc files where each file corresponds to a
             neuron in the ground truth.
-        labels : numpy.ndarray or tensorstore.TensorStore
-            Predicted segmentation mask.
         anisotropy : list[float], optional
             Image to real-world coordinates scaling factors applied to swc
             files. The default is [1.0, 1.0, 1.0]
-        black_holes_xyz_id : list
+        black_holes_xyz_id : list, optional
             ...
-        black_hole_radius : float
+        black_hole_radius : float, optional
             ...
         equivalent_ids : ...
             ...
-        valid_ids : set
+        ignore_boundary_mistakes : bool, optional
+            Indication of whether to ignore mistakes near boundary of bounding
+            box. The default is False.
+        output_dir : str, optional
+            Path to directory that each mistake site is written to. The default
+            is None.
+        pred_on_cloud : bool, optional
+            Indication of whether predicted swc files in "pred_swc_paths" are
+            on the cloud in a GCS bucket. The default is False.
+        valid_size_threshold : float, optional
             ...
+        write_to_swc : bool, optional
+            Indication of whether to write mistake sites to an swc file. The
+            default is False.
 
         Returns
         -------
         None.
 
         """
-        # Store label options
-        self.valid_ids = valid_ids
-        self.label_mask = pred_labels
-
+        # Store options
         self.anisotropy = anisotropy
         self.ignore_boundary_mistakes = ignore_boundary_mistakes
+        self.output_dir = output_dir
+        self.write_to_swc = write_to_swc
+
         self.init_black_holes(black_holes_xyz_id)
         self.black_hole_radius = black_hole_radius
 
-        self.write_to_swc = write_to_swc
-        self.output_dir = output_dir
-
         # Build Graphs
+        self.label_mask = pred_labels
+        self.pred_swc_paths = pred_swc_paths
+        self.init_valid_labels(valid_size_threshold)
+
         self.target_graphs = self.init_graphs(target_swc_paths, anisotropy)
         self.labeled_target_graphs = self.init_labeled_target_graphs()
-        self.pred_swc_paths = pred_swc_paths
 
         # Build kdtree
         self.init_xyz_to_id_node()
         self.init_kdtree()
         self.rm_spurious_intersections()
 
     # -- Initialize and Label Graphs --
+    def init_valid_labels(self, valid_size_threshold):
+        self.valid_labels = set()
+        for path in self.pred_swc_paths:
+            contents = swc_utils.read(path)
+            if len(contents) > valid_size_threshold:
+                self.valid_labels.add(utils.get_swc_id(path))
+
     def init_graphs(self, paths, anisotropy):
         """
         Initializes "self.target_graphs" by iterating over "paths" which
@@ -247,7 +271,7 @@ def __read_label(self, coord):
     def is_valid(self, label):
         """
         Validates label by checking whether it is contained in
-        "self.valid_ids".
+        "self.valid_labels".
 
         Parameters
         ----------
@@ -258,12 +282,12 @@ def is_valid(self, label):
         -------
         label : int
             There are two possibilities: (1) original label if either "label"
-            is contained in "self.valid_ids" or "self.valid_labels" is None,
-            or (2) 0 if "label" is not contained in self.valid_ids.
+            is contained in "self.valid_labels" or "self.valid_labels" is
+            None, or (2) 0 if "label" is not contained in self.valid_labels.
 
         """
-        if self.valid_ids:
-            if label not in self.valid_ids:
+        if self.valid_labels:
+            if label not in self.valid_labels:
                 return 0
         return label
 

src/segmentation_skeleton_metrics/swc_utils.py

@@ -13,100 +13,31 @@
 from segmentation_skeleton_metrics import utils
 
 
-def make_entry(node_id, parent_id, xyz):
-    """
-    Makes an entry to be written in an swc file.
+def read(path, cloud_read=False):
+    return read_from_cloud(path) if cloud_read else read_from_local(path)
 
-    Parameters
-    ----------
-    graph : networkx.Graph
-        Graph that "node_id" and "parent_id" belong to.
-    node_id : int
-        Node that entry corresponds to.
-    parent_id : int
-         Parent of node "node_id".
-    xyz : ...
-        xyz coordinate of "node_id".
 
-    Returns
-    -------
-    entry : str
-        Entry to be written in an swc file.
-
-    """
-    x, y, z = tuple(xyz)
-    entry = f"{node_id} 2 {x} {y} {z} 8 {parent_id}"
-    return entry
+def read_from_cloud(path):
+    pass
 
 
-def save(path, xyz_1, xyz_2, color=None):
+def read_from_local(path):
     """
-    Writes an swc file.
+    Reads swc file stored at "path" on local machine.
 
     Parameters
     ----------
-    path : str
-        Path on local machine that swc file will be written to.
-    entry_list : list[list]
-        List of entries that will be written to an swc file.
-    color : str, optional
-        Color of nodes. The default is None.
-
-    Returns
-    -------
-    None.
-
-    """
-    with open(path, "w") as f:
-        # Preamble
-        if color is not None:
-            f.write("# COLOR " + color)
-        else:
-            f.write("# id, type, z, y, x, r, pid")
-        f.write("\n")
-
-        # Entries
-        f.write(make_entry(1, -1, xyz_1))
-        f.write("\n")
-        f.write(make_entry(2, 1, xyz_2))
-
-
-def to_graph(path, anisotropy=[1.0, 1.0, 1.0]):
-    """
-    Reads an swc file and builds an undirected graph from it.
-
-    Parameters
-    ----------
-    path : str
+    Path : str
         Path to swc file to be read.
-    anisotropy : list[float], optional
-        Image to real-world coordinates scaling factors for (x, y, z) that is
-        applied to swc files. The default is [1.0, 1.0, 1.0].
 
     Returns
     -------
-    networkx.Graph
-        Graph constructed from an swc file.
+    list
+        List such that each entry is a line from the swc file.
 
     """
-    graph = nx.Graph(swc_id=utils.get_swc_id(path))
-    with open(path, "r") as f:
-        offset = [0, 0, 0]
-        for line in f.readlines():
-            if line.startswith("# OFFSET"):
-                parts = line.split()
-                offset = read_xyz(parts[2:5])
-            if not line.startswith("#"):
-                parts = line.split()
-                child = int(parts[0])
-                parent = int(parts[-1])
-                xyz = read_xyz(
-                    parts[2:5], anisotropy=anisotropy, offset=offset
-                )
-                graph.add_node(child, xyz=xyz)
-                if parent != -1:
-                    graph.add_edge(parent, child)
-    return graph
+    with open(path, "r") as file:
+        return file.readlines()
 
 
 def get_xyz_coords(path, anisotropy=[1.0, 1.0, 1.0]):
@@ -167,3 +98,96 @@ def read_xyz(xyz, anisotropy=[1.0, 1.0, 1.0], offset=[0, 0, 0]):
     """
     xyz = [float(xyz[i]) + offset[i] for i in range(3)]
     return np.array([xyz[i] / anisotropy[i] for i in range(3)], dtype=int)
+
+
+def save(path, xyz_1, xyz_2, color=None):
+    """
+    Writes an swc file.
+
+    Parameters
+    ----------
+    path : str
+        Path on local machine that swc file will be written to.
+    entry_list : list[list]
+        List of entries that will be written to an swc file.
+    color : str, optional
+        Color of nodes. The default is None.
+
+    Returns
+    -------
+    None.
+
+    """
+    with open(path, "w") as f:
+        # Preamble
+        if color is not None:
+            f.write("# COLOR " + color)
+        else:
+            f.write("# id, type, z, y, x, r, pid")
+        f.write("\n")
+
+        # Entries
+        f.write(make_entry(1, -1, xyz_1))
+        f.write("\n")
+        f.write(make_entry(2, 1, xyz_2))
+
+
+def make_entry(node_id, parent_id, xyz):
+    """
+    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.
+    parent_id : int
+         Parent of node "node_id".
+    xyz : ...
+        xyz coordinate of "node_id".
+
+    Returns
+    -------
+    entry : str
+        Entry to be written in an swc file.
+
+    """
+    x, y, z = tuple(xyz)
+    entry = f"{node_id} 2 {x} {y} {z} 8 {parent_id}"
+    return entry
+
+
+def to_graph(path, anisotropy=[1.0, 1.0, 1.0]):
+    """
+    Reads an swc file and builds an undirected graph from it.
+
+    Parameters
+    ----------
+    path : str
+        Path to swc file to be read.
+    anisotropy : list[float], optional
+        Image to real-world coordinates scaling factors for (x, y, z) that is
+        applied to swc files. The default is [1.0, 1.0, 1.0].
+
+    Returns
+    -------
+    networkx.Graph
+        Graph built from an swc file.
+
+    """
+    graph = nx.Graph(swc_id=utils.get_swc_id(path))
+    offset = [0, 0, 0]    
+    for line in read(path):
+        if line.startswith("# OFFSET"):
+            parts = line.split()
+            offset = read_xyz(parts[2:5])
+        if not line.startswith("#"):
+            parts = line.split()
+            child = int(parts[0])
+            parent = int(parts[-1])
+            xyz = read_xyz(parts[2:5], anisotropy=anisotropy, offset=offset)
+            graph.add_node(child, xyz=xyz)
+            if parent != -1:
+                graph.add_edge(parent, child)
+    return graph