octopi 1.0__py3-none-any.whl

octopi/models/AttentionUnet.py

@@ -0,0 +1,56 @@
+from monai.networks.nets import AttentionUnet
+import torch.nn as nn
+import torch
+
+class myAttentionUnet:
+    def __init__(self):
+
+        # Placeholder for the model and config
+        self.model = None
+        self.config = None
+
+    def build_model( self, config: dict ):
+        """Creates the AttentionUnet model based on provided parameters."""
+        
+        self.model = AttentionUnet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=config['num_classes'],
+            channels=config['channels'],
+            strides=config['strides'],
+            dropout=config['dropout']
+        )
+        
+        return self.model
+    
+    def bayesian_search(self, trial, num_classes: int):
+        """Defines the Bayesian optimization search space and builds the model with suggested parameters."""
+        
+        # Define the search space
+        num_layers = trial.suggest_int("num_layers", 3, 5)
+        hidden_layers = trial.suggest_int("hidden_layers", 1, 3)
+        base_channel = trial.suggest_categorical("base_channel", [8, 16, 32])
+        
+        # Create channel sizes and strides
+        downsampling_channels = [base_channel * (2 ** i) for i in range(num_layers)]
+        hidden_channels = [downsampling_channels[-1]] * hidden_layers
+        channels = downsampling_channels + hidden_channels
+        strides = [2] * (num_layers - 1) + [1] * hidden_layers
+        dropout = trial.suggest_float("dropout", 0.0, 0.2)
+        
+        config = {
+            'architecture': 'AttentionUnet',
+            'num_classes': num_classes,
+            'channels': channels,
+            'strides': strides,
+            'dropout': dropout
+        }
+
+        return self.build_model(config)
+
+    def get_model_parameters(self):
+        """Retrieve stored model parameters."""
+        if self.model is None:
+            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")
+        
+        return self.config

octopi/models/MedNeXt.py

@@ -0,0 +1,111 @@
+from monai.networks.nets import MedNeXt
+import torch.nn as nn
+import torch
+
+class myMedNeXt:
+    def __init__(self ):
+        # Placeholder for the model and config
+        self.model = None
+        self.config = None
+
+    def build_model(
+        self,
+        init_filters=32,
+        encoder_expansion_ratio=2,
+        decoder_expansion_ratio=2,
+        bottleneck_expansion_ratio=2,
+        kernel_size=7,
+        deep_supervision=False,
+        use_residual_connection=False,
+        norm_type="group",
+        global_resp_norm=False,
+        blocks_down=(2, 2, 2, 2),
+        blocks_bottleneck=2,
+        blocks_up=(2, 2, 2, 2),
+    ):
+        """
+        Create the MedNeXt model with the specified hyperparameters.
+        Note: For cryoET with small objects, a shallower network might help preserve details.
+        """
+        self.model = MedNeXt(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=self.num_classes,
+            init_filters=init_filters,
+            encoder_expansion_ratio=encoder_expansion_ratio,
+            decoder_expansion_ratio=decoder_expansion_ratio,
+            bottleneck_expansion_ratio=bottleneck_expansion_ratio,
+            kernel_size=kernel_size,
+            deep_supervision=deep_supervision,
+            use_residual_connection=use_residual_connection,
+            norm_type=norm_type,
+            global_resp_norm=global_resp_norm,
+            blocks_down=blocks_down,
+            blocks_bottleneck=blocks_bottleneck,
+            blocks_up=blocks_up,
+        ).to(self.device)
+
+    def bayesian_search(self, trial):
+        """
+        Defines the Bayesian optimization search space and builds the model with suggested parameters.
+        The search space has been adapted for cryoET applications:
+          - Small kernel sizes (3 or 5) to capture fine details.
+          - Choice of a shallower vs. deeper architecture to balance resolution and feature extraction.
+          - Robust normalization options for low signal-to-noise data.
+        """
+        # Core hyperparameters
+        init_filters = trial.suggest_categorical("init_filters", [16, 32])
+        encoder_expansion_ratio = trial.suggest_int("encoder_expansion_ratio", 1, 3)
+        decoder_expansion_ratio = trial.suggest_int("decoder_expansion_ratio", 1, 3)
+        bottleneck_expansion_ratio = trial.suggest_int("bottleneck_expansion_ratio", 1, 4)
+        kernel_size = trial.suggest_categorical("kernel_size", [3, 5])
+        deep_supervision = trial.suggest_categorical("deep_supervision", [True, False])
+        norm_type = trial.suggest_categorical("norm_type", ["group", "instance"])
+        # For extremely low SNR, you might opt to disable global response normalization
+        global_resp_norm = trial.suggest_categorical("global_resp_norm", [False])
+        
+        # Architecture: shallow vs. deep.
+        # For small objects, a shallower network (fewer downsampling stages) may preserve spatial detail.
+        architecture = trial.suggest_categorical("architecture", ["shallow", "deep"])
+        if architecture == "shallow":
+            blocks_down = (2, 2, 2)         # 3 downsampling stages
+            blocks_bottleneck = 2
+            blocks_up = (2, 2, 2)
+        else:
+            blocks_down = (2, 2, 2, 2)        # 4 downsampling stages
+            blocks_bottleneck = 2
+            blocks_up = (2, 2, 2, 2)
+
+        self.build_model(
+            init_filters=init_filters,
+            encoder_expansion_ratio=encoder_expansion_ratio,
+            decoder_expansion_ratio=decoder_expansion_ratio,
+            bottleneck_expansion_ratio=bottleneck_expansion_ratio,
+            kernel_size=kernel_size,
+            deep_supervision=deep_supervision,
+            use_residual_connection=True,
+            norm_type=norm_type,
+            global_resp_norm=False,
+            blocks_down=blocks_down,
+            blocks_bottleneck=blocks_bottleneck,
+            blocks_up=blocks_up,
+        )
+
+    def get_model_parameters(self):
+        """Retrieve stored model parameters."""
+        if self.model is None:
+            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")
+        
+        return {
+            'architecture': 'MedNeXt',
+            'num_classes': self.num_classes,
+            'init_filters': self.model.init_filters,
+            'encoder_expansion_ratio': self.model.encoder_expansion_ratio,
+            'decoder_expansion_ratio': self.model.decoder_expansion_ratio,
+            'bottleneck_expansion_ratio': self.model.bottleneck_expansion_ratio,
+            'kernel_size': self.model.kernel_size,
+            'deep_supervision': self.model.do_ds,
+            'use_residual_connection': self.model.use_residual_connection,
+            'norm_type': self.model.norm_type,
+            'global_resp_norm': self.model.global_resp_norm,
+        }

octopi/models/ModelTemplate.py

@@ -0,0 +1,36 @@
+import torch.nn as nn
+import torch
+
+class myModelTemplate:
+    def __init__(self):
+        """
+        Initialize the model template.
+        """
+        # Placeholder for the model and config
+        self.model = None
+        self.config = None
+
+    def build_model(self, config: dict):
+        """
+        Build the model based on provided parameters in a config dictionary.
+        """
+        pass
+
+    def bayesian_search(self, trial, num_classes: int):
+        """
+        Define the hyperparameter search space for Bayesian optimization and build the model.
+
+        The search space below is just an example and can be customized.
+        
+        Args:
+            trial (optuna.trial.Trial): Optuna trial object.
+            num_classes (int): Number of classes in the dataset.
+        """
+        pass
+    
+    def get_model_parameters(self):
+        """Retrieve stored model parameters."""
+        if self.model is None:
+            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")
+        
+        return self.config

octopi/models/SegResNet.py

@@ -0,0 +1,92 @@
+from monai.networks.nets import SegResNetDS
+import torch.nn as nn
+import torch
+
+class mySegResNet:
+    def __init__(self):
+        # Placeholder for the model and config
+        self.model = None
+        self.config = None
+
+    def build_model(
+        self,
+        config: dict
+    ):
+        """
+        Creates the SegResNetDS model based on provided parameters.
+        
+        Args:
+            init_filters (int): Number of output channels for the initial convolution.
+            blocks_down (tuple): Tuple defining the number of blocks at each downsampling stage.
+            dsdepth (int): Depth for deep supervision (number of output scales).
+            act (str): Activation type.
+            norm (str): Normalization type.
+            blocks_up (tuple or None): Number of upsample blocks (if None, uses default behavior).
+            upsample_mode (str): Upsampling mode, e.g. 'deconv' or 'trilinear'.
+            resolution (optional): If provided, adjusts non-isotropic kernels to isotropic spacing.
+            preprocess (callable or None): Optional preprocessing function for the input.
+        
+        Returns:
+            torch.nn.Module: The instantiated SegResNetDS model on the specified device.
+        """
+        self.model = SegResNetDS(
+            spatial_dims=3,
+            init_filters=config['init_filters'],
+            in_channels=1,
+            out_channels=config['num_classes'],
+            act=config['act'],
+            norm=config['norm'],
+            blocks_down=config['blocks_down'],
+            blocks_up=config['blocks_up'],
+            dsdepth=config['dsdepth'],
+            preprocess=config['preprocess'],
+            upsample_mode=config['upsample_mode'],
+            resolution=config['resolution']
+        )
+        return self.model.to(self.device)
+    
+    def bayesian_search(self, trial):
+        """
+        Defines the Bayesian optimization search space and builds the model with suggested parameters.
+        
+        Args:
+            trial (optuna.trial.Trial): An Optuna trial object.
+        
+        Returns:
+            torch.nn.Module: The model built with hyperparameters suggested by the trial.
+        """
+        # Define search space parameters
+        init_filters = trial.suggest_categorical("init_filters", [16, 32, 64])
+        dsdepth = trial.suggest_int("dsdepth", 1, 3)
+        blocks_down = trial.suggest_categorical("blocks_down", [(1, 2, 2, 4), (1, 2, 2, 2), (1, 1, 2, 2)])
+        act = trial.suggest_categorical("act", ['relu', 'leaky_relu', "LeakyReLU", "PReLU", "GELU", "ELU"])
+        norm = trial.suggest_categorical("norm", ['batch', 'instance'])
+        upsample_mode = trial.suggest_categorical("upsample_mode", ['deconv', 'trilinear'])
+        
+        self.config = {
+            'init_filters': init_filters,
+            'blocks_down': blocks_down,
+            'dsdepth': dsdepth,
+            'act': act,
+            'norm': norm,
+            'blocks_up': None,  # using default upsampling blocks
+            'upsample_mode': upsample_mode,
+            'resolution': None,
+            'preprocess': None
+        }
+        return self.build_model(self.config)
+
+    def get_model_parameters(self):
+        """
+        Retrieve stored model parameters.
+        
+        Returns:
+            dict: A dictionary of key model parameters.
+        
+        Raises:
+            ValueError: If the model has not been built yet.
+        """
+        if self.model is None:
+            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")
+        
+        return self.config

octopi/models/Unet.py

@@ -0,0 +1,59 @@
+from monai.networks.nets import UNet
+import torch.nn as nn
+import torch
+
+class myUNet:
+    def __init__(self):
+        # Placeholder for the model and config
+        self.model = None
+        self.config = None
+
+    def build_model( self, config: dict ):
+        """Creates the Unet model based on provided parameters."""
+        
+        self.config = config
+        self.model = UNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=config['num_classes'],
+            channels=config['channels'],
+            strides=config['strides'],
+            num_res_units=config['num_res_units'],
+            dropout=config['dropout']
+        )
+        return self.model
+    
+    def bayesian_search(self, trial, num_classes: int):
+        """Defines the Bayesian optimization search space and builds the model with suggested parameters."""
+        
+        # Define the search space
+        num_layers = trial.suggest_int("num_layers", 3, 5)
+        hidden_layers = trial.suggest_int("hidden_layers", 1, 3)
+        base_channel = trial.suggest_categorical("base_channel", [8, 16, 32])
+        num_res_units = trial.suggest_int("num_res_units", 1, 3)
+        dropout = trial.suggest_float("dropout", 0.0, 0.3)
+        
+        # Create channel sizes and strides
+        downsampling_channels = [base_channel * (2 ** i) for i in range(num_layers)]
+        hidden_channels = [downsampling_channels[-1]] * hidden_layers
+        channels = downsampling_channels + hidden_channels
+        strides = [2] * (num_layers - 1) + [1] * hidden_layers
+
+        # Create config dictionary
+        self.config = {
+            'architecture': 'Unet',
+            'num_classes': num_classes,
+            'channels': channels,
+            'strides': strides,
+            'num_res_units': num_res_units,
+            'dropout': dropout
+        }
+        
+        return self.build_model(self.config)
+
+    def get_model_parameters(self):
+        """Retrieve stored model parameters."""
+        if self.model is None:
+            raise ValueError("Model has not been initialized yet. Call build_model() or bayesian_search() first.")
+        
+        return self.config

octopi/models/UnetPlusPlus.py

@@ -0,0 +1,47 @@
+from monai.networks.nets import  BasicUNetPlusPlus
+
+class myUNetPlusPlus:
+    def __init__(self, num_classes, device):
+        self.device = device
+        self.num_classes = num_classes
+        
+
+    def build_model(
+        self, 
+        features=(16, 16, 32, 64, 128, 16),
+        dropout=0.2,
+        upsample='deconv',
+        activation='relu'
+    ):
+
+        model = BasicUNetPlusPlus(
+                spatial_dims=3,
+                in_channels=1,
+                out_channels=n_classes,
+                deep_supervision=True,
+                features=features,  # Halve the features
+                dropout=dropout,
+                upsample=upsample,
+                act=activation
+            )  
+        
+        return model.to(self.device)
+    
+    def bayesian_search(self, trial):
+        
+        act = trial.suggest_categorical("activation", ["LeakyReLU", "PReLU", "GELU", "ELU"])   
+        dropout_prob = trial.suggest_float("dropout", 0.0, 0.5)
+        upsample = trial.suggest_categorical("upsample", ["deconv", "pixelshuffle", "nontrainable"])        
+        model_parameters = {"activation": act,'dropout': dropout_prob, 'upsample': upsample} 
+
+        model = self.build_model(features, dropout_prob, upsample, act)
+        return model
+    
+    def get_model_parameters(self):
+        return {
+            'model_name': 'UnetPlusPlus',
+            'features': self.features,
+            'dropout': self.dropout,
+            'upsample': self.upsample,
+            'activation': self.activation
+        }

octopi/models/common.py

@@ -0,0 +1,62 @@
+from monai.losses import FocalLoss, TverskyLoss
+from octopi import losses
+from octopi.models import (
+    Unet, AttentionUnet, MedNeXt, SegResNet
+)
+
+def get_model(architecture):
+
+    # Initialize model based on architecture
+    if architecture == "Unet":
+        model = Unet.myUNet()
+    elif architecture == "AttentionUnet":
+        model = AttentionUnet.myAttentionUnet()
+    elif architecture == "MedNeXt":
+        model = MedNeXt.myMedNeXt()
+    elif architecture == "SegResNet":
+        model = SegResNet.mySegResNet()
+    else:
+        raise ValueError(f"Model type {architecture} not supported!\nPlease use one of the following: Unet, AttentionUnet, MedNeXt, SegResNet")
+
+    return model
+
+def get_loss_function(trial, loss_name = None):
+
+    # Loss function selection
+    if loss_name is None:
+        loss_name = trial.suggest_categorical(
+            "loss_function", 
+            ["FocalLoss", "WeightedFocalTverskyLoss", 'FocalTverskyLoss'])
+
+    if loss_name == "FocalLoss":
+        gamma = round(trial.suggest_float("gamma", 0.1, 2), 3)
+        loss_function = FocalLoss(include_background=True, to_onehot_y=True, use_softmax=True, gamma=gamma)
+
+    elif loss_name == "TverskyLoss":
+        alpha = round(trial.suggest_float("alpha", 0.1, 0.5), 3)
+        beta = 1.0 - alpha
+        loss_function = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True, alpha=alpha, beta=beta)
+
+    elif loss_name == 'WeightedFocalTverskyLoss':
+        gamma = round(trial.suggest_float("gamma", 0.1, 2), 3)
+        alpha = round(trial.suggest_float("alpha", 0.1, 0.5), 3)
+        beta = 1.0 - alpha
+        weight_tversky = round(trial.suggest_float("weight_tversky", 0.1, 0.9), 3)
+        weight_focal = 1.0 - weight_tversky
+        loss_function = losses.WeightedFocalTverskyLoss(
+            gamma=gamma, alpha=alpha, beta=beta,
+            weight_tversky=weight_tversky, weight_focal=weight_focal
+        )
+
+    elif loss_name == 'FocalTverskyLoss':
+        gamma = round(trial.suggest_float("gamma", 0.1, 2), 3)
+        alpha = round(trial.suggest_float("alpha", 0.1, 0.5), 3)
+        beta = 1.0 - alpha
+        loss_function = losses.FocalTverskyLoss(gamma=gamma, alpha=alpha, beta=beta)
+
+    return loss_function
+
+
+#### TODO : Models to try Adding? 
+# 1. Swin UNETR 
+# 2. Swin-Conv-UNet

octopi/processing/create_targets_from_picks.py

@@ -0,0 +1,106 @@
+from octopi.processing.segmentation_from_picks import from_picks
+import octopi.processing.writers as write
+from octopi import io 
+from typing import List
+from tqdm import tqdm
+import numpy as np
+
+def generate_targets(
+    root,
+    train_targets: dict,
+    voxel_size: float = 10,
+    tomo_algorithm: str = 'wbp',
+    radius_scale: float = 0.8,    
+    target_segmentation_name: str = 'targets',
+    target_user_name: str = 'monai',
+    target_session_id: str = '1',
+    run_ids: List[str] = None,
+    ):
+    """
+    Generate segmentation targets from picks in CoPick configuration.
+
+    Args:
+        copick_config_path (str): Path to CoPick configuration file.
+        picks_user_id (str): User ID associated with picks.
+        picks_session_id (str): Session ID associated with picks.
+        target_segmentation_name (str): Name for the target segmentation.
+        target_user_name (str): User name associated with target segmentation.
+        target_session_id (str): Session ID for the target segmentation.
+        voxel_size (float): Voxel size for tomogram reconstruction.
+        tomo_algorithm (str): Tomogram reconstruction algorithm.
+        radius_scale (float): Scale factor for target object radius.
+    """
+
+    # Default session ID to 1 if not provided
+    if target_session_id is None:
+        target_session_id = '1'
+
+    print('Creating Targets for the following objects:', ', '.join(train_targets.keys()))
+
+    # Get Target Names
+    target_names = list(train_targets.keys())
+
+    # If runIDs are not provided, load all runs
+    if run_ids is None:
+        run_ids = [run.name for run in root.runs]
+
+    # Iterate Over All Runs
+    for runID in tqdm(run_ids):
+
+        # Get Run
+        numPicks = 0
+        run = root.get_run(runID)
+
+        # Get Tomogram 
+        tomo = io.get_tomogram_array(run, voxel_size, tomo_algorithm)
+        
+        # Initialize Target Volume
+        target = np.zeros(tomo.shape, dtype=np.uint8)
+
+        # Generate Targets
+        # Applicable segmentations
+        query_seg = []
+        for target_name in target_names:
+            if not train_targets[target_name]["is_particle_target"]:            
+                query_seg += run.get_segmentations(
+                    name=target_name,
+                    user_id=train_targets[target_name]["user_id"],
+                    session_id=train_targets[target_name]["session_id"],
+                    voxel_size=voxel_size
+                )     
+
+        # Add Segmentations to Target
+        for seg in query_seg:
+            classLabel = root.get_object(seg.name).label
+            segvol = seg.numpy()
+            # Set all non-zero values to the class label
+            segvol[segvol > 0] = classLabel
+            target[:] = segvol 
+
+        # Applicable picks
+        query = []
+        for target_name in target_names:
+            if train_targets[target_name]["is_particle_target"]:
+                query += run.get_picks(
+                    object_name=target_name,
+                    user_id=train_targets[target_name]["user_id"],
+                    session_id=train_targets[target_name]["session_id"],
+                )
+
+        # Add Picks to Target  
+        for pick in query:
+            numPicks += len(pick.points)
+            target = from_picks(pick, 
+                                target, 
+                                train_targets[pick.pickable_object_name]['radius'] * radius_scale,
+                                train_targets[pick.pickable_object_name]['label'],
+                                voxel_size
+                                )
+
+        # Write Segmentation for non-empty targets
+        if target.max() > 0 and numPicks > 0:
+            tqdm.write(f'Annotating {numPicks} picks in {runID}...')    
+            write.segmentation(run, target, target_user_name, 
+                               name = target_segmentation_name, session_id= target_session_id, 
+                               voxel_size = voxel_size)
+    print('Creation of targets complete!')