From a75e11c7f73a6654e7cf6def71395156fef10484 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Tue, 15 Oct 2024 21:17:33 +0000
Subject: [PATCH] refactor: updated reporting results

---
 src/deep_neurographs/inference.py  | 16 ++++++++++------
 src/deep_neurographs/neurograph.py |  3 +++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/deep_neurographs/inference.py b/src/deep_neurographs/inference.py
index 900ccd6..32f59d6 100644
--- a/src/deep_neurographs/inference.py
+++ b/src/deep_neurographs/inference.py
@@ -169,7 +169,6 @@ def run(self, fragments_pointer):
         # Finish
         self.report("Final Graph...")
         self.report_graph()
-        self.report("\n")
 
         t, unit = util.time_writer(time() - t0)
         self.report(f"Total Runtime: {round(t, 4)} {unit}\n")
@@ -231,12 +230,17 @@ def build_graph(self, fragments_pointer):
         # Save valid labels and current graph
         swcs_path = os.path.join(self.output_dir, "processed-swcs.zip")
         labels_path = os.path.join(self.output_dir, "valid_labels.txt")
-        self.graph.to_zipped_swcs(swcs_path)
+        n_saved = self.graph.to_zipped_swcs(swcs_path)
         self.graph.save_labels(labels_path)
+        self.report(f"# SWCs Saved: {n_saved}")
 
+        # Report runtime
         t, unit = util.time_writer(time() - t0)
+        self.report(f"Module Runtime: {round(t, 4)} {unit}")
+
+        # Report graph overview 
+        self.report("\nInitial Graph...")
         self.report_graph()
-        self.report(f"Module Runtime: {round(t, 4)} {unit}\n")
 
     def generate_proposals(self, radius=None):
         """
@@ -410,7 +414,7 @@ def report_graph(self):
         self.report(f"# Connected Components: {n_components}")
         self.report(f"# Nodes: {n_nodes}")
         self.report(f"# Edges: {n_edges}")
-        self.report(f"Memory Consumption: {usage} GBs")
+        self.report(f"Memory Consumption: {usage} GBs\n")
 
 
 class InferenceEngine:
@@ -571,9 +575,9 @@ def get_batch_dataset(self, neurograph, batch):
         ...
 
         """
-        t0 = time()
+        #t0 = time()
         features = self.feature_generator.run(neurograph, batch, self.radius)
-        print("Feature Generation:", time() - t0)
+        #print("Feature Generation:", time() - t0)
         computation_graph = batch["graph"] if type(batch) is dict else None
         dataset = ml_util.init_dataset(
             neurograph,
diff --git a/src/deep_neurographs/neurograph.py b/src/deep_neurographs/neurograph.py
index 50dda00..09d5047 100644
--- a/src/deep_neurographs/neurograph.py
+++ b/src/deep_neurographs/neurograph.py
@@ -955,10 +955,13 @@ def component_path_length(self, root):
     # --- write graph to swcs ---
     def to_zipped_swcs(self, zip_path, color=None, min_size=100):
         with zipfile.ZipFile(zip_path, "w") as zip_writer:
+            cnt = 0
             for nodes in nx.connected_components(self):
                 root = util.sample_once(nodes)
                 if self.component_path_length(root) > min_size:
                     self.to_zipped_swc(zip_writer, nodes, color)
+                    cnt += 1
+            return cnt
 
     def to_zipped_swc(self, zip_writer, nodes, color):
         with StringIO() as text_buffer: