""" 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." )