Refactor datasets models

src/deep_neurographs/intake.py

@@ -16,7 +16,6 @@
 from deep_neurographs.utils import graph_util as gutil
 from deep_neurographs.utils import img_util, swc_util, util
 
-
 MIN_SIZE = 30
 NODE_SPACING = 2
 SMOOTH_BOOL = True
@@ -125,7 +124,7 @@ def run(self, swc_pointer):
             # Graph size
             n_components = util.reformat_number(len(irreducibles))
             print("\nGraph Overview...")
-            print("# connected components:", n_connected_components)
+            print("# connected components:", n_components)
             print("# nodes:", util.reformat_number(n_nodes))
             print("# edges:", util.reformat_number(n_edges))
 
@@ -155,15 +154,17 @@ def get_irreducibles(self, swc_dicts):
                 i += 1
 
             # Store results
+            desc = "Extract Graphs"
             irreducibles = []
             n_nodes, n_edges = 0, 0
-            for process in tqdm(as_completed(processes), desc="Extract Graphs"):
+            for process in tqdm(as_completed(processes), desc=desc):
                 irreducibles_i = process.result()
                 irreducibles.extend(irreducibles_i)
                 n_nodes += count_nodes(irreducibles_i)
                 n_edges += count_edges(irreducibles_i)
         return irreducibles, n_nodes, n_edges
 
+
 # --- utils ---
 def count_nodes(irreducibles):
     """

src/deep_neurographs/machine_learning/graph_datasets.py

@@ -15,7 +15,7 @@
 import torch
 from torch_geometric.data import Data as GraphData
 
-from deep_neurographs.machine_learning import feature_generation, datasets
+from deep_neurographs.machine_learning import datasets, feature_generation
 from deep_neurographs.utils import gnn_util
 
 

src/deep_neurographs/machine_learning/heterograph_datasets.py

@@ -17,7 +17,7 @@
 import torch
 from torch_geometric.data import HeteroData as HeteroGraphData
 
-from deep_neurographs.machine_learning import feature_generation, datasets
+from deep_neurographs.machine_learning import datasets, feature_generation
 from deep_neurographs.utils import gnn_util
 
 DTYPE = torch.float32

src/deep_neurographs/neurograph.py

@@ -19,7 +19,9 @@
 
 from deep_neurographs import generate_proposals, geometry
 from deep_neurographs.geometry import dist as get_dist
-from deep_neurographs.machine_learning.groundtruth_generation import init_targets
+from deep_neurographs.machine_learning.groundtruth_generation import (
+    init_targets,
+)
 from deep_neurographs.utils import graph_util as gutil
 from deep_neurographs.utils import img_util, swc_util, util
 

src/deep_neurographs/utils/ml_util.py

@@ -16,8 +16,8 @@
 from sklearn.ensemble import AdaBoostClassifier, RandomForestClassifier
 
 from deep_neurographs.machine_learning import (
-    feature_generation,
     datasets,
+    feature_generation,
     graph_datasets,
     heterograph_datasets,
 )

src/deep_neurographs/utils/swc_util.py

@@ -9,15 +9,20 @@
 
 """
 
-from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
-from google.cloud import storage
+import os
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    ThreadPoolExecutor,
+    as_completed,
+)
 from io import BytesIO
 from tqdm import tqdm
 from zipfile import ZipFile
 
 import networkx as nx
 import numpy as np
-import os
+from google.cloud import storage
+from tqdm import tqdm
 
 from deep_neurographs.utils import util