From 551e4121afaf74eccc19a5b8dbcf7cb5495625eb Mon Sep 17 00:00:00 2001
From: Anna Grim <108307071+anna-grim@users.noreply.github.com>
Date: Tue, 9 Apr 2024 14:48:34 -0700
Subject: [PATCH] upds (#109)

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>
---
 src/deep_neurographs/intake.py                |  4 +---
 .../machine_learning/inference.py             |  6 +++---
 src/deep_neurographs/reconstruction.py        | 19 ++++++-------------
 3 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py
index ca00b72..2b10d59 100644
--- a/src/deep_neurographs/intake.py
+++ b/src/deep_neurographs/intake.py
@@ -243,7 +243,7 @@ def download_gcs_zips(bucket_name, gcs_path, min_size, anisotropy):
     # Initializations
     bucket = storage.Client().bucket(bucket_name)
     zip_paths = utils.list_gcs_filenames(bucket, gcs_path, ".zip")
-    chunk_size = int(len(zip_paths) * 0.05)
+    chunk_size = int(len(zip_paths) * 0.02)
 
     # Parse
     cnt = 1
@@ -259,8 +259,6 @@ def download_gcs_zips(bucket_name, gcs_path, min_size, anisotropy):
             cnt, t1 = report_progress(
                 i, len(zip_paths), chunk_size, cnt, t0, t1
             )
-            break
-
     return swc_dicts
 
 
diff --git a/src/deep_neurographs/machine_learning/inference.py b/src/deep_neurographs/machine_learning/inference.py
index 10077fb..40fc886 100644
--- a/src/deep_neurographs/machine_learning/inference.py
+++ b/src/deep_neurographs/machine_learning/inference.py
@@ -113,7 +113,7 @@ def run_without_seeds(
     for i, batch in enumerate(batches):
         # Prediction
         proposals_i = [proposals[j] for j in batch]
-        accepts_i = predict(
+        accepts_i, graph = predict(
             neurograph,
             graph,
             img_path,
@@ -160,7 +160,7 @@ def predict(
 
     # Run model
     proposal_probs = run_model(dataset, model, model_type)
-    accepts = build.get_accepted_proposals(
+    accepts, graph = build.get_accepted_proposals(
         neurograph,
         graph,
         proposal_probs,
@@ -168,7 +168,7 @@ def predict(
         high_threshold=0.95,
         low_threshold=confidence_threshold,
     )
-    return accepts
+    return accepts, graph
 
 
 def build_from_soma(
diff --git a/src/deep_neurographs/reconstruction.py b/src/deep_neurographs/reconstruction.py
index 5f1edef..9f487e9 100644
--- a/src/deep_neurographs/reconstruction.py
+++ b/src/deep_neurographs/reconstruction.py
@@ -115,7 +115,7 @@ def get_structure_aware_accepts(
 ):
     # Add best preds
     best_preds, best_probs = get_best_preds(neurograph, preds, high_threshold)
-    accepts = check_cycles_sequential(graph, best_preds, best_probs)
+    accepts, graph = check_cycles_sequential(graph, best_preds, best_probs)
     if len(best_preds) == len(preds.keys()):
         return accepts
 
@@ -128,9 +128,9 @@ def get_structure_aware_accepts(
             good_preds.append(edge)
             good_probs.append(prob)
 
-    more_accepts = check_cycles_sequential(graph, good_preds, good_probs)
-    accepts.extend(more_accepts)
-    return accepts
+    more_accepts, graph = check_cycles_sequential(graph, good_preds, good_probs)
+    accepts.extend(more_accepts)    
+    return accepts, graph
 
 
 def get_subgraphs(graph, edge):
@@ -203,7 +203,7 @@ def check_cycles_sequential(graph, edges, probs):
         if not created_cycle:
             graph.add_edges_from([tuple(edges[i])])
             accepts.append(edges[i])
-    return accepts
+    return accepts, graph
 
 
 def get_best_preds(neurograph, preds, threshold):
@@ -220,14 +220,7 @@ def fuse_branches(neurograph, edges):
     simple_cnt = 0
     complex_cnt = 0
     for edge in edges:
-        if neurograph.is_simple(edge):
-            simple_cnt += 1
-            neurograph.merge_proposal(edge)
-        else:
-            complex_cnt += 1
-            #print("merge not implemented for complex")
-    print("# simple:", simple_cnt)
-    print("# complex:", complex_cnt)
+        neurograph.merge_proposal(edge)
     return neurograph