octopi 1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

octopi/utils/config.py

@@ -0,0 +1,57 @@
+"""
+Configuration utilities for MLflow setup and reproducibility.
+"""
+
+from dotenv import load_dotenv
+import torch, numpy as np
+import os, random
+import octopi
+
+
+def mlflow_setup():
+    """
+    Set up MLflow configuration from environment variables.
+    """
+    module_root = os.path.dirname(octopi.__file__)
+    dotenv_path = module_root.replace('src/octopi','') + '.env'
+    load_dotenv(dotenv_path=dotenv_path)
+
+    # MLflow setup
+    username = os.getenv('MLFLOW_TRACKING_USERNAME')
+    password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    if not password or not username:
+        print("Password not found in environment, loading from .env file...")
+        load_dotenv()  # Loads environment variables from a .env file
+        username = os.getenv('MLFLOW_TRACKING_USERNAME')
+        password = os.getenv('MLFLOW_TRACKING_PASSWORD')
+    
+    # Check again after loading .env file
+    if not password:
+        raise ValueError("Password is not set in environment variables or .env file!")
+    else:
+        print("Password loaded successfully")
+        os.environ['MLFLOW_TRACKING_USERNAME'] = username
+        os.environ['MLFLOW_TRACKING_PASSWORD'] = password    
+
+    return os.getenv('MLFLOW_TRACKING_URI')
+
+
+def set_seed(seed):
+    """
+    Set random seeds for reproducibility across Python, NumPy, and PyTorch.
+    """
+    # Set the seed for Python's random module
+    random.seed(seed)
+
+    # Set the seed for NumPy
+    np.random.seed(seed)
+
+    # Set the seed for PyTorch (both CPU and GPU)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # If using multi-GPU
+
+    # Ensure reproducibility of operations by disabling certain optimizations
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False 

octopi/utils/io.py

@@ -0,0 +1,128 @@
+"""
+File I/O utilities for YAML and JSON operations.
+"""
+
+import os, json, yaml
+
+
+# Create a custom dumper that uses flow style for lists only.
+class InlineListDumper(yaml.SafeDumper):
+    def represent_list(self, data):
+        node = super().represent_list(data)
+        node.flow_style = True  # Use inline style for lists
+        return node
+
+
+def save_parameters_yaml(params: dict, output_path: str):
+    """
+    Save parameters to a YAML file.
+    """
+    InlineListDumper.add_representer(list, InlineListDumper.represent_list)
+    with open(output_path, 'w') as f:
+        yaml.dump(params, f, Dumper=InlineListDumper, default_flow_style=False, sort_keys=False)
+
+
+def load_yaml(path: str) -> dict:
+    """
+    Load a YAML file and return the contents as a dictionary.
+    """
+    if os.path.exists(path):
+        with open(path, 'r') as f:
+            return yaml.safe_load(f)
+    else:
+        raise FileNotFoundError(f"File not found: {path}")
+    
+
+def save_results_to_json(results, filename: str):
+    """
+    Save training results to a JSON file.
+    """
+    results = prepare_inline_results_json(results)
+    with open(os.path.join(filename), "w") as json_file:
+        json.dump( results, json_file, indent=4 )
+    print(f"Training Results saved to {filename}")
+
+
+def prepare_inline_results_json(results):
+    """
+    Prepare results for inline JSON formatting.
+    """
+    # Traverse the dictionary and format lists of lists as inline JSON
+    for key, value in results.items():
+        # Check if the value is a list of lists (like [[epoch, value], ...])
+        if isinstance(value, list) and all(isinstance(item, list) and len(item) == 2 for item in value):
+            # Format the list of lists as a single-line JSON string
+            results[key] = json.dumps(value)
+    return results 
+
+def get_optimizer_parameters(trainer):
+    """
+    Extract optimizer parameters from a trainer object.
+    """
+    optimizer_parameters = {
+        'my_num_samples': trainer.num_samples,  
+        'val_interval': trainer.val_interval,
+        'lr': trainer.optimizer.param_groups[0]['lr'],
+        'optimizer': trainer.optimizer.__class__.__name__,
+        'metrics_function': trainer.metrics_function.__class__.__name__,
+        'loss_function': trainer.loss_function.__class__.__name__,
+    }
+
+    # Log Tversky Loss Parameters
+    if trainer.loss_function.__class__.__name__ == 'TverskyLoss':
+        optimizer_parameters['alpha'] = trainer.loss_function.alpha
+    elif trainer.loss_function.__class__.__name__ == 'FocalLoss':
+        optimizer_parameters['gamma'] = trainer.loss_function.gamma
+    elif trainer.loss_function.__class__.__name__ == 'WeightedFocalTverskyLoss':
+        optimizer_parameters['alpha'] = trainer.loss_function.alpha
+        optimizer_parameters['gamma'] = trainer.loss_function.gamma
+        optimizer_parameters['weight_tversky'] = trainer.loss_function.weight_tversky
+    elif trainer.loss_function.__class__.__name__ == 'FocalTverskyLoss':
+        optimizer_parameters['alpha'] = trainer.loss_function.alpha
+        optimizer_parameters['gamma'] = trainer.loss_function.gamma
+
+    return optimizer_parameters
+
+
+def save_parameters_to_yaml(model, trainer, dataloader, filename: str):
+    """
+    Save training parameters to a YAML file.
+    """
+    
+    parameters = {
+        'model': model.get_model_parameters(),
+        'optimizer': get_optimizer_parameters(trainer),
+        'dataloader': dataloader.get_dataloader_parameters()
+    }
+
+    save_parameters_yaml(parameters, filename)
+    print(f"Training Parameters saved to {filename}") 
+
+def flatten_params(params, parent_key=''):
+    """
+    Helper function to flatten and serialize nested parameters.
+    """
+    flattened = {}
+    for key, value in params.items():
+        new_key = f"{parent_key}.{key}" if parent_key else key
+        if isinstance(value, dict):
+            flattened.update(flatten_params(value, new_key))
+        elif isinstance(value, list):
+            flattened[new_key] = ', '.join(map(str, value))  # Convert list to a comma-separated string
+        else:
+            flattened[new_key] = value
+    return flattened
+
+
+def prepare_for_inline_json(data):
+    """
+    Manually join specific lists into strings for inline display.
+    """
+    for key in ["trainRunIDs", "valRunIDs", "testRunIDs"]:
+        if key in data['dataloader']:
+            data['dataloader'][key] = f"[{', '.join(map(repr, data['dataloader'][key]))}]"
+
+    for key in ['channels', 'strides']:
+        if key in data['model']:
+                data['model'][key] = f"[{', '.join(map(repr, data['model'][key]))}]"
+    return data

octopi/{utils.py → utils/parsers.py}

@@ -1,58 +1,12 @@
-from monai.networks.nets import UNet, AttentionUnet
+"""
+Argument parsing and configuration utilities.
+"""
+
+import argparse, os, random
+import torch, numpy as np
 from typing import List, Tuple, Union
 from dotenv import load_dotenv
-import argparse, octopi
-import torch, random, os, yaml
-from typing import List
-import numpy as np
-
-##############################################################################################################################
-
-def mlflow_setup():
-
-    module_root = os.path.dirname(octopi.__file__)
-    dotenv_path = module_root.replace('src/octopi','') + '.env'
-    load_dotenv(dotenv_path=dotenv_path)
-
-    # MLflow setup
-    username = os.getenv('MLFLOW_TRACKING_USERNAME')
-    password = os.getenv('MLFLOW_TRACKING_PASSWORD')
-    if not password or not username:
-        print("Password not found in environment, loading from .env file...")
-        load_dotenv()  # Loads environment variables from a .env file
-        username = os.getenv('MLFLOW_TRACKING_USERNAME')
-        password = os.getenv('MLFLOW_TRACKING_PASSWORD')
-    
-    # Check again after loading .env file
-    if not password:
-        raise ValueError("Password is not set in environment variables or .env file!")
-    else:
-        print("Password loaded successfully")
-        os.environ['MLFLOW_TRACKING_USERNAME'] = username
-        os.environ['MLFLOW_TRACKING_PASSWORD'] = password    
-
-    return os.getenv('MLFLOW_TRACKING_URI')
-
-##############################################################################################################################
-
-def set_seed(seed):
-    # Set the seed for Python's random module
-    random.seed(seed)
-
-    # Set the seed for NumPy
-    np.random.seed(seed)
-
-    # Set the seed for PyTorch (both CPU and GPU)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)  # If using multi-GPU
-
-    # Ensure reproducibility of operations by disabling certain optimizations
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-    
-###############################################################################################################################
+import octopi
 
 def parse_list(value: str) -> List[str]:
     """
@@ -62,7 +16,6 @@ def parse_list(value: str) -> List[str]:
     value = value.strip("[]")  # Remove brackets if present
     return [x.strip() for x in value.split(",")]
 
-###############################################################################################################################
 
 def parse_int_list(value: str) -> List[int]:
     """
@@ -71,7 +24,6 @@ def parse_int_list(value: str) -> List[int]:
     """
     return [int(x) for x in parse_list(value)]
 
-###############################################################################################################################
 
 def string2bool(value: str):
     """
@@ -86,7 +38,6 @@ def string2bool(value: str):
     else:
         raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")
 
-###############################################################################################################################
 
 def parse_target(value: str) -> Tuple[str, Union[str, None], Union[str, None]]:
     """
@@ -130,6 +81,7 @@ def parse_seg_target(value: str) -> List[Tuple[str, Union[str, None], Union[str,
             )
     return targets
 
+
 def parse_copick_configs(config_entries: List[str]):
     """
     Parse a string representing a list of CoPick configuration file paths.
@@ -172,6 +124,7 @@ def parse_copick_configs(config_entries: List[str]):
 
     return copick_configs
 
+
 def parse_data_split(value: str) -> Tuple[float, float, float]:
     """
     Parse data split ratios from string input.
@@ -208,31 +161,4 @@ def parse_data_split(value: str) -> Tuple[float, float, float]:
     if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
         raise ValueError(f"Ratios must sum to 1.0, got {train_ratio + val_ratio + test_ratio}")
     
-    return round(train_ratio, 2), round(val_ratio, 2), round(test_ratio, 2)
-
-##############################################################################################################################
-
-# Create a custom dumper that uses flow style for lists only.
-class InlineListDumper(yaml.SafeDumper):
-    def represent_list(self, data):
-        node = super().represent_list(data)
-        node.flow_style = True  # Use inline style for lists
-        return node
-
-def save_parameters_yaml(params: dict, output_path: str):
-    """
-    Save parameters to a YAML file.
-    """
-    InlineListDumper.add_representer(list, InlineListDumper.represent_list)
-    with open(output_path, 'w') as f:
-        yaml.dump(params, f, Dumper=InlineListDumper, default_flow_style=False, sort_keys=False)
-
-def load_yaml(path: str) -> dict:
-    """
-    Load a YAML file and return the contents as a dictionary.
-    """
-    if os.path.exists(path):
-        with open(path, 'r') as f:
-            return yaml.safe_load(f)
-    else:
-        raise FileNotFoundError(f"File not found: {path}")
+    return round(train_ratio, 2), round(val_ratio, 2), round(test_ratio, 2) 

octopi/{stopping_criteria.py → utils/stopping_criteria.py}

@@ -7,9 +7,9 @@ class EarlyStoppingChecker:
     
     def __init__(self, 
                  max_nan_epochs=15, 
-                 plateau_patience=20, 
-                 plateau_min_delta=0.001,
-                 stagnation_patience=50, 
+                 plateau_patience=10, 
+                 plateau_min_delta=0.005,
+                 stagnation_patience=30, 
                  convergence_window=5, 
                  convergence_threshold=0.005,
                  val_interval=15,

octopi/{visualization_tools.py → utils/visualization_tools.py}

@@ -1,7 +1,7 @@
+from copick_utils.io import readers
 import matplotlib.colors as mcolors
 from typing import Optional, List
 import matplotlib.pyplot as plt
-from octopi import io
 import numpy as np
 
 # Define the plotting function
@@ -76,7 +76,7 @@ def show_tomo_points(tomo, run, objects, user_id, vol_slice, session_id = None,
 
     for name,_,_ in objects:
         try:    
-            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id, session_id=session_id)
+            coordinates = readers.coordinates(run, name=name, user_id=user_id, session_id=session_id)
             close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
             plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)
         except:
@@ -94,7 +94,7 @@ def compare_tomo_points(tomo, run, objects, vol_slice, user_id1, user_id2,
 
     for name,_,_ in objects:
         try:
-            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id1, session_id=session_id1)
+            coordinates = readers.coordinates(run, name=name, user_id=user_id1, session_id=session_id1)
             close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
             plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)    
         except: 
@@ -106,7 +106,7 @@ def compare_tomo_points(tomo, run, objects, vol_slice, user_id1, user_id2,
     
     for name,_,_ in objects:
         try:
-            coordinates = io.get_copick_coordinates(run, name=name, user_id=user_id2, session_id=session_id2)
+            coordinates = readers.coordinates(run, name=name, user_id=user_id2, session_id=session_id2)
             close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
             plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)
         except:

octopi/workflows.py

@@ -0,0 +1,236 @@
+from octopi.extract.localize import process_localization
+import octopi.processing.evaluate as octopi_evaluate
+from monai.metrics import ConfusionMatrixMetric
+from octopi.models import common as builder
+from octopi.pytorch import segmentation
+from octopi.datasets import generators
+from octopi.pytorch import trainer 
+import multiprocess as mp
+import copick, torch, os
+from octopi.utils import io
+from tqdm import tqdm
+    
+def train(config, target_info, tomo_algorithm, voxel_size, loss_function,
+          model_config = None, model_weights = None, trainRunIDs = None, validateRunIDs = None,
+          model_save_path = 'results', best_metric = 'fBeta2', num_epochs = 1000, use_ema = True):
+    """
+    Train a UNet Model for Segmentation
+
+    Args:
+        config (str): Path to the Copick Config File
+        target_info (list): List containing the target user ID, target session ID, and target algorithm
+        tomo_algorithm (str): The tomographic algorithm to use for segmentation
+        voxel_size (float): The voxel size of the data
+        loss_function (str): The loss function to use for training
+        model_config (dict): The model configuration
+        model_weights (str): The path to the model weights
+        trainRunIDs (list): The list of run IDs to use for training
+        validateRunIDs (list): The list of run IDs to use for validation
+        model_save_path (str): The path to save the model
+        best_metric (str): The metric to use for early stopping
+        num_epochs (int): The number of epochs to train for
+    """
+
+    # If No Model Configuration is Provided, Use the Default Configuration
+    if model_config is None:
+        root = copick.from_file(config)
+        model_config = {
+            'architecture': 'Unet',
+            'num_classes': root.pickable_objects[-1].label + 1,
+            'dim_in': 80,
+            'strides': [2, 2, 1],
+            'channels': [48, 64, 80, 80],
+            'dropout': 0.0, 'num_res_units': 1,
+        }
+        print('No Model Configuration Provided, Using Default Configuration')
+        print(model_config)
+    
+    data_generator = generators.TrainLoaderManager(
+            config, 
+            target_info[0], 
+            target_session_id = target_info[2],
+            target_user_id = target_info[1],
+            tomo_algorithm = tomo_algorithm,
+            voxel_size = voxel_size,
+            Nclasses = model_config['num_classes'],
+            tomo_batch_size = 15 ) 
+
+    data_generator.get_data_splits(
+        trainRunIDs = trainRunIDs,
+        validateRunIDs = validateRunIDs,
+        train_ratio = 0.9, val_ratio = 0.1, test_ratio = 0.0,
+        create_test_dataset = False)
+
+    # Get the reload frequency
+    data_generator.get_reload_frequency(num_epochs)
+
+    # Monai Functions
+    metrics_function = ConfusionMatrixMetric(include_background=False, metric_name=["recall",'precision','f1 score'], reduction="none")
+    
+    # Build the Model
+    model_builder = builder.get_model(model_config['architecture'])
+    model = model_builder.build_model(model_config)
+    
+    # Load the Model Weights if Provided 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if model_weights: 
+        state_dict = torch.load(model_weights, map_location=device, weights_only=True)
+        model.load_state_dict(state_dict)     
+    model.to(device) 
+
+    # Optimizer
+    optimizer = torch.optim.Adam(model.parameters(), 1e-3)
+
+    # Create UNet-Trainer
+    model_trainer = trainer.ModelTrainer(
+        model, device, loss_function, metrics_function, optimizer,
+        use_ema = use_ema
+    )
+
+    results = model_trainer.train(
+        data_generator, model_save_path, max_epochs=num_epochs,
+        crop_size=model_config['dim_in'], my_num_samples=16,
+        val_interval=10, best_metric=best_metric, verbose=True
+    )
+    
+    # Save parameters and results
+    parameters_save_name = os.path.join(model_save_path, "model_config.yaml")
+    io.save_parameters_to_yaml(model_builder, model_trainer, data_generator, parameters_save_name)
+
+    # TODO: Write Results to Zarr or Another File Format? 
+    results_save_name = os.path.join(model_save_path, "results.json")
+    io.save_results_to_json(results, results_save_name)
+
+def segment(config, tomo_algorithm, voxel_size, model_weights, model_config, 
+            seg_info = ['predict', 'octopi', '1'], use_tta = False, run_ids = None):
+    """
+    Segment a Dataset using a Trained Model or Ensemble of Models
+
+    Args:
+        config (str): Path to the Copick Config File
+        tomo_algorithm (str): The tomographic algorithm to use for segmentation
+        voxel_size (float): The voxel size of the data
+        model_weights (str, list): The path to the model weights or a list of paths to the model weights
+        model_config (str, list): The model configuration or a list of model configurations
+        seg_info (list): The segmentation information
+        use_tta (bool): Whether to use test time augmentation
+        run_ids (list): The list of run IDs to use for segmentation
+    """
+
+    # Initialize the Predictor
+    predict = segmentation.Predictor(
+        config,
+        model_config,
+        model_weights,
+        apply_tta = use_tta
+    )
+
+    # Run batch prediction
+    predict.batch_predict(
+        runIDs=run_ids,
+        num_tomos_per_batch=15,
+        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]
+    )
+
+def localize(config, voxel_size, seg_info, pick_user_id, pick_session_id, n_procs = 16,
+            method = 'watershed', filter_size = 10, radius_min_scale = 0.4, radius_max_scale = 1.0,
+            run_ids = None):
+    """
+    Extract 3D Coordinates from the Segmentation Maps
+
+    Args:
+        config (str): Path to the Copick Config File
+        voxel_size (float): The voxel size of the data
+        seg_info (list): The segmentation information
+        pick_user_id (str): The user ID of the pick
+        pick_session_id (str): The session ID of the pick
+        n_procs (int): The number of processes to use for parallelization
+        method (str): The method to use for localization
+        filter_size (int): The filter size to use for localization
+        radius_min_scale (float): The minimum radius scale to use for localization
+        radius_max_scale (float): The maximum radius scale to use for localization
+        run_ids (list): The list of run IDs to use for localization
+    """
+
+    # Load the Copick Config
+    root = copick.from_file(config) 
+
+    # Get objects that can be Picked
+    objects = [(obj.name, obj.label, obj.radius) for obj in root.pickable_objects if obj.is_particle]    
+
+    # Get all RunIDs
+    if run_ids is None:
+        run_ids = [run.name for run in root.runs]
+    n_run_ids = len(run_ids)
+
+     # Run Localization - Main Parallelization Loop
+    print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")
+    with mp.Pool(processes=n_procs) as pool:
+        with tqdm(total=n_run_ids, desc="Localization", unit="run") as pbar:
+            worker_func = lambda run_id: process_localization(
+                root.get_run(run_id),  
+                objects, 
+                seg_info,
+                method, 
+                voxel_size,
+                filter_size,
+                radius_min_scale, 
+                radius_max_scale,
+                pick_session_id,
+                pick_user_id
+            )
+
+            for _ in pool.imap_unordered(worker_func, run_ids, chunksize=1):
+                pbar.update(1)
+
+    print('Localization Complete!')
+    
+
+def evaluate(config, 
+             gt_user_id, gt_session_id,
+             pred_user_id, pred_session_id,
+             run_ids = None, distance_threshold = 0.5, save_path = None):
+    """
+    Evaluate the Localization on a Dataset
+
+    Args:
+        config (str): Path to the Copick Config File
+        gt_user_id (str): The user ID of the ground truth
+        gt_session_id (str): The session ID of the ground truth
+        pred_user_id (str): The user ID of the predicted coordinates
+        pred_session_id (str): The session ID of the predicted coordinates
+        run_ids (list): The list of run IDs to use for evaluation
+        distance_threshold (float): The distance threshold to use for evaluation
+        save_path (str): The path to save the evaluation results
+    """
+             
+    print('Running Evaluation on the Following Query:')
+    print(f'Distance Threshold: {distance_threshold}')
+    print(f'GT User ID: {gt_user_id}, GT Session ID: {gt_session_id}')
+    print(f'Pred User ID: {pred_user_id}, Pred Session ID: {pred_session_id}')
+    print(f'Run IDs: {run_ids}')
+    
+    # Load the Copick Config
+    root = copick.from_file(config) 
+
+    # For Now Lets Assume Object Names are None..
+    object_names = None
+
+    # Run Evaluation
+    eval = octopi_evaluate.evaluator(
+        config,
+        gt_user_id,
+        gt_session_id,
+        pred_user_id,
+        pred_session_id, 
+        object_names=object_names
+    )
+
+    eval.run(
+        distance_threshold_scale=distance_threshold, 
+        runIDs=run_ids, save_path=save_path
+    )