Refactor integrate upds

src/deep_neurographs/geometry.py

@@ -20,21 +20,38 @@ def get_directional(neurograph, i, origin, window_size):
     branches = neurograph.get_branches(i)
     branches = translate_branches(branches, origin)
     if len(branches) == 1:
-        return compute_tangent(get_sub_branch(branches[0], window_size))
+        return compute_tangent(get_subarray(branches[0], window_size))
     elif len(branches) == 2:
-        branch_1 = get_sub_branch(branches[0], window_size)
-        branch_2 = get_sub_branch(branches[1], window_size)
+        branch_1 = get_subarray(branches[0], window_size)
+        branch_2 = get_subarray(branches[1], window_size)
         branch = np.concatenate((branch_1, branch_2))
         return compute_tangent(branch)
     else:
         return np.array([0, 0, 0])
 
 
-def get_sub_branch(branch, window_size):
-    if branch.shape[0] < window_size:
-        return branch
+def get_subarray(arr, window_size):
+    """
+    Extracts a sub-array of a specified window size from a given input array.
+
+    Parameters
+    ----------
+    branch : numpy.ndarray
+        Array from which the sub-branch will be extracted.
+    window_size : int
+        Size of the window to extract from "arr".
+
+    Returns
+    -------
+    numpy.ndarray
+        A sub-array of the specified window size. If the input array is
+        smaller than the window size, the entire branch array is returned.
+
+    """
+    if arr.shape[0] < window_size:
+        return arr
     else:
-        return branch[0:window_size, :]
+        return arr[0:window_size, :]
 
 
 def compute_svd(xyz):
@@ -65,6 +82,22 @@ def compute_svd(xyz):
 
 
 def compute_tangent(xyz):
+    """
+    Computes the tangent vector at a given point or along a curve defined by
+    an array of points.
+
+    Parameters
+    ----------
+    xyz : numpy.ndarray
+        Array containing either two xyz coordinates or an arbitrary number of
+        defining a curve.
+
+    Returns
+    -------
+    numpy.ndarray
+        Tangent vector at the specified point or along the curve.
+
+    """
     if xyz.shape[0] == 2:
         tangent = (xyz[1] - xyz[0]) / dist(xyz[1], xyz[0])
     else:
@@ -74,6 +107,21 @@ def compute_tangent(xyz):
 
 
 def compute_normal(xyz):
+    """
+    Computes the normal vector of a plane defined by an array of xyz
+    coordinates using Singular Value Decomposition (SVD).
+
+    Parameters
+    ----------
+    xyz : numpy.ndarray
+        An array of xyz coordinates that normal vector is to be computed of.
+
+    Returns
+    -------
+    numpy.ndarray
+        The normal vector of the array "xyz".
+
+    """
     U, S, VT = compute_svd(xyz)
     normal = VT[-1]
     return normal / np.linalg.norm(normal)
@@ -150,16 +198,6 @@ def fit_spline(xyz, s=None):
     return spline_x, spline_y, spline_z
 
 
-def sample_path(path, n_points):
-    if len(path) > 5:
-        t = np.linspace(0, 1, n_points)
-        spline_x, spline_y, spline_z = fit_spline(path)
-        path = np.column_stack((spline_x(t), spline_y(t), spline_z(t)))
-    else:
-        path = make_line(path[0], path[-1], 10)
-    return path.astype(int)
-
-
 # Image feature extraction
 def get_profile(img, xyz_arr, process_id=None, window=[5, 5, 5]):
     profile = []

src/deep_neurographs/machine_learning/graph_datasets.py

@@ -78,6 +78,7 @@ class GraphDataset:
     Custom dataset for homogenous graphs.
 
     """
+
     def __init__(
         self,
         neurograph,
@@ -106,6 +107,7 @@ class HeteroGraphDataset:
     Custom dataset for heterogenous graphs.
 
     """
+
     def __init__(
         self,
         neurograph,

src/deep_neurographs/machine_learning/graph_models.py

@@ -11,14 +11,14 @@
 import torch
 import torch.nn.functional as F
 from torch.nn import ELU, Linear
-from torch_geometric.nn import GCNConv
+from torch_geometric.nn import GATConv, GCNConv
 
 
 class GCN(torch.nn.Module):
     def __init__(self, input_channels):
         super().__init__()
-        self.conv1 = GCNConv(input_channels, input_channels // 2)
-        self.conv2 = GCNConv(input_channels // 2, input_channels // 2)
+        self.conv1 = GCNConv(input_channels, input_channels)
+        self.conv2 = GCNConv(input_channels, input_channels // 2)
         self.conv3 = GCNConv(input_channels // 2, 1)
         self.ELU = ELU()
 
@@ -38,11 +38,35 @@ def forward(self, x, edge_index):
         return x
 
 
+class GAT(torch.nn.Module):
+    def __init__(self, input_channels):
+        super().__init__()
+        self.conv1 = GATConv(input_channels, input_channels)
+        self.conv2 = GATConv(input_channels, input_channels // 2)
+        self.conv3 = GATConv(input_channels // 2, 1)
+        self.ELU = ELU()
+
+    def forward(self, x, edge_index):
+        # Layer 1
+        x = self.conv1(x, edge_index)
+        # x = self.ELU(x)
+        # x = F.dropout(x, p=0.25)
+
+        # Layer 2
+        # x = self.conv2(x, edge_index)
+        # x = self.ELU(x)
+        # x = F.dropout(x, p=0.25)
+
+        # Layer 3
+        # x = self.conv3(x, edge_index)
+        return x
+
+
 class MLP(torch.nn.Module):
     def __init__(self, input_channels):
         super().__init__()
-        self.linear1 = Linear(input_channels, input_channels // 2)
-        self.linear2 = Linear(input_channels // 2, input_channels // 2)
+        self.linear1 = Linear(input_channels, input_channels)
+        self.linear2 = Linear(input_channels, input_channels // 2)
         self.linear3 = Linear(input_channels // 2, 1)
         self.ELU = ELU()
 

src/deep_neurographs/machine_learning/graph_trainer.py

@@ -8,15 +8,20 @@
 
 """
 
+from copy import deepcopy
 from random import sample, shuffle
+
 import numpy as np
 import torch
-from copy import deepcopy
-from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
-from torch.nn.functional import sigmoid
+from sklearn.metrics import (
+    accuracy_score,
+    f1_score,
+    precision_score,
+    recall_score,
+)
 from torch.utils.tensorboard import SummaryWriter
-from deep_neurographs.machine_learning import ml_utils
 
+from deep_neurographs.machine_learning import ml_utils
 
 LR = 1e-3
 N_EPOCHS = 1000
@@ -29,6 +34,7 @@ class GraphTrainer:
     Custom class that trains graph neural networks.
 
     """
+
     def __init__(
         self,
         model,
@@ -98,7 +104,7 @@ def run_on_graphs(self, graph_datasets):
                 y_i, hat_y_i = self.train(graph_datasets[graph_id].data, epoch)
                 y.extend(toCPU(y_i))
                 hat_y.extend(toCPU(hat_y_i))
-            train_score = self.compute_metrics(y, hat_y, "train", epoch)
+            self.compute_metrics(y, hat_y, "train", epoch)
 
             # Test
             if epoch % 10 == 0:
@@ -236,11 +242,11 @@ def compute_metrics(self, y, hat_y, prefix, epoch):
         f1 = f1_score(y, hat_y)
 
         # Log
-        self.writer.add_scalar(prefix + '_accuracy:', accuracy, epoch)
-        self.writer.add_scalar(prefix + '_accuracy_df:', accuracy_dif, epoch)
-        self.writer.add_scalar(prefix + '_precision:', precision, epoch)
-        self.writer.add_scalar(prefix + '_recall:', recall, epoch)
-        self.writer.add_scalar(prefix + '_f1:', f1, epoch)
+        self.writer.add_scalar(prefix + "_accuracy:", accuracy, epoch)
+        self.writer.add_scalar(prefix + "_accuracy_df:", accuracy_dif, epoch)
+        self.writer.add_scalar(prefix + "_precision:", precision, epoch)
+        self.writer.add_scalar(prefix + "_recall:", recall, epoch)
+        self.writer.add_scalar(prefix + "_f1:", f1, epoch)
         return f1
 
 
@@ -349,7 +355,7 @@ def truncate(hat_y, y):
         Truncated "hat_y".
 
     """
-    return hat_y[0: y.size(0), 0]
+    return hat_y[: y.size(0), 0]
 
 
 def get_predictions(hat_y, threshold=0.5):

src/deep_neurographs/machine_learning/inference.py

@@ -9,7 +9,6 @@
 """
 
 from copy import deepcopy
-from time import time
 
 import fastremap
 import networkx as nx
@@ -268,7 +267,7 @@ def run_model(dataset, model, model_type):
 
             # Postprocess
             hat_y_i = np.array(hat_y_i)
-            hat_y.extend(hat_y_i.tolist())
+            hat_y.extend(hat_y_i[:, 0].tolist())
     else:
         data = dataset["dataset"]
         hat_y = model.predict_proba(data["inputs"])[:, 1]

src/deep_neurographs/machine_learning/ml_utils.py

@@ -195,4 +195,4 @@ def sigmoid(x):
     Sigmoid applied to "x".
 
     """
-    return 1.0/(1.0 + np.exp(-x))
+    return 1.0 / (1.0 + np.exp(-x))

src/deep_neurographs/machine_learning/trainer.py

@@ -87,9 +87,9 @@ def fit_deep_model(
 
     # Fit model
     pylightning_trainer = pl.Trainer(
-        accelerator="gpu",
+        # accelerator="gpu",
         callbacks=[ckpt_callback],
-        devices=1,
+        # devices=1,
         enable_model_summary=False,
         enable_progress_bar=False,
         logger=logger,

src/deep_neurographs/reconstruction.py

@@ -30,23 +30,26 @@ def get_accepted_propoals_blocks(
 ):
     accepts = dict()
     for block_id in blocks:
+        # Threshold prediction
+        preds_upd = threshold_preds(
+            preds,
+            idx_to_edge,
+            low_threshold,
+            valid_idxs=block_to_idxs[block_id],
+        )
+
         # Get accepts
         if structure_aware:
             graph = neurographs[block_id].copy()
-            accepts[block_id] = get_structure_aware_accepts(
+            accepts[block_id], _ = get_structure_aware_accepts(
                 neurographs[block_id],
                 graph,
-                preds,
+                preds_upd,
                 high_threshold=high_threshold,
                 low_threshold=low_threshold,
             )
         else:
-            preds = threshold_preds(
-                preds,
-                idx_to_edge,
-                low_threshold,
-                valid_idxs=block_to_idxs[block_id],
-            )
+
             accepts[block_id] = preds.keys()
     return accepts
 
@@ -114,7 +117,7 @@ def get_structure_aware_accepts(
     best_preds, best_probs = get_best_preds(neurograph, preds, high_threshold)
     accepts, graph = check_cycles_sequential(graph, best_preds, best_probs)
     if len(best_preds) == len(preds.keys()):
-        return accepts
+        return accepts, graph
 
     # Add remaining preds
     best_preds = set(best_preds)
@@ -156,7 +159,7 @@ def get_subgraphs(graph, edge):
         subgraph = nx.union(subgraph_1, subgraph_2)
         return subgraph
     except:
-        return False 
+        return False
 
 
 def check_cycles_parallelized(graph, edge_list):
@@ -201,6 +204,7 @@ def check_cycles_parallelized(graph, edge_list):
 def check_cycles_sequential(graph, edges, probs):
     accepts = []
     for i in np.argsort(probs):
+        print(i, edges)
         subgraph = get_subgraphs(graph, edges[i])
         if subgraph:
             created_cycle, _ = gutils.creates_cycle(subgraph, tuple(edges[i]))
@@ -241,6 +245,7 @@ def save_prediction(neurograph, accepted_proposals, output_dir):
     save_corrections(neurograph, accepted_proposals, corrections_dir)
     save_connections(neurograph, accepted_proposals, connections_path)
 
+
 def save_corrections(neurograph, accepted_proposals, output_dir):
     for cnt, (i, j) in enumerate(accepted_proposals):
         # Info

src/deep_neurographs/swc_utils.py

@@ -359,7 +359,7 @@ def make_entry(graph, i, parent, node_to_idx):
     ...
 
     """
-    r = set_radius(graph, i)
+    r = graph[i]["radius"]
     x, y, z = tuple(graph.nodes[i]["xyz"])
     node_to_idx[i] = len(node_to_idx)
     entry = f"{node_to_idx[i]} 2 {x} {y} {z} {r} {node_to_idx[parent]}"