octopi 1.4.0__py3-none-any.whl

octopi/entry_points/run_optuna.py

@@ -0,0 +1,161 @@
+from octopi.entry_points import common
+from octopi.utils import parsers
+import rich_click as click
+
+def save_parameters(config: tuple,
+                    target_info: tuple,
+                    tomo_alg: str,
+                    voxel_size: float,
+                    model_type: str,
+                    mlflow_experiment_name: str,
+                    random_seed: int,
+                    num_trials: int,
+                    best_metric: str,
+                    num_epochs: int,
+                    tomo_batch_size: int,
+                    trainRunIDs: list,
+                    validateRunIDs: list,
+                    data_split: str,
+                    output_path: str):
+    """
+    Save the Optuna search parameters to a YAML file.
+    """
+    import octopi.utils.io as io
+    import pprint
+
+    # Organize parameters into categories
+    params = {
+        "input": {
+            "copick_config": config,
+            "target_info": target_info,
+            "tomo_algorithm": tomo_alg,
+            "voxel_size": voxel_size,          
+        },
+        "optimization": {
+            "model_type": model_type,
+            "mlflow_experiment_name": mlflow_experiment_name,
+            "random_seed": random_seed,
+            "num_trials": num_trials,
+            "best_metric": best_metric
+        },
+        "training": {
+            "num_epochs": num_epochs,            
+            "tomo_batch_size": tomo_batch_size,
+            "trainRunIDs": trainRunIDs,
+            "validateRunIDs": validateRunIDs,
+            "data_split": data_split
+        }
+    }
+
+    # Print the parameters
+    print(f"\nParameters for Model Architecture Search:")
+    pprint.pprint(params); print()
+
+    # Save to YAML file
+    io.save_parameters_yaml(params, output_path)
+
+
+@click.command('model-explore', help="Perform model architecture search with Optuna and MLflow integration")
+# Training Arguments
+@click.option('--random-seed', type=int, default=42,
+              help="Random seed for reproducibility")
+@common.train_parameters(octopi=True)
+# Model Arguments
+@click.option('--model-type', type=click.Choice(['Unet', 'AttentionUnet', 'MedNeXt', 'SegResNet'], case_sensitive=False),
+              default='Unet',
+              help="Model type to use for training")
+# Input Arguments
+@click.option('-split', '--data-split', type=str, default='0.8',
+              help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
+@click.option('-vruns', '--validateRunIDs', type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of validation run IDs, e.g., run3,run4 or [run3,run4]")
+@click.option('-truns', '--trainRunIDs', type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of training run IDs, e.g., run1,run2 or [run1,run2]")
+@click.option('--mlflow-experiment-name', type=str, default="model-search",
+              help="Name of the MLflow experiment")
+@click.option('-alg', '--tomo-alg', type=str, default='wbp',
+              help="Tomogram algorithm used for training")
+@click.option('-tinfo', '--target-info', type=str, default="targets,octopi,1",
+              callback=lambda ctx, param, value: parsers.parse_target(value),
+              help="Target information, e.g., 'name' or 'name,user_id,session_id'")
+@common.config_parameters(single_config=False)
+def cli(config, voxel_size, target_info, tomo_alg, mlflow_experiment_name, 
+        trainrunids, validaterunids, data_split,
+        model_type,
+        num_epochs, val_interval, tomo_batch_size, best_metric, num_trials, random_seed):
+    """
+    CLI entry point for running optuna model architecture search.
+    """
+
+    run_model_explore(
+        config, voxel_size, target_info, tomo_alg, mlflow_experiment_name, 
+        trainrunids, validaterunids, data_split, model_type,
+        num_epochs, val_interval, tomo_batch_size, best_metric, num_trials, random_seed
+    )
+
+def run_model_explore(config, voxel_size, target_info, tomo_alg, mlflow_experiment_name, 
+        trainrunids, validaterunids, data_split, model_type,
+        num_epochs, val_interval, tomo_batch_size, best_metric, num_trials, random_seed):
+    """
+    Run the model exploration.
+    """
+    from octopi.pytorch.model_search_submitter import ModelSearchSubmit
+    import os
+
+    # Parse the CoPick configuration paths
+    if len(config) > 1:
+        copick_configs = parsers.parse_copick_configs(config)
+    else:
+        copick_configs = config[0]
+
+    # Create the model exploration directory
+    os.makedirs(f'explore_results_{model_type}', exist_ok=True)
+
+    # Save parameters
+    save_parameters(
+        config=config,
+        target_info=target_info,
+        tomo_alg=tomo_alg,
+        voxel_size=voxel_size,
+        model_type=model_type,
+        mlflow_experiment_name=mlflow_experiment_name,
+        random_seed=random_seed,
+        num_trials=num_trials,
+        best_metric=best_metric,
+        num_epochs=num_epochs,
+        tomo_batch_size=tomo_batch_size,
+        trainRunIDs=trainrunids,
+        validateRunIDs=validaterunids,
+        data_split=data_split,
+        output_path=f'explore_results_{model_type}/octopi.yaml'
+    )
+
+    # Call the function with parsed arguments
+    search = ModelSearchSubmit(
+        copick_config=copick_configs,
+        target_name=target_info[0],
+        target_user_id=target_info[1],
+        target_session_id=target_info[2],
+        tomo_algorithm=tomo_alg,
+        voxel_size=voxel_size,
+        model_type=model_type,
+        mlflow_experiment_name=mlflow_experiment_name,
+        random_seed=random_seed,
+        num_epochs=num_epochs,
+        num_trials=num_trials,
+        trainRunIDs=trainrunids,
+        validateRunIDs=validaterunids, 
+        tomo_batch_size=tomo_batch_size,
+        best_metric=best_metric,
+        val_interval=val_interval,
+        data_split=data_split
+    )
+
+    # Run the model search
+    search.run_model_search()
+
+
+if __name__ == "__main__":
+    cli()

octopi/entry_points/run_segment.py

@@ -0,0 +1,154 @@
+from octopi.entry_points import common
+from typing import List, Tuple
+import rich_click as click
+
+def inference(
+    copick_config_path: str,
+    model_weights: str, 
+    model_config: str,
+    seg_info: Tuple[str,str,str],
+    voxel_size: float,
+    tomo_algorithm: str,
+    tomo_batch_size: int,
+    run_ids: List[str],
+    ):
+    """
+    Perform segmentation inference using a model on provided tomograms.
+
+    Args:
+        copick_config_path (str): Path to CoPick configuration file.
+        run_ids (List[str]): List of tomogram run IDs for inference.
+        model_weights (str): Path to the trained model weights file.
+        channels (List[int]): List of channel sizes for each layer.
+        strides (List[int]): List of strides for the layers.
+        res_units (int): Number of residual units for the model.
+        voxel_size (float): Voxel size for tomogram reconstruction.
+        tomo_algorithm (str): Tomogram reconstruction algorithm to use.
+        segmentation_name (str): Name for the segmentation output.
+        segmentation_user_id (str): User ID associated with the segmentation.
+        segmentation_session_id (str): Session ID for this segmentation run.
+    """
+    from octopi.pytorch import segmentation
+    import torch
+    
+    gpu_count = torch.cuda.device_count()
+    print(f"Number of GPUs available: {gpu_count}")
+
+    if ',' in model_weights:
+        model_weights = model_weights.split(',')
+    if ',' in model_config:
+        model_config = model_config.split(',')
+    if isinstance(model_weights, list) and isinstance(model_config, list):
+        if len(model_weights) != len(model_config):
+            raise ValueError("Number of model weights and model configs must match for ensemble prediction.")
+        print("\nUsing Model Ensemble (Soup) Segmentation.")
+        print('Model Weights:', model_weights)
+        print('Model Configs:', model_config)
+    else:
+        print("Using Single Model Segmentation.")
+
+    if gpu_count > 1:
+        print("Using Multi-GPU Predictor.")
+        predict = segmentation.MultiGPUPredictor(
+            copick_config_path,
+            model_config,
+            model_weights
+        )
+
+        # Run Multi-GPU inference
+        predict.multi_gpu_inference(
+            runIDs=run_ids,
+            tomo_algorithm=tomo_algorithm,
+            voxel_spacing=voxel_size,
+            segmentation_name=seg_info[0],
+            segmentation_user_id=seg_info[1],
+            segmentation_session_id=seg_info[2],
+            save=True
+        )
+
+    else:
+        print("Using Single-GPU Predictor.")
+        predict = segmentation.Predictor(
+            copick_config_path,
+            model_config,
+            model_weights,
+        )
+
+        # Run batch prediction
+        predict.batch_predict(
+            runIDs=run_ids,
+            num_tomos_per_batch=tomo_batch_size,
+            tomo_algorithm=tomo_algorithm,
+            voxel_spacing=voxel_size,
+            segmentation_name=seg_info[0],
+            segmentation_user_id=seg_info[1],
+            segmentation_session_id=seg_info[2]
+        )
+
+    print("Inference completed successfully.")
+
+
+@click.command('segment')
+# Inference Arguments
+@common.inference_parameters()
+# Model Arguments
+@common.inference_model_parameters()
+# Input Arguments
+@common.config_parameters(single_config=True)
+def cli(config, voxel_size,
+        model_config, model_weights,
+        tomo_alg, seg_info, tomo_batch_size, run_ids):
+    """
+    Segment volumes using trained neural network models.
+    
+    It supports both single model inference and model ensembles 
+    (model soups) for improved accuracy. Multi-GPU inference is automatically enabled when 
+    multiple GPUs are available.
+    
+    The segmentation masks are saved as zarr arrays in your copick project, organized by 
+    segmentation name, user ID, and session ID for easy tracking and comparison.
+    
+    \b
+    Examples:
+      # Segment with a single model
+      octopi segment -c config.json \\
+        --model-config model.yaml --model-weights model.pth \\
+        --seg-info predictions,octopi,1
+    
+    \b
+      # Segment with model ensemble (comma-separated)
+      octopi segment -c config.json \\
+        --model-config model1.yaml,model2.yaml \\
+        --model-weights model1.pth,model2.pth \\
+        --seg-info ensemble,octopi,1
+    
+    \b
+      # Segment specific runs only
+      octopi segment -c config.json \\
+        --model-config model.yaml --model-weights model.pth \\
+        --run-ids TS_001,TS_002,TS_003 \\
+        --tomo-batch-size 10
+    """
+    
+    # Set default values if not provided
+    seg_info = list(seg_info)  # Convert tuple to list
+    if seg_info[1] is None:
+        seg_info[1] = "octopi"
+    if seg_info[2] is None:
+        seg_info[2] = "1"
+
+    # Call the inference function with parsed arguments
+    inference(
+        copick_config_path=config,
+        model_weights=model_weights,
+        model_config=model_config,
+        seg_info=seg_info,
+        voxel_size=voxel_size,
+        tomo_algorithm=tomo_alg,
+        tomo_batch_size=tomo_batch_size,
+        run_ids=run_ids,
+    )
+
+
+if __name__ == "__main__":
+    cli()

octopi/entry_points/run_train.py

@@ -0,0 +1,189 @@
+from typing import List, Optional, Tuple
+from octopi.entry_points import common
+from octopi.utils import parsers
+from octopi import cli_context
+import rich_click as click
+
+# Configure rich-click
+click.rich_click.USE_RICH_MARKUP = True
+click.rich_click.SHOW_ARGUMENTS = True
+click.rich_click.GROUP_ARGUMENTS_OPTIONS = True
+
+def train_model(
+    copick_config_path: str,
+    target_info: Tuple[str, str, str],
+    tomo_algorithm: str = 'wbp',
+    voxel_size: float = 10,
+    trainRunIDs: List[str] = None,
+    validateRunIDs: List[str] = None,    
+    model_config: str = None,
+    model_weights: Optional[str] = None,
+    model_save_path: str = 'results',
+    num_tomo_crops: int = 16,
+    tomo_batch_size: int = 15,
+    lr: float = 1e-3,
+    tversky_alpha: float = 0.5,
+    num_epochs: int = 100,  
+    val_interval: int = 5,
+    best_metric: str = 'avg_f1',
+    data_split: str = '0.8'
+    ):
+    """
+    Train a 3D U-Net model using the specified CoPick configuration and target information.
+    """
+
+    import matplotlib
+    # Force a headless-safe backend everywhere (must be BEFORE pyplot import)
+    matplotlib.use("Agg", force=True)
+
+    from octopi.datasets import generators, multi_config_generator
+    from monai.losses import TverskyLoss
+    from octopi.utils import parsers, io
+    from octopi.workflows import train
+
+    # Initialize the data generator to manage training and validation datasets
+    print(f'Training with {copick_config_path}\n')
+
+    # Multi-config training
+    if isinstance(copick_config_path, dict):
+        data_generator = multi_config_generator.MultiConfigTrainLoaderManager(
+            copick_config_path, 
+            target_info[0], 
+            target_session_id = target_info[2],
+            target_user_id = target_info[1],
+            tomo_algorithm = tomo_algorithm,
+            voxel_size = voxel_size,
+            tomo_batch_size = tomo_batch_size )
+    else:  # Single-config training
+        data_generator = generators.TrainLoaderManager(
+            copick_config_path, 
+            target_info[0], 
+            target_session_id = target_info[2],
+            target_user_id = target_info[1],
+            tomo_algorithm = tomo_algorithm,
+            voxel_size = voxel_size,
+            tomo_batch_size = tomo_batch_size )
+
+    # Get the data splits
+    ratios = parsers.parse_data_split(data_split)
+    data_generator.get_data_splits(
+        trainRunIDs = trainRunIDs,
+        validateRunIDs = validateRunIDs,
+        train_ratio = ratios[0], val_ratio = ratios[1], test_ratio = ratios[2],
+        create_test_dataset = False)
+    
+    # Get the reload frequency
+    data_generator.get_reload_frequency(num_epochs)
+    model_config['num_classes'] = data_generator.Nclasses
+
+    # Monai Functions
+    alpha = tversky_alpha
+    beta = 1 - alpha
+    loss_function = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True, alpha=alpha, beta=beta)  
+    
+    # Train the Model
+    train(
+        data_generator, loss_function, 
+        model_config = model_config, model_weights = model_weights,
+        best_metric = best_metric, num_epochs = num_epochs,
+        model_save_path = model_save_path, lr0 = lr
+    )
+
+def get_model_config(channels, strides, res_units, dim_in):
+    """
+    Create a model configuration dictionary if no model configuration file is provided.
+    """
+    model_config = {
+        'architecture': 'Unet',
+        'channels': channels,
+        'strides': strides,
+        'num_res_units': res_units, 
+        'dropout': 0.1,
+        'dim_in': dim_in
+    }
+    return model_config
+
+@click.command('train', help="Train 3D CNN U-Net models")
+# Training Arguments (applied in reverse order)
+@common.train_parameters(octopi=False)
+# UNet-Model Arguments
+@common.model_parameters(octopi=False)
+# Fine-Tuning Arguments
+@click.option('-mw', '--model-weights', type=click.Path(exists=True), default=None,
+              help="Path to the model weights file (typically used for fine-tuning)")
+@click.option('-mc', '--model-config', type=click.Path(exists=True), default=None,
+              help="Path to the model configuration file (typically used for fine-tuning)")
+# Input Arguments
+@click.option('-split', '--data-split', type=str, default='0.8',
+              help="Data split ratios. Either a single value (e.g., '0.8' for 80/20/0 split) or two comma-separated values (e.g., '0.7,0.1' for 70/10/20 split)")
+@click.option('-vruns', "--validateRunIDs", type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of validation run IDs, e.g., run4,run5,run6")
+@click.option('-truns', "--trainRunIDs", type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of training run IDs, e.g., run1,run2,run3")
+@click.option('-alg',"--tomo-alg", type=str, default='wbp',
+              help="Tomogram algorithm used for training")
+@click.option('-tinfo', "--target-info", type=str, default="targets,octopi,1",
+              callback=lambda ctx, param, value: parsers.parse_target(value),
+              help="Target information, e.g., 'name' or 'name,user_id,session_id'. Default is 'targets,octopi,1'.")
+@common.config_parameters(single_config=False)
+def cli(config, voxel_size, target_info, tomo_alg, trainrunids, validaterunids, data_split,
+        model_config, model_weights,
+        channels, strides, res_units, dim_in,
+        num_epochs, val_interval, tomo_batch_size, best_metric, 
+        num_tomo_crops, lr, tversky_alpha, model_save_path):
+    """
+    CLI entry point for training models where results can either be saved to a local directory or a server with MLFlow.
+    """
+
+    run_train(config, voxel_size, target_info, tomo_alg, trainrunids, validaterunids, data_split,
+        model_config, model_weights,
+        channels, strides, res_units, dim_in,
+        num_epochs, val_interval, tomo_batch_size, best_metric, 
+        num_tomo_crops, lr, tversky_alpha, model_save_path)
+
+def run_train(config, voxel_size, target_info, tomo_alg, trainrunids, validaterunids, data_split,
+        model_config, model_weights,
+        channels, strides, res_units, dim_in,
+        num_epochs, val_interval, tomo_batch_size, best_metric, 
+        num_tomo_crops, lr, tversky_alpha, model_save_path):
+    """
+    Run the training model.
+    """
+    import octopi.utils.io as io
+
+    # Parse the CoPick configuration paths
+    if len(config) > 1:
+        copick_configs = parsers.parse_copick_configs(config)
+    else:
+        copick_configs = config[0]
+    
+    if model_config:
+        model_config_dict = io.load_yaml(model_config)
+    else:
+        model_config_dict = get_model_config(channels, strides, res_units, dim_in)
+
+    # Call the training function
+    train_model(
+        copick_config_path=copick_configs, 
+        target_info=target_info,
+        tomo_algorithm=tomo_alg,
+        voxel_size=voxel_size,
+        model_config=model_config_dict,
+        model_weights=model_weights,
+        model_save_path=model_save_path,
+        num_tomo_crops=num_tomo_crops,
+        tomo_batch_size=tomo_batch_size,
+        lr=lr,
+        tversky_alpha=tversky_alpha,
+        num_epochs=num_epochs,
+        val_interval=val_interval,
+        best_metric=best_metric,
+        trainRunIDs=trainrunids,
+        validateRunIDs=validaterunids,
+        data_split=data_split
+    )
+
+if __name__ == '__main__':
+    cli()