octopi 1.4.0__py3-none-any.whl

octopi/entry_points/create_slurm_submission.py

@@ -0,0 +1,251 @@
+from octopi.entry_points import run_train, run_segment_predict, run_localize, run_optuna
+from octopi.utils.submit_slurm import create_shellsubmit, create_multiconfig_shellsubmit
+from octopi.processing.importers import cli_mrcs_parser, cli_dataportal_parser
+from octopi.entry_points import common 
+from octopi import utils
+import argparse
+
+def create_train_script(args):
+    """
+    Create a SLURM script for training 3D CNN models
+    """
+
+    strconfigs = ""
+    for config in args.config:
+        strconfigs += f"--config {config}"
+
+    command = f"""
+octopi train \\
+    {strconfigs} \\
+    --model-save-path {args.model_save_path} \\
+    --target-info {','.join(args.target_info)} \\
+    --voxel-size {args.voxel_size} --tomo-alg {args.tomo_alg} --Nclass {args.Nclass} \\
+    --tomo-batch-size {args.tomo_batch_size} --num-tomo-crops {args.num_tomo_crops} \\
+    --best-metric {args.best_metric} --num-epochs {args.num_epochs} --val-interval {args.val_interval} \\
+    """
+
+    # If a model config is provided, use it to build the model
+    if args.model_config is not None:
+        command += f" --model-config {args.model_config}"
+    else:
+        channels = ",".join(map(str, args.channels))
+        strides = ",".join(map(str, args.strides))
+        command += (
+            f" --tversky-alpha {args.tversky_alpha}"
+            f" --channels {channels}"
+            f" --strides {strides}"
+            f" --dim-in {args.dim_in}"
+            f" --res-units {args.res_units}"
+        )
+
+    # If Model Weights are provided, use them to initialize the model
+    if args.model_weights is not None and args.model_config is not None:
+        command += f" --model-weights {args.model_weights}"
+
+    create_shellsubmit(
+        job_name = args.job_name,
+        output_file = 'trainer.log',
+        shell_name = 'train_octopi.sh',
+        conda_path = args.conda_env,
+        command = command,
+        num_gpus = 1,
+        gpu_constraint = args.gpu_constraint
+    )
+
+def train_model_slurm():
+    """
+    Create a SLURM script for training 3D CNN models
+    """
+    parser_description = "Create a SLURM script for training 3D CNN models"
+    args = run_train.train_model_parser(parser_description, add_slurm=True)
+    create_train_script(args)   
+
+def create_model_explore_script(args):
+    """
+    Create a SLURM script for running bayesian optimization on 3D CNN models
+    """
+    strconfigs = ""
+    for config in args.config:
+        strconfigs += f"--config {config}"
+
+    command = f"""
+octopi model-explore \\
+    --model-type {args.model_type} --model-save-path {args.model_save_path} \\
+    --voxel-size {args.voxel_size} --tomo-alg {args.tomo_alg} --Nclass {args.Nclass} \\
+    --val-interval {args.val_interval} --num-epochs {args.num_epochs} --num-trials {args.num_trials} \\
+    --best-metric {args.best_metric} --mlflow-experiment-name {args.mlflow_experiment_name} \\
+    --target-name {args.target_name} --target-session-id {args.target_session_id} --target-user-id {args.target_user_id} \\
+    {strconfigs}
+"""
+
+    create_shellsubmit(
+        job_name = args.job_name,
+        output_file = 'optuna.log',
+        shell_name = 'model_explore.sh',
+        conda_path = args.conda_env,
+        command = command,
+        num_gpus = 1,
+        gpu_constraint = args.gpu_constraint
+    )
+
+def model_explore_slurm():
+    """
+    Create a SLURM script for running bayesian optimization on 3D CNN models
+    """
+    parser_description = "Create a SLURM script for running bayesian optimization on 3D CNN models"
+    args = run_optuna.optuna_parser(parser_description, add_slurm=True)
+    create_model_explore_script(args)   
+
+def create_inference_script(args):
+    """
+    Create a SLURM script for running inference on 3D CNN models
+    """
+
+    if len(args.config.split(',')) > 1:
+
+        create_multiconfig_shellsubmit(
+            job_name = args.job_name,
+            output_file = 'predict.log',
+            shell_name = 'segment.sh',
+            conda_path = args.conda_env,
+            base_inputs = args.base_inputs,
+            config_inputs = args.config_inputs,
+            command = args.command,
+            num_gpus = args.num_gpus,
+            gpu_constraint = args.gpu_constraint
+        )
+    else:
+
+        command = f"""
+octopi inference \\
+    --config {args.config} \\
+    --seg-info  {",".join(args.seg_info)} \\
+    --model-weights {args.model_weights} \\
+    --dim-in {args.dim_in} --res-units {args.res_units} \\
+    --model-type {args.model_type} --channels {",".join(map(str, args.channels))} --strides {",".join(map(str, args.strides))} \\
+    --voxel-size {args.voxel_size} --tomo-alg {args.tomo_alg} --Nclass {args.Nclass}
+"""
+
+        create_shellsubmit(
+            job_name = args.job_name,
+            output_file = 'predict.log',
+            shell_name = 'segment.sh',
+            conda_path = args.conda_env,
+            command = command,
+            num_gpus = 1,
+            gpu_constraint = args.gpu_constraint
+        )
+
+def inference_slurm():
+    """
+    Create a SLURM script for running segmentation predictions with a specified model and configuration on CryoET Tomograms.
+    """
+    parser_description = "Create a SLURM script for running segmentation predictions with a specified model and configuration on CryoET Tomograms."
+    args = run_segment_predict.inference_parser(parser_description, add_slurm=True)
+    create_inference_script(args)
+
+def create_localize_script(args):
+    """"
+    Create a SLURM script for running localization on predicted segmentation masks
+    """
+    if len(args.config.split(',')) > 1:
+    
+        create_multiconfig_shellsubmit(
+            job_name = args.job_name,
+            output_file = args.output,
+            shell_name = args.output_script,
+            conda_path = args.conda_env,
+            base_inputs = args.base_inputs,
+            config_inputs = args.config_inputs,
+            command = args.command
+        )
+    else:
+
+        command = f"""
+octopi localize \\
+    --config {args.config} \\
+    --voxel-size {args.voxel_size} --pick-session-id {args.pick_session_id} --pick-user-id {args.pick_user_id} \\
+    --method {args.method}  --seg-info {",".join(args.seg_info)} \\
+"""
+        if args.pick_objects is not None:
+            command += f" --pick-objects {args.pick_objects}"
+
+        create_shellsubmit(
+            job_name = args.job_name,
+            output_file = 'localize.log',
+            shell_name = 'localize.sh',
+            conda_path = args.conda_env,
+            command = command,
+            num_gpus = 0
+        )
+
+def localize_slurm():
+    """
+    Create a SLURM script for running localization on predicted segmentation masks
+    """
+    parser_description = "Create a SLURM script for localization on predicted segmentation masks"
+    args = run_localize.localize_parser(parser_description, add_slurm=True)
+    create_localize_script(args)
+        
+def create_extract_mb_picks_script(args):
+    pass
+
+def extract_mb_picks_slurm():
+    pass
+
+
+def create_import_mrc_script(args):
+    """
+    Create a SLURM script for importing mrc volumes and potentialy downsampling
+    """
+    command = f"""
+octopi import-mrc-volumes \\
+    --mrcs-path {args.mrcs_path} \\
+    --config {args.config} --target-tomo-type {args.target_tomo_type} \\
+    --input-voxel-size {args.input_voxel_size} --output-voxel-size {args.output_voxel_size}
+"""
+
+    create_shellsubmit(
+        job_name = args.job_name,
+        output_file = 'importer.log',
+        shell_name = 'mrc_importer.sh',
+        conda_path = args.conda_env,
+        command = command
+    )
+
+def import_mrc_slurm():
+    """
+    Create a SLURM script for importing mrc volumes and potentialy downsampling
+    """
+    parser_description = "Create a SLURM script for importing mrc volumes and potentialy downsampling"
+    args = cli_mrcs_parser(parser_description, add_slurm=True)
+    create_import_mrc_script(args)
+
+
+def create_download_dataportal_script(args):
+    """
+    Create a SLURM script for downloading tomograms from the Dataportal
+    """
+    command = f"""
+octopi download-dataportal \\
+    --config {args.config} --datasetID {args.datasetID} \\
+    --overlay-path {args.overlay_path} 
+    --dataportal-name {args.dataportal_name} --target-tomo-type {args.target_tomo_type} \\
+    --input-voxel-size {args.input_voxel_size} --output-voxel-size {args.output_voxel_size}
+"""
+
+    create_shellsubmit(
+        job_name = args.job_name,
+        output_file = 'importer.log',
+        shell_name = 'dataportal_importer.sh',
+        conda_path = args.conda_env,
+        command = command
+    )
+
+def download_dataportal_slurm():
+    """
+    Create a SLURM script for downloading tomograms from the Dataportal
+    """
+    parser_description = "Create a SLURM script for downloading tomograms from the Dataportal"
+    args = cli_dataportal_parser(parser_description, add_slurm=True)
+    create_download_dataportal_script(args)

octopi/entry_points/groups.py

@@ -0,0 +1,152 @@
+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
+
+click.rich_click.COMMAND_GROUPS = {
+    "routines": [
+        {
+            "name": "Pre-Processing",
+            "commands": ["download", "import", "create-targets"]
+        },
+        {
+            "name": "Training",
+            "commands": ["train", "model-explore"]
+        },
+        {
+            "name": "Inference",
+            "commands": ["segment", "localize", "membrane-extract", "evaluate"]
+        }
+    ]
+}
+
+# Define option groups for all subcommands
+# Key format: "parent_command_name subcommand_name" or just "subcommand_name"
+click.rich_click.OPTION_GROUPS = {
+    "routines train": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--voxel-size", "--target-info", "--tomo-alg", 
+                       "--trainRunIDs", "--validateRunIDs", "--data-split"]
+        },
+        {
+            "name": "Fine-Tuning Arguments",
+            "options": ["--model-config", "--model-weights"]
+        },
+        {
+            "name": "Training Arguments",
+            "options": ["--num-epochs", "--val-interval", "--tomo-batch-size", "--best-metric", 
+                       "--num-tomo-crops", "--lr", "--tversky-alpha", "--model-save-path"]
+        },
+        {
+            "name": "UNet-Model Arguments",
+            "options": ["--Nclass", "--channels", "--strides", "--res-units", "--dim-in"]
+        }        
+    ],
+    "routines create-targets": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--target", "--picks-session-id", "--picks-user-id", 
+                       "--seg-target", "--run-ids"]
+        },
+        {
+            "name": "Parameters",
+            "options": ["--tomo-alg", "--radius-scale", "--voxel-size"]
+        },
+        {
+            "name": "Output Arguments",
+            "options": ["--target-segmentation-name", "--target-user-id", "--target-session-id"]
+        }
+    ],
+    "routines segment": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--voxel-size", "--tomo-alg"]
+        },
+        {
+            "name": "Model Arguments",
+            "options": ["--model-config", "--model-weights"]
+        },
+        {
+            "name": "Inference Arguments",
+            "options": ["--seg-info", "--tomo-batch-size", "--run-ids"]
+        }
+    ],
+    "routines localize": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--method", "--seg-info", "--voxel-size", "--runIDs"]
+        },
+        {
+            "name": "Localize Arguments",
+            "options": ["--radius-min-scale", "--radius-max-scale", "--filter-size", 
+                       "--pick-objects", "--n-procs"]
+        },
+        {
+            "name": "Output Arguments",
+            "options": ["--pick-session-id", "--pick-user-id"]
+        }
+    ],
+    "routines model-explore": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--voxel-size", "--target-info", "--tomo-alg", 
+                       "--mlflow-experiment-name", "--trainRunIDs", "--validateRunIDs", "--data-split"]
+        },
+        {
+            "name": "Model Arguments",
+            "options": ["--model-type"]
+        },
+        {
+            "name": "Training Arguments",
+            "options": ["--num-epochs", "--val-interval", "--tomo-batch-size", "--best-metric", 
+                       "--num-trials", "--random-seed"]
+        }
+    ],
+    "routines evaluate": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--ground-truth-user-id", "--ground-truth-session-id",
+                       "--predict-user-id", "--predict-session-id", "--run-ids"]
+        },
+        {
+            "name": "Evaluation Parameters",
+            "options": ["--distance-threshold-scale", "--object-names"]
+        },
+        {
+            "name": "Output Arguments",
+            "options": ["--save-path"]
+        }
+    ],
+    "routines membrane-extract": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--voxel-size", "--picks-info", "--membrane-info",
+                       "--organelle-info", "--runIDs"]
+        },
+        {
+            "name": "Parameters",
+            "options": ["--distance-threshold", "--n-procs"]
+        },
+        {
+            "name": "Output Arguments",
+            "options": ["--save-user-id", "--save-session-id"]
+        }
+    ],
+    "routines download": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--datasetID", "--overlay-path"]
+        },
+        {
+            "name": "Tomogram Settings",
+            "options": ["--dataportal-name", "--target-tomo-type"]
+        },
+        {
+            "name": "Voxel Settings",
+            "options": ["--input-voxel-size", "--output-voxel-size"]
+        }
+    ],    
+}

octopi/entry_points/run_create_targets.py

@@ -0,0 +1,234 @@
+from typing import List, Tuple, Union
+from collections import defaultdict
+from octopi.utils import parsers
+import rich_click as click
+
+def create_sub_train_targets(
+    config: str,
+    pick_targets: List[Tuple[str, Union[str, None], Union[str, None]]],
+    seg_targets: List[Tuple[str, Union[str, None], Union[str, None]]],
+    voxel_size: float,
+    radius_scale: float,    
+    tomogram_algorithm: str,
+    target_segmentation_name: str,
+    target_user_id: str,
+    target_session_id: str,
+    run_ids: List[str],    
+    ):
+    import octopi.processing.create_targets_from_picks as create_targets
+    import copick
+
+    # Load Copick Project 
+    root = copick.from_file(config)
+
+    # Create empty dictionary for all targets
+    train_targets = defaultdict(dict)
+
+    # Create dictionary for particle targets
+    value = 1 
+    for t in pick_targets:
+        # Parse the target
+        obj_name, user_id, session_id = t 
+        obj = root.get_object(obj_name)
+
+        # Check if the object is valid
+        if obj is None:
+            print(f'Warning - Skipping Particle Target: "{obj_name}", as it is not a valid name in the config file.')
+            continue
+
+        if obj_name in train_targets:
+            print(f'Warning - Skipping Particle Target: "{obj_name}, {user_id}, {session_id}", as it has already been added to the target list.')
+            continue
+        
+        # Get the label and radius of the object
+        label = value # Assign labels sequentially
+        info = {
+            "label": label,
+            "user_id": user_id,
+            "session_id": session_id,
+            "is_particle_target": True,
+            "radius": root.get_object(obj_name).radius,
+        }
+        train_targets[obj_name] = info
+        value += 1
+
+    # Create dictionary for segmentation targets
+    train_targets = add_segmentation_targets(root, seg_targets, train_targets, value)   
+
+    create_targets.generate_targets(
+        config, train_targets, voxel_size, tomogram_algorithm, radius_scale,
+        target_segmentation_name, target_user_id, 
+        target_session_id, run_ids
+    )
+
+
+def create_all_train_targets(
+    config: str,
+    seg_targets: List[List[Tuple[str, Union[str, None], Union[str, None]]]],
+    picks_session_id: str,
+    picks_user_id: str,
+    voxel_size: float,
+    radius_scale: float,
+    tomogram_algorithm: str,
+    target_segmentation_name: str,
+    target_user_id: str,
+    target_session_id: str,
+    run_ids: List[str],    
+    ):     
+    import octopi.processing.create_targets_from_picks as create_targets
+    import copick
+
+    # Load Copick Project 
+    root = copick.from_file(config)
+
+    # Create empty dictionary for all targets
+    target_objects = defaultdict(dict)
+
+    # Create dictionary for particle targets
+    for object in root.pickable_objects:
+        info = {
+            "label": object.label,
+            "radius": object.radius,
+            "user_id": picks_user_id,
+            "session_id": picks_session_id,
+            "is_particle_target": True,
+        }
+        target_objects[object.name] = info
+
+    # Create dictionary for segmentation targets
+    target_objects = add_segmentation_targets(root, seg_targets, target_objects)
+
+    create_targets.generate_targets(
+        config, target_objects, voxel_size, tomogram_algorithm, 
+        radius_scale, target_segmentation_name, target_user_id, 
+        target_session_id, run_ids 
+    )
+
+def add_segmentation_targets(
+    root,
+    seg_targets,
+    train_targets: dict,
+    start_value: int = -1
+    ):
+
+    # Create dictionary for segmentation targets
+    for s in seg_targets:
+
+        # Parse Segmentation Target
+        obj_name, user_id, session_id = s
+
+        # Add Segmentation Target
+        if start_value > 0:
+            value = start_value
+            start_value += 1
+        else:
+            value = root.get_object(obj_name).label
+
+        try:
+            info = {
+                "label": value,
+                "user_id": user_id,
+                "session_id": session_id,
+                "is_particle_target": False,                 
+                "radius": None,    
+            }
+            train_targets[obj_name] = info
+
+        # If Segmentation Target is not found, print warning
+        except:
+            print(f'Warning - Skipping Segmentation Name: "{obj_name}", as it is not a valid object in the Copick project.')
+
+    return train_targets    
+
+
+@click.command('create-targets')
+# Output Arguments
+@click.option('-sid', '--target-session-id', type=str, default="1",
+              help="Session ID for the target segmentation")
+@click.option('-uid','--target-user-id', type=str, default="octopi",
+              help="User ID associated with the target segmentation")
+@click.option('-name', '--target-segmentation-name', type=str, default='targets',
+              help="Name for the target segmentation")
+# Parameters
+@click.option('-vs', '--voxel-size', type=float, default=10,
+              help="Voxel size for tomogram reconstruction")
+@click.option('-rs', '--radius-scale', type=float, default=0.7,
+              help="Scale factor for object radius")
+@click.option('-alg', '--tomo-alg', type=str, default="wbp",
+              help="Tomogram reconstruction algorithm")
+# Input Arguments
+@click.option('--run-ids', type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of run IDs")
+@click.option('--seg-target', type=str, multiple=True,
+              callback=lambda ctx, param, value: [parsers.parse_target(v) for v in value] if value else [],
+              help='Segmentation targets: "name" or "name,user_id,session_id"')
+@click.option('--picks-user-id', type=str, default=None,
+              help="User ID associated with the picks")
+@click.option('--picks-session-id', type=str, default=None,
+              help="Session ID for the picks")
+@click.option('-t', '--target', type=str, multiple=True,
+              callback=lambda ctx, param, value: [parsers.parse_target(v) for v in value] if value else None,
+              help='Target specifications: "name" or "name,user_id,session_id"')
+@click.option('-c', '--config', type=click.Path(exists=True), required=True,
+              help="Path to the CoPick configuration file")
+def cli(config, target, picks_session_id, picks_user_id, seg_target, run_ids,
+        tomo_alg, radius_scale, voxel_size,
+        target_segmentation_name, target_user_id, target_session_id):
+    """
+    Generate segmentation targets from CoPick configurations.
+
+    This tool allows users to specify target labels for training in two ways:
+
+    1. Manual Specification: Define a subset of pickable objects using --target name or --target name,user_id,session_id
+
+    2. Automated Query: Provide --picks-session-id and/or --picks-user-id to automatically retrieve all pickable objects
+
+    Example Usage:
+
+        Manual: octopi create-targets --config config.json --target ribosome --target apoferritin,123,456
+
+        Automated: octopi create-targets --config config.json --picks-session-id 123 --picks-user-id 456
+    """
+
+    # Print Summary To User
+    print('\nGenerating Target Segmentation Masks from the Following Copick-Query:')
+    if target is not None and len(target) > 0:
+        print(f'    - Targets: {target}\n')
+    else:
+        print(f'    -  UserID: {picks_user_id} -- SessionID: {picks_session_id} \n')
+
+    # Check if either target or seg_target is provided
+    if (target is not None and len(target) > 0) or seg_target:
+        # If at least one --target is provided, call create_sub_train_targets
+        create_sub_train_targets(
+            config=config,
+            pick_targets=target if target else [],
+            seg_targets=seg_target,
+            voxel_size=voxel_size,
+            radius_scale=radius_scale,
+            tomogram_algorithm=tomo_alg,
+            target_segmentation_name=target_segmentation_name,
+            target_user_id=target_user_id,
+            target_session_id=target_session_id,
+            run_ids=run_ids,
+        )
+    else:
+        # If no --target is provided, call create_all_train_targets
+        create_all_train_targets(
+            config=config,
+            seg_targets=seg_target,
+            picks_session_id=picks_session_id,
+            picks_user_id=picks_user_id,
+            voxel_size=voxel_size,
+            radius_scale=radius_scale,
+            tomogram_algorithm=tomo_alg,
+            target_segmentation_name=target_segmentation_name,
+            target_user_id=target_user_id,
+            target_session_id=target_session_id,
+            run_ids=run_ids,
+        )
+
+
+if __name__ == "__main__":
+    cli()