Bug val dataset (#558)

* bug: val dataset length

* bug: save mistake mips multimodal

* bug: query error

* refactor: mkdir

* refactor: updated merge inference

---------

Co-authored-by: anna-grim <[email protected]>

Author: anna-grim

src/neuron_proofreader/merge_proofreading/merge_inference.py

@@ -19,7 +19,9 @@
 import numpy as np
 import torch
 
-from neuron_proofreader import exp
+from neuron_proofreader.machine_learning.point_cloud_models import (
+    subgraph_to_point_cloud,
+)
 from neuron_proofreader.utils import img_util, ml_util, util
 
 
@@ -35,6 +37,7 @@ def __init__(
         anisotropy=(1.0, 1.0, 1.0),
         batch_size=32,
         device="cuda",
+        is_multimodal=False,
         min_size=0,
         prefetch=128,
         remove_detected_sites=False,
@@ -62,6 +65,7 @@ def __init__(
             patch_shape,
             anisotropy=anisotropy,
             batch_size=batch_size,
+            is_multimodal=is_multimodal,
             min_size=min_size,
             prefetch=prefetch,
             step_size=step_size,
@@ -173,6 +177,7 @@ def __init__(
         patch_shape,
         anisotropy=(1.0, 1.0, 1.0),
         batch_size=16,
+        is_multimodal=False,
         min_size=0,
         prefetch=128,
         step_size=10,
@@ -186,14 +191,15 @@ def __init__(
         self.batch_size = batch_size
         self.distance_traversed = 0
         self.graph = graph
+        self.is_multimodal = is_multimodal
         self.min_size = min_size
         self.patch_shape = patch_shape
         self.prefetch = prefetch
         self.step_size = step_size
         self.subgraph_radius = subgraph_radius
 
         # Image reader
-        self.img_reader = img_util.init_reader(img_path)
+        self.img_reader = img_util.TensorStoreReader(img_path)
 
     # --- Core routines ---
     def __iter__(self):
@@ -222,7 +228,14 @@ def submit_thread():
                         # Process completed thread
                         nodes, patch_centers = pending.pop(thread)
                         img, offset = thread.result()
-                        yield self.get_multimodal_batch(img, offset, patch_centers, nodes)
+                        if self.is_multimodal:
+                            yield self.get_multimodal_batch(
+                                img, offset, patch_centers, nodes
+                            )
+                        else:
+                            yield self.get_batch(
+                                img, offset, patch_centers, nodes
+                            )
 
                         # Continue submitting threads
                         submit_thread()
@@ -338,10 +351,10 @@ def get_multimodal_batch(self, img, offset, patch_centers, nodes):
             patches[i, 1, ...] = label_mask[s]
 
             subgraph = self.graph.get_rooted_subgraph(node, self.subgraph_radius)
-            point_clouds[i] = exp.subgraph_to_point_cloud(subgraph)
+            point_clouds[i] = subgraph_to_point_cloud(subgraph)
 
         # Compile batch dictionary
-        batch = exp.TensorDict({
+        batch = ml_util.TensorDict({
             "img": ml_util.to_tensor(patches),
             "point_cloud": ml_util.to_tensor(point_clouds)
         })

src/neuron_proofreader/utils/util.py

@@ -40,14 +40,15 @@ def listdir(path, extension=None):
 
     Returns
     -------
-    List[str]
+    filenames : List[str]
         Filenames in directory with extension "extension" if provided.
         Otherwise, list of all files in directory.
     """
-    if extension is None:
-        return [f for f in os.listdir(path)]
+    filenames = [f for f in os.listdir(path) if not f.startswith(".")]
+    if extension:
+        return [f for f in filenames if f.endswith(extension)]
     else:
-        return [f for f in os.listdir(path) if f.endswith(extension)]
+        return filenames
 
 
 def list_files_in_zip(zip_content):
@@ -138,8 +139,7 @@ def mkdir(path, delete=False):
     if delete:
         rmdir(path)
 
-    if not os.path.exists(path):
-        os.mkdir(path)
+    os.makedirs(path, exist_ok=True)
 
 
 def rmdir(path):