utils/preprocessing.py

"""
utils/preprocessing.py
───────────────────────
Image preprocessing pipeline that converts a raw histopathology image
(file path, PIL Image, or numpy array) into a normalised tensor of
shape (1, 3, 224, 224) ready for model inference.

Normalization follows ImageNet standards as required by the spec:
    Mean : [0.485, 0.456, 0.406]
    Std  : [0.229, 0.224, 0.225]
"""

from __future__ import annotations

from pathlib import Path
from typing import Union

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

# ── ImageNet normalization constants ────────────────────────────────────────
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

# ── Target tensor shape ─────────────────────────────────────────────────────
TARGET_SIZE = (224, 224)


# ── Transform pipeline ───────────────────────────────────────────────────────
def build_inference_transform() -> transforms.Compose:
    """
    Returns the deterministic inference transform pipeline.

    Steps
    ─────
    1. Resize shortest edge to 256 px (preserves aspect ratio).
    2. Centre-crop to 224 × 224.
    3. Convert PIL image to float32 tensor in [0, 1].
    4. Normalize with ImageNet mean / std.
    """
    return transforms.Compose([
        transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.CenterCrop(TARGET_SIZE),
        transforms.ToTensor(),                          # → [0, 1]
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])




# ── StainJitter — Fix 1 ──────────────────────────────────────────────────────
class StainJitter:
    """
    Randomly perturb H&E stain concentrations in HED colour space.

    Why this works
    ──────────────
    H&E-stained slides vary in colour between labs due to differences in
    staining batches, fixation protocols, and scanner calibrations.
    Standard RGB colour jitter doesn't model this — it shifts all three
    channels independently. StainJitter works in HED space (Haematoxylin,
    Eosin, DAB), which directly corresponds to the actual stains in the tissue.
    Perturbing HED channels simulates real-world staining variation without
    needing a reference image or external library.

    Implementation
    ──────────────
    Uses the Ruifrok & Johnston HED deconvolution matrix to decompose
    RGB into stain concentrations, perturbs each channel with a random
    scale (alpha) and shift (beta), then reconstructs the RGB image.
    Pure NumPy — no external dependencies beyond what is already installed.

    Parameters
    ----------
    strength : float
        Controls the magnitude of perturbation.
        0.05 = ±5% scale variation + ±5% shift variation per channel.
        Typical values: 0.03 (subtle) to 0.10 (aggressive).
    p : float
        Probability of applying the transform. Default 0.5.
    """

    # Ruifrok & Johnston HED deconvolution matrix
    # Rows = [Haematoxylin, Eosin, DAB] stain absorption vectors
    HED = np.array([
        [0.6500286,  0.7044536,  0.2860126],
        [0.7044522,  0.4956977,  0.5079795],
        [0.2860126,  0.5079795,  0.8128560],
    ], dtype=np.float64)

    # Pre-compute inverse once at class level
    HED_INV = np.linalg.inv(HED)

    def __init__(self, strength: float = 0.05, p: float = 0.5) -> None:
        self.strength = strength
        self.p        = p

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

        # PIL → float64 numpy in [0, 1]
        rgb = np.array(img, dtype=np.float64) / 255.0

        # Convert to optical density — Beer-Lambert law
        # Clamp to avoid log(0)
        od = -np.log(np.clip(rgb, 1e-6, 1.0))

        # Decompose into HED stain concentrations
        # od = concentrations @ HED  →  concentrations = od @ HED_INV
        hed = od @ self.HED_INV                      # (H, W, 3) HED concentrations

        # Perturb each stain channel independently
        alpha = np.random.uniform(
            1.0 - self.strength,
            1.0 + self.strength,
            size=(1, 1, 3),
        )
        beta = np.random.uniform(
            -self.strength,
            +self.strength,
            size=(1, 1, 3),
        )
        hed_perturbed = hed * alpha + beta

        # Reconstruct optical density then RGB
        od_reconstructed = hed_perturbed @ self.HED
        rgb_out = np.exp(-od_reconstructed)
        rgb_out = np.clip(rgb_out, 0.0, 1.0)

        return Image.fromarray((rgb_out * 255).astype(np.uint8), mode="RGB")

def build_training_transform() -> transforms.Compose:
    """
    Augmentation pipeline for fine-tuning.
    Included for completeness; inference always uses build_inference_transform().
    """
    return transforms.Compose([
        StainJitter(strength=0.05, p=0.5),           # Fix 1: H&E stain augmentation
        transforms.RandomResizedCrop(TARGET_SIZE, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2,
                               saturation=0.1),
        transforms.RandomRotation(degrees=15),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])


# ── Preprocessing entry point ────────────────────────────────────────────────
class ImagePreprocessor:
    """
    Accepts multiple input types and returns a (1, 3, 224, 224) tensor.

    Supported inputs
    ────────────────
    - str / pathlib.Path  : local file path to PNG / JPG / TIFF
    - PIL.Image.Image     : already-loaded PIL image
    - np.ndarray          : HxWx3 uint8 or float32 array
    - torch.Tensor        : CxHxW or 1xCxHxW (skips PIL stage)
    """

    def __init__(self) -> None:
        self._transform = build_inference_transform()

    def __call__(
        self,
        image: Union[str, Path, "Image.Image", np.ndarray, torch.Tensor],
    ) -> torch.Tensor:
        """
        Parameters
        ----------
        image : see supported inputs above

        Returns
        -------
        torch.Tensor
            Shape (1, 3, 224, 224), dtype float32, ImageNet-normalised.
        """
        pil_image = self._to_pil(image)
        tensor    = self._transform(pil_image)   # (3, 224, 224)
        return tensor.unsqueeze(0)               # (1, 3, 224, 224)

    # ── Type dispatch helpers ────────────────────────────────────────────────
    @staticmethod
    def _to_pil(image) -> "Image.Image":
        if isinstance(image, (str, Path)):
            return Image.open(image).convert("RGB")

        if isinstance(image, Image.Image):
            return image.convert("RGB")

        if isinstance(image, np.ndarray):
            if image.dtype != np.uint8:
                image = (np.clip(image, 0, 1) * 255).astype(np.uint8)
            if image.ndim == 2:
                image = np.stack([image] * 3, axis=-1)  # grayscale → RGB
            return Image.fromarray(image, mode="RGB")

        if isinstance(image, torch.Tensor):
            t   = image.squeeze(0) if image.ndim == 4 else image
            arr = (t.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            return Image.fromarray(arr, mode="RGB")

        raise TypeError(
            f"Unsupported image type: {type(image)}. "
            "Expected str, Path, PIL.Image, np.ndarray, or torch.Tensor."
        )