octopi 1.4.0__py3-none-any.whl

octopi/utils/io.py

@@ -0,0 +1,215 @@
+"""
+File I/O utilities for YAML and JSON operations.
+"""
+
+import os, json, yaml, copick, glob
+import pandas as pd
+
+
+# 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_csv(results, filename: str):
+    """Save training results to a CSV file (aligned to validation steps)."""
+    data = {}
+    
+    # Get validation steps (these are the steps we want to keep)
+    val_steps = set()
+    for key, value in results.items():
+        if isinstance(value, list) and value and isinstance(value[0], tuple):
+            if key.startswith(('val_', 'avg_', 'f1_', 'recall_', 'precision_', 'fbeta_')):
+                val_steps.update([item[0] for item in value])
+                break
+    
+    val_steps = sorted(val_steps)
+    data['step'] = val_steps
+    
+    # Extract all metrics, filtering to validation steps only
+    for key, value in results.items():
+        if isinstance(value, list) and value and isinstance(value[0], tuple):
+            step_to_value = {item[0]: item[1] for item in value}
+            data[key] = [step_to_value.get(step, None) for step in val_steps]
+    
+    df = pd.DataFrame(data)
+    df.to_csv(filename, index=False)
+    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.
+    """
+
+    # Check for the target configuration file for model labels
+    target_config = check_target_config_path(dataloader)
+    
+    # Extract and flatten parameters
+    parameters = {
+        'model': model.get_model_parameters(),
+        'labels': target_config['input']['labels'],
+        '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
+
+def check_target_config_path(data_generator):
+    """
+    Check for the target configuration file in the CoPick overlay or static root directories.
+    If the session_id is not provided, search for the most recent file matching the target name
+    """
+    
+    # Open the Copick Project for MultiConfig or SingleConfig Workflow
+    if isinstance(data_generator.config, dict):
+        sessions = list(data_generator.config.keys())
+        config_path = data_generator.config[sessions[0]]
+    else:
+        config_path = data_generator.config
+
+    # Get the Target Config File
+    return get_config(
+        config_path, data_generator.target_name, 'targets',
+        data_generator.target_user_id, data_generator.target_session_id
+    )
+
+def get_config(config_path, name, process, user_id=None, session_id=None):
+    """
+    Get the configuration for a specific process and target name.
+    """
+
+    # Get the Overlay and Static Roots
+    root = copick.from_file(config_path)
+
+    # Remove the local:// prefix from static_root if it exists  
+    overlay_root = remove_prefix(root.config.overlay_root)
+    try:  # Check if static_root is available
+        static_root = remove_prefix(root.config.static_root)
+    except: # If not, set it to None
+        static_root = None
+    
+    # Two Search Patterns, Either only a name provided or name, user_id, session_id
+    if session_id is None:
+        pattern = glob.glob(os.path.join(overlay_root, 'logs', f"{process}_*{name}.yaml"))
+        if len(pattern) == 0 and static_root is not None:
+            pattern = glob.glob(os.path.join(static_root, 'logs', f"{process}_*{name}.yaml"))
+        fname = pattern[-1]
+    else:
+        fname = f"{process}-{user_id}_{session_id}_{name}.yaml"
+
+    # The Target Config File Should Either in be the Overlay or Static Root
+    if os.path.exists(os.path.join(overlay_root, 'logs', fname)):
+        path = os.path.join(overlay_root, 'logs', fname)
+    elif static_root is not None and os.path.exists(os.path.join(static_root, 'logs', fname)):
+        path = os.path.join(static_root, 'logs', fname)
+    else:
+        raise FileNotFoundError(f"Target config file not found: {fname}")
+
+    # Load the Target Config File
+    with open(path, 'r') as f:
+        target_config = yaml.safe_load(f)
+    return target_config
+
+def remove_prefix(text: str) -> str:
+    """
+    Remove a prefix from a string if it exists.
+    """
+    # Check if the text is None
+    if text is None:
+        return None
+    elif text[:8] == 'local://':
+        text = text[8:]
+    return text

octopi/utils/losses.py

@@ -0,0 +1,86 @@
+from monai.losses import FocalLoss, TverskyLoss
+import torch
+
+class WeightedFocalTverskyLoss(torch.nn.Module):
+    def __init__(
+        self, gamma=1.0, alpha=0.7, beta=0.3, 
+        weight_tversky=0.5, weight_focal=0.5, 
+        smooth=1e-5, **kwargs ):
+        """
+        Weighted combination of Focal and Tversky loss.
+
+        Args:
+            gamma (float): Focus parameter for Focal Loss.
+            alpha (float): Weight for false positives in Tversky Loss.
+            beta (float): Weight for false negatives in Tversky Loss.
+            weight_tversky (float): Weight of Tversky loss in the combination.
+            weight_focal (float): Weight of Focal loss in the combination.
+            smooth (float): Smoothing factor to avoid division by zero.
+        """
+        super().__init__()
+        self.tversky_loss = TverskyLoss(
+            alpha=alpha, beta=beta, include_background=True, 
+            to_onehot_y=True, softmax=True,
+            smooth_nr=smooth, smooth_dr=smooth, **kwargs
+        )
+        self.focal_loss = FocalLoss(
+            include_background=True, to_onehot_y=True, 
+            use_softmax=True, gamma=gamma
+        )
+        self.alpha = alpha
+        self.beta = beta
+        self.gamma = gamma
+        self.weight_tversky = weight_tversky
+        self.weight_focal = weight_focal
+
+    def forward(self, y_pred, y_true):
+        """
+        Compute the combined loss.
+
+        Args:
+            y_pred (Tensor): Predicted probabilities (B, C, ...).
+            y_true (Tensor): Ground truth labels (B, C, ...), one-hot encoded.
+
+        Returns:
+            Tensor: Weighted combination of Tversky and Focal losses.
+        """
+        tversky = self.tversky_loss(y_pred, y_true)
+        focal = self.focal_loss(y_pred, y_true)
+        return self.weight_tversky * tversky + self.weight_focal * focal
+    
+class FocalTverskyLoss(TverskyLoss):
+    def __init__(
+        self, 
+        alpha=0.7, beta=0.3, gamma=1.0, smooth=1e-5, **kwargs):
+        """
+        Focal Tversky Loss with an additional power term for harder samples.
+
+        From https://arxiv.org/abs/1810.07842
+
+        Args:
+            alpha (float): Weight for false positives.
+            beta (float): Weight for false negatives.
+            gamma (float): Focus parameter (like Focal Loss).
+            smooth (float): Smoothing factor to avoid division by zero.
+        """
+        super().__init__(
+            alpha=alpha, beta=beta, 
+            include_background=True, 
+            to_onehot_y=True, softmax=True, 
+            smooth_nr=smooth, smooth_dr=smooth, **kwargs)
+        self.gamma = gamma
+        self.alpha = alpha
+        self.beta = beta
+
+    def forward(self, y_pred, y_true):
+        """
+        Args:
+            y_pred (Tensor): Predicted probabilities (B, C, ...).
+            y_true (Tensor): Ground truth labels (B, C, ...), one-hot encoded.
+
+        Returns:
+            Tensor: Loss value.
+        """
+        tversky_loss = super().forward(y_pred, y_true)
+        modified_loss = torch.pow(tversky_loss, 1 / self.gamma)
+        return modified_loss

octopi/utils/parsers.py

@@ -0,0 +1,162 @@
+"""
+Argument parsing and configuration utilities.
+"""
+
+from typing import List, Tuple, Union
+from dotenv import load_dotenv
+import argparse
+
+def parse_list(value: str) -> List[str]:
+    """
+    Parse a string representing a list of items.
+    Supports formats like '[item1,item2,item3]' or 'item1,item2,item3'.
+    """
+    value = value.strip("[]")  # Remove brackets if present
+    return [x.strip() for x in value.split(",")]
+
+
+def parse_int_list(value: str) -> List[int]:
+    """
+    Parse a string representing a list of integers.
+    Supports formats like '[1,2,3]' or '1,2,3'.
+    """
+    return [int(x) for x in parse_list(value)]
+
+
+def string2bool(value: str):
+    """
+    Custom function to convert string values to boolean.
+    """
+    if isinstance(value, bool):
+        return value
+    if value.lower() in {'True', 'true', 't', '1', 'yes'}:
+        return True
+    elif value.lower() in {'False', 'false', 'f', '0', 'no'}:
+        return False
+    else:
+        raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")
+
+
+def parse_target(value: str) -> Tuple[str, Union[str, None], Union[str, None]]:
+    """
+    Parse a single target string.
+    Expected formats:
+      - "name"
+      - "name,user_id,session_id"
+    """
+    parts = value.split(',')
+    if len(parts) == 1:
+        obj_name = parts[0]
+        return obj_name, None, None
+    elif len(parts) == 3:
+        obj_name, user_id, session_id = parts
+        return obj_name, user_id, session_id
+    else:
+        raise argparse.ArgumentTypeError(
+            f"Invalid target format: '{value}'. Expected 'name' or 'name,user_id,session_id'."
+        )
+
+
+def parse_seg_target(value: str) -> List[Tuple[str, Union[str, None], Union[str, None]]]:
+    """
+    Parse segmentation targets. Each target can have the format:
+      - "name"
+      - "name,user_id,session_id"
+    Multiple targets can be comma-separated.
+    """
+    targets = []
+    for target in value.split(';'):  # Use ';' as a separator for multiple targets
+        parts = target.split(',')
+        if len(parts) == 1:
+            name = parts[0]
+            targets.append((name, None, None))
+        elif len(parts) == 3:
+            name, user_id, session_id = parts
+            targets.append((name, user_id, session_id))
+        else:
+            raise argparse.ArgumentTypeError(
+                f"Invalid seg-target format: '{target}'. Expected 'name' or 'name,user_id,session_id'."
+            )
+    return targets
+
+
+def parse_copick_configs(config_entries: List[str]):
+    """
+    Parse a string representing a list of CoPick configuration file paths.
+    """
+    # Process the --config arguments into a dictionary
+    copick_configs = {}
+
+    for config_entry in config_entries:
+        if ',' in config_entry:
+            # Entry has a session name and a config path
+            try:
+                session_name, config_path = config_entry.split(',', 1)
+                copick_configs[session_name] = config_path
+            except ValueError:
+                raise argparse.ArgumentTypeError(
+                    f"Invalid format for --config entry: '{config_entry}'. Expected 'session_name,/path/to/config.json'."
+                )
+        else:
+            # Single configuration path without a session name
+            # if "default" in copick_configs:
+            #     raise argparse.ArgumentTypeError(
+            #         f"Only one single-path --config entry is allowed when using default configurations. "
+            #         f"Detected duplicate: {config_entry}"
+            #     )
+            # copick_configs["default"] = config_entry
+            copick_configs = config_entry
+
+        # if ',' in config_entry:
+        #     parts = config_entry.split(',')
+        #     if len(parts) == 2:
+        #         # Entry with session name and config path
+        #         session_name, config_path = parts
+        #         copick_configs[session_name] = {"path": config_path, "algorithm": None}
+        #     elif len(parts) == 3:
+        #         # Entry with session name, config path, and algorithm
+        #         session_name, config_path, algorithm = parts
+        #         copick_configs[session_name] = {"path": config_path, "algorithm": algorithm}    
+        # else:
+        #     copick_configs = config_entry
+
+    return copick_configs
+
+
+def parse_data_split(value: str) -> Tuple[float, float, float]:
+    """
+    Parse data split ratios from string input.
+    
+    Args:
+        value: Either a single float (e.g., "0.8") or two comma-separated floats (e.g., "0.7,0.1")
+    
+    Returns:
+        Tuple of (train_ratio, val_ratio, test_ratio)
+    
+    Examples:
+        "0.8" -> (0.8, 0.2, 0.0)
+        "0.7,0.1" -> (0.7, 0.1, 0.2)
+    """
+    parts = value.split(',')
+    
+    if len(parts) == 1:
+        # Single value provided - use it as train ratio
+        train_ratio = float(parts[0])
+        val_ratio = 1.0 - train_ratio
+        test_ratio = 0.0
+    elif len(parts) == 2:
+        # Two values provided - use as train and val ratios
+        train_ratio = float(parts[0])
+        val_ratio = float(parts[1])
+        test_ratio = 1.0 - train_ratio - val_ratio
+    else:
+        raise ValueError("Data split must be either a single value or two comma-separated values")
+    
+    # Validate ratios
+    if train_ratio < 0 or val_ratio < 0 or test_ratio < 0:
+        raise ValueError("All ratios must be non-negative")
+    
+    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) 

octopi/utils/progress.py

@@ -0,0 +1,78 @@
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    BarColumn,
+    MofNCompleteColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+import json
+
+# Minimal helper to get a shared Console without top-level Rich dependency.
+def get_console():
+    return Console()
+
+def _progress(iterable, description="Processing"):
+    """
+    Wrap an iterable with a Rich progress bar.
+
+    Args:
+        iterable: Any iterable object (e.g., list, generator).
+        description: Text label to display above the progress bar.
+
+    Yields:
+        Each item from the iterable, while updating the progress bar.
+
+    Example:
+        for x in _progress(range(10), "Doing work"):
+            time.sleep(0.5)
+    """
+
+    console = Console()
+
+    # The generator itself yields items while advancing the progress bar
+    with Progress(
+        SpinnerColumn(),
+        TextColumn(f"[bold blue]{description}"),
+        BarColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        transient=False,
+        console=console,
+    ) as progress:
+        task = progress.add_task(description, total=len(iterable))
+        for item in iterable:
+            yield item
+            progress.advance(task)
+
+def print_summary(process: str, **kwargs):
+    """
+    Pretty-print download parameters using Rich in a clean table with green highlights.
+    """
+
+    console = Console()
+
+    # ---- Header ----
+    console.rule(f"[bold green]{process} Parameters Summary[/bold green]")
+
+    # ---- Table (no outer box) ----
+    table = Table(
+        show_header=True,
+        header_style="bold magenta",
+        expand=False,
+        border_style="green",   # table borders only
+    )
+    table.add_column("Parameter", style="cyan", no_wrap=True)
+    table.add_column("Value", style="white")
+
+    for key, value in kwargs.items():
+        if isinstance(value, (dict, list)):
+            value = json.dumps(value, indent=2)
+        table.add_row(str(key), str(value))
+
+    console.print(table)   # Print table directly, NO panel wrapper

octopi/utils/stopping_criteria.py

@@ -0,0 +1,143 @@
+import numpy as np
+
+class EarlyStoppingChecker:
+    """
+    A class to manage various early stopping criteria for model training.
+    """
+    
+    def __init__(self, 
+                 max_nan_epochs=15, 
+                 plateau_patience=20, 
+                 plateau_min_delta=0.005,
+                 stagnation_patience=30, 
+                 convergence_window=5, 
+                 convergence_threshold=0.005,
+                 val_interval=15,
+                 monitor_metric='avg_fbeta'):
+        """
+        Initialize early stopping parameters.
+        
+        Args:
+            max_nan_epochs: Maximum number of epochs with NaN loss before stopping
+            plateau_patience: Number of validation checks to wait for plateau detection
+            plateau_min_delta: Minimum change to qualify as improvement
+            stagnation_patience: Number of validation intervals to wait for best metric improvement
+            convergence_window: Window size for calculating improvement rate
+            convergence_threshold: Minimum improvement rate threshold
+            val_interval: Number of epochs between validation runs
+            monitor_metric: Primary metric to monitor for early stopping criteria
+        """
+        self.max_nan_epochs = max_nan_epochs
+        self.plateau_patience = plateau_patience
+        self.plateau_min_delta = plateau_min_delta
+        self.stagnation_patience = stagnation_patience
+        self.convergence_window = convergence_window
+        self.convergence_threshold = convergence_threshold
+        self.val_interval = val_interval
+        self.monitor_metric = monitor_metric
+        
+        # Counters
+        self.nan_counter = 0
+        
+        # Flags for detailed reporting
+        self.stopped_reason = None
+
+    def check_for_nan(self, epoch_loss):
+        """Check for NaN in the loss."""
+        if np.isnan(epoch_loss):
+            self.nan_counter += 1
+            if self.nan_counter > self.max_nan_epochs:
+                self.stopped_reason = f"NaN values in loss for more than {self.max_nan_epochs} epochs"
+                return True
+        else:
+            self.nan_counter = 0  # Reset the counter if loss is valid
+        return False
+
+    def check_for_plateau(self, results):
+        """Detect plateaus in validation metrics."""
+        if len(results[self.monitor_metric]) < self.plateau_patience + 1:
+            return False
+            
+        # Get the last 'patience' number of validation points
+        recent_values = [x[1] for x in results[self.monitor_metric][-self.plateau_patience:]]
+        # Find the max value in the window
+        max_value = max(recent_values)
+        # Find the min value in the window
+        min_value = min(recent_values)
+        
+        # If the range of values is small, consider it a plateau
+        if max_value - min_value < self.plateau_min_delta:
+            self.stopped_reason = f"{self.monitor_metric} plateaued for {self.plateau_patience} validations"
+            return True
+        
+        return False
+
+    def check_best_metric_stagnation(self, results):
+        """Stop if best metric hasn't improved for a number of validation intervals."""
+        if "best_metric_epoch" not in results or len(results[self.monitor_metric]) < self.stagnation_patience + 1:
+            return False
+            
+        # Get epoch of the best metric so far
+        best_epoch = results["best_metric_epoch"]
+        current_epoch = results[self.monitor_metric][-1][0]
+        
+        # Check if it's been more than 'patience' validation intervals
+        if (current_epoch - best_epoch) >= (self.stagnation_patience * self.val_interval):
+            self.stopped_reason = f"No improvement for {self.stagnation_patience} validation intervals"
+            return True
+            
+        return False
+
+    # def check_convergence_rate(self, results):
+    #     """Stop when improvement rate slows below threshold."""
+    #     if len(results[self.monitor_metric]) < self.convergence_window + 1:
+    #         return False
+        
+    #     # Calculate average improvement rate over window
+    #     recent_values = [x[1] for x in results[self.monitor_metric][-(self.convergence_window+1):]]
+    #     improvements = [recent_values[i+1] - recent_values[i] for i in range(self.convergence_window)]
+    #     avg_improvement = sum(improvements) / self.convergence_window
+        
+    #     if avg_improvement < self.convergence_threshold and avg_improvement > 0:
+    #         self.stopped_reason = f"Convergence rate ({avg_improvement:.6f}) below threshold"
+    #         return True
+            
+    #     return False
+
+    def should_stop_training(self, epoch_loss, results=None, check_metrics=False):
+        """
+        Comprehensive check for whether training should stop.
+        
+        Args:
+            epoch_loss: Current epoch's loss value
+            results: Dictionary containing training metrics history
+            check_metrics: Whether to also check validation metrics-based criteria
+            
+        Returns:
+            bool: True if training should stop, False otherwise
+        """
+        # Check for NaN in loss (can be done every epoch)
+        if self.check_for_nan(epoch_loss):
+            return True
+            
+        # Only check metric-based criteria if requested and results are provided
+        if check_metrics and results:
+            # Check for plateau in validation metrics
+            if self.check_for_plateau(results):
+                return True
+                
+            # Check if best metric hasn't improved for a while
+            if self.check_best_metric_stagnation(results):
+                return True
+                
+            # # Check if convergence rate has slowed down
+            # if self.check_convergence_rate(results):
+            #     return True
+                
+        return False
+    
+    def get_stopped_reason(self):
+        """Get the reason for stopping, if any."""
+        if self.stopped_reason:
+            return f"Early stopping triggered: {self.stopped_reason}"
+        return "No early stopping criteria met."