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/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")