fastMONAI 0.5.4__py3-none-any.whl → 0.6.0__py3-none-any.whl

fastMONAI/utils.py

@@ -2,7 +2,7 @@
 
 # %% auto 0
 __all__ = ['store_variables', 'load_variables', 'store_patch_variables', 'load_patch_variables', 'print_colab_gpu_info',
-           'ModelTrackingCallback', 'MLflowUIManager']
+           'ModelTrackingCallback', 'create_mlflow_callback', 'MLflowUIManager']
 
 # %% ../nbs/07_utils.ipynb 1
 import pickle
@@ -13,6 +13,7 @@ import mlflow.pytorch
 import os
 import tempfile
 import json
+from datetime import datetime
 from fastai.callback.core import Callback
 from fastcore.foundation import L
 from typing import Any
@@ -136,12 +137,147 @@ def print_colab_gpu_info():
     else: print(colab_gpu_msg)
 
 # %% ../nbs/07_utils.ipynb 9
+def _detect_patch_workflow(dls) -> bool:
+    """Detect if DataLoaders are patch-based (MedPatchDataLoaders).
+    
+    Args:
+        dls: DataLoaders instance
+        
+    Returns:
+        True if dls is a MedPatchDataLoaders instance
+    """
+    return hasattr(dls, 'patch_config') or hasattr(dls, '_patch_config')
+
+
+def _extract_size_from_transforms(tfms) -> list | None:
+    """Extract target size from PadOrCrop transform if present.
+    
+    Args:
+        tfms: List of transforms
+        
+    Returns:
+        Target size as list, or None if not found
+    """
+    if tfms is None:
+        return None
+    for tfm in tfms:
+        if hasattr(tfm, 'pad_or_crop') and hasattr(tfm.pad_or_crop, 'target_shape'):
+            return list(tfm.pad_or_crop.target_shape)
+    return None
+
+
+def _extract_standard_config(learn) -> dict:
+    """Extract config from standard MedDataBlock workflow.
+    
+    Args:
+        learn: fastai Learner instance
+        
+    Returns:
+        Dictionary with extracted configuration
+    """
+    from fastMONAI.vision_core import MedBase
+    dls = learn.dls
+
+    # Get preprocessing from MedBase class attributes
+    apply_reorder = MedBase.apply_reorder
+    target_spacing = MedBase.target_spacing
+
+    # Extract item_tfms from DataLoaders pipeline
+    item_tfms = []
+    if hasattr(dls, 'after_item') and dls.after_item:
+        item_tfms = list(dls.after_item.fs)
+
+    # Extract size from PadOrCrop transform
+    size = _extract_size_from_transforms(item_tfms)
+
+    return {
+        'apply_reorder': apply_reorder,
+        'target_spacing': target_spacing,
+        'size': size,
+        'item_tfms': item_tfms,
+        'batch_size': dls.bs,
+        'patch_config': None,
+    }
+
+
+def _extract_patch_config(learn) -> dict:
+    """Extract config from MedPatchDataLoaders workflow.
+    
+    Args:
+        learn: fastai Learner instance
+        
+    Returns:
+        Dictionary with extracted configuration including patch-specific params
+    """
+    dls = learn.dls
+    patch_config = getattr(dls, '_patch_config', None) or getattr(dls, 'patch_config', None)
+
+    config = {
+        'apply_reorder': getattr(dls, '_apply_reorder', patch_config.apply_reorder if patch_config else False),
+        'target_spacing': getattr(dls, '_target_spacing', patch_config.target_spacing if patch_config else None),
+        'size': patch_config.patch_size if patch_config else None,
+        'item_tfms': getattr(dls, '_pre_patch_tfms', []) or [],
+        'batch_size': dls.bs,
+    }
+
+    # Add patch-specific params for logging
+    if patch_config:
+        config['patch_config'] = {
+            'patch_size': patch_config.patch_size,
+            'patch_overlap': patch_config.patch_overlap,
+            'samples_per_volume': patch_config.samples_per_volume,
+            'sampler_type': patch_config.sampler_type,
+            'label_probabilities': str(patch_config.label_probabilities) if patch_config.label_probabilities else None,
+            'queue_length': patch_config.queue_length,
+            'aggregation_mode': patch_config.aggregation_mode,
+            'padding_mode': patch_config.padding_mode,
+            'keep_largest_component': patch_config.keep_largest_component,
+        }
+    else:
+        config['patch_config'] = None
+
+    return config
+
+
+def _extract_loss_name(learn) -> str:
+    """Extract loss function name from Learner.
+    
+    Args:
+        learn: fastai Learner instance
+        
+    Returns:
+        Name of the loss function
+    """
+    loss_func = learn.loss_func
+    # Handle CustomLoss wrapper
+    if hasattr(loss_func, 'loss_func'):
+        inner = loss_func.loss_func
+        return inner._get_name() if hasattr(inner, '_get_name') else inner.__class__.__name__
+    return loss_func._get_name() if hasattr(loss_func, '_get_name') else loss_func.__class__.__name__
+
+
+def _extract_model_name(learn) -> str:
+    """Extract model architecture name from Learner.
+    
+    Args:
+        learn: fastai Learner instance
+        
+    Returns:
+        Name of the model architecture
+    """
+    model = learn.model
+    return model._get_name() if hasattr(model, '_get_name') else model.__class__.__name__
+
+# %% ../nbs/07_utils.ipynb 10
 class ModelTrackingCallback(Callback):
     """
     A FastAI callback for comprehensive MLflow experiment tracking.
     
     This callback automatically logs hyperparameters, metrics, model artifacts,
-    and configuration to MLflow during training.
+    and configuration to MLflow during training. If a SaveModelCallback is present,
+    the best model checkpoint will also be logged as an artifact.
+    
+    Supports auto-managed runs when created via `create_mlflow_callback()`.
     """
     
     def __init__(
@@ -151,7 +287,13 @@ class ModelTrackingCallback(Callback):
         item_tfms: list[Any],
         size: list[int], 
         target_spacing: list[float], 
-        apply_reorder: bool
+        apply_reorder: bool,
+        experiment_name: str = None,
+        run_name: str = None,
+        auto_start: bool = False,
+        patch_config: dict = None,
+        extra_params: dict = None,
+        extra_tags: dict = None
     ):
         """
         Initialize the MLflow tracking callback.
@@ -159,9 +301,16 @@ class ModelTrackingCallback(Callback):
         Args:
             model_name: Name of the model architecture for registration
             loss_function: Name of the loss function being used
+            item_tfms: List of item transforms
             size: Model input dimensions
             target_spacing: Resampling dimensions
             apply_reorder: Whether reordering augmentation is applied
+            experiment_name: MLflow experiment name (used with auto_start)
+            run_name: MLflow run name (auto-generated if None)
+            auto_start: If True, auto-starts/stops MLflow run
+            patch_config: Patch configuration dict for logging (from MedPatchDataLoaders)
+            extra_params: Additional parameters to log
+            extra_tags: MLflow tags to set on the run
         """
         self.model_name = model_name
         self.loss_function = loss_function
@@ -170,6 +319,15 @@ class ModelTrackingCallback(Callback):
         self.target_spacing = target_spacing
         self.apply_reorder = apply_reorder
         
+        # New auto-management fields
+        self.experiment_name = experiment_name
+        self.run_name = run_name
+        self.auto_start = auto_start
+        self.patch_config = patch_config
+        self.extra_params = extra_params or {}
+        self.extra_tags = extra_tags or {}
+        self._auto_started = False
+        
         self.config = self._build_config()
         
     def extract_all_params(self, tfm):
@@ -235,6 +393,12 @@ class ModelTrackingCallback(Callback):
             separators=(',', ': ')
         )
         
+        # Add patch-specific params if present
+        if self.patch_config:
+            for key, value in self.patch_config.items():
+                if value is not None:
+                    params[f"patch_{key}"] = value if isinstance(value, (int, float, str, bool)) else str(value)
+        
         return params
     
     def _extract_epoch_metrics(self) -> dict[str, float]:
@@ -284,6 +448,23 @@ class ModelTrackingCallback(Callback):
         if os.path.exists(weights_file):
             mlflow.log_artifact(weights_file, "model")
         
+        # Auto-detect SaveModelCallback and log best model
+        from fastai.callback.tracker import SaveModelCallback
+        best_model_cb = None
+        for cb in self.learn.cbs:
+            if isinstance(cb, SaveModelCallback):
+                best_model_path = self.learn.path/self.learn.model_dir/f'{cb.fname}.pth'
+                if best_model_path.exists():
+                    mlflow.log_artifact(str(best_model_path), "model")
+                    print(f"Logged best model artifact: {cb.fname}.pth")
+                    best_model_cb = cb
+                break
+
+        # Load best model weights before export if SaveModelCallback was used
+        if best_model_cb is not None:
+            self.learn.load(best_model_cb.fname)
+            print(f"Loaded best model weights ({best_model_cb.fname}) for learner export")
+
         # Remove MLflow callbacks before exporting learner for inference
         # This prevents the callback from being triggered during inference
         original_cbs = self.learn.cbs.copy()  # Save original callbacks
@@ -312,8 +493,22 @@ class ModelTrackingCallback(Callback):
         )
     
     def before_fit(self) -> None:
-        """Log hyperparameters before training starts."""
+        """Log hyperparameters before training starts. Auto-start run if configured."""
+        # Auto-start run if requested
+        if self.auto_start:
+            if self.experiment_name:
+                mlflow.set_experiment(self.experiment_name)
+            self.run_name = self.run_name or f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            mlflow.start_run(run_name=self.run_name)
+            self._auto_started = True
+        
+        # Set tags
+        if self.extra_tags:
+            mlflow.set_tags(self.extra_tags)
+        
+        # Log params
         params = self._extract_training_params()
+        params.update(self.extra_params)
         mlflow.log_params(params)
     
     def after_epoch(self) -> None:
@@ -332,10 +527,104 @@ class ModelTrackingCallback(Callback):
             self._save_model_artifacts(temp_path)
             
             self._register_pytorch_model()
-            
-        print(f"MLflow run completed. Run ID: {mlflow.active_run().info.run_id}")
+        
+        run_id = mlflow.active_run().info.run_id
+        print(f"MLflow run completed. Run ID: {run_id}")
+        
+        # End run if auto-started
+        if self._auto_started:
+            mlflow.end_run()
+            self._auto_started = False
 
-# %% ../nbs/07_utils.ipynb 10
+# %% ../nbs/07_utils.ipynb 11
+def create_mlflow_callback(
+    learn,
+    experiment_name: str = None,
+    run_name: str = None,
+    auto_start: bool = True,
+    model_name: str = None,
+    extra_params: dict = None,
+    extra_tags: dict = None,
+) -> ModelTrackingCallback:
+    """Create MLflow tracking callback with auto-extracted configuration.
+    
+    This factory function automatically extracts configuration from the Learner,
+    eliminating the need to manually specify parameters like size, transforms,
+    loss function, etc.
+    
+    Auto-extracts from Learner:
+    - Preprocessing: apply_reorder, target_spacing, size/patch_size
+    - Transforms: item_tfms or pre_patch_tfms
+    - Training: loss_func, model architecture
+    
+    Args:
+        learn: fastai Learner instance
+        experiment_name: MLflow experiment name. If None, uses model name.
+        run_name: MLflow run name. If None, auto-generates with timestamp.
+        auto_start: If True, auto-starts/stops MLflow run in before_fit/after_fit.
+        model_name: Override auto-extracted model name for registration.
+        extra_params: Additional parameters to log (e.g., {'dropout': 0.5}).
+        extra_tags: MLflow tags to set on the run.
+    
+    Returns:
+        ModelTrackingCallback ready to use with learn.fit()
+    
+    Example:
+        >>> # Instead of this (6 manual params):
+        >>> # mlflow_callback = ModelTrackingCallback(
+        >>> #     model_name=f"{task}_{model._get_name()}",
+        >>> #     loss_function=loss_func.loss_func._get_name(),
+        >>> #     item_tfms=item_tfms,
+        >>> #     size=size,
+        >>> #     target_spacing=target_spacing,
+        >>> #     apply_reorder=True,
+        >>> # )
+        >>> # with mlflow.start_run(run_name="training"):
+        >>> #     learn.fit_one_cycle(30, lr, cbs=[mlflow_callback])
+        >>> 
+        >>> # Do this (zero manual params):
+        >>> callback = create_mlflow_callback(learn, experiment_name="Task02_Heart")
+        >>> learn.fit_one_cycle(30, lr, cbs=[callback, save_best])
+    """
+    # Detect workflow and extract config
+    if _detect_patch_workflow(learn.dls):
+        config = _extract_patch_config(learn)
+    else:
+        config = _extract_standard_config(learn)
+    
+    # Extract model/loss info
+    _model_name = model_name or _extract_model_name(learn)
+    _loss_name = _extract_loss_name(learn)
+    
+    # Set experiment name
+    if experiment_name is None:
+        experiment_name = _model_name
+    
+    # Validate size was extracted
+    if config['size'] is None:
+        raise ValueError(
+            "Could not auto-extract 'size'. Either:\n"
+            "1. Add PadOrCrop to item_tfms, or\n"
+            "2. Use MedPatchDataLoaders with PatchConfig, or\n"
+            "3. Use ModelTrackingCallback directly with manual params"
+        )
+    
+    return ModelTrackingCallback(
+        model_name=_model_name,
+        loss_function=_loss_name,
+        item_tfms=config['item_tfms'],
+        size=config['size'],
+        target_spacing=config['target_spacing'],
+        apply_reorder=config['apply_reorder'],
+        experiment_name=experiment_name,
+        run_name=run_name,
+        auto_start=auto_start,
+        patch_config=config['patch_config'],
+        extra_params=extra_params,
+        extra_tags=extra_tags,
+    )
+
+# %% ../nbs/07_utils.ipynb 13
 import subprocess
 import threading
 import time

fastMONAI/vision_augmentation.py

@@ -3,8 +3,8 @@
 # %% auto 0
 __all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'ZNormalization', 'RescaleIntensity', 'NormalizeIntensity',
            'BraTSMaskConverter', 'BinaryConverter', 'RandomGhosting', 'RandomSpike', 'RandomNoise', 'RandomBiasField',
-           'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion', 'RandomElasticDeformation',
-           'RandomAffine', 'RandomFlip', 'OneOf']
+           'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion', 'RandomCutout',
+           'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf']
 
 # %% ../nbs/03_vision_augment.ipynb 2
 from fastai.data.all import *
@@ -425,7 +425,139 @@ class RandomMotion(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
+# %% ../nbs/03_vision_augment.ipynb 22
+def _create_ellipsoid_mask(shape, center, radii):
+    """Create a 3D ellipsoid mask.
+
+    Args:
+        shape: (D, H, W) shape of the volume
+        center: (z, y, x) center of ellipsoid
+        radii: (rz, ry, rx) radii along each axis
+
+    Returns:
+        Boolean mask where True = inside ellipsoid
+    """
+    z, y, x = torch.meshgrid(
+        torch.arange(shape[0]),
+        torch.arange(shape[1]),
+        torch.arange(shape[2]),
+        indexing='ij'
+    )
+    dist = ((z - center[0]) / radii[0]) ** 2 + \
+           ((y - center[1]) / radii[1]) ** 2 + \
+           ((x - center[2]) / radii[2]) ** 2
+    return dist <= 1.0
+
 # %% ../nbs/03_vision_augment.ipynb 23
+class _TioRandomCutout(tio.IntensityTransform):
+    """TorchIO-compatible RandomCutout for patch-based workflows."""
+
+    def __init__(self, holes=1, spatial_size=8, fill_value=None,
+                 max_holes=None, max_spatial_size=None, p=0.2, **kwargs):
+        super().__init__(p=p, **kwargs)
+        self.holes = holes
+        self.spatial_size = spatial_size
+        self.fill_value = fill_value
+        self.max_holes = max_holes
+        self.max_spatial_size = max_spatial_size
+
+    def _apply_cutout(self, data, fill_val):
+        """Apply spherical cutout(s) to a tensor."""
+        result = data.clone()
+        n_holes = torch.randint(self.holes, (self.max_holes or self.holes) + 1, (1,)).item()
+
+        spatial_shape = data.shape[1:]  # (D, H, W)
+        min_size = self.spatial_size if isinstance(self.spatial_size, int) else self.spatial_size[0]
+        max_size = self.max_spatial_size or self.spatial_size
+        max_size = max_size if isinstance(max_size, int) else max_size[0]
+
+        for _ in range(n_holes):
+            # Random size for this hole
+            size = torch.randint(min_size, max_size + 1, (3,))
+            radii = size.float() / 2
+
+            # Random center (ensure hole fits in volume)
+            center = [
+                torch.randint(int(radii[i].item()),
+                              max(spatial_shape[i] - int(radii[i].item()), int(radii[i].item()) + 1),
+                              (1,)).item()
+                for i in range(3)
+            ]
+
+            mask = _create_ellipsoid_mask(spatial_shape, center, radii)
+            result[:, mask] = fill_val
+
+        return result
+
+    def apply_transform(self, subject):
+        for image in self.get_images(subject):
+            data = image.data
+            fill_val = self.fill_value if self.fill_value is not None else float(data.min())
+            result = self._apply_cutout(data, fill_val)
+            image.set_data(result)
+        return subject
+
+# %% ../nbs/03_vision_augment.ipynb 24
+class RandomCutout(DisplayedTransform):
+    """Randomly erase spherical regions in 3D medical images.
+
+    Simulates post-operative surgical cavities by filling random ellipsoid
+    volumes with specified values. Useful for training on pre-op images
+    to generalize to post-op scans.
+
+    Args:
+        holes: Minimum number of cutout regions. Default: 1.
+        max_holes: Maximum number of regions. Default: 3.
+        spatial_size: Minimum cutout diameter in voxels. Default: 8.
+        max_spatial_size: Maximum cutout diameter. Default: 16.
+        fill: Fill value - 'min', 'mean', 'random', or float. Default: 'min'.
+        p: Probability of applying transform. Default: 0.2.
+
+    Example:
+        >>> # Simulate post-op cavities with dark spherical voids
+        >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10,
+        ...                    max_spatial_size=25, fill='min', p=0.2)
+    """
+
+    split_idx, order = 0, 1
+
+    def __init__(self, holes=1, max_holes=3, spatial_size=8,
+                 max_spatial_size=16, fill='min', p=0.2):
+        self.holes = holes
+        self.max_holes = max_holes
+        self.spatial_size = spatial_size
+        self.max_spatial_size = max_spatial_size
+        self.fill = fill
+        self.p = p
+
+        self._tio_cutout = _TioRandomCutout(
+            holes=holes, spatial_size=spatial_size,
+            fill_value=None if isinstance(fill, str) else fill,
+            max_holes=max_holes, max_spatial_size=max_spatial_size, p=p
+        )
+
+    @property
+    def tio_transform(self):
+        """Return TorchIO-compatible transform for patch-based workflows."""
+        return self._tio_cutout
+
+    def _get_fill_value(self, tensor):
+        if self.fill == 'min': return float(tensor.min())
+        elif self.fill == 'mean': return float(tensor.mean())
+        elif self.fill == 'random':
+            return torch.empty(1).uniform_(float(tensor.min()), float(tensor.max())).item()
+        else: return self.fill
+
+    def encodes(self, o: MedImage):
+        if torch.rand(1).item() > self.p: return o
+        fill_val = self._get_fill_value(o)
+        result = self._tio_cutout._apply_cutout(o.clone(), fill_val)
+        return MedImage.create(result)
+
+    def encodes(self, o: MedMask):
+        return o
+
+# %% ../nbs/03_vision_augment.ipynb 26
 class RandomElasticDeformation(CustomDictTransform):
     """Apply TorchIO `RandomElasticDeformation`."""
 
@@ -438,7 +570,7 @@ class RandomElasticDeformation(CustomDictTransform):
             image_interpolation=image_interpolation,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 24
+# %% ../nbs/03_vision_augment.ipynb 27
 class RandomAffine(CustomDictTransform):
     """Apply TorchIO `RandomAffine`."""
 
@@ -454,14 +586,14 @@ class RandomAffine(CustomDictTransform):
             default_pad_value=default_pad_value,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 25
+# %% ../nbs/03_vision_augment.ipynb 28
 class RandomFlip(CustomDictTransform):
     """Apply TorchIO `RandomFlip`."""
 
     def __init__(self, axes='LR', p=0.5):
         super().__init__(tio.RandomFlip(axes=axes, flip_probability=p))
 
-# %% ../nbs/03_vision_augment.ipynb 26
+# %% ../nbs/03_vision_augment.ipynb 29
 class OneOf(CustomDictTransform):
     """Apply only one of the given transforms using TorchIO `OneOf`."""