octopi 1.0__py3-none-any.whl

octopi/entry_points/run_localize.py

@@ -0,0 +1,222 @@
+from octopi.entry_points import common
+from octopi.extract import localize
+from octopi import utils
+import copick, argparse, pprint
+from typing import List, Tuple
+import multiprocess as mp
+from tqdm import tqdm
+
+def pick_particles(
+    copick_config_path: str,
+    method: str,
+    seg_info: Tuple[str, str, str],
+    voxel_size: float,
+    pick_session_id: str,
+    pick_user_id: str,
+    radius_min_scale: float,
+    radius_max_scale: float,
+    filter_size: float,
+    pick_objects: List[str],
+    runIDs: List[str],
+    n_procs: int,
+    ):
+
+    # Load the Copick Project
+    root = copick.from_file(copick_config_path)    
+
+    # Get objects that can be Picked
+    objects = [(obj.name, obj.label, obj.radius) for obj in root.pickable_objects if obj.is_particle]
+
+     # Verify each object has the required attributes
+    for obj in objects:
+        if len(obj) < 3 or not isinstance(obj[2], (float, int)):
+            raise ValueError(f"Invalid object format: {obj}. Expected a tuple with (name, label, radius).")
+
+    # Filter elements
+    if pick_objects is not None:
+        objects = [obj for obj in objects if obj[0] in pick_objects]
+
+    print(f'Running Localization on the Following Objects: ')
+    print(', '.join([f'{obj[0]} (Label: {obj[1]})' for obj in objects]) + '\n')
+
+    # Either Specify Input RunIDs or Run on All RunIDs
+    if runIDs:  print('Running Localization on the Following RunIDs: ' + ', '.join(runIDs) + '\n')
+    run_ids = runIDs if runIDs else [run.name for run in root.runs]
+    n_run_ids = len(run_ids)
+
+    # Determine the number of processes to use
+    if n_procs is None:
+        n_procs = min(int(mp.cpu_count()//4), n_run_ids)
+    print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")
+
+    # Initialize tqdm progress bar
+    with tqdm(total=n_run_ids, desc="Localization", unit="run") as pbar:
+        for _iz in range(0, n_run_ids, n_procs):
+
+            start_idx = _iz
+            end_idx = min(_iz + n_procs, n_run_ids)  # Ensure end_idx does not exceed n_run_ids
+            print(f"\nProcessing runIDs from {start_idx} -> {end_idx } (out of {n_run_ids})")
+
+            processes = []                
+            for _in in range(n_procs):
+                _iz_this = _iz + _in
+                if _iz_this >= n_run_ids:
+                    break
+                run_id = run_ids[_iz_this]
+                run = root.get_run(run_id)
+                p = mp.Process(
+                    target=localize.processs_localization,
+                    args=(run,  
+                          objects, 
+                          seg_info,
+                          method, 
+                          voxel_size,
+                          filter_size,
+                          radius_min_scale, 
+                          radius_max_scale,
+                          pick_session_id,
+                          pick_user_id),
+                )
+                processes.append(p)
+
+            for p in processes:
+                p.start()
+
+            for p in processes:
+                p.join()
+
+            for p in processes:
+                p.close()
+
+            # Update tqdm progress bar
+            pbar.update(len(processes))
+
+    print('Localization Complete!')
+
+def localize_parser(parser_description, add_slurm: bool = False):
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    input_group = parser.add_argument_group("Input Arguments")
+    input_group.add_argument("--config", type=str, required=True, help="Path to the CoPick configuration file.")
+    input_group.add_argument("--method", type=str, choices=['watershed', 'com'], default='watershed', required=False, help="Localization method to use.")
+    input_group.add_argument('--seg-info', type=utils.parse_target, required=True, help='Query for the organelles segmentations (e.g., "name" or "name,user_id,session_id".).')
+    input_group.add_argument("--voxel-size", type=float, default=10, required=False, help="Voxel size for localization.")
+    input_group.add_argument("--runIDs", type=utils.parse_list, default = None, required=False, help="List of runIDs to run inference on, e.g., run1,run2,run3 or [run1,run2,run3].")
+
+    localize_group = parser.add_argument_group("Localize Arguments")
+    localize_group.add_argument("--radius-min-scale", type=float, default=0.5, required=False, help="Minimum radius scale for particles.")
+    localize_group.add_argument("--radius-max-scale", type=float, default=1.0, required=False, help="Maximum radius scale for particles.")
+    localize_group.add_argument("--filter-size", type=int, default=10, required=False, help="Filter size for localization.")
+    localize_group.add_argument("--pick-objects", type=utils.parse_list, default=None, required=False, help="Specific Objects to Find Picks for.")
+    localize_group.add_argument("--n-procs", type=int, default=None, required=False, help="Number of CPU processes to parallelize runs across. Defaults to the max number of cores available or available runs.")
+
+    output_group = parser.add_argument_group("Output Arguments")
+    output_group.add_argument("--pick-session-id", type=str, default='1', required=False, help="Session ID for the particle picks.")
+    output_group.add_argument("--pick-user-id", type=str, default='monai', required=False, help="User ID for the particle picks.")
+
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'localize', gpus = 0)
+
+    args = parser.parse_args()
+    return args
+
+# Entry point with argparse
+def cli():
+    
+    parser_description = "Localized particles in tomograms using multiprocessing."
+    args = localize_parser(parser_description)
+
+    # Save JSON with Parameters
+    output_yaml = f'localize_{args.pick_user_id}_{args.pick_session_id}.yaml'    
+    save_parameters(args, output_yaml)    
+
+    # Set multiprocessing start method
+    mp.set_start_method("spawn")
+    
+    pick_particles(
+        copick_config_path=args.config,
+        method=args.method,
+        seg_info=args.seg_info,
+        voxel_size=args.voxel_size,
+        pick_session_id=args.pick_session_id,
+        pick_user_id=args.pick_user_id,
+        radius_min_scale=args.radius_min_scale,
+        radius_max_scale=args.radius_max_scale,
+        filter_size=args.filter_size,
+        runIDs=args.runIDs,
+        pick_objects=args.pick_objects,
+        n_procs=args.n_procs,
+    )
+
+def save_parameters(args: argparse.Namespace, 
+                    output_path: str):
+
+    # Organize parameters into categories
+    params = {
+        "input": {
+            "config": args.config,
+            "seg_name": args.seg_info[0],
+            "seg_user_id": args.seg_info[1],
+            "seg_session_id": args.seg_info[2],
+            "voxel_size": args.voxel_size
+        },
+        "output": {
+            "pick_session_id": args.pick_session_id,
+            "pick_user_id": args.pick_user_id
+        },
+        "parameters": {
+            "method": args.method,
+            "radius_min_scale": args.radius_min_scale,
+            "radius_max_scale": args.radius_max_scale,
+            "filter_size": args.filter_size,
+            "runIDs": args.runIDs
+        }
+    }
+
+    # Print the parameters
+    print(f"\nParameters for Localization:")
+    pprint.pprint(params); print()
+
+    # Save to YAML file
+    utils.save_parameters_yaml(params, output_path)
+
+if __name__ == "__main__":
+    cli()
+
+# def time_pick_particles():
+#     import json, time
+
+#     # Set multiprocessing start method
+#     mp.set_start_method("spawn")
+
+#     copick_config_path = "/mnt/simulations/ml_challenge/ml_config.json"  # Replace with your actual path
+#     n_procs_list = [1, 4, 8, 16, 32]  # Adjust based on your needs
+#     n_procs_list = [32, 16, 8, 4, 1]
+#     timing_results = {}
+
+#     session_id = 1
+#     for n_procs in n_procs_list:
+#         print(f"Testing with {n_procs} processes...")
+#         start_time = time.time()
+#         pick_particles(
+#             copick_config_path=copick_config_path,
+#             pick_session_id=str(session_id),
+#             n_procs=n_procs
+#         )
+#         elapsed_time = time.time() - start_time
+#         timing_results[n_procs] = elapsed_time
+#         print(f"Elapsed time with {n_procs} processes: {elapsed_time:.2f} seconds")
+
+#         session_id +=1 
+
+#     # Save timing results to a JSON file
+#     with open("timing_results.json", "w") as f:
+#         json.dump(timing_results, f, indent=4)
+
+#     print("Timing results saved to 'timing_results.json'")
+
+# if __name__ == "__main__":
+#     time_pick_particles()

octopi/entry_points/run_optuna.py

@@ -0,0 +1,139 @@
+from octopi.pytorch.model_search_submitter import ModelSearchSubmit
+from octopi.entry_points import common
+import argparse, os, pprint
+from octopi import utils
+
+def optuna_parser(parser_description, add_slurm: bool = False):
+    """
+    Create an argument parser for model architecture search using Optuna.
+    
+    Args:
+        parser_description (str): Description of the parser
+        add_slurm (bool): Whether to add SLURM-specific arguments
+    """
+    
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    # Input Arguments
+    input_group = parser.add_argument_group("Input Arguments")
+    common.add_config(input_group, single_config=False)
+    input_group.add_argument("--target-info", type=utils.parse_target, default="targets,octopi,1", 
+                             help="Target information, e.g., 'name' or 'name,user_id,session_id'")    
+    input_group.add_argument("--tomo-alg", default='wbp', 
+                             help="Tomogram algorithm used for training")
+    input_group.add_argument("--mlflow-experiment-name", type=str, default="model-search", required=False, 
+                             help="Name of the MLflow experiment (default: 'model-search').")
+    input_group.add_argument("--trainRunIDs", type=utils.parse_list, default=None, required=False, 
+                             help="List of training run IDs, e.g., run1,run2 or [run1,run2].")
+    input_group.add_argument("--validateRunIDs", type=utils.parse_list, default=None, required=False, 
+                             help="List of validation run IDs, e.g., run3,run4 or [run3,run4].")    
+    input_group.add_argument('--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)")
+
+    model_group = parser.add_argument_group("Model Arguments")
+    model_group.add_argument("--model-type", type=str, default='Unet', required=False, 
+                             choices=['Unet', 'AttentionUnet'],
+                             help="Model type to use for training")
+    model_group.add_argument("--Nclass", type=int, default=3, required=False, help="Number of prediction classes in the model")
+
+    train_group = parser.add_argument_group("Training Arguments")
+    common.add_train_parameters(train_group, octopi = True)
+    train_group.add_argument("--random-seed", type=int, default=42, required=False, 
+                             help="Random seed for reproducibility (default: 42).")
+
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'optuna')
+
+    args = parser.parse_args()
+    return args
+
+# Entry point with argparse
+def cli():
+    """
+    CLI entry point for running optuna model archetecture search.
+    """
+
+    description="Perform model architecture search with Optuna and MLflow integration."
+    args = optuna_parser(description)
+
+    # Parse the CoPick configuration paths
+    if len(args.config) > 1:    copick_configs = utils.parse_copick_configs(args.config)
+    else:                       copick_configs = args.config[0]
+
+    # Create the model exploration directory
+    os.makedirs(f'explore_results_{args.model_type}', exist_ok=True)
+
+    # Save JSON with Parameters
+    save_parameters(args, f'explore_results_{args.model_type}/octopi.yaml')
+
+    # Call the function with parsed arguments
+    search = ModelSearchSubmit(
+        copick_config=copick_configs,
+        target_name=args.target_info[0],
+        target_user_id=args.target_info[1],
+        target_session_id=args.target_info[2],
+        tomo_algorithm=args.tomo_alg,
+        voxel_size=args.voxel_size,
+        Nclass=args.Nclass,
+        model_type=args.model_type,
+        mlflow_experiment_name=args.mlflow_experiment_name,
+        random_seed=args.random_seed,
+        num_epochs=args.num_epochs,
+        num_trials=args.num_trials,
+        trainRunIDs=args.trainRunIDs,
+        validateRunIDs=args.validateRunIDs, 
+        tomo_batch_size=args.tomo_batch_size,
+        best_metric=args.best_metric,
+        val_interval=args.val_interval,
+        data_split=args.data_split
+    )
+
+    # Run the model search
+    search.run_model_search()
+
+def save_parameters(args: argparse.Namespace, 
+                    output_path: str):
+    """
+    Save the Optuna search parameters to a JSON file.
+    Args:
+        args: Parsed arguments from argparse.
+        output_path: Path to save the JSON file.
+    """
+    # Organize parameters into categories
+    params = {
+        "input": {
+            "copick_config": args.config,
+            "target_info": args.target_info,
+            "tomo_algorithm": args.tomo_alg,
+            "voxel_size": args.voxel_size,
+            "Nclass": args.Nclass,            
+        },
+        "optimization": {
+            "model_type": args.model_type,
+            "mlflow_experiment_name": args.mlflow_experiment_name,
+            "random_seed": args.random_seed,
+            "num_trials": args.num_trials,
+            "best_metric": args.best_metric
+        },
+        "training": {
+            "num_epochs": args.num_epochs,            
+            "tomo_batch_size": args.tomo_batch_size,
+            "trainRunIDs": args.trainRunIDs,
+            "validateRunIDs": args.validateRunIDs,
+            "data_split": args.data_split
+        }
+    }
+
+    # Print the parameters
+    print(f"\nParameters for Model Architecture Search:")
+    pprint.pprint(params); print()
+
+    # Save to YAML file
+    utils.save_parameters_yaml(params, output_path)
+
+if __name__ == "__main__":
+    cli()

octopi/entry_points/run_segment_predict.py

@@ -0,0 +1,166 @@
+from octopi.pytorch import segmentation
+from octopi.entry_points import common
+import torch, argparse, json, pprint, yaml, os
+from octopi import utils
+from typing import List, Tuple
+
+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.
+    """
+    
+    gpu_count = torch.cuda.device_count()
+    print(f"Number of GPUs available: {gpu_count}")
+
+    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.")
+
+def inference_parser(parser_description, add_slurm: bool = False):
+    """
+    Parse the arguments for the inference
+    """
+    parser = argparse.ArgumentParser(
+        description=parser_description,
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    input_group = parser.add_argument_group("Input Arguments")
+    common.add_config(input_group, single_config=True)
+
+    model_group = parser.add_argument_group("Model Arguments")
+    common.inference_model_parameters(model_group)
+
+    inference_group = parser.add_argument_group("Inference Arguments")
+    common.add_inference_parameters(inference_group)
+
+    if add_slurm:
+        slurm_group = parser.add_argument_group("SLURM Arguments")
+        common.add_slurm_parameters(slurm_group, 'segment_predict', gpus = 2)
+
+    args = parser.parse_args()
+    return args
+
+# Entry point with argparse
+def cli():
+    """
+    CLI entry point for running inference.
+    """
+    
+    # Parse the arguments
+    parser_description = "Run segmentation predictions with a specified model and configuration on CryoET Tomograms."
+    args = inference_parser(parser_description)
+
+    # Set default values if not provided
+    args.seg_info = list(args.seg_info)  # Convert tuple to list
+    if args.seg_info[1] is None:
+        args.seg_info[1] = "octopi"
+
+    if args.seg_info[2] is None:
+        args.seg_info[2] = "1"
+
+    # Save JSON with Parameters
+    output_json = f'segment-predict_{args.seg_info[1]}_{args.seg_info[2]}_{args.seg_info[0]}.yaml'
+    save_parameters(args, output_json)
+
+    # Call the inference function with parsed arguments
+    inference(
+        copick_config_path=args.config,
+        model_weights=args.model_weights,
+        model_config=args.model_config,
+        seg_info=args.seg_info,
+        voxel_size=args.voxel_size,
+        tomo_algorithm=args.tomo_alg,
+        tomo_batch_size=args.tomo_batch_size,
+        run_ids=args.run_ids,
+    )
+
+def save_parameters(args: argparse.Namespace, 
+                    output_path: str):  
+
+    # Load the model config
+    model_config = utils.load_yaml(args.model_config)
+
+    # Create parameters dictionary
+    params = {
+        "inputs": {
+            "config": args.config,
+            "model_config": args.model_config,
+            "model_weights": args.model_weights,
+            "tomo_algorithm": args.tomo_alg,
+            "voxel_size": args.voxel_size
+        },
+        "outputs": {
+            "segmentation_name": args.seg_info[0],
+            "segmentation_user_id": args.seg_info[1],
+            "segmentation_session_id": args.seg_info[2]
+        },
+        'model': model_config['model']
+    }            
+
+    # Print the parameters
+    print(f"\nParameters for Inference (Segment Prediction):")
+    pprint.pprint(params); print()
+
+    # Save to YAML file
+    utils.save_parameters_yaml(params, output_path)
+
+if __name__ == "__main__":
+    cli()