refactor: feature generation (#262)

Co-authored-by: anna-grim <[email protected]>
AllenInstitute · Oct 8, 2024 · c772340 · c772340
1 parent 2a2a8a8
commit c772340
Show file tree

Hide file tree

Showing 11 changed files with 866 additions and 1,018 deletions.
diff --git a/src/deep_neurographs/config.py b/src/deep_neurographs/config.py
@@ -97,6 +97,7 @@ class MLConfig:
     threshold: float = 0.6
     model_type: str = "GraphNeuralNet"
     n_epochs: int = 1000
+    use_img_embedding: bool = False
     validation_split: float = 0.15
     weight_decay: float = 1e-3
 

diff --git a/src/deep_neurographs/generate_proposals.py b/src/deep_neurographs/generate_proposals.py
@@ -319,8 +319,8 @@ def run_trimming(neurograph, proposals, radius, progress_bar):
 
 def trim_endpoints(neurograph, i, j, radius):
     # Initializations
-    branch_i = neurograph.get_branch(i)
-    branch_j = neurograph.get_branch(j)
+    branch_i = neurograph.branch(i)
+    branch_j = neurograph.branch(j)
 
     # Check both orderings
     idx_i, idx_j = trim_endpoints_ordered(branch_i, branch_j)
@@ -394,8 +394,8 @@ def trim_to_idx(neurograph, i, idx):
 
     """
     # Update node
-    branch_xyz = neurograph.get_branch(i, key="xyz")
-    branch_radii = neurograph.get_branch(i, key="radius")
+    branch_xyz = neurograph.branch(i, key="xyz")
+    branch_radii = neurograph.branch(i, key="radius")
     neurograph.nodes[i]["xyz"] = branch_xyz[idx]
     neurograph.nodes[i]["radius"] = branch_radii[idx]