octopi 1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

octopi/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.2.0"

octopi/datasets/cached_datset.py

@@ -1,7 +1,7 @@
 from typing import List, Tuple, Callable, Optional, Dict, Any
 from monai.transforms import Compose
 from monai.data import CacheDataset
-from octopi import io
+from octopi.datasets import io
 from tqdm import tqdm
 import os, sys
 

octopi/datasets/generators.py

@@ -1,7 +1,7 @@
 from octopi.datasets import dataset, augment, cached_datset
 from monai.data import DataLoader, SmartCacheDataset, CacheDataset, Dataset
 from typing import List, Optional
-from octopi import io
+from octopi.datasets import io
 import torch, os, random, gc
 import multiprocess as mp
 

octopi/datasets/io.py

@@ -0,0 +1,200 @@
+"""
+Data loading, processing, and dataset operations for the datasets module.
+"""
+
+from monai.data import DataLoader, CacheDataset, Dataset
+from monai.transforms import (
+    Compose, 
+    NormalizeIntensityd,
+    EnsureChannelFirstd,  
+)
+from sklearn.model_selection import train_test_split
+from collections import defaultdict
+from copick_utils.io import readers
+import copick, torch, os, random
+from typing import List
+from tqdm import tqdm
+
+
+def load_training_data(root, 
+                       runIDs: List[str],
+                       voxel_spacing: float, 
+                       tomo_algorithm: str, 
+                       segmenation_name: str,
+                       segmentation_session_id: str = None,
+                       segmentation_user_id: str = None,
+                       progress_update: bool = True):
+    """
+    Load training data from CoPick runs.
+    """
+    data_dicts = []
+    # Use tqdm for progress tracking only if progress_update is True
+    iterable = tqdm(runIDs, desc="Loading Training Data") if progress_update else runIDs
+    for runID in iterable:
+        run = root.get_run(str(runID))
+        tomogram = readers.tomogram(run, voxel_spacing, tomo_algorithm)
+        segmentation = readers.segmentation(run, 
+                                              voxel_spacing,
+                                              segmenation_name,
+                                              segmentation_session_id, 
+                                              segmentation_user_id)
+        data_dicts.append({"image": tomogram, "label": segmentation})
+
+    return data_dicts 
+
+
+def load_predict_data(root, 
+                      runIDs: List[str],
+                      voxel_spacing: float, 
+                      tomo_algorithm: str):
+    """
+    Load prediction data from CoPick runs.
+    """
+    data_dicts = []
+    for runID in tqdm(runIDs):
+        run = root.get_run(str(runID))
+        tomogram = readers.tomogram(run, voxel_spacing, tomo_algorithm)
+        data_dicts.append({"image": tomogram})
+
+    return data_dicts 
+
+
+def create_predict_dataloader(
+    root,
+    voxel_spacing: float, 
+    tomo_algorithm: str,       
+    runIDs: str = None,       
+    ): 
+    """
+    Create a dataloader for prediction data.
+    """
+    # define pre transforms
+    pre_transforms = Compose(
+        [   EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"),
+            NormalizeIntensityd(keys=["image"]),
+    ])
+
+    # Split trainRunIDs, validateRunIDs, testRunIDs
+    if runIDs is None:
+        runIDs = [run.name for run in root.runs]
+    test_files = load_predict_data(root, runIDs, voxel_spacing, tomo_algorithm) 
+
+    bs = min( len(test_files), 4)
+    test_ds = CacheDataset(data=test_files, transform=pre_transforms)
+    test_loader = DataLoader(test_ds, 
+                            batch_size=bs, 
+                            shuffle=False, 
+                            num_workers=4, 
+                            pin_memory=torch.cuda.is_available())
+    return test_loader, test_ds
+
+
+def adjust_to_multiple(value, multiple = 16):
+    """
+    Adjust a value to be a multiple of the specified number.
+    """
+    return int((value // multiple) * multiple)
+
+
+def get_input_dimensions(dataset, crop_size: int):
+    """
+    Get input dimensions for the dataset.
+    """
+    nx = dataset[0]['image'].shape[1]
+    if crop_size > nx:
+        first_dim = adjust_to_multiple(nx/2)
+        return first_dim, crop_size, crop_size
+    else:
+        return crop_size, crop_size, crop_size
+
+
+def get_num_classes(copick_config_path: str):
+    """
+    Get the number of classes from a CoPick configuration.
+    """
+    root = copick.from_file(copick_config_path)
+    return len(root.pickable_objects) + 1
+
+
+def split_multiclass_dataset(runIDs, 
+                             train_ratio: float = 0.7, 
+                             val_ratio: float = 0.15, 
+                             test_ratio: float = 0.15, 
+                             return_test_dataset: bool = True,
+                             random_state: int = 42):
+    """
+    Splits a given dataset into three subsets: training, validation, and testing. If the dataset
+    has categories (as tuples), splits are balanced across all categories. If the dataset is a 1D
+    list, it is split without categorization.
+
+    Parameters:
+    - runIDs: A list of items to split. It can be a 1D list or a list of tuples (category, value).
+    - train_ratio: Proportion of the dataset for training.
+    - val_ratio: Proportion of the dataset for validation.
+    - test_ratio: Proportion of the dataset for testing.
+    - return_test_dataset: Whether to return the test dataset.
+    - random_state: Random state for reproducibility.
+
+    Returns:
+    - trainRunIDs: Training subset.
+    - valRunIDs: Validation subset.
+    - testRunIDs: Testing subset (if return_test_dataset is True, otherwise None).
+    """
+
+    # Ensure the ratios add up to 1
+    assert train_ratio + val_ratio + test_ratio == 1.0, "Ratios must add up to 1.0"
+
+    # Check if the dataset has categories
+    if isinstance(runIDs[0], tuple) and len(runIDs[0]) == 2:
+        # Group by category
+        grouped = defaultdict(list)
+        for item in runIDs:
+            grouped[item[0]].append(item)
+
+        # Split each category
+        trainRunIDs, valRunIDs, testRunIDs = [], [], []
+        for category, items in grouped.items():
+            # Shuffle for randomness
+            random.shuffle(items)
+            # Split into train and remaining
+            train_items, remaining = train_test_split(items, test_size=(1 - train_ratio), random_state=random_state)
+            trainRunIDs.extend(train_items)
+
+            if return_test_dataset:
+                # Split remaining into validation and test
+                val_items, test_items = train_test_split(
+                    remaining,
+                    test_size=(test_ratio / (val_ratio + test_ratio)),
+                    random_state=random_state,
+                )
+                valRunIDs.extend(val_items)
+                testRunIDs.extend(test_items)
+            else:
+                valRunIDs.extend(remaining)
+                testRunIDs = []
+    else:
+        # If no categories, split as a 1D list
+        trainRunIDs, remaining = train_test_split(runIDs, test_size=(1 - train_ratio), random_state=random_state)
+        if return_test_dataset:
+            valRunIDs, testRunIDs = train_test_split(
+                remaining,
+                test_size=(test_ratio / (val_ratio + test_ratio)),
+                random_state=random_state,
+            )
+        else:
+            valRunIDs = remaining
+            testRunIDs = []
+
+    return trainRunIDs, valRunIDs, testRunIDs    
+
+
+def load_copick_config(path: str):
+    """
+    Load a CoPick configuration from file.
+    """
+    if os.path.isfile(path):
+        root = copick.from_file(path)
+    else:
+        raise FileNotFoundError(f"Copick Config Path does not exist: {path}")
+    
+    return root 

octopi/datasets/multi_config_generator.py

@@ -1,7 +1,7 @@
 from octopi.datasets import dataset, augment, cached_datset
 from octopi.datasets.generators import TrainLoaderManager
 from monai.data import DataLoader, SmartCacheDataset, CacheDataset, Dataset
-from octopi import io
+from octopi.datasets import io
 import multiprocess as mp
 from typing import List
 from tqdm import tqdm

octopi/entry_points/common.py

@@ -1,4 +1,4 @@
-from octopi import utils
+from octopi.utils import parsers
 import argparse
 
 def add_model_parameters(parser, octopi = False):
@@ -8,8 +8,8 @@ def add_model_parameters(parser, octopi = False):
     
     # Add U-Net model parameters
     parser.add_argument("--Nclass", type=int, required=False, default=3, help="Number of prediction classes in the model")
-    parser.add_argument("--channels", type=utils.parse_int_list, required=False, default='32,64,96,96', help="List of channel sizes")
-    parser.add_argument("--strides", type=utils.parse_int_list, required=False, default='2,2,1', help="List of stride sizes")
+    parser.add_argument("--channels", type=parsers.parse_int_list, required=False, default='32,64,96,96', help="List of channel sizes")
+    parser.add_argument("--strides", type=parsers.parse_int_list, required=False, default='2,2,1', help="List of stride sizes")
     parser.add_argument("--res-units", type=int, required=False, default=1, help="Number of residual units in the UNet")
     parser.add_argument("--dim-in", type=int, required=False, default=96, help="Input dimension for the UNet model")
 
@@ -52,11 +52,11 @@ def add_inference_parameters(parser):
 
     parser.add_argument("--tomo-alg", required=False, default = 'wbp', 
                         help="Tomogram algorithm used for produces segmentation prediction masks.")    
-    parser.add_argument("--seg-info", type=utils.parse_target, required=False, 
+    parser.add_argument("--seg-info", type=parsers.parse_target, required=False, 
                         default='predict,octopi,1', help='Information Query to save Segmentation predictions under, e.g., (e.g., "name" or "name,user_id,session_id" - Default UserID is octopi and SessionID is 1')
     parser.add_argument("--tomo-batch-size", type=int, default=25, required=False, 
                         help="Batch size for tomogram processing.")
-    parser.add_argument("--run-ids", type=utils.parse_list, default=None, required=False, 
+    parser.add_argument("--run-ids", type=parsers.parse_list, default=None, required=False, 
                         help="List of run IDs for prediction, e.g., run1,run2 or [run1,run2]. If not provided, all available runs will be processed.")   
     
 def add_localize_parameters(parser):

octopi/entry_points/create_slurm_submission.py

@@ -1,5 +1,5 @@
 from octopi.entry_points import run_train, run_segment_predict, run_localize, run_optuna
-from octopi.submit_slurm import create_shellsubmit, create_multiconfig_shellsubmit
+from octopi.utils.submit_slurm import create_shellsubmit, create_multiconfig_shellsubmit
 from octopi.processing.importers import cli_mrcs_parser, cli_dataportal_parser
 from octopi.entry_points import common 
 from octopi import utils

octopi/entry_points/run_create_targets.py

@@ -1,8 +1,8 @@
 import octopi.processing.create_targets_from_picks as create_targets
 from typing import List, Tuple, Union
+from octopi.utils import io, parsers
 from collections import defaultdict
 import argparse, copick, yaml, os
-from octopi import utils, io 
 from tqdm import tqdm
 import numpy as np
 
@@ -160,16 +160,16 @@ def parse_args():
 
     input_group = parser.add_argument_group("Input Arguments")
     input_group.add_argument("--config", type=str, required=True, help="Path to the CoPick configuration file.")
-    input_group.add_argument("--target", type=utils.parse_target, action="append", default=None, help='Target specifications: "name" or "name,user_id,session_id".')
+    input_group.add_argument("--target", type=parsers.parse_target, action="append", default=None, help='Target specifications: "name" or "name,user_id,session_id".')
     input_group.add_argument("--picks-session-id", type=str, default=None, help="Session ID for the picks.")
     input_group.add_argument("--picks-user-id", type=str, default=None, help="User ID associated with the picks.")
-    input_group.add_argument("--seg-target", type=utils.parse_target, action="append", default=[], help='Segmentation targets: "name" or "name,user_id,session_id".')
-    input_group.add_argument("--run-ids", type=utils.parse_list, default=None, help="List of run IDs.")
+    input_group.add_argument("--seg-target", type=parsers.parse_target, action="append", default=[], help='Segmentation targets: "name" or "name,user_id,session_id".')
+    input_group.add_argument("--run-ids", type=parsers.parse_list, default=None, help="List of run IDs.")
 
     # Parameters
     parameters_group = parser.add_argument_group("Parameters")
     parameters_group.add_argument("--tomo-alg", type=str, default="wbp", help="Tomogram reconstruction algorithm.")    
-    parameters_group.add_argument("--radius-scale", type=float, default=0.8, help="Scale factor for object radius.")
+    parameters_group.add_argument("--radius-scale", type=float, default=0.7, help="Scale factor for object radius.")
     parameters_group.add_argument("--voxel-size", type=float, default=10, help="Voxel size for tomogram reconstruction.")    
     
     output_group = parser.add_argument_group("Output Arguments")
@@ -275,7 +275,7 @@ def save_parameters(args, output_path: str):
     existing_data[input_key] = new_entry[input_key]
 
     # Save back to the YAML file
-    utils.save_parameters_yaml(existing_data, output_path)
+    io.save_parameters_yaml(existing_data, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_evaluate.py

@@ -1,5 +1,5 @@
 import octopi.processing.evaluate as evaluate
-import octopi.utils as utils
+from octopi.utils import parsers
 from typing import List
 import argparse
 
@@ -31,6 +31,7 @@ def cli():
     """
     CLI entry point for running evaluation.
     """
+
     parser = argparse.ArgumentParser(
         description='Run evaluation on pick and place predictions.',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
@@ -43,8 +44,8 @@ def cli():
     parser.add_argument('--predict-session-id', type=str, required=False, default= None, help='Session ID for prediction data')
     parser.add_argument('--save-path', type=str, required=False, default= None, help='Path to save evaluation results')
     parser.add_argument('--distance-threshold-scale', type=float, required=False, default = 0.8, help='Compute Distance Threshold Based on Particle Radius')
-    parser.add_argument('--object-names', type=utils.parse_list, default=None, required=False, help='Optional list of object names to evaluate, e.g., ribosome,apoferritin or [ribosome,apoferritin].')
-    parser.add_argument('--run-ids', type=utils.parse_list, default=None, required=False, help='Optional list of run IDs to evaluate, e.g., run1,run2,run3 or [run1,run2,run3].')
+    parser.add_argument('--object-names', type=parsers.parse_list, default=None, required=False, help='Optional list of object names to evaluate, e.g., ribosome,apoferritin or [ribosome,apoferritin].')
+    parser.add_argument('--run-ids', type=parsers.parse_list, default=None, required=False, help='Optional list of run IDs to evaluate, e.g., run1,run2,run3 or [run1,run2,run3].')
 
     args = parser.parse_args()
 

octopi/entry_points/run_extract_mb_picks.py

@@ -1,5 +1,5 @@
 from octopi.extract import membranebound_extract as extract
-from octopi import utils, io
+from octopi.utils import parsers
 import argparse, json, pprint, copick, json
 from typing import List, Tuple, Optional
 import multiprocess as mp
@@ -58,12 +58,12 @@ def cli():
     parser.add_argument('--config', type=str, required=True, help='Path to the configuration file.')
     parser.add_argument('--voxel-size', type=float, required=False, default=10, help='Voxel size.')
     parser.add_argument('--distance-threshold', type=float, required=False, default=10, help='Distance threshold.')
-    parser.add_argument('--picks-info', type=utils.parse_target, required=True, help='Query for the picks (e.g., "name" or "name,user_id,session_id".).')
-    parser.add_argument('--membrane-info', type=utils.parse_target, required=False, help='Query for the membrane segmentation (e.g., "name" or "name,user_id,session_id".).')
-    parser.add_argument('--organelle-info', type=utils.parse_target, required=False, help='Query for the organelles segmentations (e.g., "name" or "name,user_id,session_id".).')
+    parser.add_argument('--picks-info', type=parsers.parse_target, required=True, help='Query for the picks (e.g., "name" or "name,user_id,session_id".).')
+    parser.add_argument('--membrane-info', type=parsers.parse_target, required=False, help='Query for the membrane segmentation (e.g., "name" or "name,user_id,session_id".).')
+    parser.add_argument('--organelle-info', type=parsers.parse_target, required=False, help='Query for the organelles segmentations (e.g., "name" or "name,user_id,session_id".).')
     parser.add_argument('--save-user-id', type=str, required=False, default=None, help='User ID to save the new picks.')
     parser.add_argument('--save-session-id', type=str, required=True, help='Session ID to save the new picks.')
-    parser.add_argument('--runIDs', type=utils.parse_list, required=False, help='List of run IDs to process.')
+    parser.add_argument('--runIDs', type=parsers.parse_list, required=False, help='List of run IDs to process.')
     parser.add_argument('--n-procs', type=int, required=False, default=None, help='Number of processes to use.')
 
     args = parser.parse_args()

octopi/entry_points/run_localize.py

@@ -1,11 +1,10 @@
 from octopi.entry_points import common
+from octopi.utils import parsers, io
 from octopi.extract import localize
-from octopi import utils
 import copick, argparse, pprint
 from typing import List, Tuple
 import multiprocess as mp
 from tqdm import tqdm
-import os
 
 def pick_particles(
     copick_config_path: str,
@@ -53,13 +52,13 @@ def pick_particles(
 
     # Nprocesses shouldnt exceed computation resource or number of available runs
     n_run_ids = len(run_ids)
-    n_procs = min(mp.mp.cpu_count(), n_procs, n_run_ids)
+    n_procs = min(mp.cpu_count(), n_procs, n_run_ids)
 
     # Run Localization - Main Parallelization Loop
     print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")
     with mp.Pool(processes=n_procs) as pool:
         with tqdm(total=n_run_ids, desc="Localization", unit="run") as pbar:
-            worker_func = lambda run_id: localize.processs_localization(
+            worker_func = lambda run_id: localize.process_localization(
                 root.get_run(run_id),  
                 objects, 
                 seg_info,
@@ -85,20 +84,20 @@ def localize_parser(parser_description, add_slurm: bool = False):
     input_group = parser.add_argument_group("Input Arguments")
     input_group.add_argument("--config", type=str, required=True, help="Path to the CoPick configuration file.")
     input_group.add_argument("--method", type=str, choices=['watershed', 'com'], default='watershed', required=False, help="Localization method to use.")
-    input_group.add_argument('--seg-info', type=utils.parse_target, required=True, help='Query for the organelles segmentations (e.g., "name" or "name,user_id,session_id".).')
+    input_group.add_argument('--seg-info', type=parsers.parse_target, required=False, default='predict,octopi,1', help='Query for the organelles segmentations (e.g., "name" or "name,user_id,session_id".).')
     input_group.add_argument("--voxel-size", type=float, default=10, required=False, help="Voxel size for localization.")
-    input_group.add_argument("--runIDs", type=utils.parse_list, default = None, required=False, help="List of runIDs to run inference on, e.g., run1,run2,run3 or [run1,run2,run3].")
+    input_group.add_argument("--runIDs", type=parsers.parse_list, default = None, required=False, help="List of runIDs to run inference on, e.g., run1,run2,run3 or [run1,run2,run3].")
 
     localize_group = parser.add_argument_group("Localize Arguments")
     localize_group.add_argument("--radius-min-scale", type=float, default=0.5, required=False, help="Minimum radius scale for particles.")
     localize_group.add_argument("--radius-max-scale", type=float, default=1.0, required=False, help="Maximum radius scale for particles.")
     localize_group.add_argument("--filter-size", type=int, default=10, required=False, help="Filter size for localization.")
-    localize_group.add_argument("--pick-objects", type=utils.parse_list, default=None, required=False, help="Specific Objects to Find Picks for.")
+    localize_group.add_argument("--pick-objects", type=parsers.parse_list, default=None, required=False, help="Specific Objects to Find Picks for.")
     localize_group.add_argument("--n-procs", type=int, default=8, required=False, help="Number of CPU processes to parallelize runs across. Defaults to the max number of cores available or available runs.")
 
     output_group = parser.add_argument_group("Output Arguments")
     output_group.add_argument("--pick-session-id", type=str, default='1', required=False, help="Session ID for the particle picks.")
-    output_group.add_argument("--pick-user-id", type=str, default='monai', required=False, help="User ID for the particle picks.")
+    output_group.add_argument("--pick-user-id", type=str, default='octopi', required=False, help="User ID for the particle picks.")
 
     if add_slurm:
         slurm_group = parser.add_argument_group("SLURM Arguments")
@@ -165,7 +164,7 @@ def save_parameters(args: argparse.Namespace,
     pprint.pprint(params); print()
 
     # Save to YAML file
-    utils.save_parameters_yaml(params, output_path)
+    io.save_parameters_yaml(params, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_optuna.py

@@ -1,7 +1,7 @@
 from octopi.pytorch.model_search_submitter import ModelSearchSubmit
 from octopi.entry_points import common
+from octopi.utils import parsers, io
 import argparse, os, pprint
-from octopi import utils
 
 def optuna_parser(parser_description, add_slurm: bool = False):
     """
@@ -20,22 +20,22 @@ def optuna_parser(parser_description, add_slurm: bool = False):
     # Input Arguments
     input_group = parser.add_argument_group("Input Arguments")
     common.add_config(input_group, single_config=False)
-    input_group.add_argument("--target-info", type=utils.parse_target, default="targets,octopi,1", 
+    input_group.add_argument("--target-info", type=parsers.parse_target, default="targets,octopi,1", 
                              help="Target information, e.g., 'name' or 'name,user_id,session_id'")    
     input_group.add_argument("--tomo-alg", default='wbp', 
                              help="Tomogram algorithm used for training")
     input_group.add_argument("--mlflow-experiment-name", type=str, default="model-search", required=False, 
                              help="Name of the MLflow experiment (default: 'model-search').")
-    input_group.add_argument("--trainRunIDs", type=utils.parse_list, default=None, required=False, 
+    input_group.add_argument("--trainRunIDs", type=parsers.parse_list, default=None, required=False, 
                              help="List of training run IDs, e.g., run1,run2 or [run1,run2].")
-    input_group.add_argument("--validateRunIDs", type=utils.parse_list, default=None, required=False, 
+    input_group.add_argument("--validateRunIDs", type=parsers.parse_list, default=None, required=False, 
                              help="List of validation run IDs, e.g., run3,run4 or [run3,run4].")    
     input_group.add_argument('--data-split', type=str, default='0.8', help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) "
                                 "or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
 
     model_group = parser.add_argument_group("Model Arguments")
     model_group.add_argument("--model-type", type=str, default='Unet', required=False, 
-                             choices=['Unet', 'AttentionUnet'],
+                             choices=['Unet', 'AttentionUnet', 'MedNeXt', 'SegResNet'],
                              help="Model type to use for training")
     model_group.add_argument("--Nclass", type=int, default=3, required=False, help="Number of prediction classes in the model")
 
@@ -61,7 +61,7 @@ def cli():
     args = optuna_parser(description)
 
     # Parse the CoPick configuration paths
-    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    if len(args.config) > 1:    copick_configs = parsers.parse_copick_configs(args.config)
     else:                       copick_configs = args.config[0]
 
     # Create the model exploration directory
@@ -133,7 +133,7 @@ def save_parameters(args: argparse.Namespace,
     pprint.pprint(params); print()
 
     # Save to YAML file
-    utils.save_parameters_yaml(params, output_path)
+    io.save_parameters_yaml(params, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_segment_predict.py

@@ -1,8 +1,8 @@
+import torch, argparse, json, pprint, yaml, os
 from octopi.pytorch import segmentation
 from octopi.entry_points import common
-import torch, argparse, json, pprint, yaml, os
-from octopi import utils
 from typing import List, Tuple
+from octopi.utils import io
 
 def inference(
     copick_config_path: str,
@@ -136,7 +136,7 @@ def save_parameters(args: argparse.Namespace,
                     output_path: str):  
 
     # Load the model config
-    model_config = utils.load_yaml(args.model_config)
+    model_config = io.load_yaml(args.model_config)
 
     # Create parameters dictionary
     params = {
@@ -160,7 +160,7 @@ def save_parameters(args: argparse.Namespace,
     pprint.pprint(params); print()
 
     # Save to YAML file
-    utils.save_parameters_yaml(params, output_path)
+    io.save_parameters_yaml(params, output_path)
 
 if __name__ == "__main__":
     cli()

octopi/entry_points/run_train.py

@@ -2,12 +2,11 @@ from octopi.datasets import generators, multi_config_generator
 from monai.losses import DiceLoss, FocalLoss, TverskyLoss
 from octopi.models import common as builder
 from monai.metrics import ConfusionMatrixMetric
+from octopi.utils import parsers, io
 from octopi.entry_points import common 
 from octopi.pytorch import trainer 
-from octopi import io, utils
 import torch, os, argparse
 from typing import List, Optional, Tuple
-import pprint
 
 def train_model(
     copick_config_path: str,
@@ -56,7 +55,7 @@ def train_model(
     
 
     # Get the data splits
-    ratios = utils.parse_data_split(data_split)
+    ratios = parsers.parse_data_split(data_split)
     data_generator.get_data_splits(
         trainRunIDs = trainRunIDs,
         validateRunIDs = validateRunIDs,
@@ -114,11 +113,11 @@ def train_model_parser(parser_description, add_slurm: bool = False):
     # Input Arguments
     input_group = parser.add_argument_group("Input Arguments")
     common.add_config(input_group, single_config=False)
-    input_group.add_argument("--target-info", type=utils.parse_target, default="targets,octopi,1", 
+    input_group.add_argument("--target-info", type=parsers.parse_target, default="targets,octopi,1", 
                              help="Target information, e.g., 'name' or 'name,user_id,session_id'. Default is 'targets,octopi,1'.")
     input_group.add_argument("--tomo-alg", default='wbp', help="Tomogram algorithm used for training")
-    input_group.add_argument("--trainRunIDs", type=utils.parse_list, help="List of training run IDs, e.g., run1,run2,run3")
-    input_group.add_argument("--validateRunIDs", type=utils.parse_list, help="List of validation run IDs, e.g., run4,run5,run6")
+    input_group.add_argument("--trainRunIDs", type=parsers.parse_list, help="List of training run IDs, e.g., run1,run2,run3")
+    input_group.add_argument("--validateRunIDs", type=parsers.parse_list, help="List of validation run IDs, e.g., run4,run5,run6")
     input_group.add_argument('--data-split', type=str, default='0.8', help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) "
                                 "or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
     
@@ -153,11 +152,11 @@ def cli():
     args = train_model_parser(parser_description)
 
     # Parse the CoPick configuration paths
-    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    if len(args.config) > 1:    copick_configs = parsers.parse_copick_configs(args.config)
     else:                       copick_configs = args.config[0]
 
     if args.model_config:
-        model_config = utils.load_yaml(args.model_config)
+        model_config = io.load_yaml(args.model_config)
     else:
         model_config = get_model_config(args.channels, args.strides, args.res_units, args.Nclass, args.dim_in)
 

octopi/extract/localize.py

@@ -3,15 +3,15 @@ from scipy.cluster.hierarchy import fcluster, linkage
 from skimage.segmentation import watershed
 from typing import List, Optional, Tuple
 from skimage.measure import regionprops
+from copick_utils.io import readers
 from scipy.spatial import distance
 from dataclasses import dataclass
-from octopi import io
 import scipy.ndimage as ndi
 from tqdm import tqdm
 import numpy as np
 import gc
 
-def processs_localization(run,  
+def process_localization(run,  
                           objects, 
                           seg_info: Tuple[str, str, str],
                           method: str = 'com', 
@@ -27,12 +27,12 @@ def processs_localization(run,
         raise ValueError(f"Invalid method '{method}'. Expected 'watershed' or 'com'.")
 
     # Get Segmentation
-    seg = io.get_segmentation_array(run, 
-                                    voxel_size, 
-                                    seg_info[0], 
-                                    user_id=seg_info[1], 
-                                    session_id=seg_info[2],
-                                    raise_error=False)
+    seg = readers.segmentation(
+        run, voxel_size, 
+        seg_info[0], 
+        user_id=seg_info[1], 
+        session_id=seg_info[2],
+        raise_error=False)
 
     # Preprocess Segmentation
     # seg = preprocess_segmentation(seg, voxel_size, objects)
@@ -99,8 +99,8 @@ def extract_particle_centroids_via_watershed(
         max_particle_size (int): Maximum size threshold for particles.
     """
 
-    if maxima_filter_size is None or maxima_filter_size < 0:
-        AssertionError('Enter a Non-Zero Filter Size!')
+    if maxima_filter_size is None or maxima_filter_size <= 0:
+        raise ValueError('Enter a Non-Zero Filter Size!')
 
     # Calculate minimum and maximum particle volumes based on the given radii
     min_particle_size = (4 / 3) * np.pi * (min_particle_radius ** 3) 
@@ -117,12 +117,8 @@ def extract_particle_centroids_via_watershed(
     # Structuring element for erosion and dilation
     struct_elem = ball(1)
     eroded = binary_erosion(binary_mask, struct_elem)
-    del binary_mask
-    gc.collect()
 
     dilated = binary_dilation(eroded, struct_elem)
-    del eroded
-    gc.collect()
 
     # Distance transform and local maxima detection
     distance = ndi.distance_transform_edt(dilated)
@@ -131,12 +127,11 @@ def extract_particle_centroids_via_watershed(
     # Watershed segmentation
     markers, _ = ndi.label(local_max)
     del local_max
-    markers = markers.astype(np.uint8)
     gc.collect()
 
     watershed_labels = watershed(-distance, markers, mask=dilated)
+    distance, markers, dilated = None, None, None
     del distance, markers, dilated
-    watershed_labels = watershed_labels.astype(np.uint8)
     gc.collect()
 
     # Extract region properties and filter based on particle size
@@ -147,9 +142,6 @@ def extract_particle_centroids_via_watershed(
             # Option 1: Use all centroids
             all_centroids.append(region.centroid)
 
-    del watershed_labels
-    gc.collect()
-
     return all_centroids
 
 def extract_particle_centroids_via_com(