utils/mammogram_preprocessing.py

"""
utils/mammogram_preprocessing.py
──────────────────────────────────
Preprocessing pipeline for full-field digital mammography (FFDM).

Handles both DICOM files and standard image formats (PNG, JPG).
Applies mammogram-specific preprocessing:
  - VOI LUT windowing for DICOM files
  - Breast region normalisation
  - Grayscale to 3-channel RGB conversion
  - Mammogram-appropriate augmentations (no stain jitter)

Install
───────
  pip install pydicom pylibjpeg python-gdcm
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional, Union

import numpy as np
from PIL import Image
from torchvision import transforms

# DICOM support — optional import so the rest of the codebase doesn't break
# if pydicom is not installed
try:
    import pydicom
    from pydicom.pixel_data_handlers.util import apply_voi_lut
    PYDICOM_AVAILABLE = True
except ImportError:
    PYDICOM_AVAILABLE = False

# ── Constants ──────────────────────────────────────────────────────────────────
MAMMOGRAM_SIZE    = 512            # EfficientNet-B4 input resolution
IMAGENET_MEAN     = [0.485, 0.456, 0.406]
IMAGENET_STD      = [0.229, 0.224, 0.225]


# ── DICOM loader ───────────────────────────────────────────────────────────────

def load_dicom(path: Union[str, Path]) -> Image.Image:
    """
    Load a DICOM mammogram file and convert to an 8-bit RGB PIL Image.

    Applies VOI LUT windowing if available in the DICOM metadata,
    otherwise falls back to min-max normalisation.

    MONOCHROME1 images (where high pixel = dark) are inverted so
    the tissue appears bright on a dark background, matching the
    visual convention used during model training.

    Parameters
    ----------
    path : str | Path
        Path to a .dcm DICOM file.

    Returns
    -------
    PIL.Image.Image  — RGB image ready for preprocessing.

    Raises
    ------
    ImportError  — if pydicom is not installed.
    FileNotFoundError — if the file does not exist.
    """
    if not PYDICOM_AVAILABLE:
        raise ImportError(
            "pydicom not installed. Run: pip install pydicom pylibjpeg python-gdcm"
        )

    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"DICOM file not found: {path}")

    dcm = pydicom.dcmread(str(path))

    # Apply VOI LUT windowing (converts to display-ready values)
    try:
        pixel_array = apply_voi_lut(dcm.pixel_array, dcm)
    except Exception:
        pixel_array = dcm.pixel_array.astype(np.float32)

    pixel_array = pixel_array.astype(np.float32)

    # Normalise to [0, 255]
    p_min, p_max = pixel_array.min(), pixel_array.max()
    if p_max > p_min:
        pixel_array = (pixel_array - p_min) / (p_max - p_min) * 255.0
    else:
        pixel_array = np.zeros_like(pixel_array)

    pixel_array = pixel_array.astype(np.uint8)

    # MONOCHROME1: invert so tissue is bright
    if hasattr(dcm, "PhotometricInterpretation"):
        if dcm.PhotometricInterpretation == "MONOCHROME1":
            pixel_array = 255 - pixel_array

    # Convert grayscale → RGB (EfficientNet expects 3 channels)
    if pixel_array.ndim == 2:
        rgb = np.stack([pixel_array, pixel_array, pixel_array], axis=-1)
    else:
        rgb = pixel_array

    return Image.fromarray(rgb, mode="RGB")


def load_mammogram(path: Union[str, Path]) -> Image.Image:
    """
    Load a mammogram from either DICOM or standard image format.

    Automatically detects format by file extension.

    Parameters
    ----------
    path : str | Path
        Path to DICOM (.dcm) or image file (.png, .jpg, .tiff).

    Returns
    -------
    PIL.Image.Image — RGB image.
    """
    path = Path(path)
    if path.suffix.lower() in {".dcm", ".dicom"}:
        return load_dicom(path)

    # Standard image format
    img = Image.open(path).convert("RGB")

    # If grayscale was saved as single-channel PNG, already converted above.
    # But if it's a true grayscale mammogram saved as PNG:
    if img.mode == "L":
        arr = np.array(img)
        rgb = np.stack([arr, arr, arr], axis=-1)
        return Image.fromarray(rgb, mode="RGB")

    return img


# ── Mammogram-specific augmentation transforms ──────────────────────────────────

class BreastRegionEnhancer:
    """
    Enhances breast tissue contrast using adaptive histogram equalisation.

    Applied per-channel to improve visibility of masses and calcifications
    without affecting the overall image structure.

    Parameters
    ----------
    clip_limit : float
        CLAHE clip limit. Higher = more aggressive enhancement.
    p : float
        Probability of applying this transform.
    """

    def __init__(self, clip_limit: float = 2.0, p: float = 0.5) -> None:
        self.clip_limit = clip_limit
        self.p          = p

    def __call__(self, img: Image.Image) -> Image.Image:
        if np.random.random() > self.p:
            return img

        try:
            import cv2
            arr = np.array(img)
            clahe = cv2.createCLAHE(
                clipLimit    = self.clip_limit,
                tileGridSize = (8, 8),
            )
            # Apply CLAHE to each channel
            enhanced = np.stack(
                [clahe.apply(arr[:, :, c]) for c in range(arr.shape[2])],
                axis=-1,
            )
            return Image.fromarray(enhanced.astype(np.uint8), mode="RGB")
        except ImportError:
            return img   # OpenCV not available — skip silently


class RandomElasticDeformation:
    """
    Random elastic deformation — simulates tissue compression variation
    from different mammography unit pressures.

    Parameters
    ----------
    alpha : float
        Strength of deformation.
    sigma : float
        Smoothness of deformation field.
    p : float
        Probability of applying.
    """

    def __init__(
        self,
        alpha: float = 34.0,
        sigma: float = 4.0,
        p:     float = 0.3,
    ) -> None:
        self.alpha = alpha
        self.sigma = sigma
        self.p     = p

    def __call__(self, img: Image.Image) -> Image.Image:
        if np.random.random() > self.p:
            return img

        try:
            from scipy.ndimage import gaussian_filter, map_coordinates

            arr    = np.array(img, dtype=np.float32)
            h, w   = arr.shape[:2]
            dx     = gaussian_filter(
                (np.random.rand(h, w) * 2 - 1) * self.alpha, self.sigma
            )
            dy     = gaussian_filter(
                (np.random.rand(h, w) * 2 - 1) * self.alpha, self.sigma
            )
            x, y   = np.meshgrid(np.arange(w), np.arange(h))
            coords = [
                np.clip(y + dy, 0, h - 1).ravel(),
                np.clip(x + dx, 0, w - 1).ravel(),
            ]

            result = np.stack([
                map_coordinates(arr[:, :, c], coords, order=1).reshape(h, w)
                for c in range(arr.shape[2])
            ], axis=-1)

            return Image.fromarray(result.clip(0, 255).astype(np.uint8), mode="RGB")
        except ImportError:
            return img   # scipy not available — skip


def build_mammogram_train_transform() -> transforms.Compose:
    """
    Training augmentation pipeline for mammograms.

    Mammogram-appropriate augmentations — NO stain jitter (mammograms
    are X-ray images, not H&E-stained tissue). Augmentations simulate:
    - Different patient positioning (flips, rotation)
    - Tissue compression variation (elastic deformation)
    - Scanner variation (brightness/contrast)
    - Tissue contrast differences (CLAHE)
    """
    return transforms.Compose([
        BreastRegionEnhancer(clip_limit=2.0, p=0.4),
        RandomElasticDeformation(alpha=34.0, sigma=4.0, p=0.3),
        transforms.Resize((MAMMOGRAM_SIZE, MAMMOGRAM_SIZE)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.2),
        transforms.RandomRotation(degrees=10),
        transforms.ColorJitter(
            brightness = 0.15,
            contrast   = 0.15,
        ),
        transforms.RandomAffine(
            degrees    = 0,
            translate  = (0.05, 0.05),
            scale      = (0.95, 1.05),
        ),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])


def build_mammogram_inference_transform() -> transforms.Compose:
    """
    Inference transform — resize and normalise only. No augmentation."""
    return transforms.Compose([
        transforms.Resize((MAMMOGRAM_SIZE, MAMMOGRAM_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])