black reformat

src/deep_neurographs/deep_learning/datasets.py

@@ -245,15 +245,15 @@ def __init__(self):
         None
 
         """
-        self.blur = tio.RandomBlur(std=(0, 0.4))
-        self.noise = tio.RandomNoise(std=(0, 0.03))
-        self.apply_geometric = tio.Compose(
-            {
-                # tio.RandomFlip(axes=(0, 1, 2)),
-                tio.RandomAffine(
-                    degrees=20, scales=(0.8, 1), image_interpolation="nearest"
-                )
-            }
+        self.transform = tio.Compose(
+            [
+                tio.RandomBlur(std=(0, 0.4)),
+                tio.RandomNoise(std=(0, 0.03)),
+                tio.RandomFlip(axes=(0, 1, 2)),
+                # tio.RandomAffine(
+                # degrees=20, scales=(0.8, 1), image_interpolation="nearest"
+                # )
+            ]
         )
 
     def run(self, arr):
@@ -271,10 +271,7 @@ def run(self, arr):
             Transformed array after being run through augmentation pipeline.
 
         """
-        arr = self.blur(arr)
-        arr = self.noise(arr)
-        arr = self.apply_geometric(arr)
-        return arr
+        return self.transform(arr)
 
 
 def reformat(arr):

src/deep_neurographs/deep_learning/train.py

@@ -17,6 +17,7 @@
 import torch.nn.functional as F
 import torch.utils.data as torch_data
 from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.profilers import PyTorchProfiler
 from sklearn.ensemble import AdaBoostClassifier, RandomForestClassifier
 from torch.utils.data import DataLoader
 from torcheval.metrics.functional import (
@@ -26,15 +27,14 @@
     binary_recall,
 )
 
-from deep_neurographs import utils
 from deep_neurographs.deep_learning import datasets as ds
 from deep_neurographs.deep_learning import models
 
 logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
 
 
 BATCH_SIZE = 32
-NUM_WORKERS = 1
+NUM_WORKERS = 0
 SHUFFLE = True
 SUPPORTED_CLFS = [
     "AdaBoost",
@@ -90,6 +90,7 @@ def train_network(
     logger=True,
     max_epochs=100,
     model_summary=True,
+    profile=False,
     progress_bar=True,
 ):
     # Load data
@@ -101,14 +102,20 @@ def train_network(
         shuffle=SHUFFLE,
     )
     valid_loader = DataLoader(
-        valid_set, num_workers=NUM_WORKERS, batch_size=BATCH_SIZE
+        valid_set,
+        batch_size=BATCH_SIZE,
+        num_workers=NUM_WORKERS,
+        pin_memory=True,
     )
 
-    # Fit model
+    # Configure trainer
     model = LitNeuralNet(net)
     checkpoint_callback = ModelCheckpoint(
         save_top_k=1, monitor="val_f1", mode="max"
     )
+    profiler = PyTorchProfiler() if profile else None
+
+    # Fit model
     trainer = pl.Trainer(
         accelerator="gpu",
         callbacks=[checkpoint_callback],
@@ -117,6 +124,7 @@ def train_network(
         enable_progress_bar=progress_bar,
         logger=logger,
         max_epochs=max_epochs,
+        profiler=profiler,
     )
     trainer.fit(model, train_loader, valid_loader)
     return model

src/deep_neurographs/densegraph.py

@@ -1,3 +1,13 @@
+"""
+Created on Sat November 04 15:30:00 2023
+
+@author: Anna Grim
+@email: [email protected]
+
+Class of graphs that are built from swc files.
+
+"""
+
 import os
 
 import networkx as nx
@@ -6,7 +16,7 @@
 from scipy.spatial import KDTree
 
 from deep_neurographs import swc_utils, utils
-from deep_neurographs.geometry_utils import dist, make_line
+from deep_neurographs.geometry_utils import dist
 
 
 class DenseGraph:
@@ -69,27 +79,17 @@ def check_aligned(self, pred_xyz_i, pred_xyz_j):
         if self.xyz_to_swc[xyz_i] != self.xyz_to_swc[xyz_j]:
             return False
 
-        # Compare pred and target distances
+        # Compute distances
         pred_xyz_i = np.array(pred_xyz_i)
         pred_xyz_j = np.array(pred_xyz_j)
         pred_dist = dist(pred_xyz_i, pred_xyz_j)
 
         target_path, target_dist = self.connect_nodes(graph_id, xyz_i, xyz_j)
         target_dist = max(target_dist, 1)
 
+        # Check criteria
         ratio = min(pred_dist, target_dist) / max(pred_dist, target_dist)
         if ratio < 0.5 and pred_dist > 15:
             return False
-        # elif ratio < 0.2:
-        #    return False
-
-        # Compare projected predicted path
-        proj_dists = []
-        proj_nodes = set()
-        for xyz in make_line(pred_xyz_i, pred_xyz_j, len(target_path)):
-            proj_xyz, proj_d = self.get_projection(xyz)
-            swc = self.xyz_to_swc[tuple(proj_xyz)]
-            proj_nodes.add(self.xyz_to_node[swc][tuple(proj_xyz)])
-            proj_dists.append(proj_d)
-
-        return True
+        else:
+            return True

src/deep_neurographs/feature_extraction.py

@@ -52,14 +52,14 @@ def generate_mutable_features(
 
     """
     features = {"skel": generate_mutable_skel_features(neurograph)}
-    if img_path and img_profile:
-        features["img"] = generate_mutable_img_profile_features(
-            neurograph, img_path, anisotropy=anisotropy
-        )
-    elif img_path and not img_profile:
+    if img_path and labels_path:
         features["img"] = generate_mutable_img_chunk_features(
             neurograph, img_path, labels_path, anisotropy=anisotropy
         )
+    elif img_path and img_profile:
+        features["img"] = generate_mutable_img_profile_features(
+            neurograph, img_path, anisotropy=anisotropy
+        )
     return features
 
 

src/deep_neurographs/geometry_utils.py

@@ -143,7 +143,7 @@ def get_profile(
 def fill_path(img, path, val=-1):
     for xyz in path:
         x, y, z = tuple(np.round(xyz).astype(int))
-        img[x - 1 : x + 1, y - 1 : y + 1, z - 1 : z + 1] = val
+        img[(x - 1) : x + 1, (y - 1) : y + 1, (z - 1) : z + 1] = val
     return img
 
 

src/deep_neurographs/intake.py

@@ -10,9 +10,6 @@
 
 import os
 
-import torch
-from torch_geometric.data import Data
-
 from deep_neurographs import s3_utils, swc_utils, utils
 from deep_neurographs.neurograph import NeuroGraph
 

src/deep_neurographs/neurograph.py

@@ -308,7 +308,7 @@ def init_targets(self, target_neurograph):
             proj_xyz_j, d_j = target_neurograph.get_projection(xyz_j)
 
             # Check criteria
-            if d_i > 9 or d_j > 9:
+            if d_i > 8 or d_j > 8:
                 continue
             elif self.check_cycle((i, j)):
                 continue

src/deep_neurographs/utils.py

@@ -133,18 +133,18 @@ def open_tensorstore(path):
 
 def read_img_chunk(img, xyz, shape):
     return img[
-        xyz[2] - shape[2] // 2 : xyz[2] + shape[2] // 2,
-        xyz[1] - shape[1] // 2 : xyz[1] + shape[1] // 2,
-        xyz[0] - shape[0] // 2 : xyz[0] + shape[0] // 2,
+        (xyz[2] - shape[2] // 2) : xyz[2] + shape[2] // 2,
+        (xyz[1] - shape[1] // 2) : xyz[1] + shape[1] // 2,
+        (xyz[0] - shape[0] // 2) : xyz[0] + shape[0] // 2,
     ].transpose(2, 1, 0)
 
 
 def read_tensorstore(ts_arr, xyz, shape):
     arr = (
         ts_arr[
-            xyz[0] - shape[0] // 2 : xyz[0] + shape[0] // 2,
-            xyz[1] - shape[1] // 2 : xyz[1] + shape[1] // 2,
-            xyz[2] - shape[2] // 2 : xyz[2] + shape[2] // 2,
+            (xyz[0] - shape[0] // 2) : xyz[0] + shape[0] // 2,
+            (xyz[1] - shape[1] // 2) : xyz[1] + shape[1] // 2,
+            (xyz[2] - shape[2] // 2) : xyz[2] + shape[2] // 2,
         ]
         .read()
         .result()