octopi 1.0__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,9 +8,9 @@ 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,128,128', 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("--res-units", type=int, required=False, default=2, help="Number of residual units in the UNet")
+    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")
 
 def inference_model_parameters(parser):
@@ -24,7 +24,7 @@ def add_train_parameters(parser, octopi = False):
     """
     Add training parameters to the parser.
     """
-    parser.add_argument("--num-epochs", type=int, required=False, default=100, help="Number of training epochs")
+    parser.add_argument("--num-epochs", type=int, required=False, default=1000, help="Number of training epochs")
     parser.add_argument("--val-interval", type=int, required=False, default=10, help="Interval for validation metric calculations")
     parser.add_argument("--tomo-batch-size", type=int, required=False, default=15, help="Number of tomograms to load per epoch for training")
     parser.add_argument("--best-metric", type=str, default='avg_f1', required=False, help="Metric to Monitor for Determining Best Model. To track fBetaN, use fBetaN with N as the beta-value.")
@@ -32,8 +32,8 @@ def add_train_parameters(parser, octopi = False):
     if not octopi:
         parser.add_argument("--num-tomo-crops", type=int, required=False, default=16, help="Number of tomogram crops to use per patch")    
         parser.add_argument("--lr", type=float, required=False, default=1e-3, help="Learning rate for the optimizer")
-        parser.add_argument("--tversky-alpha", type=float, required=False, default=0.5, help="Alpha parameter for the Tversky loss")
-        parser.add_argument("--model-save-path", required=True, help="Path to model save directory")
+        parser.add_argument("--tversky-alpha", type=float, required=False, default=0.3, help="Alpha parameter for the Tversky loss")
+        parser.add_argument("--model-save-path", required=False, default='results', help="Path to model save directory")
     else:
         parser.add_argument("--num-trials", type=int, default=10, required=False, help="Number of trials for architecture search (default: 10).")
 
@@ -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
@@ -16,19 +16,27 @@ def create_train_script(args):
 
     command = f"""
 octopi train \\
+    {strconfigs} \\
     --model-save-path {args.model_save_path} \\
-    --target-info {args.target_info} \\
-    --voxel-size {args.voxel_size} --tomo-algorithm {args.tomo_algorithm} --Nclass {args.Nclass} \\
-    --best-metric {args.best_metric} --num-epochs {args.num_epochs} --val-interval {args.val_interval} \\
+    --target-info {','.join(args.target_info)} \\
+    --voxel-size {args.voxel_size} --tomo-alg {args.tomo_alg} --Nclass {args.Nclass} \\
     --tomo-batch-size {args.tomo_batch_size} --num-tomo-crops {args.num_tomo_crops} \\
-    {strconfigs}
-"""
+    --best-metric {args.best_metric} --num-epochs {args.num_epochs} --val-interval {args.val_interval} \\
+    """
 
     # If a model config is provided, use it to build the model
     if args.model_config is not None:
         command += f" --model-config {args.model_config}"
     else:
-        command += f" --tversky-alpha {args.tversky_alpha} --channels {args.channels} --strides {args.strides} --dim-in {args.dim_in} --res-units {args.res_units}"
+        channels = ",".join(map(str, args.channels))
+        strides = ",".join(map(str, args.strides))
+        command += (
+            f" --tversky-alpha {args.tversky_alpha}"
+            f" --channels {channels}"
+            f" --strides {strides}"
+            f" --dim-in {args.dim_in}"
+            f" --res-units {args.res_units}"
+        )
 
     # If Model Weights are provided, use them to initialize the model
     if args.model_weights is not None and args.model_config is not None:
@@ -240,4 +248,4 @@ def download_dataportal_slurm():
     """
     parser_description = "Create a SLURM script for downloading tomograms from the Dataportal"
     args = cli_dataportal_parser(parser_description, add_slurm=True)
-    create_download_dataportal_script(args)
+    create_download_dataportal_script(args)

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
@@ -30,46 +30,23 @@ def extract_membrane_bound_picks(
     if n_procs is None:
         n_procs = min(mp.cpu_count(), n_run_ids)
     print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")   
-
-    # Initialize tqdm progress bar
-    with tqdm(total=n_run_ids, desc="Membrane-Protein Isolation", unit="run") as pbar:
-        for _iz in range(0, n_run_ids, n_procs):
-
-            start_idx = _iz
-            end_idx = min(_iz + n_procs, n_run_ids)  # Ensure end_idx does not exceed n_run_ids
-            print(f"\nProcessing runIDs from {start_idx} -> {end_idx } (out of {n_run_ids})")
-
-            processes = []                
-            for _in in range(n_procs):
-                _iz_this = _iz + _in
-                if _iz_this >= n_run_ids:
-                    break
-                run_id = run_ids[_iz_this]
-                run = root.get_run(run_id)
-                p = mp.Process(
-                    target=extract.process_membrane_bound_extract,
-                    args=(run,  
-                          voxel_size, 
-                          picks_info, 
-                          membrane_info,
-                          organelle_info,
-                          save_user_id, 
-                          save_session_id,
-                          distance_threshold),
-                )
-                processes.append(p)
-
-            for p in processes:
-                p.start()
-
-            for p in processes:
-                p.join()
-
-            for p in processes:
-                p.close()
-
-            # Update tqdm progress bar
-            pbar.update(len(processes))
+    
+    # Run Membrane-Protein Isolation - Main Parallelization Loop
+    with mp.Pool(processes=n_procs) as pool:
+        with tqdm(total=n_run_ids, desc="Membrane-Protein Isolation", unit="run") as pbar:
+            worker_func = lambda run_id: extract.process_membrane_bound_extract(
+                root.get_run(run_id),  
+                voxel_size, 
+                picks_info, 
+                membrane_info,
+                organelle_info,
+                save_user_id, 
+                save_session_id,
+                distance_threshold
+            )
+
+            for _ in pool.imap_unordered(worker_func, run_ids, chunksize=1):
+                pbar.update(1)
 
     print('Extraction of Membrane-Bound Proteins Complete!')
 
@@ -81,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,6 +1,6 @@
 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
@@ -40,56 +40,39 @@ def pick_particles(
     print(', '.join([f'{obj[0]} (Label: {obj[1]})' for obj in objects]) + '\n')
 
     # Either Specify Input RunIDs or Run on All RunIDs
-    if runIDs:  print('Running Localization on the Following RunIDs: ' + ', '.join(runIDs) + '\n')
-    run_ids = runIDs if runIDs else [run.name for run in root.runs]
+    if runIDs:
+        print('Running Localization on the Following RunIDs: ' + ', '.join(runIDs) + '\n')
+        run_ids = runIDs
+    else:
+        run_ids = [run.name for run in root.runs if run.get_voxel_spacing(voxel_size) is not None]
+        skipped_run_ids = [run.name for run in root.runs if run.get_voxel_spacing(voxel_size) is None]
+
+        if skipped_run_ids:
+            print(f"Warning: skipping runs with no voxel spacing {voxel_size}: {skipped_run_ids}")
+
+    # Nprocesses shouldnt exceed computation resource or number of available runs
     n_run_ids = len(run_ids)
+    n_procs = min(mp.cpu_count(), n_procs, n_run_ids)
 
-    # Determine the number of processes to use
-    if n_procs is None:
-        n_procs = min(int(mp.cpu_count()//4), n_run_ids)
+    # Run Localization - Main Parallelization Loop
     print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")
-
-    # Initialize tqdm progress bar
-    with tqdm(total=n_run_ids, desc="Localization", unit="run") as pbar:
-        for _iz in range(0, n_run_ids, n_procs):
-
-            start_idx = _iz
-            end_idx = min(_iz + n_procs, n_run_ids)  # Ensure end_idx does not exceed n_run_ids
-            print(f"\nProcessing runIDs from {start_idx} -> {end_idx } (out of {n_run_ids})")
-
-            processes = []                
-            for _in in range(n_procs):
-                _iz_this = _iz + _in
-                if _iz_this >= n_run_ids:
-                    break
-                run_id = run_ids[_iz_this]
-                run = root.get_run(run_id)
-                p = mp.Process(
-                    target=localize.processs_localization,
-                    args=(run,  
-                          objects, 
-                          seg_info,
-                          method, 
-                          voxel_size,
-                          filter_size,
-                          radius_min_scale, 
-                          radius_max_scale,
-                          pick_session_id,
-                          pick_user_id),
-                )
-                processes.append(p)
-
-            for p in processes:
-                p.start()
-
-            for p in processes:
-                p.join()
-
-            for p in processes:
-                p.close()
-
-            # Update tqdm progress bar
-            pbar.update(len(processes))
+    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.process_localization(
+                root.get_run(run_id),  
+                objects, 
+                seg_info,
+                method, 
+                voxel_size,
+                filter_size,
+                radius_min_scale, 
+                radius_max_scale,
+                pick_session_id,
+                pick_user_id
+            )
+
+            for _ in pool.imap_unordered(worker_func, run_ids, chunksize=1):
+                pbar.update(1)
 
     print('Localization Complete!')
 
@@ -101,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("--n-procs", type=int, default=None, required=False, help="Number of CPU processes to parallelize runs across. Defaults to the max number of cores available or available runs.")
+    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")
@@ -181,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()