octopi 1.0__py3-none-any.whl

octopi/pytorch/model_search_submitter.py

@@ -0,0 +1,290 @@
+from octopi.datasets import generators, multi_config_generator
+from octopi.pytorch import hyper_search
+import torch, mlflow, optuna
+from octopi import utils
+from typing import List
+import pandas as pd
+
+class ModelSearchSubmit:
+    def __init__(
+        self,
+        copick_config: str,
+        target_name: str,
+        target_user_id: str,
+        target_session_id: str,
+        tomo_algorithm: str,
+        voxel_size: float,
+        Nclass: int,
+        model_type: str,
+        mlflow_experiment_name: str,
+        random_seed: int,
+        num_epochs: int,
+        num_trials: int,
+        tomo_batch_size: int,
+        best_metric: str,
+        val_interval: int,
+        trainRunIDs: List[str],
+        validateRunIDs: List[str],
+        data_split: str
+    ):
+        """
+        Initialize the ModelSearch class for architecture search with Optuna.
+
+        Args:
+            copick_config (str or dict): Path to the CoPick configuration file or a dictionary for multi-config training.
+            target_name (str): Name of the target for segmentation.
+            target_user_id (str): Optional user ID for tracking.
+            target_session_id (str): Optional session ID for tracking.
+            tomo_algorithm (str): Tomogram algorithm to use.
+            voxel_size (float): Voxel size for tomograms.
+            Nclass (int): Number of prediction classes.
+            model_type (str): Type of model to use.
+            mlflow_experiment_name (str): MLflow experiment name.
+            random_seed (int): Seed for reproducibility.
+            num_epochs (int): Number of epochs per trial.
+            num_trials (int): Number of trials for hyperparameter optimization.
+            tomo_batch_size (int): Batch size for tomogram loading.
+            best_metric (str): Metric to optimize.
+            val_interval (int): Validation interval.
+            trainRunIDs (List[str]): List of training run IDs.
+            validateRunIDs (List[str]): List of validation run IDs.
+            data_split (str): Data split ratios.
+        """
+
+        # Input parameters 
+        self.copick_config = copick_config
+        self.target_name = target_name
+        self.target_user_id = target_user_id
+        self.target_session_id = target_session_id
+        self.tomo_algorithm = tomo_algorithm
+        self.voxel_size = voxel_size
+        self.Nclass = Nclass
+        self.model_type = model_type
+        self.mlflow_experiment_name = mlflow_experiment_name
+        self.random_seed = random_seed
+        self.num_epochs = num_epochs
+        self.num_trials = num_trials
+        self.tomo_batch_size = tomo_batch_size
+        self.best_metric = best_metric
+        self.val_interval = val_interval
+        self.trainRunIDs = trainRunIDs
+        self.validateRunIDs = validateRunIDs
+        self.data_split = data_split
+        
+        # Data generator - will be initialized in _initialize_data_generator()
+        self.data_generator = None
+
+        # Set random seed for reproducibility
+        utils.set_seed(self.random_seed)
+
+        # Initialize dataset generator
+        self._initialize_data_generator()
+
+    def _initialize_data_generator(self):
+        """Initializes the data generator for training and validation datasets."""
+        self._print_input_configs()
+
+        if isinstance(self.copick_config, dict):
+            self.data_generator = multi_config_generator.MultiConfigTrainLoaderManager(
+                self.copick_config,
+                self.target_name,
+                target_session_id=self.target_session_id,
+                target_user_id=self.target_user_id,
+                tomo_algorithm=self.tomo_algorithm,
+                voxel_size=self.voxel_size,
+                Nclasses=self.Nclass,
+                tomo_batch_size=self.tomo_batch_size
+            )
+        else:
+            self.data_generator = generators.TrainLoaderManager(
+                self.copick_config,
+                self.target_name,
+                target_session_id=self.target_session_id,
+                target_user_id=self.target_user_id,
+                tomo_algorithm=self.tomo_algorithm,
+                voxel_size=self.voxel_size,
+                Nclasses=self.Nclass,
+                tomo_batch_size=self.tomo_batch_size
+            )
+
+        # Split datasets into training and validation
+        ratios = utils.parse_data_split(self.data_split)
+        self.data_generator.get_data_splits(
+            trainRunIDs=self.trainRunIDs,
+            validateRunIDs=self.validateRunIDs,
+            train_ratio = ratios[0], val_ratio = ratios[1], test_ratio = ratios[2],
+            create_test_dataset = False
+        )
+        
+        # Get the reload frequency
+        self.data_generator.get_reload_frequency(self.num_epochs)
+        
+    def _print_input_configs(self):
+        """Prints training configuration for debugging purposes."""
+        print(f'\nTraining with:')
+        if isinstance(self.copick_config, dict):
+            for session, config in self.copick_config.items():
+                print(f'  {session}: {config}')
+        else:
+            print(f'  {self.copick_config}')
+        print()
+
+    def run_model_search(self):
+        """Performs model architecture search using Optuna and MLflow."""
+
+        # Set up MLflow tracking
+        try:
+            tracking_uri = utils.mlflow_setup()
+            mlflow.set_tracking_uri(tracking_uri)
+        except Exception as e:
+            print(f'Failed to set up MLflow tracking: {e}')
+            pass
+
+        mlflow.set_experiment(self.mlflow_experiment_name)
+    
+        # Create a storage object with heartbeat configuration
+        storage_url = f"sqlite:///explore_results_{self.model_type}/trials.db"
+        self.storage = optuna.storages.RDBStorage(
+            url=storage_url,
+            heartbeat_interval=60,  # Record heartbeat every minute
+            grace_period=600,       # 10 minutes grace period
+            failed_trial_callback=optuna.storages.RetryFailedTrialCallback(max_retry=1)
+        )
+
+        # Detect GPU availability
+        gpu_count = torch.cuda.device_count()
+        print(f'Running Architecture Search Over {gpu_count} GPUs\n')
+
+        # Initialize model search object
+        if gpu_count > 1:
+            self._multi_gpu_optuna(gpu_count)
+        else:
+            model_search = hyper_search.BayesianModelSearch(self.data_generator, self.model_type)
+            self._single_gpu_optuna(model_search)
+
+    def _single_gpu_optuna(self, model_search):
+        """Runs Optuna optimization on a single GPU."""
+        
+        with mlflow.start_run(nested=False) as parent_run:
+
+            model_search.parent_run_id = parent_run.info.run_id
+            model_search.parent_run_name = parent_run.info.run_name
+
+            # Log the experiment parameters
+            mlflow.log_params({"random_seed": self.random_seed})
+            mlflow.log_params(self.data_generator.get_dataloader_parameters())
+            mlflow.log_params({"parent_run_name": parent_run.info.run_name})
+
+            # Determine device
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+            # Create the study and run the optimization
+            study = self.get_optuna_study()
+            study.optimize(
+                lambda trial: model_search.objective(
+                    trial, self.num_epochs, device,
+                    val_interval=self.val_interval,
+                    best_metric=self.best_metric,
+                ),
+                n_trials=self.num_trials
+            )
+
+            # Save contour plot
+            self.save_contour_plot_as_png(study)
+
+            print(f"Best trial: {study.best_trial.value}")
+            print(f"Best params: {study.best_params}")
+
+    def _multi_gpu_optuna(self, gpu_count):
+        """Runs Optuna optimization on multiple GPUs."""
+        with mlflow.start_run() as parent_run:
+
+            # Log the experiment parameters
+            mlflow.log_params({"random_seed": self.random_seed})
+            mlflow.log_params(self.data_generator.get_dataloader_parameters())
+            mlflow.log_params({"parent_run_name": parent_run.info.run_name})
+
+            # Run multi-GPU optimization
+            study = self.get_optuna_study()
+            study.optimize(
+                lambda trial: BayesianModelSearch(self.data_generator, self.model_type).multi_gpu_objective(
+                    parent_run, trial,
+                    self.num_epochs,
+                    best_metric=self.best_metric,
+                    val_interval=self.val_interval,
+                    gpu_count=gpu_count
+                ),
+                n_trials=self.num_trials,
+                n_jobs=gpu_count
+            )
+
+            # Save contour Plot
+            self.save_contour_plot_as_png(study)
+
+        print(f"Best trial: {study.best_trial.value}")
+        print(f"Best params: {study.best_params}")
+
+    def _get_optuna_sampler(self):
+        """Returns Optuna's TPE sampler with default settings."""
+        return optuna.samplers.TPESampler(
+            n_startup_trials=10,
+            n_ei_candidates=24,
+            multivariate=True
+        )
+        # return optuna.samplers.BoTorchSampler(
+        #     n_startup_trials=10,
+        #     multivariate=True
+        # )
+
+    def get_optuna_study(self):
+        """Returns the Optuna study object."""
+        return optuna.create_study(
+                storage=self.storage,
+                direction="maximize",
+                sampler=self._get_optuna_sampler(),
+                load_if_exists=True,
+                pruner=self._get_optuna_pruner()
+            )
+
+    def _get_optuna_pruner(self):
+        """Returns Optuna's pruning strategy."""
+        return optuna.pruners.MedianPruner()
+
+    def save_contour_plot_as_png(self, study):
+        """
+        Save the contour plot of hyperparameter interactions as a PNG, 
+        automatically extracting parameter names from the study object.
+
+        Args:
+            study: The Optuna study object.
+            output_path: Path to save the PNG file.
+        """
+        # Extract all parameter names from the study trials
+        all_params = set()
+        for trial in study.trials:
+            all_params.update(trial.params.keys())
+        all_params = list(all_params)  # Convert to a sorted list for consistency
+
+        # Generate the contour plot
+        fig = optuna.visualization.plot_contour(study, params=all_params) 
+
+        # Adjust figure size and font size
+        fig.update_layout(
+            width=6000, height=6000,  # Large figure size
+            font=dict(size=40)  # Increase font size for better readability
+        )
+
+        # Save the plot as a PNG file
+        fig.write_image(f'explore_results_{self.model_type}/contour_plot.png', scale=1)  
+
+        # Extract trial data
+        trials = [
+            {**trial.params, 'objective_value': trial.value}
+            for trial in study.trials if trial.state == optuna.trial.TrialState.COMPLETE
+        ]
+
+        # Convert to DataFrame
+        df = pd.DataFrame(trials)
+
+        # Save to CSV
+        df.to_csv(f"explore_results_{self.model_type}/optuna_results.csv", index=False)

octopi/pytorch/segmentation.py

@@ -0,0 +1,317 @@
+from monai.inferers import sliding_window_inference
+from monai.data import decollate_batch
+from torch.multiprocessing import Pool
+from monai.data import MetaTensor
+from monai.transforms import (
+    Compose, AsDiscrete, Activations
+)
+import octopi.processing.writers as write
+from octopi.models import common
+from typing import List, Optional
+import torch, copick, gc, os
+from octopi import io, utils
+from tqdm import tqdm
+import numpy as np
+
+class Predictor:
+
+    def __init__(self, 
+                 config: str,
+                 model_config: str,
+                 model_weights: str,
+                 apply_tta: bool = True,
+                 device: Optional[str] = None):
+
+        self.config = config
+        self.root = copick.from_file(config)
+
+        # Load the model config
+        model_config = utils.load_yaml(model_config)
+
+        self.Nclass = model_config['model']['num_classes']     
+        self.dim_in = model_config['model']['dim_in']
+        self.input_dim = None
+        
+        # Get the number of GPUs available
+        num_gpus = torch.cuda.device_count()
+        if num_gpus == 0:
+            raise RuntimeError("No GPUs available.")
+
+        # Set the device
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = device
+        print('Running Inference On: ', self.device)
+
+        # Check to see if the model weights file exists
+        if not os.path.exists(model_weights):
+            raise ValueError(f"Model weights file does not exist: {model_weights}")
+
+        # Load the model weights
+        model_builder = common.get_model(model_config['model']['architecture'])
+        model_builder.build_model(model_config['model'])
+        self.model = model_builder.model
+        state_dict = torch.load(model_weights, map_location=self.device, weights_only=True)
+        self.model.load_state_dict(state_dict)
+        self.model.to(self.device)
+        self.model.eval()
+
+       # Initialize TTA if enabled
+        self.apply_tta = apply_tta
+        if self.apply_tta: 
+            self.create_tta_augmentations() 
+            # self.post_transforms = Compose([
+            #     Activations(softmax=True)  # Keep probability output
+            # ])
+            self.softmax_transform = Compose([
+                Activations(softmax=True)  # Keep probability output
+            ])
+            
+            # Create the final discretization transform
+            self.discretize_transform = AsDiscrete(argmax=True)
+        else:
+            # Define the post-processing transforms
+            self.post_transforms = Compose([
+                Activations(softmax=True),
+                AsDiscrete(argmax=True)
+            ])
+            
+        
+    def _run_inference(self, input):
+        """Apply sliding window inference to the input."""
+        with torch.no_grad():
+            predictions = sliding_window_inference(
+                inputs=input,
+                roi_size=(self.dim_in, self.dim_in, self.dim_in),
+                sw_batch_size=4,  # one window is proecessed at a time
+                predictor=self.model,
+                overlap=0.5,
+            )
+            return [self.post_transforms(i) for i in decollate_batch(predictions)]
+
+    def _run_inference_tta(self, input_data):
+        """Memory-efficient TTA implementation that returns proper discrete segmentation maps."""
+        
+        batch_size = input_data.shape[0]
+        results = []
+        
+        # Process one sample at a time
+        for sample_idx in range(batch_size):
+            # Extract single sample
+            single_sample = input_data[sample_idx:sample_idx+1]
+            
+            # Initialize probability accumulator for this sample
+            # Shape: [1, Nclass, Z, Y, X]
+            acc_probs = torch.zeros(
+                (1, self.Nclass, *single_sample.shape[2:]), 
+                dtype=torch.float32, device=self.device
+            )
+            
+            # Process each augmentation
+            with torch.no_grad():
+                for tta_transform, inverse_transform in zip(self.tta_transforms, self.inverse_tta_transforms):
+                    # Apply transform to single sample
+                    aug_sample = tta_transform(single_sample)
+                    
+                    # Free memory
+                    torch.cuda.empty_cache()
+                    
+                    # Run inference (one sample at a time)
+                    predictions = sliding_window_inference(
+                        inputs=aug_sample,
+                        roi_size=(self.dim_in, self.dim_in, self.dim_in),
+                        sw_batch_size=4,  # Process one window at a time
+                        predictor=self.model,
+                        overlap=0.5,
+                    )
+                    
+                    # Get softmax probabilities
+                    probs = self.softmax_transform(predictions[0])  # Get first (only) item
+                    
+                    # Apply inverse transform with correct dimensions
+                    inv_probs = inverse_transform(probs)
+                    
+                    # Accumulate probabilities
+                    acc_probs[0] += inv_probs
+                    
+                    # Clear memory
+                    del predictions, probs, inv_probs, aug_sample
+            
+            # Average accumulated probabilities
+            acc_probs = acc_probs / len(self.tta_transforms)
+            
+            # Convert to discrete prediction - get argmax along class dimension
+            # This gives us a tensor of shape [1, Z, Y, X] with discrete class indices
+            discrete_pred = torch.argmax(acc_probs, dim=1)
+            
+            # Add to results - keeping only the spatial dimensions [Z, Y, X]
+            results.append(discrete_pred[0])
+            
+            # Clear memory
+            del acc_probs, discrete_pred
+            torch.cuda.empty_cache()
+        
+        return results
+        
+    def predict_on_gpu(self, 
+                        runIDs: List[str],
+                        voxel_spacing: float,
+                        tomo_algorithm: str ):
+
+        # Load data for the current batch
+        test_loader, test_dataset = io.create_predict_dataloader(
+            self.root,
+            voxel_spacing, tomo_algorithm,
+            runIDs)
+        
+        # Determine Input Crop Size.
+        if self.input_dim is None:
+            self.input_dim = io.get_input_dimensions(test_dataset, self.dim_in)
+        
+        predictions = []
+        with torch.no_grad():
+            for data in tqdm(test_loader):
+                tomogram = data['image'].to(self.device)
+                if self.apply_tta: data['pred'] = self._run_inference_tta(tomogram)
+                else:              data['pred']  = self._run_inference(tomogram)
+                for idx in range(len(data['image'])):
+                    predictions.append(data['pred'][idx].squeeze(0).numpy(force=True)) 
+
+        return predictions
+
+    def batch_predict(self, 
+                      num_tomos_per_batch = 15, 
+                      runIDs: Optional[str] = None,
+                      voxel_spacing: float = 10,
+                      tomo_algorithm: str = 'denoised', 
+                      segmentation_name: str = 'prediction',
+                      segmentation_user_id: str = 'octopi',
+                      segmentation_session_id: str = '0'):
+
+        """Run inference on tomograms in batches."""                          
+        
+        # If runIDs are not provided, load all runs
+        if runIDs is None:
+            runIDs = [run.name for run in self.root.runs]
+        
+        # Iterate over batches of runIDs
+        for i in range(0, len(runIDs), num_tomos_per_batch):
+
+            # Get a batch of runIDs
+            batch_ids = runIDs[i:i + num_tomos_per_batch]  
+            print('Running Inference on the Follow RunIDs: ', batch_ids)
+
+            predictions = self.predict_on_gpu(batch_ids, voxel_spacing, tomo_algorithm)
+
+            # Save Predictions to Corresponding RunID
+            for ind in range(len(batch_ids)):
+                run = self.root.get_run(batch_ids[ind])
+                seg = predictions[ind]
+                write.segmentation(run, seg, segmentation_user_id, segmentation_name, 
+                                   segmentation_session_id, voxel_spacing)
+
+            # After processing and saving predictions for a batch:
+            del predictions  # Remove reference to the list holding prediction arrays
+            torch.cuda.empty_cache()  # Clear unused GPU memory
+            gc.collect()  # Trigger garbage collection for CPU memory
+
+        print('Predictions Complete!')
+
+    def create_tta_augmentations(self):
+        """Define TTA augmentations and inverse transforms."""
+
+        # Instead of Flip lets rotate around the first axis 3 times (90,180,270)
+        self.tta_transforms = [
+            lambda x: x,                    # Identity (no augmentation)
+            lambda x: torch.rot90(x, k=1, dims=(3, 4)),  # 90° rotation
+            lambda x: torch.rot90(x, k=2, dims=(3, 4)),  # 180° rotation
+            lambda x: torch.rot90(x, k=3, dims=(3, 4)),  # 270° rotation
+            # Flip(spatial_axis=0),         # Flip along x-axis (depth)
+            # Flip(spatial_axis=1),         # Flip along y-axis (height)
+            # Flip(spatial_axis=2),         # Flip along z-axis (width)
+        ]
+
+        # Define inverse transformations (flip back to original orientation)
+        self.inverse_tta_transforms = [
+            lambda x: x,                           # Identity (no transformation needed)
+            lambda x: torch.rot90(x, k=-1, dims=(2, 3)),  # Inverse of 90° (i.e. -90°)
+            lambda x: torch.rot90(x, k=-2, dims=(2, 3)),  # Inverse of 180° (i.e. -180°)
+            lambda x: torch.rot90(x, k=-3, dims=(2, 3)),  # Inverse of 270° (i.e. -270°)
+            # Flip(spatial_axis=0),         # Undo Flip along x-axis
+            # Flip(spatial_axis=1),         # Undo Flip along y-axis
+            # Flip(spatial_axis=2),         # Undo Flip along z-axis
+        ]
+
+###################################################################################################################################################
+
+class MultiGPUPredictor(Predictor):
+
+    def __init__(self, 
+                 config: str,
+                 model_config: str,
+                 model_weights: str):
+        super().__init__(config, model_config, model_weights)
+        self.num_gpus = torch.cuda.device_count()
+        if self.num_gpus < 2:
+            raise RuntimeError("MultiGPUPredictor requires at least 2 GPUs.")
+    
+    def predict_on_gpu(self, gpu_id: int, batch_ids: List[str], voxel_spacing: float, tomo_algorithm: str) -> List[np.ndarray]:
+        """Helper function to run inference on a single GPU."""
+        device = torch.device(f'cuda:{gpu_id}')
+        self.model.to(device)
+        
+        # Load data specific to the batch assigned to this GPU
+        test_loader = io.load_predict_data(self.root, batch_ids, voxel_spacing, tomo_algorithm)
+        predictions = []
+        
+        with torch.no_grad():
+            for data in tqdm(test_loader, desc=f"GPU {gpu_id}"):
+                tomogram = data['image'].to(device)
+                data["prediction"] = self.run_inference(tomogram)
+                data = [self.post_processing(i) for i in decollate_batch(data)]
+                for b in data:
+                    predictions.append(b['prediction'].squeeze(0).cpu().numpy())
+                    
+        return predictions
+
+    def multi_gpu_inference(self, 
+                            num_tomos_per_batch: int = 15, 
+                            runIDs: Optional[List[str]] = None,
+                            voxel_spacing: float = 10,
+                            tomo_algorithm: str = 'denoised', 
+                            save: bool = False,
+                            segmentation_name: str = 'prediction',
+                            segmentation_user_id: str = 'monai',
+                            segmentation_session_id: str = '0') -> Optional[List[np.ndarray]]:
+        """Run inference across multiple GPUs, optionally saving results or returning predictions."""
+        
+        runIDs = runIDs or [run.name for run in self.root.runs]
+        all_predictions = []
+
+        # Divide runIDs into batches for each GPU
+        batches = [runIDs[i:i + num_tomos_per_batch] for i in range(0, len(run_ids), num_tomos_per_batch)]
+        
+        # Run inference in parallel across GPUs
+        for i in range(0, len(batches), self.num_gpus):
+            gpu_batches = batches[i:i + self.num_gpus]
+            with Pool(processes=self.num_gpus) as pool:
+                results = pool.starmap(
+                    self.predict_on_gpu,
+                    [(gpu_id, gpu_batches[gpu_id], voxel_spacing, tomo_algorithm) for gpu_id in range(len(gpu_batches))]
+                )
+
+            # Collect and save results
+            for gpu_id, predictions in enumerate(results):
+                if save:
+                    for idx, run_id in enumerate(gpu_batches[gpu_id]):
+                        run = self.root.get_run(run_id)
+                        segmentation = predictions[idx]
+                        write.segmentation(run, segmentation, segmentation_user_id, segmentation_name, 
+                                           segmentation_session_id, voxel_spacing)
+                else:
+                    all_predictions.extend(predictions)
+
+        print('Multi-GPU predictions complete.')
+        
+        return None if save else all_predictions