octopi 1.4.0__py3-none-any.whl

octopi/datasets/io.py

@@ -0,0 +1,200 @@
+"""
+Data loading, processing, and dataset operations for the datasets module.
+"""
+
+from monai.data import DataLoader, CacheDataset
+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), os.cpu_count() or 4)
+    test_ds = CacheDataset(data=test_files, transform=pre_transforms)
+    test_loader = DataLoader(test_ds, 
+                            batch_size=bs, 
+                            shuffle=False, 
+                            num_workers=bs, 
+                            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/mixup.py

@@ -0,0 +1,49 @@
+from monai.transforms import Transform
+from torch.distributions import Beta
+from torch import nn
+import numpy as np
+import torch
+
+class MixupTransformd(Transform):
+    """
+    A dictionary-based wrapper for Mixup augmentation that applies to batches.
+    This needs to be applied after batching, typically directly in the training loop.
+    """
+    def __init__(self, keys=("image", "label"), mix_beta=0.2, mixadd=False, prob=0.5):
+        self.keys = keys
+        self.mixup = Mixup(mix_beta=mix_beta, mixadd=mixadd)
+        self.prob = prob
+        
+    def __call__(self, data):
+        d = dict(data)
+        if np.random.random() < self.prob:  # Apply with probability
+            d[self.keys[0]], d[self.keys[1]] = self.mixup(d[self.keys[0]], d[self.keys[1]])
+        return d
+
+class Mixup(nn.Module):
+    def __init__(self, mix_beta, mixadd=False):
+
+        super(Mixup, self).__init__()
+        self.beta_distribution = Beta(mix_beta, mix_beta)
+        self.mixadd = mixadd
+
+    def forward(self, X, Y, Z=None):
+
+        bs = X.shape[0]
+        n_dims = len(X.shape)
+        perm = torch.randperm(bs)
+        coeffs = self.beta_distribution.rsample(torch.Size((bs,))).to(X.device)
+        X_coeffs = coeffs.view((-1,) + (1,)*(X.ndim-1))
+        Y_coeffs = coeffs.view((-1,) + (1,)*(Y.ndim-1))
+        
+        X = X_coeffs * X + (1-X_coeffs) * X[perm]
+
+        if self.mixadd:
+            Y = (Y + Y[perm]).clip(0, 1)
+        else:
+            Y = Y_coeffs * Y + (1 - Y_coeffs) * Y[perm]
+                
+        if Z:
+            return X, Y, Z
+
+        return X, Y

octopi/datasets/multi_config_generator.py

@@ -0,0 +1,252 @@
+from octopi.datasets import dataset, augment, cached_datset
+from octopi.datasets.generators import TrainLoaderManager
+from monai.data import DataLoader, SmartCacheDataset, CacheDataset, Dataset
+from octopi.datasets import io
+import multiprocess as mp
+from typing import List
+from tqdm import tqdm
+import torch, gc
+
+class MultiConfigTrainLoaderManager(TrainLoaderManager):
+
+    def __init__(self, 
+                 configs: dict,  # Dictionary of session names and config paths
+                 target_name: str,
+                 target_session_id: str = None,
+                 target_user_id: str = None,
+                 voxel_size: float = 10, 
+                 tomo_algorithm: List[str] = ['wbp'], 
+                 tomo_batch_size: int = 15
+                 ):
+        """
+        Initialize MultiConfigTrainLoaderManager with multiple configs.
+
+        Args:
+            configs (list): List of config file paths.
+            Other arguments are inherited from TrainLoaderManager.
+        """
+
+        # Initialize shared attributes manually (skip super().__init__ to avoid invalid config handling)
+        self.config = configs
+        self.roots = {name: io.load_copick_config(path) for name, path in configs.items()}
+
+        # Target and algorithm parameters
+        self.target_name = target_name
+        self.target_session_id = target_session_id
+        self.target_user_id = target_user_id
+        self.voxel_size = voxel_size
+        self.tomo_algorithm = tomo_algorithm
+
+        # Data management parameters
+        self.tomo_batch_size = tomo_batch_size
+        self.reload_training_dataset = True
+        self.reload_validation_dataset = True
+        self.val_loader = None
+        self.train_loader = None
+
+        # Initialize Run IDs placeholder
+        self.myRunIDs = {}
+
+        # Initialize the input dimensions   
+        self.nx = None
+        self.ny = None
+        self.nz = None        
+
+    def get_available_runIDs(self):
+        """
+        Identify and return a combined list of run IDs with available segmentations 
+        across all configured CoPick projects.
+
+        Returns:
+            List of tuples: [(session_name, run_name), ...]
+        """
+        available_runIDs = []  
+        for name, root in self.roots.items():
+            runIDs = [run.name for run in root.runs]                       
+            for run in runIDs:
+                run = root.get_run(run)
+                seg = run.get_segmentations(
+                    name=self.target_name, 
+                    session_id=self.target_session_id, 
+                    user_id=self.target_user_id,
+                    voxel_size=float(self.voxel_size)
+                )
+                if len(seg) > 0:
+                    available_runIDs.append((name, run.name))  # Include session name for disambiguation
+
+        # If No Segmentations are Found, Inform the User
+        if len(available_runIDs) == 0:
+            print(
+                f"[Error] No segmentations found for the target query:\n"
+                f"TargetName: {self.target_name}, UserID: {self.target_user_id}, "
+                f"SessionID: {self.target_session_id}\n"
+                f"Please check the target name, user ID, and session ID.\n"
+            )
+            exit()
+
+        return available_runIDs        
+
+    def get_data_splits(self, 
+                        trainRunIDs: str = None,
+                        validateRunIDs: str = None,
+                        train_ratio: float = 0.8,
+                        val_ratio: float = 0.1,
+                        test_ratio: float = 0.1,
+                        create_test_dataset: bool = True):
+        """
+        Override to handle run IDs as (session_name, run_name) tuples.
+        """
+        # Use the get_available_runIDs method to handle multiple projects
+        runIDs = self.get_available_runIDs()
+        return super().get_data_splits(trainRunIDs = runIDs, 
+                                       train_ratio = train_ratio, 
+                                       val_ratio = val_ratio, 
+                                       test_ratio = test_ratio, 
+                                       create_test_dataset = create_test_dataset)
+
+    def _initialize_train_iterators(self):
+        """
+        Initialize the training data iterators with multi-config support.
+        """
+        self.padded_train_list = self._get_padded_list(self.myRunIDs['train'], self.train_batch_size)
+        self.train_data_iter = iter(self._get_data_batches(self.padded_train_list, self.train_batch_size))
+
+    def _initialize_val_iterators(self):
+        """
+        Initialize the validation data iterators with multi-config support.
+        """
+        self.padded_val_list = self._get_padded_list(self.myRunIDs['validate'], self.val_batch_size)
+        self.val_data_iter = iter(self._get_data_batches(self.padded_val_list, self.val_batch_size))
+
+    def _load_data(self, runIDs):
+        """
+        Load data from multiple CoPick projects for given run IDs.
+
+        Args:
+            runIDs (list): List of (session_name, run_name) tuples.
+
+        Returns:
+            List: Combined dataset for the specified run IDs.
+        """
+
+        data = []
+        for session_name, run_name in tqdm(runIDs):
+            root = self.roots[session_name]
+            data.extend(io.load_training_data(
+                root, [run_name], self.voxel_size, self.tomo_algorithm, 
+                self.target_name, self.target_session_id, self.target_user_id, 
+                progress_update=False ))
+        self._check_max_label_value(data)
+        return data
+
+    def create_train_dataloaders(self, *args, **kwargs):
+        """
+        Override data loading to fetch from multiple projects.
+        """
+        my_crop_size = kwargs.get("crop_size", 96)
+        my_num_samples = kwargs.get("num_samples", 128)
+
+        # If reloads are disabled and loaders already exist, reuse them
+        if self.reload_frequency < 0 and (self.train_loader is not None) and (self.val_loader is not None):
+            return self.train_loader, self.val_loader        
+
+        # Estimate Max Number of Threads with mp.cpu_count
+        n_procs = min(mp.cpu_count(), 4)               
+
+        if self.train_loader is None: 
+            # Fetch the next batch of run IDs
+            trainRunIDs = self._extract_run_ids('train_data_iter', self._initialize_train_iterators)
+            train_files = self._load_data(trainRunIDs)
+
+            # # Create the cached dataset with non-random transforms
+            train_ds = SmartCacheDataset(data=train_files, transform=augment.get_transforms(), cache_rate=0.5)
+
+            # # Delete the training files to free memory
+            train_files = None
+            gc.collect()
+
+            # Create the cached dataset with non-random transforms
+            # train_ds = cached_datset.MultiConfigCacheDataset(
+            #     self, trainRunIDs, transform=augment.get_transforms(), cache_rate=1.0
+            # )
+
+            # I need to read (nx,ny,nz) and scale the crop size to make sure it isnt larger than nx.
+            if self.nx is None: (self.nx,self.ny,self.nz) = train_ds[0]['image'].shape[1:]
+            self.input_dim = io.get_input_dimensions(train_ds, my_crop_size)
+
+            # Wrap the cached dataset to apply random transforms during iteration
+            self.dynamic_train_dataset = dataset.DynamicDataset(
+                data=train_ds, 
+                transform=augment.get_random_transforms(self.input_dim, my_num_samples, self.Nclasses)
+            ) 
+
+            self.train_loader = DataLoader(
+                self.dynamic_train_dataset,
+                batch_size=1,
+                shuffle=True,
+                num_workers=n_procs,
+                pin_memory=torch.cuda.is_available(),             
+            )
+
+        else:
+            # Fetch the next batch of run IDs
+            trainRunIDs = self._extract_run_ids('train_data_iter', self._initialize_train_iterators)
+            train_files = self._load_data(trainRunIDs)
+            train_ds = CacheDataset(data=train_files, transform=augment.get_transforms(), cache_rate=1.0)
+            self.dynamic_train_dataset.update_data(train_ds)
+
+        # We Only Need to Reload the Validation Dataset if the Total Number of Runs is larger than 
+        # the tomo batch size
+        if self.val_loader is None: 
+
+            validateRunIDs = self._extract_run_ids('val_data_iter', self._initialize_val_iterators)             
+            val_files  = self._load_data(validateRunIDs)    
+
+            # # Create validation dataset
+            val_ds = SmartCacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0)            
+
+            # # Delete the validation files to free memory
+            val_files = None
+            gc.collect()
+
+            # Create the cached dataset with non-random transforms
+            # val_ds = cached_datset.MultiConfigCacheDataset(
+            #     self, validateRunIDs, transform=augment.get_transforms(), cache_rate=1.0
+            # )
+
+            # # I need to read (nx,ny,nz) and scale the crop size to make sure it isnt larger than nx.
+            # if self.nx is None:
+            #     (self.nx,self.ny,self.nz) = val_ds[0]['image'].shape[1:]
+
+            # if crop_size > self.nx: self.input_dim = (self.nx, crop_size, crop_size)
+            # else:                   self.input_dim = (crop_size, crop_size, crop_size)
+
+            # Wrap the cached dataset to apply random transforms during iteration
+            self.dynamic_validation_dataset = dataset.DynamicDataset( data=val_ds )
+
+            # Create validation DataLoader
+            self.val_loader  = DataLoader(
+                self.dynamic_validation_dataset,
+                batch_size=1,
+                num_workers=n_procs,
+                pin_memory=torch.cuda.is_available(),
+                shuffle=False,  # Ensure the data order remains consistent            
+            )
+        else:
+            validateRunIDs = self._extract_run_ids('val_data_iter', self._initialize_val_iterators)             
+            val_files   = self._load_data(validateRunIDs)
+
+            val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0)
+            self.dynamic_validation_dataset.update_data(val_ds)
+            
+        return self.train_loader, self.val_loader
+
+
+    def tmp_return_datasets(self):
+        trainRunIDs = self._extract_run_ids('train_data_iter', self._initialize_train_iterators)
+        train_files = self._load_data(trainRunIDs)   
+
+        validateRunIDs = self._extract_run_ids('val_data_iter', self._initialize_val_iterators)
+        val_files = self._load_data(validateRunIDs)
+
+        return train_files, val_files        

octopi/entry_points/common.py

@@ -0,0 +1,119 @@
+from octopi.utils import parsers
+import rich_click as click
+from typing import List, Tuple
+
+def model_parameters(octopi: bool = False):
+    """Decorator for adding model parameters"""
+    def decorator(f):
+        f = click.option("-dim", "--dim-in", type=int, default=96, 
+                        help="Input dimension for the UNet model")(f)
+        f = click.option("--res-units", type=int, default=1,
+                        help="Number of residual units in the UNet")(f)
+        f = click.option("-s", "--strides", type=str, default='2,2,1',
+                        callback=lambda ctx, param, value: parsers.parse_int_list(value) if value else value,
+                        help="List of stride sizes")(f)
+        f = click.option("-ch", "--channels", type=str, default='32,64,96,96',
+                        callback=lambda ctx, param, value: parsers.parse_int_list(value) if value else value,
+                        help="List of channel sizes")(f)
+        return f
+    return decorator
+
+def inference_model_parameters():
+    """Decorator for adding inference model parameters"""
+    def decorator(f):
+        f = click.option("-mw", "--model-weights", type=click.Path(exists=True), required=True,
+                        help="Path to the model weights file")(f)
+        f = click.option("-mc", "--model-config", type=click.Path(exists=True), required=True,
+                        help="Path to the model configuration file")(f)
+        return f
+    return decorator
+
+def train_parameters(octopi: bool = False):
+    """Decorator for adding training parameters"""
+    def decorator(f):
+        if octopi:
+            f = click.option("-nt", "--num-trials", type=int, default=10,
+                            help="Number of trials for architecture search")(f)
+        else:
+            f = click.option("-o", "--model-save-path", type=click.Path(), default='results',
+                            help="Path to model save directory")(f)
+            f = click.option("--tversky-alpha", type=float, default=0.3,
+                            help="Alpha parameter for the Tversky loss")(f)
+            f = click.option("-lr", "--lr", type=float, default=1e-3,
+                            help="Learning rate for the optimizer")(f)
+            f = click.option('-ncrops', "--num-tomo-crops", type=int, default=16,
+                            help="Number of tomogram crops to use per patch")(f)
+        
+        f = click.option("--best-metric", type=str, default='avg_f1',
+                        help="Metric to Monitor for Determining Best Model. To track fBetaN, use fBetaN with N as the beta-value.")(f)
+        f = click.option('-ntomos', "--tomo-batch-size", type=int, default=15,
+                        help="Number of tomograms to load per epoch for training")(f)
+        f = click.option("--val-interval", type=int, default=10,
+                        help="Interval for validation metric calculations")(f)
+        f = click.option('-nepochs', "--num-epochs", type=int, default=1000,
+                        help="Number of training epochs")(f)
+        return f
+    return decorator
+
+def config_parameters(single_config: bool):
+    """Decorator for adding config parameters"""
+    def decorator(f):
+        f = click.option("-vs", "--voxel-size", type=float, default=10,
+                        help="Voxel size of tomograms used")(f)
+        if single_config:
+            f = click.option("-c", "--config", type=click.Path(exists=True), required=True,
+                            help="Path to the configuration file")(f)
+        else:
+            f = click.option("-c", "--config", type=str, multiple=True, required=True,
+                            help="Specify a single configuration path (/path/to/config.json) or multiple entries in the format session_name,/path/to/config.json. Use multiple --config entries for multiple sessions.")(f)
+        return f
+    return decorator
+
+def inference_parameters():
+    """Decorator for adding inference parameters"""
+    def decorator(f):
+        f = click.option('-runs', "--run-ids", type=str, default=None,
+                        callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+                        help="List of run IDs for prediction, e.g., run1,run2 or [run1,run2]. If not provided, all available runs will be processed.")(f)
+        f = click.option('-ntomos', "--tomo-batch-size", type=int, default=25,
+                        help="Batch size for tomogram processing")(f)
+        f = click.option('-seginfo', "--seg-info", type=str, default='predict,octopi,1',
+                        callback=lambda ctx, param, value: parsers.parse_target(value) if value else value,
+                        help='Information Query to save Segmentation predictions under (e.g., "name" or "name,user_id,session_id" - Default UserID is octopi and SessionID is 1')(f)
+        f = click.option('-alg', "--tomo-alg", type=str, default='wbp',
+                        help="Tomogram algorithm used for produces segmentation prediction masks")(f)
+        return f
+    return decorator
+
+def localize_parameters():
+    """Decorator for adding localization parameters"""
+    def decorator(f):
+        f = click.option('-seginfo', "--seg-info", type=str, required=True,
+                        help="Segmentation info")(f)
+        f = click.option("--pick-objects", type=str, required=True,
+                        help="Pick objects")(f)
+        f = click.option("--pick-session-id", type=str, default="1",
+                        help="Pick session ID")(f)
+        f = click.option('-m', "--method", type=str, default='watershed',
+                        help="Localization method")(f)
+        f = click.option('-vs', "--voxel-size", type=int, default=10,
+                        help="Voxel size")(f)
+        return f
+    return decorator
+
+def slurm_parameters(base_job_name: str, gpus: int = 1):
+    """Decorator for adding SLURM parameters"""
+    def decorator(f):
+        if gpus > 1:
+            f = click.option("--num-gpus", type=int, default=1,
+                            help="Number of GPUs to use")(f)
+        if gpus > 0:
+            f = click.option("--gpu-constraint", type=click.Choice(['a6000', 'a100', 'h100', 'h200'], case_sensitive=False),
+                            default='h100',
+                            help="GPU constraint")(f)
+        f = click.option("--job-name", type=str, default=base_job_name,
+                        help="Job name for SLURM job")(f)
+        f = click.option("--conda-env", type=click.Path(), default='/hpc/projects/group.czii/conda_environments/pyUNET/',
+                        help="Path to Conda environment")(f)
+        return f
+    return decorator