feat: save fragments (#163)

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

Author: anna-grim

…_metrics/data_handling/skeleton_graph.py …n_metrics/data_handling/graph_classes.pysrc/segmentation_skeleton_metrics/data_handling/skeleton_graph.py renamed to src/segmentation_skeleton_metrics/data_handling/graph_classes.py src/segmentation_skeleton_metrics/data_handling/skeleton_graph.py renamed to src/segmentation_skeleton_metrics/data_handling/graph_classes.py

@@ -19,7 +19,7 @@
 import numpy as np
 
 from segmentation_skeleton_metrics.data_handling import swc_loading
-from segmentation_skeleton_metrics.data_handling.skeleton_graph import (
+from segmentation_skeleton_metrics.data_handling.graph_classes import (
     FragmentGraph, LabeledGraph
 )
 from segmentation_skeleton_metrics.utils import util

src/segmentation_skeleton_metrics/data_handling/graph_loading.py

@@ -275,8 +275,52 @@ def report_summary(self, results):
 
     # --- Writers ---
     def save_fragments(self, gt_graphs, fragment_graphs):
-        assert fragment_graphs is not None
-        pass
+        """
+        Saves ground-truth graphs and their intersecting fragment graphs to
+        zipped SWC files.
+
+        Parameters
+        ----------
+        gt_graphs : Dict[str, LabeledGraph]
+            Graphs built from ground truth SWC files.
+        fragment_graphs : Dict[str, FragmentsGraph]
+            Graphs built from skeletons obtained from a segmentation.
+        """
+        # Initializations
+        fragments_dir = os.path.join(self.output_dir, "fragments")
+        util.mkdir(fragments_dir, delete=True)
+
+        # Main
+        for key, graph in gt_graphs.items():
+            # Create zip writer
+            zip_path = os.path.join(fragments_dir, f"{graph.name}.zip")
+            zip_writer = ZipFile(zip_path, "a")
+
+            # Save skeletons
+            graph.to_zipped_swc(zip_writer)
+            self.save_intersecting_fragments(
+                graph, fragment_graphs, zip_writer
+            )
+
+    @staticmethod
+    def save_intersecting_fragments(gt_graph, fragment_graphs, zip_writer):
+        """
+        Saves SWC files for all fragment graphs whose label intersects with
+        the given ground-truth graph.
+
+        Parameters
+        ----------
+        gt_graph : LabeledGraph
+            Graphs built from ground truth SWC files.
+        fragment_graphs : Dict[str, FragmentGraph]
+            Graphs built from skeletons obtained from a segmentation.
+        zip_writer : zipfile.ZipFile
+            Open ZIP file handle used to write fragments.
+        """
+        intersecting_labels = gt_graph.get_node_labels()
+        for key, graph in fragment_graphs.items():
+            if graph.label in intersecting_labels:
+                graph.to_zipped_swc(zip_writer)
 
     def save_merge_results(self, gt_graphs, fragment_graphs, output_dir):
         """
@@ -313,7 +357,7 @@ def save_merge_sites(self, zip_writer):
         Parameters
         ----------
         zip_writer : zipfile.ZipFile
-            Open ZIP file handle used to store merge site data.
+            Open ZIP file handle used to write merge site data.
         """
         merge_sites = self.metrics["# Merges"].merge_sites
         for i in range(len(merge_sites)):

src/segmentation_skeleton_metrics/evaluate.py

@@ -484,9 +484,9 @@ def find_merge_site(self, gt_graph, fragment_graph, source, visited):
 
     def verify_site(self, gt_graph, fragment_graph, gt_node, fragment_node):
         """
-        Verifies whether a given site in a fragment graph corresponds to a merge
-        mistake relative to the ground truth graph. If so, the site is saved
-        in an internal data structure.
+        Verifies whether a given site in a fragment graph corresponds to a
+        merge mistake relative to the ground truth graph. If so, the site is
+        saved in an internal data structure.
 
         Parameters
         ----------