octopi 1.4.0__py3-none-any.whl

octopi/main.py

@@ -0,0 +1,33 @@
+import rich_click as click
+from octopi import cli_context
+from octopi.entry_points import groups 
+from octopi.processing.downloader import cli as download
+from octopi.processing.importers import cli as import_tomograms
+from octopi.entry_points.run_train import cli as train_model
+from octopi.entry_points.run_optuna import cli as model_explore
+from octopi.entry_points.run_create_targets import cli as create_targets
+from octopi.entry_points.run_segment import cli as inference
+from octopi.entry_points.run_localize import cli as localize
+from octopi.entry_points.run_evaluate import cli as evaluate
+from octopi.entry_points.run_extract_mb_picks import cli as mb_extract
+
+@click.group(context_settings=cli_context)
+def routines():
+    """Octopi 🐙: 🛠️ Tools for Finding Proteins in 🧊 cryo-ET data"""
+    pass
+
+routines.add_command(download)
+routines.add_command(import_tomograms)
+routines.add_command(train_model)
+routines.add_command(create_targets)
+routines.add_command(inference)
+routines.add_command(localize)
+routines.add_command(model_explore)
+routines.add_command(evaluate)
+routines.add_command(mb_extract)
+
+@click.group(context_settings=cli_context)
+def slurm_routines():
+    """Slurm-Octopi 🐙: 🛠️ Tools for Finding Proteins in 🧊 cryo-ET data"""
+    pass
+

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,72 @@
+from monai.losses import FocalLoss, TverskyLoss
+from octopi.utils 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
+
+def get_default_unet_params():
+
+    model_config = {
+            'architecture': 'Unet',
+            'dim_in': 80,
+            'strides': [2, 2, 1],
+            'channels': [48, 64, 80, 80],
+            'dropout': 0.0, 'num_res_units': 1,
+        }
+    return model_config
+
+#### TODO : Models to try Adding? 
+# 1. Swin UNETR 
+# 2. Swin-Conv-UNet