"""Lesion feature extraction: runs Phase 4's models as inference engines over images that have no lesion ground truth of their own (APTOS), to build the per-image feature vector Phase 5's fusion head concats with the grading model's CNN embedding. **Cross-dataset domain shift, faced honestly, not hidden**: every model used here was trained and validated exclusively on IDRiD (2848x4288, one capture device/protocol). APTOS images range from 1050x1050 to 3216x2136 across several different capture devices, with no resolution normalisation applied before inference -- a cheap sanity check (the hard-exudate model against 4 images per DR grade) showed mean predicted probability rising with grade even so (grade 0: 0.0000-0.0025 positive-pixel fraction at threshold 0.5; grade 3: 0.023-0.114), which is why this was judged worth building rather than assumed broken -- but it was never validated against real APTOS lesion ground truth (APTOS ships none), so treat these as a genuine, undemonstrated hypothesis under the fusion head's own McNemar test, not a calibrated measurement the way IDRiD-evaluated numbers elsewhere in this project are. IDRiD's own grading images ("B. Disease Grading") are byte-identical to its localisation images ("C. Localization") -- real OD/fovea/lesion ground truth would be available there. Deliberately not used: docs/04_ROADMAP.md Phase 8 reserves IDRiD (with Messidor-2) as the locked, evaluate-once final test set, and using it to build or tune the fusion head now would contaminate it. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path import cv2 import numpy as np import torch from drdetect.segmentation.localization import ( LocalizationModule, heatmap_argmax, working_res_od_diameter, ) from drdetect.segmentation.microaneurysms import ClassifierModule, PatchClassifier, score_candidates from drdetect.segmentation.model import build_segmentation_model from drdetect.segmentation.module import SegmentationModule from drdetect.segmentation.quadrants import assign_quadrant __all__ = [ "LesionModels", "load_lesion_models", "extract_lesion_features", "lesion_centroids", "FEATURE_NAMES", "QUADRANT_ORDER", ] LOCALIZATION_SIZE = (512, 768) # (height, width), matches training SEGMENTATION_PATCH = 512 MA_TUNED_THRESHOLD = 0.95 # docs/15's val-tuned operating point, frozen there _MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) _STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) QUADRANT_ORDER = [ "superior-foveal-side", "inferior-foveal-side", "superior-disc-side", "inferior-disc-side", ] FEATURE_NAMES = [ "hard_exudate_mean_prob", "soft_exudate_mean_prob", "haemorrhage_mean_prob", "microaneurysm_count", "distance_to_fovea_mean", *[f"quadrant_frac_{q}" for q in QUADRANT_ORDER], ] @dataclass class LesionModels: localization: LocalizationModule hard_exudates: SegmentationModule soft_exudates: SegmentationModule haemorrhages: SegmentationModule microaneurysm_classifier: ClassifierModule device: torch.device def _load_segmentation(checkpoint: str, device: torch.device) -> SegmentationModule: model = build_segmentation_model("resnet34", pretrained=False, classes=1) module = SegmentationModule.load_from_checkpoint(checkpoint, model=model, map_location=device) return module.eval().to(device) def load_lesion_models( *, localization_checkpoint: str, hard_exudate_checkpoint: str, soft_exudate_checkpoint: str, haemorrhage_checkpoint: str, microaneurysm_checkpoint: str, device: torch.device, ) -> LesionModels: loc_model = build_segmentation_model("resnet34", pretrained=False, classes=2) localization = ( LocalizationModule.load_from_checkpoint( localization_checkpoint, model=loc_model, map_location=device ) .eval() .to(device) ) ma_model = PatchClassifier() microaneurysm_classifier = ( ClassifierModule.load_from_checkpoint( microaneurysm_checkpoint, model=ma_model, map_location=device ) .eval() .to(device) ) return LesionModels( localization=localization, hard_exudates=_load_segmentation(hard_exudate_checkpoint, device), soft_exudates=_load_segmentation(soft_exudate_checkpoint, device), haemorrhages=_load_segmentation(haemorrhage_checkpoint, device), microaneurysm_classifier=microaneurysm_classifier, device=device, ) def _predict_od_fovea( image_rgb: np.ndarray, models: LesionModels ) -> tuple[tuple[float, float], tuple[float, float], float]: """Returns (od_xy, fovea_xy, od_diameter), all in LOCALIZATION_SIZE's resized pixel space (not the input image's own resolution).""" resized = cv2.resize(image_rgb, (LOCALIZATION_SIZE[1], LOCALIZATION_SIZE[0])) rgb = (resized.astype(np.float32) / 255.0 - _MEAN) / _STD tensor = torch.from_numpy(rgb.transpose(2, 0, 1)).float().unsqueeze(0).to(models.device) with torch.no_grad(): logits = models.localization(tensor) probs = torch.sigmoid(logits).squeeze(0).cpu().numpy() od_xy = heatmap_argmax(probs[0]) fovea_xy = heatmap_argmax(probs[1]) diameter = working_res_od_diameter(LOCALIZATION_SIZE) return od_xy, fovea_xy, diameter def _tile_predict( module: SegmentationModule, image_rgb: np.ndarray, device: torch.device ) -> np.ndarray: """Same non-overlapping reflect-padded tiling as scripts/evaluate_segmentation.py's tile_predict, inlined here rather than imported from a script module so library code doesn't depend on scripts/.""" patch_size = SEGMENTATION_PATCH h, w = image_rgb.shape[:2] ph = (patch_size - h % patch_size) % patch_size pw = (patch_size - w % patch_size) % patch_size padded = cv2.copyMakeBorder(image_rgb, 0, ph, 0, pw, cv2.BORDER_REFLECT_101) ph_full, pw_full = padded.shape[:2] prob_map = np.zeros((ph_full, pw_full), dtype=np.float32) tiles, positions = [], [] for top in range(0, ph_full, patch_size): for left in range(0, pw_full, patch_size): tiles.append(padded[top : top + patch_size, left : left + patch_size]) positions.append((top, left)) batch_size = 4 with torch.no_grad(): for start in range(0, len(tiles), batch_size): chunk = tiles[start : start + batch_size] arrs = [ ((t.astype(np.float32) / 255.0 - _MEAN) / _STD).transpose(2, 0, 1) for t in chunk ] batch = torch.from_numpy(np.stack(arrs)).float().to(device) logits = module(batch) probs = torch.sigmoid(logits).squeeze(1).cpu().numpy() for (top, left), prob in zip(positions[start : start + batch_size], probs, strict=True): prob_map[top : top + patch_size, left : left + patch_size] = prob return prob_map[:h, :w] def lesion_centroids(prob_map: np.ndarray, threshold: float = 0.5) -> list[tuple[float, float]]: """Connected-component centroids of a thresholded probability map -- a handful of per-lesion centroids is what quadrant/distance features need, not per-pixel detail.""" binary = (prob_map > threshold).astype(np.uint8) if binary.sum() == 0: return [] n, _labels, stats, centroids = cv2.connectedComponentsWithStats(binary, connectivity=8) return [tuple(centroids[i]) for i in range(1, n) if stats[i, cv2.CC_STAT_AREA] >= 3] def extract_lesion_features(image_path: str | Path, models: LesionModels) -> dict[str, float]: """One image -> the flat feature dict Phase 5's fusion head concats with the grading CNN's embedding. Every model runs at its own native working convention (full-resolution tiling for the three DeepLabV3+ segmentation models, a fixed 512x768 resize for localisation, full-resolution candidate generation for microaneurysms) -- see this module's own docstring for the cross-dataset resolution caveat that follows from that. """ image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR) image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) od_xy, fovea_xy, od_diameter = _predict_od_fovea(image_rgb, models) he_prob = _tile_predict(models.hard_exudates, image_rgb, models.device) se_prob = _tile_predict(models.soft_exudates, image_rgb, models.device) ha_prob = _tile_predict(models.haemorrhages, image_rgb, models.device) # Segmentation ran at the image's OWN resolution; localisation predicts in # LOCALIZATION_SIZE's resized space. Rescale lesion centroids into that # same space before computing distance/quadrant features, since od_xy/ # fovea_xy live there. scale_x = LOCALIZATION_SIZE[1] / image_rgb.shape[1] scale_y = LOCALIZATION_SIZE[0] / image_rgb.shape[0] all_centroids: list[tuple[float, float]] = [] for prob_map in (he_prob, se_prob, ha_prob): for cx, cy in lesion_centroids(prob_map): all_centroids.append((cx * scale_x, cy * scale_y)) # score_candidates expects a ground-truth mask (it's normally used for # scored evaluation, see evaluate_microaneurysms.py) purely to compute # labels/recall bookkeeping this call never reads -- an all-zero dummy # mask makes that bookkeeping cheap and correct-by-construction (0 # recovered of 0 true instances) without a separate inference-only code # path for a function that otherwise does exactly what's needed here. dummy_mask = np.zeros(image_bgr.shape[:2], dtype=np.uint8) ma_result = score_candidates( image_bgr, dummy_mask, models.microaneurysm_classifier.model, device=str(models.device) ) accepted = ma_result["scores"] > MA_TUNED_THRESHOLD ma_count = int(accepted.sum()) for c, keep in zip(ma_result["candidates"], accepted, strict=True): if keep: all_centroids.append((c.cx * scale_x, c.cy * scale_y)) if all_centroids: distances = [ float(np.linalg.norm(np.array(c) - np.array(fovea_xy))) / od_diameter for c in all_centroids ] distance_to_fovea_mean = float(np.mean(distances)) quadrant_counts = dict.fromkeys(QUADRANT_ORDER, 0) for c in all_centroids: quadrant_counts[assign_quadrant(c, od_xy, fovea_xy)] += 1 total = len(all_centroids) quadrant_fracs = {q: quadrant_counts[q] / total for q in QUADRANT_ORDER} else: distance_to_fovea_mean = float("nan") quadrant_fracs = dict.fromkeys(QUADRANT_ORDER, 0.0) return { "hard_exudate_mean_prob": float(he_prob.mean()), "soft_exudate_mean_prob": float(se_prob.mean()), "haemorrhage_mean_prob": float(ha_prob.mean()), "microaneurysm_count": float(ma_count), "distance_to_fovea_mean": distance_to_fovea_mean, **{f"quadrant_frac_{q}": quadrant_fracs[q] for q in QUADRANT_ORDER}, }