octopi 1.0__py3-none-any.whl

octopi/processing/segmentation_from_picks.py

@@ -0,0 +1,167 @@
+# This code is adapted from the copick-utils project,
+# originally available at: https://github.com/copick/copick-utils/blob/main/src/copick_utils/segmentation/segmentation_from_picks.py
+# Licensed under the MIT License.
+
+# Copyright (c) 2023 The copick-utils authors
+
+import numpy as np
+import zarr
+from scipy.ndimage import zoom
+import copick
+
+def from_picks(pick,
+               seg_volume,
+               radius: float = 10.0,
+               label_value: int = 1,
+               voxel_spacing: float = 10):
+    """
+    Paints picks into a segmentation volume as spheres.
+
+    Parameters:
+    -----------
+    pick : copick.models.CopickPicks
+        Copick object containing `points`, where each point has a `location` attribute with `x`, `y`, `z` coordinates.
+    seg_volume : numpy.ndarray
+        3D segmentation volume (numpy array) where the spheres are painted. Shape should be (Z, Y, X).
+    radius : float, optional
+        The radius of the spheres to be inserted in physical units (not voxel units). Default is 10.0.
+    label_value : int, optional
+        The integer value used to label the sphere regions in the segmentation volume. Default is 1.
+    voxel_spacing : float, optional
+        The spacing of voxels in the segmentation volume, used to scale the radius of the spheres. Default is 10.
+    Returns:
+    --------
+    numpy.ndarray
+        The modified segmentation volume with spheres inserted at pick locations.    
+    """
+    def create_sphere(shape, center, radius, val):
+        zc, yc, xc = center
+        z, y, x = np.indices(shape)
+        distance_sq = (x - xc)**2 + (y - yc)**2 + (z - zc)**2
+        sphere = np.zeros(shape, dtype=np.float32)
+        sphere[distance_sq <= radius**2] = val
+        return sphere
+
+    def get_relative_target_coordinates(center, delta, shape):
+        low = max(int(np.floor(center - delta)), 0)
+        high = min(int(np.ceil(center + delta + 1)), shape)
+        return low, high
+
+    # Adjust radius for voxel spacing
+    radius_voxel = max(radius / voxel_spacing, 1)
+    delta = int(np.ceil(radius_voxel))
+
+    # Paint each pick as a sphere
+    for point in pick.points:
+        # Convert the pick's location from angstroms to voxel units
+        cx, cy, cz = point.location.x / voxel_spacing, point.location.y / voxel_spacing, point.location.z / voxel_spacing
+
+        # Calculate subarray bounds
+        xLow, xHigh = get_relative_target_coordinates(cx, delta, seg_volume.shape[2])
+        yLow, yHigh = get_relative_target_coordinates(cy, delta, seg_volume.shape[1])
+        zLow, zHigh = get_relative_target_coordinates(cz, delta, seg_volume.shape[0])
+
+        # Subarray shape
+        subarray_shape = (zHigh - zLow, yHigh - yLow, xHigh - xLow)
+        if any(dim <= 0 for dim in subarray_shape):
+            continue
+
+        # Compute the local center of the sphere within the subarray
+        local_center = (cz - zLow, cy - yLow, cx - xLow)
+        sphere = create_sphere(subarray_shape, local_center, radius_voxel, label_value)
+
+        # Assign Sphere to Segmentation Target Volume
+        seg_volume[zLow:zHigh, yLow:yHigh, xLow:xHigh] = np.maximum(seg_volume[zLow:zHigh, yLow:yHigh, xLow:xHigh], sphere)
+
+    return seg_volume
+
+
+def downsample_to_exact_shape(array, target_shape):
+    """
+    Downsamples a 3D array to match the target shape using nearest-neighbor interpolation.
+    Ensures that the resulting array has the exact target shape.
+    """
+    zoom_factors = [t / s for t, s in zip(target_shape, array.shape)]
+    return zoom(array, zoom_factors, order=0)
+
+
+def segmentation_from_picks(radius, painting_segmentation_name, run, voxel_spacing, tomo_type, pickable_object, pick_set, user_id="paintedPicks", session_id="0"):
+    """
+    Paints picks from a run into a multiscale segmentation array, representing them as spheres in 3D space.
+
+    Parameters:
+    -----------
+    radius : float
+        Radius of the spheres in physical units.
+    painting_segmentation_name : str
+        The name of the segmentation dataset to be created or modified.
+    run : copick.Run
+        The current Copick run object.
+    voxel_spacing : float
+        The spacing of the voxels in the tomogram data.
+    tomo_type : str
+        The type of tomogram to retrieve.
+    pickable_object : copick.models.CopickObject
+        The object that defines the label value to be used in segmentation.
+    pick_set : copick.models.CopickPicks
+        The set of picks containing the locations to paint spheres.
+    user_id : str, optional
+        The ID of the user creating the segmentation. Default is "paintedPicks".
+    session_id : str, optional
+        The session ID for this segmentation. Default is "0".
+
+    Returns:
+    --------
+    copick.Segmentation
+        The created or modified segmentation object.
+    """
+    # Fetch the tomogram and determine its multiscale structure
+    tomogram = run.get_voxel_spacing(voxel_spacing).get_tomogram(tomo_type)
+    if not tomogram:
+        raise ValueError("Tomogram not found for the given parameters.")
+
+    # Use copick to create a new segmentation if one does not exist
+    segs = run.get_segmentations(user_id=user_id, session_id=session_id, is_multilabel=True, name=painting_segmentation_name, voxel_size=voxel_spacing)
+    if len(segs) == 0:
+        seg = run.new_segmentation(voxel_spacing, painting_segmentation_name, session_id, True, user_id=user_id)
+    else:
+        seg = segs[0]
+
+    segmentation_group = zarr.open(seg.zarr(), mode="a")
+    highest_res_name = "0"
+
+    # Get the highest resolution dimensions and create a new array if necessary
+    tomogram_zarr = zarr.open(tomogram.zarr(), "r")
+
+    highest_res_shape = tomogram_zarr[highest_res_name].shape
+    if highest_res_name not in segmentation_group:
+        segmentation_group.create(highest_res_name, shape=highest_res_shape, dtype=np.uint16, overwrite=True)
+
+    # Initialize or load the highest resolution array
+    highest_res_seg = segmentation_group[highest_res_name][:]
+    highest_res_seg.fill(0)
+
+    # Paint picks into the highest resolution array
+    highest_res_seg = from_picks(pick_set, highest_res_seg, radius, pickable_object.label, voxel_spacing)
+
+    # Write back the highest resolution data
+    segmentation_group[highest_res_name][:] = highest_res_seg
+
+    # Downsample to create lower resolution scales
+    multiscale_metadata = tomogram_zarr.attrs.get('multiscales', [{}])[0].get('datasets', [])
+    for level_index, level_metadata in enumerate(multiscale_metadata):
+        if level_index == 0:
+            continue
+
+        level_name = level_metadata.get("path", str(level_index))
+        expected_shape = tuple(tomogram_zarr[level_name].shape)
+
+        # Compute scaling factors relative to the highest resolution shape
+        scaled_array = downsample_to_exact_shape(highest_res_seg, expected_shape)
+
+        # Create/overwrite the Zarr array for this level
+        segmentation_group.create_dataset(level_name, shape=expected_shape, data=scaled_array, dtype=np.uint16, overwrite=True)
+
+        segmentation_group[level_name][:] = scaled_array
+
+    return seg

octopi/processing/writers.py

@@ -0,0 +1,102 @@
+# This code is adapted from the copick-utils project,
+# originally available at: https://github.com/copick/copick-utils/blob/main/src/copick_utils/writers/write.py
+# Licensed under the MIT License.
+
+# Copyright (c) 2023 The copick-utils authors
+
+from typing import Any, Dict, List
+import numpy as np
+
+def tomogram(
+    run,
+    input_volume,
+    voxel_size=10,
+    algorithm="wbp"
+):
+    """
+    Writes a volumetric tomogram into an OME-Zarr format within a Copick directory.
+
+    Parameters:
+    -----------
+    run : copick.Run
+        The current Copick run object.
+    input_volume : np.ndarray
+        The volumetric tomogram data to be written.
+    voxel_size : float, optional
+        The size of the voxels in physical units. Default is 10.
+    algorithm : str, optional
+        The tomographic reconstruction algorithm to use. Default is 'wbp'.
+
+    Returns:
+    --------
+    copick.Tomogram
+        The created or modified tomogram object.
+    """
+    
+    # Retrieve or create voxel spacing
+    voxel_spacing = run.get_voxel_spacing(voxel_size)
+    if voxel_spacing is None:
+        voxel_spacing = run.new_voxel_spacing(voxel_size=voxel_size)
+    
+    # Check if We Need to Create a New Tomogram for Given Algorithm
+    tomogram = voxel_spacing.get_tomogram(algorithm)
+    if tomogram is None:
+        tomogram = voxel_spacing.new_tomogram(tomo_type=algorithm)
+    
+    # Write the tomogram data
+    tomogram.from_numpy(input_volume)
+
+
+def segmentation(
+    run,
+    segmentation_volume,
+    user_id,
+    name="segmentation",
+    session_id="0",
+    voxel_size=10,
+    multilabel=True
+):
+    """
+    Writes a segmentation into an OME-Zarr format within a Copick directory.
+
+    Parameters:
+    -----------
+    run : copick.Run
+        The current Copick run object.
+    segmentation_volume : np.ndarray
+        The segmentation data to be written.
+    user_id : str
+        The ID of the user creating the segmentation.
+    name : str, optional
+        The name of the segmentation dataset to be created or modified. Default is 'segmentation'.
+    session_id : str, optional
+        The session ID for this segmentation. Default is '0'.
+    voxel_size : float, optional
+        The size of the voxels in physical units. Default is 10.
+    multilabel : bool, optional
+        Whether the segmentation is a multilabel segmentation. Default is True.
+
+    Returns:
+    --------
+    copick.Segmentation
+        The created or modified segmentation object.
+    """
+    
+    # Retrieve or create a segmentation
+    segmentations = run.get_segmentations(name=name, user_id=user_id, session_id=session_id)
+
+    # If no segmentation exists or no segmentation at the given voxel size, create a new one
+    if len(segmentations) == 0 or any(seg.voxel_size != voxel_size for seg in segmentations):
+        segmentation = run.new_segmentation(
+            voxel_size=voxel_size,
+            name=name,
+            session_id=session_id,
+            is_multilabel=multilabel,
+            user_id=user_id
+        )
+    else:
+        # Overwrite the current segmentation at the specified voxel size if it exists
+        segmentation = next(seg for seg in segmentations if seg.voxel_size == voxel_size)
+    
+    # Write the segmentation data
+    segmentation.from_numpy(segmentation_volume, dtype=np.uint8)

octopi/pytorch/hyper_search.py

@@ -0,0 +1,243 @@
+from monai.losses import FocalLoss, TverskyLoss
+from monai.metrics import ConfusionMatrixMetric
+from octopi.pytorch import trainer
+from mlflow.tracking import MlflowClient
+from octopi.models import common
+from octopi import io, losses
+import torch, mlflow, optuna, gc
+
+class BayesianModelSearch:
+
+    def __init__(self, data_generator, model_type="Unet", parent_run_id=None, parent_run_name=None):
+        """
+        Class to handle model creation, training, and optimization.
+
+        Args:
+            data_generator (object): Data generator object containing dataset properties.
+            model_type (str): Type of model to build ("UNet", "AttentionUnet").
+        """
+        self.data_generator = data_generator
+        self.Nclasses = data_generator.Nclasses
+        self.device = None
+        self.model_type = model_type
+        self.model = None
+        self.loss_function = None
+        self.metrics_function = None
+        self.sampling = None
+        self.parent_run_id = parent_run_id
+        
+        # Define results directory path
+        self.results_dir = f'explore_results_{self.model_type}'            
+
+    def my_build_model(self, trial):
+        """Builds and initializes a model based on Optuna-suggested parameters."""
+       
+        # Build the model
+        self.model_builder = common.get_model(self.model_type)
+        self.model_builder.bayesian_search(trial, self.Nclasses)
+        self.model = self.model_builder.model.to(self.device)
+        self.config = self.model_builder.config
+
+        # Define loss function
+        self.loss_function = common.get_loss_function(trial)
+
+        # Define metrics
+        self.metrics_function = ConfusionMatrixMetric(
+            include_background=False,
+            metric_name=["recall", "precision", "f1 score"],
+            reduction="none"
+        )
+
+        # Sample crop size and num_samples
+        self.sampling = {
+            'crop_size': trial.suggest_int("crop_size", 48, 160, step=16),
+            'num_samples': 8
+        }
+        self.config['dim_in'] = self.sampling['crop_size']
+
+    def _define_optimizer(self):
+        # Define optimizer
+        lr0 = 1e-3
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr0, weight_decay=1e-5)
+
+    def _train_model(self, trial, model_trainer, epochs, val_interval, crop_size, num_samples, best_metric):
+        """Handles model training and error handling."""
+        try:
+            results = model_trainer.train(
+                self.data_generator,
+                model_save_path=None,
+                crop_size=crop_size,
+                max_epochs=epochs,
+                val_interval=val_interval,
+                my_num_samples=num_samples,
+                best_metric=best_metric,
+                use_mlflow=True,
+                verbose=False
+            )
+            return results['best_metric']
+
+        except torch.cuda.OutOfMemoryError:
+            print(f"[Trial Failed] OOM Error for model={model_trainer.model}, crop_size={crop_size}, num_samples={num_samples}")
+            trial.set_user_attr("out_of_memory", True)
+            raise optuna.TrialPruned()
+
+        except Exception as e:
+            print(f"[Trial Failed] Unexpected error: {e}")
+            trial.set_user_attr("error", str(e))
+            raise optuna.TrialPruned()        
+
+    def objective(self, trial, epochs, device, val_interval=15, best_metric="avg_f1"):
+        """Runs the full training process for a given trial."""
+
+        # Set device
+        self.device = device
+        
+        # Set a unique run name for each trial
+        trial_num = f"trial_{trial.number}"
+
+        # Start MLflow run
+        with mlflow.start_run(run_name=trial_num, nested=True):
+            
+            # Build model
+            self.my_build_model(trial)
+
+            # Create trainer
+            self._define_optimizer()
+            model_trainer = trainer.ModelTrainer(self.model, self.device, self.loss_function, self.metrics_function, self.optimizer)
+
+            # Train model and evaluate score
+            score = self._train_model(
+                trial, model_trainer, epochs, val_interval, 
+                self.sampling['crop_size'], self.sampling['num_samples'], 
+                best_metric)               
+
+            # Log parameters and metrics
+            params = {
+                'model': self.model_builder.get_model_parameters(),
+                'optimizer': io.get_optimizer_parameters(model_trainer)
+            }
+            model_trainer.my_log_params(io.flatten_params(params))
+
+            # Explicitly set the parent run ID
+            mlflow.log_param("parent_run_id", self.parent_run_id)
+            mlflow.log_param("parent_run_name", self.parent_run_name)         
+
+            # Save best model
+            self._save_best_model(trial, model_trainer, score)
+
+            # Cleanup
+            self.cleanup(model_trainer, self.optimizer)
+
+        return score        
+
+    def _setup_parallel_trial_run(self, trial, parent_run=None, gpu_count=1):
+        """Set up parallel MLflow runs and assign GPU for the trial."""
+        
+        trial_num = f"trial_{trial.number}"
+
+        # Create a child run under the parent MLflow experiment
+        mlflow_client = MlflowClient()
+        self.trial_run = mlflow_client.create_run(
+            experiment_id=mlflow_client.get_run(parent_run.info.run_id).info.experiment_id,
+            tags={"mlflow.parentRunId": parent_run.info.run_id},
+            run_name=trial_num
+        )
+        target_run_id = self.trial_run.info.run_id
+        print(f"Logging trial {trial.number} data to MLflow run: {target_run_id}")
+
+        # Assign GPU device
+        if gpu_count > 1:
+            gpu_id = trial.number % gpu_count
+            device = torch.device(f"cuda:{gpu_id}")
+            torch.cuda.set_device(device)
+        else:
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        return device, mlflow_client, target_run_id
+    
+    def multi_gpu_objective(self, parent_run, trial, epochs, val_interval=5, best_metric="avg_f1", gpu_count=1):
+        """
+        Trains model on multiple GPUs using a parent MLflow run.
+        """
+        self.device, self.client, self.target_run_id = self._setup_parallel_trial_run(trial, parent_run, gpu_count)
+        
+        # Build model
+        self.my_build_model(trial)
+
+        # Create trainer
+        self._define_optimizer()
+        model_trainer = trainer.ModelTrainer(self.model, self.device, self.loss_function, self.metrics_function, self.optimizer)
+        model_trainer.set_parallel_mlflow(self.client, self.target_run_id)
+
+        # # Train model
+        # score = self._train_model(
+        #     trial, model_trainer, epochs, val_interval, 
+        #     self.sampling['crop_size'], self.sampling['num_samples'], 
+        #     best_metric)
+
+        # Train Model, with error handling
+        try:
+            score = self._train_model(
+                trial, model_trainer, epochs, val_interval, 
+                self.sampling['crop_size'], self.sampling['num_samples'], 
+                best_metric)
+        except Exception as e:
+            print(f"[Trial Failed] Unexpected error: {e}")
+            trial.set_user_attr("error", str(e))
+            raise optuna.TrialPruned()
+        
+        # Log training parameters
+        params = {
+            'model': self.model_builder.get_model_parameters(),
+            'optimizer': io.get_optimizer_parameters(model_trainer)
+        }    
+        model_trainer.my_log_params(io.flatten_params(params))
+        model_trainer.my_log_params({"parent_run_name": parent_run.info.run_name})
+
+        # Save best model
+        self._save_best_model(trial, model_trainer, score)
+
+        # Cleanup
+        self.cleanup(model_trainer, self.optimizer)
+        return score
+
+    def _save_best_model(self, trial, model_trainer, score):
+        """Saves the best model if it improves upon previous scores."""
+        best_score_so_far = self.get_best_score(trial)
+        if score > best_score_so_far:
+            torch.save(model_trainer.model_weights, f'{self.results_dir}/best_model.pth')
+            io.save_parameters_to_yaml(self.model_builder, model_trainer, self.data_generator, 
+                                    f'{self.results_dir}/best_model_config.yaml')
+
+    def get_best_score(self, trial):
+        """Retrieve the best score from the trial."""
+        try:
+            return trial.study.best_value
+        except ValueError:
+            return 0
+
+    def cleanup(self, model_trainer, optimizer):
+        """Handles cleanup of resources."""
+
+        # Delete the trainer and optimizer objects
+        del model_trainer, optimizer
+
+        # If the model object holds GPU memory, delete it explicitly and set it to None
+        if hasattr(self, "model"):
+            del self.model
+            self.model = None
+
+        # Optional: If your model_builder or other objects hold GPU references, delete them too
+        if hasattr(self, "model_builder"):
+            del self.model_builder
+            self.model_builder = None
+
+        # Clear the CUDA cache and force garbage collection
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        # Try Terminating Multi-GPU Runs, if not run current run (works for single GPU runs)
+        try:
+            self.client.set_terminated(self.target_run_id, status="FINISHED")
+        except:
+            mlflow.end_run()