Feat upd cnn3d

src/neuron_proofreader/merge_proofreading/merge_datamodules.py

@@ -38,10 +38,10 @@
 from neuron_proofreader.machine_learning.image_dataloader import (
     DetectionPatchLoader as PatchLoader,
 )
-from neuron_proofreader.machine_learning.geometric_gnn_models import (
+from neuron_proofreader.models.geometric_gnn_models import (
     subgraph_to_data,
 )
-from neuron_proofreader.machine_learning.point_cloud_models import (
+from neuron_proofreader.models.point_cloud_models import (
     subgraph_to_point_cloud,
 )
 from neuron_proofreader.skeleton_graph import SkeletonGraph

src/neuron_proofreader/models/geometric_gnn_models.py

@@ -15,7 +15,7 @@
 import torch
 import torch.nn.functional as F
 
-from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.models.vision_models import CNN3D
 from neuron_proofreader.utils.ml_util import FeedForwardNet
 
 

src/neuron_proofreader/models/gnn_models.py

@@ -15,7 +15,7 @@
 import torch
 import torch.nn.init as init
 
-from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.models.vision_models import CNN3D
 from neuron_proofreader.split_proofreading import split_feature_extraction
 from neuron_proofreader.utils.ml_util import FeedForwardNet
 

src/neuron_proofreader/models/new_vision_models.py

@@ -22,30 +22,27 @@ class NewCNN3D(nn.Module):
 
     def __init__(
         self,
-        patch_shape,
-        in_channels=2,
-        first_out_channels=16,
-        output_dim=1,
-        channel_growth=2,
+        input_shape,
+        base_channels=16,
+        channel_multiplier=2,
+        depth=5,
         dropout=0.1,
         max_channels=256,
-        num_blocks=5,
         num_single_blocks=2,
+        output_dim=1,
         use_double=True,
     ):
         """
         Instantiates a CNN3D object.
 
         Parameters
         ----------
-        patch_shape : Tuple[int]
-            Shape of input image patch.
         output_dim : int, optional
             Dimension of output. Default is 1.
         dropout : float, optional
             Fraction of values to randomly drop during training. Default is
             0.1.
-        num_blocks : int, optional
+        depth : int, optional
             Number of convolutional blocks. Default is 5.
         use_double : bool, optional
             Indication of whether to use double convolution. Default is True.
@@ -55,13 +52,11 @@ def __init__(
 
         # Instance attributes
         self.dropout = dropout
-        self.patch_shape = patch_shape
-
         self.encode = Encoder3D(
-            in_channels,
-            first_out_channels,
-            num_blocks,
-            channel_growth=channel_growth,
+            input_shape[0],
+            base_channels,
+            depth,
+            channel_multiplier=channel_multiplier,
             use_double=use_double,
             max_channels=max_channels,
             num_single_blocks=num_single_blocks,
@@ -121,8 +116,8 @@ def __init__(
         self,
         in_channels,
         out_channels,
-        num_blocks,
-        channel_growth=2,
+        depth,
+        channel_multiplier=2,
         use_double=True,
         max_channels=256,
         num_single_blocks=2,
@@ -136,9 +131,9 @@ def __init__(
             Number of channels input to the first block.
         out_channels : int
             Number of channels output by the first block.
-        num_blocks : int
+        depth : int
             Number of conv blocks in the encoder.
-        channel_growth : float, optional
+        channel_multiplier : float, optional
             Multiplicative channel growth factor per layer. Default is 2.
         use_double : bool, optional
             Indication of whether blocks use double convolution. Default is True.
@@ -150,14 +145,14 @@ def __init__(
 
         # Create encoding blocks
         blocks = list()
-        for i in range(num_blocks):
+        for i in range(depth):
             use_double = i > num_single_blocks
             block = ConvBlock3D(
                 in_channels, out_channels, use_double=use_double
             )
             blocks.append(block)
             in_channels = block.out_channels
-            out_channels = int(min(out_channels * channel_growth, max_channels))
+            out_channels = int(min(out_channels * channel_multiplier, max_channels))
 
         # Instance attributes
         self.blocks = nn.ModuleList(blocks)
@@ -176,7 +171,7 @@ def __init__(
     ):
         """
         Instantiates a ConvBlock3D object.
- 
+
         Parameters
         ----------
         in_channels : int

src/neuron_proofreader/models/point_cloud_models.py

@@ -14,7 +14,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.models.vision_models import CNN3D
 from neuron_proofreader.utils.ml_util import FeedForwardNet
 
 

src/neuron_proofreader/utils/swc_util.py

@@ -48,6 +48,7 @@ class Reader:
     """
 
     gcs_client = None
+    s3_client = None
 
     def __init__(
         self, anisotropy=(1.0, 1.0, 1.0), min_swc_pts=1, verbose=True
@@ -71,9 +72,22 @@ def __init__(
 
     @classmethod
     def _get_gcs_client(cls):
-        if cls._gcs_client is None:
-            cls._gcs_client = storage.Client()
-        return cls._gcs_client
+        if cls.gcs_client is None:
+            cls.gcs_client = storage.Client()
+        return cls.gcs_client
+
+    @classmethod
+    def _get_s3_client(cls):
+        # TEMP
+        return cls.s3_client
+
+    @classmethod
+    def _get_client(cls, path):
+        if util.is_gcs_path(path):
+            return cls._get_gcs_client()
+        elif util.is_s3_path(path):
+            return cls._get_s3_client()
+        return None
 
     # --- Read Data ---
     def __call__(self, swc_pointer):
@@ -156,7 +170,8 @@ def read_swc(self, path):
             Dictionary whose keys and values are the attribute names and
             values from an SWC file.
         """
-        content = util.read_txt(path).splitlines()
+        client = self._get_client(path)
+        content = util.read_txt(path, client=client).splitlines()
         filename = os.path.basename(path)
         return self.parse(content, filename)
 
@@ -310,31 +325,6 @@ def read_from_cloud(self, path):
         else:
             return list()
 
-    def read_gcs_swc(self, path):
-        """
-        Reads a single SWC file stored in a GCS bucket.
-
-        Parameters
-        ----------
-        path : List[str]
-            Path to SWC file to be read.
-
-        Returns
-        -------
-        Deque[dict]
-            Dictionaries whose keys and values are the attribute names and
-            values from an SWC file.
-        """
-        # Initialize cloud reader
-        bucket_name, key = util.parse_cloud_path(path)
-        bucket = self._get_gcs_client().bucket(bucket_name)
-        blob = bucket.blob(key)
-
-        # Parse swc contents
-        content = blob.download_as_text().splitlines()
-        filename = os.path.basename(key)
-        return self.parse(content, filename)
-
     def read_gcs_zip(self, path):
         """
         Reads SWC files stored in a ZIP archive downloaded from a GCS

src/neuron_proofreader/utils/util.py

@@ -221,7 +221,7 @@ def read_json(path):
         return json.load(f)
 
 
-def read_txt(path):
+def read_txt(path, client=None):
     """
     Reads txt file at the given path.
 
@@ -236,9 +236,9 @@ def read_txt(path):
         Text from the txt file.
     """
     if is_s3_path(path):
-        return read_s3_txt(path)
+        return read_s3_txt(path, client=client)
     elif is_gcs_path(path):
-        return read_gcs_txt(path)
+        return read_gcs_txt(path, client=client)
     else:
         with open(path, "r") as f:
             return f.read()
@@ -520,7 +520,7 @@ def list_gcs_subprefixes(path):
     return subdirs
 
 
-def read_gcs_txt(path):
+def read_gcs_txt(prefix, client=None):
     """
     Reads a txt file stored in a GCS bucket.
 
@@ -534,9 +534,10 @@ def read_gcs_txt(path):
     str
         Contents of txt file.
     """
-    bucket_name, subpath = parse_cloud_path(path)
-    bucket = storage.Client().bucket(bucket_name)
-    return bucket.blob(subpath).download_as_text()
+    bucket_name, subprefix = parse_cloud_path(prefix)
+    client = client or storage.Client()
+    bucket = client.bucket(bucket_name)
+    return bucket.blob(subprefix).download_as_text()
 
 
 def read_json_from_gcs(bucket_name, blob_path):
@@ -613,23 +614,23 @@ def list_s3_paths(bucket_name, prefix, extension=""):
     return paths
 
 
-def read_s3_txt(path):
+def read_s3_txt(prefix, client=None):
     """
     Reads a txt file stored in an S3 bucket.
 
     Parameters
     ----------
-    path : str
+    prefix : str
         Path to txt file to be read.
 
     Returns
     -------
     str
         Contents of txt file.
     """
-    bucket_name, subpath = parse_cloud_path(path)
-    s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))
-    obj = s3.get_object(Bucket=bucket_name, Key=subpath)
+    bucket_name, subprefix = parse_cloud_path(prefix)
+    client = boto3.client("s3", config=Config(signature_version=UNSIGNED))
+    obj = client.get_object(Bucket=bucket_name, Key=subprefix)
     return obj["Body"].read().decode("utf-8")