diff --git a/src/neuron_proofreader/merge_proofreading/merge_datamodules.py b/src/neuron_proofreader/merge_proofreading/merge_datamodules.py index bd76f82..4d17f93 100644 --- a/src/neuron_proofreader/merge_proofreading/merge_datamodules.py +++ b/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 diff --git a/src/neuron_proofreader/models/geometric_gnn_models.py b/src/neuron_proofreader/models/geometric_gnn_models.py index 0c68d7f..35398e1 100644 --- a/src/neuron_proofreader/models/geometric_gnn_models.py +++ b/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 diff --git a/src/neuron_proofreader/models/gnn_models.py b/src/neuron_proofreader/models/gnn_models.py index 95f1387..5fc454a 100644 --- a/src/neuron_proofreader/models/gnn_models.py +++ b/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 diff --git a/src/neuron_proofreader/models/new_vision_models.py b/src/neuron_proofreader/models/new_vision_models.py index 9a41377..5548e9e 100644 --- a/src/neuron_proofreader/models/new_vision_models.py +++ b/src/neuron_proofreader/models/new_vision_models.py @@ -22,15 +22,14 @@ 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, ): """ @@ -38,14 +37,12 @@ def __init__( 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 diff --git a/src/neuron_proofreader/models/point_cloud_models.py b/src/neuron_proofreader/models/point_cloud_models.py index 23d46ca..219096a 100644 --- a/src/neuron_proofreader/models/point_cloud_models.py +++ b/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 diff --git a/src/neuron_proofreader/utils/swc_util.py b/src/neuron_proofreader/utils/swc_util.py index f0ffa73..70379e9 100644 --- a/src/neuron_proofreader/utils/swc_util.py +++ b/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 diff --git a/src/neuron_proofreader/utils/util.py b/src/neuron_proofreader/utils/util.py index 4056cbf..31c65ca 100644 --- a/src/neuron_proofreader/utils/util.py +++ b/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,13 +614,13 @@ 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 @@ -627,9 +628,9 @@ def read_s3_txt(path): 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")