octopi 1.0__py3-none-any.whl

octopi/datasets/augment.py

@@ -0,0 +1,84 @@
+from monai.transforms import (
+    Compose, 
+    RandFlipd, 
+    Orientationd, 
+    RandRotate90d, 
+    NormalizeIntensityd,
+    EnsureChannelFirstd, 
+    RandCropByLabelClassesd,
+    RandScaleIntensityd,
+    RandShiftIntensityd,
+    RandAdjustContrastd,
+    RandGaussianNoised,
+    ScaleIntensityRanged,  
+    RandomOrder,
+)
+
+def get_transforms():
+    """
+    Returns non-random transforms.
+    """
+    return Compose([
+        EnsureChannelFirstd(keys=["image", "label"], channel_dim="no_channel"),
+        NormalizeIntensityd(keys="image"),
+        Orientationd(keys=["image", "label"], axcodes="RAS")
+    ])
+
+def get_random_transforms( input_dim, num_samples, Nclasses):
+    """
+    Input:
+        input_dim: tuple of (nx, ny, nz)
+        num_samples: int
+        Nclasses: int
+
+    Returns random transforms.
+    
+    For data with a missing wedge along the first axis (causing smearing in that direction),
+    we avoid rotations that would move this artifact to other axes. We only rotate around
+    the first axis (spatial_axes=[1, 2]) and avoid flipping along the first axis.
+    """
+    return Compose([
+        RandCropByLabelClassesd(
+            keys=["image", "label"],
+            label_key="label",
+            spatial_size=[input_dim[0], input_dim[1], input_dim[2]],     
+            num_classes=Nclasses,
+            num_samples=num_samples
+        ),
+        # Only rotate around the first axis (keeping the missing wedge orientation consistent)
+        RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[1, 2], max_k=3),
+        # Avoid flipping along the first axis (where the missing wedge is)
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),  # Removed
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
+        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),        
+        RandomOrder([
+            # Intensity augmentations are still appropriate
+            RandScaleIntensityd(keys="image", prob=0.5, factors=(0.85, 1.15)),
+            RandShiftIntensityd(keys="image", prob=0.5, offsets=(-0.15, 0.15)),
+            RandAdjustContrastd(keys="image", prob=0.5, gamma=(0.85, 1.15)),
+            RandGaussianNoised(keys="image", prob=0.5, mean=0.0, std=0.5),  # Reduced noise std
+        ]),
+    ]) 
+
+    # Augmentations to Explore in the Future: 
+    # Intensity-based augmentations
+    # RandHistogramShiftd(keys="image", prob=0.5, num_control_points=(3, 5))
+    # RandGaussianSmoothd(keys="image", prob=0.5, sigma_x=(0.5, 1.5), sigma_y=(0.5, 1.5), sigma_z=(0.5, 1.5)),
+
+    # Geometric Transforms
+    # RandAffined(
+    #     keys=["image", "label"],
+    #     rotate_range=(0.1, 0.1, 0.1),  # Rotation angles (radians) for x, y, z axes
+    #     scale_range=(0.1, 0.1, 0.1),   # Scale range for isotropic/anisotropic scaling
+    #     prob=0.5,                      # Probability of applying the transform
+    #     padding_mode="border"          # Handle out-of-bounds values
+    # )    
+
+def get_predict_transforms():
+    """
+    Returns predict transforms.
+    """
+    return Compose([
+        EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"),
+        NormalizeIntensityd(keys="image")
+    ])

octopi/datasets/cached_datset.py

@@ -0,0 +1,113 @@
+from typing import List, Tuple, Callable, Optional, Dict, Any
+from monai.transforms import Compose
+from monai.data import CacheDataset
+from octopi import io
+from tqdm import tqdm
+import os, sys
+
+class MultiConfigCacheDataset(CacheDataset):
+    """
+    A custom CacheDataset that loads data lazily from multiple sources
+    with consolidated loading and caching process.
+    """
+    
+    def __init__(
+        self,
+        manager,
+        run_ids: List[Tuple[str, str]],
+        transform: Optional[Callable] = None,
+        cache_rate: float = 1.0,
+        num_workers: int = 0,
+        progress: bool = True,
+        copy_cache: bool = True,
+        cache_num: int = sys.maxsize
+    ):
+        # Save reference to manager and run_ids
+        self.manager = manager
+        self.run_ids = run_ids
+        self.progress = progress
+        
+        # Prepare empty data list first - don't load immediately
+        self.data = []
+        
+        # Initialize the parent CacheDataset with an empty list
+        # We'll override the _fill_cache method to handle loading and caching in one step
+        super().__init__(
+            data=[],  # Empty list - we'll load data in _fill_cache
+            transform=transform,
+            cache_rate=cache_rate,
+            num_workers=num_workers,
+            progress=False,  # We'll handle our own progress
+            copy_cache=copy_cache,
+            cache_num=cache_num
+        )
+    
+    def _fill_cache(self):
+        """
+        Override the parent's _fill_cache method to combine loading and caching.
+        """
+        if self.progress:
+            print("Loading and caching dataset...")
+        
+        # Load and process data in a single operation
+        self.data = []
+        iterator = tqdm(self.run_ids, desc="Loading dataset") if self.progress else self.run_ids
+        
+        for session_name, run_name in iterator:
+            root = self.manager.roots[session_name]
+            batch_data = io.load_training_data(
+                root, 
+                [run_name], 
+                self.manager.voxel_size, 
+                self.manager.tomo_algorithm, 
+                self.manager.target_name, 
+                self.manager.target_session_id, 
+                self.manager.target_user_id, 
+                progress_update=False
+            )
+            
+            self.data.extend(batch_data)
+            
+            # Process and cache this batch right away
+            for i, item in enumerate(batch_data):
+                if len(self._cache) < self.cache_num and self.cache_rate > 0.0:
+                    if np.random.random() < self.cache_rate:
+                        self._cache.append(self._transform(item))
+        
+        # Check max label value if needed
+        if hasattr(self.manager, '_check_max_label_value'):
+            self.manager._check_max_label_value(self.data)
+        
+        # Update the _data attribute to match the loaded data
+        self._data = self.data
+    
+    def __len__(self):
+        """
+        Return the length of the dataset.
+        """
+        if not self.data:
+            self._fill_cache()  # Load data if not loaded yet
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        """
+        Return the item at the given index.
+        """
+        if not self.data:
+            self._fill_cache()  # Load data if not loaded yet
+        
+        # Use parent's logic for cached items
+        if index < len(self._cache):
+            return self._cache[index]
+        
+        # Otherwise transform on-the-fly
+        return self._transform(self.data[index])
+
+# TODO: Implement Single Config Cache Dataset
+# class SingleConfigCacheDataset(CacheDataset):
+#     def __init__(self, 
+#                  root: Any, 
+#                  run_ids: List[str], 
+#                  voxel_size: float, 
+#                  tomo_algorithm: str, 
+#                  target_name: str, 

octopi/datasets/dataset.py

@@ -0,0 +1,19 @@
+from torch.utils.data import Dataset
+
+class DynamicDataset(Dataset):
+    def __init__(self, data, transform=None):
+        self.data = data
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        sample = self.data[idx]
+        if self.transform:
+            sample = self.transform(sample)
+        return sample
+
+    def update_data(self, new_data):
+        """Update the internal dataset with new data."""
+        self.data = new_data

octopi/datasets/generators.py

@@ -0,0 +1,429 @@
+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
+import torch, os, random, gc
+import multiprocess as mp
+
+class TrainLoaderManager:
+
+    def __init__(self, 
+                 config: str, 
+                 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, # Number of Tomograms to Load Per Sub-Epoch    
+                 Nclasses: int = 3): # Number of Objects + Background
+
+        # Read Copick Projectdd
+        self.config = config
+        self.root = io.load_copick_config(config)
+
+        # Copick Query for Target
+        self.target_name = target_name 
+        self.target_session_id = target_session_id
+        self.target_user_id = target_user_id
+
+        # Copick Query For Input Tomogram
+        self.voxel_size = voxel_size
+        self.tomo_algorithm = tomo_algorithm
+
+        self.Nclasses = Nclasses
+        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 the input dimensions   
+        self.nx = None
+        self.ny = None
+        self.nz = None
+
+    def get_available_runIDs(self):
+        """
+        Identify and return a list of run IDs that have segmentations available for the target.
+        
+        - Iterates through all runs in the project to check for segmentations that match 
+        the specified target name, session ID, and user ID.
+        - Only includes runs that have at least one matching segmentation.
+
+        Returns:
+            available_runIDs (list): List of run IDs with available segmentations.
+        """        
+        available_runIDs = []
+        runIDs = [run.name for run in self.root.runs]
+        for run in runIDs:
+            run = self.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(run.name)
+
+        # 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):
+        """
+        Split the available data into training, validation, and testing sets based on input parameters.
+
+        Args:
+            trainRunIDs (str): Predefined list of run IDs for training. If provided, it overrides splitting logic.
+            validateRunIDs (str): Predefined list of run IDs for validation. If provided with trainRunIDs, no splitting occurs.
+            train_ratio (float): Proportion of available data to allocate to the training set.
+            val_ratio (float): Proportion of available data to allocate to the validation set.
+            test_ratio (float): Proportion of available data to allocate to the test set.
+            create_test_dataset (bool): Whether to create a test dataset or leave it empty.
+
+        Returns:
+            myRunIDs (dict): Dictionary containing run IDs for training, validation, and testing.
+        """          
+
+        # Option 1: Only TrainRunIDs are Provided, Split into Train, Validate and Test (Optional)
+        if trainRunIDs is not None and validateRunIDs is None:
+            trainRunIDs, validateRunIDs, testRunIDs = io.split_multiclass_dataset(
+                trainRunIDs, train_ratio, val_ratio, test_ratio, 
+                return_test_dataset = create_test_dataset
+            )
+        # Option 2: TrainRunIDs and ValidateRunIDs are Provided, No Need to Split
+        elif trainRunIDs is not None and validateRunIDs is not None:
+            testRunIDs = None
+        # Option 3: Use the Entire Copick Project, Split into Train, Validate and Test
+        else:
+            runIDs = self.get_available_runIDs()          
+            trainRunIDs, validateRunIDs, testRunIDs = io.split_multiclass_dataset(
+                runIDs, train_ratio, val_ratio, test_ratio, 
+                return_test_dataset = create_test_dataset
+            )
+
+        # Swap if Test Runs is Larger than Validation Runs
+        if create_test_dataset and len(testRunIDs) > len(validateRunIDs):
+            testRunIDs, validateRunIDs = validateRunIDs, testRunIDs
+
+        # Determine if datasets fit entirely in memory based on the batch size
+        # If the validation set is smaller than the batch size, avoid reloading
+        if len(validateRunIDs) < self.tomo_batch_size:
+            self.reload_validation_dataset  = False
+
+        # If the training set is smaller than the batch size, avoid reloading
+        if len(trainRunIDs) < self.tomo_batch_size:
+            self.reload_training_dataset = False
+
+        # Store the split run IDs into a dictionary for easy access
+        self.myRunIDs = {
+            'train': trainRunIDs,
+            'validate': validateRunIDs,
+            'test': testRunIDs
+        }
+
+        print(f"Number of training samples: {len(trainRunIDs)}")
+        print(f"Number of validation samples: {len(validateRunIDs)}")
+        if testRunIDs is not None:
+            print(f'Number of test samples: {len(testRunIDs)}')    
+
+        # Define separate batch sizes
+        self.train_batch_size = min( len(self.myRunIDs['train']), self.tomo_batch_size)
+        self.val_batch_size   = min( len(self.myRunIDs['validate']), self.tomo_batch_size)
+
+        # Initialize data iterators for training and validation
+        self._initialize_val_iterators()
+        self._initialize_train_iterators()
+
+        return self.myRunIDs
+    
+    def _get_padded_list(self, data_list, batch_size):
+        # Calculate padding needed to make `data_list` a multiple of `batch_size`
+        remainder = len(data_list) % batch_size
+        if remainder > 0:
+            # Number of additional items needed to make the length a multiple of batch size
+            padding_needed = batch_size - remainder
+            # Extend `data_list` with a random subset to achieve the padding
+            data_list = data_list + random.sample(data_list, padding_needed)
+        # Shuffle the full list
+        random.shuffle(data_list)
+        return data_list        
+    
+    def _initialize_train_iterators(self):
+        # Initialize padded train and validation data lists
+        self.padded_train_list = self._get_padded_list(self.myRunIDs['train'], self.train_batch_size)
+
+        # Create iterators
+        self.train_data_iter = iter(self._get_data_batches(self.padded_train_list, self.train_batch_size))
+
+    def _initialize_val_iterators(self):     
+        # Initialize padded train and validation data lists
+        self.padded_val_list = self._get_padded_list(self.myRunIDs['validate'], self.val_batch_size)
+
+        # Create iterators
+        self.val_data_iter = iter(self._get_data_batches(self.padded_val_list, self.val_batch_size))        
+
+    def _get_data_batches(self, data_list, batch_size):
+        # Generator that yields batches of specified size
+        for i in range(0, len(data_list), batch_size):
+            yield data_list[i:i + batch_size]
+
+    def _extract_run_ids(self, data_iter_name, initialize_method):
+        # Access the instance's data iterator by name
+        data_iter = getattr(self, data_iter_name)
+        try:
+            # Attempt to get the next batch from the iterator
+            runIDs = next(data_iter)
+        except StopIteration:
+            # Reinitialize the iterator if exhausted
+            initialize_method()
+            # Update the iterator reference after reinitialization
+            data_iter = getattr(self, data_iter_name)
+            runIDs = next(data_iter)
+        # Update the instance attribute with the new iterator state
+        setattr(self, data_iter_name, data_iter)
+        return runIDs
+    
+    def create_train_dataloaders(
+        self,
+        crop_size: int = 96,
+        num_samples: int = 64):
+
+        train_batch_size = 1
+        val_batch_size = 1
+
+        # 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         
+
+        # We Only Need to Reload the Training Dataset if the Total Number of Runs is larger than 
+        # the tomo batch size
+        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 = io.load_training_data(self.root, trainRunIDs, 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(train_files)
+
+            # Create the cached dataset with non-random transforms
+            train_ds = CacheDataset(data=train_files, transform=augment.get_transforms(), cache_rate=1.0)
+
+            # Delete the training files to free memory
+            train_files = None
+            gc.collect()
+
+            # 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, 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, num_samples, self.Nclasses)
+            )
+
+            # Define the number of processes for the DataLoader
+            n_procs = min(mp.cpu_count(), 4)
+
+            # DataLoader remains the same
+            self.train_loader = DataLoader(
+                self.dynamic_train_dataset,
+                batch_size=train_batch_size,
+                shuffle=False,
+                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 = io.load_training_data(self.root, trainRunIDs, 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(train_files)
+
+            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   = io.load_training_data(self.root, validateRunIDs, 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(val_files)
+
+            # Create validation dataset
+            val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0)
+
+            # Delete the validation files to free memory
+            val_files = None
+            gc.collect()
+
+            # # 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 )
+
+            dataset_size = len(self.dynamic_validation_dataset)
+            n_procs = min(mp.cpu_count(), 8)
+
+            # Create validation DataLoader
+            self.val_loader  = DataLoader(
+                self.dynamic_validation_dataset,
+                batch_size=val_batch_size,
+                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   = io.load_training_data(self.root, validateRunIDs, 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(val_files)
+
+        return self.train_loader, self.val_loader
+    
+    def get_reload_frequency(self, num_epochs: int):
+        """
+        Automatically calculate the reload frequency for the dataset during training.
+
+        Returns:
+            int: Reload frequency (number of epochs between dataset reloads).
+        """
+        if not self.reload_training_dataset:
+            # No need to reload if all tomograms fit in memory
+            print("All training samples fit in memory. No reloading required.")
+            self.reload_frequency = -1
+
+        else:
+            # Calculate the number of segments based on total training runs and batch size
+            num_segments = (len(self.myRunIDs['train']) + self.tomo_batch_size - 1) // self.tomo_batch_size
+
+            # Calculate reload frequency to distribute reloading evenly over epochs
+            self.reload_frequency = max(num_epochs // num_segments, 1)
+
+            print(f"\nReloading {self.tomo_batch_size} tomograms every {self.reload_frequency} epochs\n")
+
+            # Warn if the number of epochs is insufficient for full dataset coverage
+            if num_epochs < num_segments:
+                print(
+                    f"Warning: Chosen number of epochs ({num_epochs}) may not be sufficient "
+                    f"to train over all training samples. Consider increasing the number of epochs "
+                    f"to at least {num_segments}\n."
+                )
+
+    def _check_max_label_value(self, train_files):
+        max_label_value = max(file['label'].max() for file in train_files)
+        if max_label_value > self.Nclasses:
+            print(f"Warning: Maximum class label value {max_label_value} exceeds the number of classes {self.Nclasses}.")
+            print("This may cause issues with the model's output layer.")
+            print("Consider adjusting the number of classes or the label values in your data.\n")
+
+    def get_dataloader_parameters(self):
+
+        parameters = {
+            'config': self.config,
+            'target_name': self.target_name,
+            'target_session_id': self.target_session_id,
+            'target_user_id': self.target_user_id,
+            'voxel_size': self.voxel_size,
+            'tomo_algorithm': self.tomo_algorithm,
+            'tomo_batch_size': self.tomo_batch_size,
+            'reload_frequency': self.reload_frequency,
+            'testRunIDs': self.myRunIDs['test'],
+            'valRunIDs': self.myRunIDs['validate'],    
+            'trainRunIDs': self.myRunIDs['train'],
+        }
+
+        return parameters                
+
+class PredictLoaderManager:
+
+    def __init__(self, 
+                 config: str, 
+                 voxel_size: float = 10, 
+                 tomo_algorithm: str = 'wbp', 
+                 tomo_batch_size: int = 15, # Number of Tomograms to Load Per Sub-Epoch    
+                 Nclasses: int = 3):
+        
+        # Read Copick Project
+        self.copick_config = config
+        self.root = io.load_copick_config(config)
+
+        # Copick Query For Input Tomogram
+        self.voxel_size = voxel_size
+        self.tomo_algorithm = tomo_algorithm
+
+        self.Nclasses = Nclasses
+        self.tomo_batch_size = tomo_batch_size 
+
+        # Initialize the input dimensions   
+        self.nx = None
+        self.ny = None
+        self.nz = None
+
+
+    def create_predict_dataloader(
+        self, 
+        voxel_spacing: float, 
+        tomo_algorithm: str,       
+        runIDs: str = None):
+
+        # Split trainRunIDs, validateRunIDs, testRunIDs
+        if runIDs is None:
+            runIDs = [run.name for run in self.root.runs]
+            
+        # Load the test data
+        test_files = io.load_predict_data(self.root, runIDs, voxel_spacing, tomo_algorithm)  
+
+        # Create the cached dataset with non-random transforms
+        test_ds = CacheDataset(data=test_files, transform=augment.get_predict_transforms())
+
+        # Read (nx,ny,nz) for input tomograms.
+        if self.nx is None:
+            (self.nx,self.ny,self.nz) = test_ds[0]['image'].shape[1:]
+
+        # Create the DataLoader
+        test_loader = DataLoader(test_ds, 
+                                batch_size=4, 
+                                shuffle=False, 
+                                num_workers=4, 
+                                pin_memory=torch.cuda.is_available())
+        return test_loader
+    
+    def get_dataloader_parameters(self):
+
+        parameters = {
+            'config': self.copick_config,
+            'voxel_size': self.voxel_size,
+            'tomo_algorithm': self.tomo_algorithm
+        }
+
+        return parameters    
+    
+