major refactor

Author: anna-grim

src/segmentation_skeleton_metrics/evaluate.py

@@ -8,10 +8,28 @@
 
 """
 
-from segmentation_skeleton_metrics.skeleton_metric import SkeletonMetric
+from zipfile import ZipFile
+
+import numpy as np
+import os
+import pandas as pd
+
+from segmentation_skeleton_metrics.skeleton_metrics import (
+    MergeCountMetric,
+    MergeRateMetric,
+    MergedEdgePercentMetric,
+    OmitEdgePercentMetric,
+    SplitEdgePercentMetric,
+    SplitCountMetric,
+    SplitRateMetric,
+    EdgeAccuracyMetric,
+    ERLMetric,
+    NormalizedERLMetric
+)
 from segmentation_skeleton_metrics.data_handling.graph_loading import (
     DataLoader, LabelHandler
 )
+from segmentation_skeleton_metrics.utils import util
 
 
 def evaluate(
@@ -21,6 +39,7 @@ def evaluate(
     anisotropy=(1.0, 1.0, 1.0),
     connections_path=None,
     fragments_pointer=None,
+    results_filename="results",
     save_merges=False,
     save_fragments=False,
     use_anisotropy=False,
@@ -56,27 +75,117 @@ def evaluate(
         anisotropy=anisotropy,
         use_anisotropy=use_anisotropy,
         verbose=verbose
-    )    
+    )
     gt_graphs = dataloader.load_groundtruth(gt_pointer, label_mask)
     fragment_graphs = dataloader.load_fragments(fragments_pointer, gt_graphs)
 
-    # Evaluator
-    skeleton_metric = SkeletonMetric(
-        gt_graphs,
-        fragment_graphs,
-        label_handler,
-        anisotropy=anisotropy,
-        output_dir=output_dir,
-        save_merges=save_merges,
-        save_fragments=save_fragments,
-    )
-    skeleton_metric.run()
+    # Run evaluation
+    evaluator = Evaluator(output_dir, results_filename, verbose)
+    evaluator.run(gt_graphs, fragment_graphs)
 
-    # Compute metrics
+    # Optional saves
+    if save_merges:
+        evaluator.save_merge_results()
 
-    # Report results
+    if save_fragments and fragment_graphs:
+        evaluator.save_fragments(gt_graphs, fragment_graphs)
 
 
 # --- Evaluator ---
 class Evaluator:
-    pass
+
+    def __init__(self, output_dir, results_filename, verbose=True):
+        # Instance attributes
+        self.output_dir = output_dir
+        self.results_filename = results_filename
+        self.verbose = verbose
+
+        # Set core metrics
+        self.metrics = {
+            "# Splits": SplitCountMetric(verbose=verbose),
+            "# Merges": MergeCountMetric(verbose=verbose),
+            "% Split Edges": SplitEdgePercentMetric(verbose=verbose),
+            "% Omit Edges": OmitEdgePercentMetric(verbose=verbose),
+            "% Merged Edges": MergedEdgePercentMetric(verbose=verbose),
+            "ERL": ERLMetric(verbose=verbose)
+        }
+
+        # Set derived metrics
+        self.derived_metrics = {
+            "Normalized ERL": NormalizedERLMetric(verbose=verbose),
+            "Edge Accuracy": EdgeAccuracyMetric(verbose=verbose),
+            "Split Rate": SplitRateMetric(verbose=verbose),
+            "Merge Rate": MergeRateMetric(verbose=verbose),
+        }
+
+    # --- Core Routines ---
+    def run(self, gt_graphs, fragment_graphs=None):
+        # Printout step
+        if self.verbose:
+            print("\n(3) Compute Metrics")
+
+        # Compute core metrics
+        results = self.init_results(gt_graphs)
+        for name, metric in self.metrics.items():
+            if name == "# Merges" and fragment_graphs:
+                results[name] = metric.compute(gt_graphs, fragment_graphs)
+            elif name != "# Merges":
+                results.update(metric.compute(gt_graphs))
+
+        # Compute derived metrics
+        for name, metric in self.derived_metrics.items():
+            if name == "Merge Rate" and fragment_graphs:
+                results[name] = metric.compute(gt_graphs, results)
+            elif name != "Merge Rate":
+                results[name] = metric.compute(gt_graphs, results)
+
+        # Save report
+        path = f"{self.output_dir}/{self.results_filename}.csv"
+        results.to_csv(path, index=False)
+        self.report_summary(results)
+
+    def init_results(self, gt_graphs):
+        cols = list(self.metrics.keys()) + list(self.derived_metrics.keys())
+        index = list(gt_graphs.keys())
+        index.sort()
+        return pd.DataFrame(np.nan, index=index, columns=cols)
+
+    def report_summary(self, results):
+        pass
+
+    # --- Writers ---
+    def save_fragments(self):
+        pass
+
+    def save_merge_results(self, gt_graphs, fragment_graphs, output_dir):
+        # Initialize a writer
+        zip_path = os.path.join(output_dir, "merged_fragments.zip")
+        util.rm_file(zip_path)
+        zip_writer = ZipFile(zip_path, "a")
+
+        # Save SWC files
+        self.save_merge_sites(zip_writer)
+        self.save_skeletons_with_merge(gt_graphs, fragment_graphs, zip_writer)
+        zip_writer.close()
+
+        # Save CSV file
+        path = os.path.join(output_dir, "merge_sites.csv")
+        self.merge_sites.to_csv(path, index=True)
+
+    def save_merge_sites(self, zip_writer):
+        merge_sites = self.metrics["# Merges"].merge_sites
+        for i in range(len(merge_sites)):
+            filename = merge_sites.index[i]
+            xyz = merge_sites["World"].iloc[i]
+            util.to_zipped_point(zip_writer, filename, xyz)
+
+    def save_skeletons_with_merge(
+        self, gt_graphs, fragment_graphs, zip_writer
+    ):
+        # Save ground truth skeletons
+        for key in self.merge_sites["GroundTruth_ID"].unique():
+            gt_graphs[key].to_zipped_swc(zip_writer)
+
+        # Save fragments
+        for key in self.fragments_with_merge:
+            fragment_graphs[key].to_zipped_swc(zip_writer)