diff --git a/README.md b/README.md index ee9743e52a39e91c8e19a41929ffd1d891f918ec..7a464b31c66b9e98778f7be4ba63b9370d953e65 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,23 @@ --- -title: Drdetect Dr Screening -emoji: πŸ’» -colorFrom: indigo -colorTo: pink +title: DR Screening Demo +emoji: πŸ‘οΈ +colorFrom: blue +colorTo: green sdk: gradio -sdk_version: 6.26.0 -python_version: '3.12' +sdk_version: "6.26.0" app_file: app.py pinned: false +license: other --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +# Diabetic Retinopathy Screening β€” Live Demo + +**Research prototype only. NOT a medical device. NOT a diagnosis.** Upload a fundus photograph to +see this project's quality gate, ICDR grading, and Grad-CAM evidence run end to end. + +Full project, every experiment, and the complete evidence trail (including this model's honest +external-validation limitations β€” read before trusting any output): +[github.com/adarshcod30/Diabetic-Retinopathy-Detection](https://github.com/adarshcod30/Diabetic-Retinopathy-Detection). + +Model weights: [huggingface.co/adarshcod30/drdetect-dr-screening](https://huggingface.co/adarshcod30/drdetect-dr-screening) +(research use only β€” see that repo's own model card for the full license terms). diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..3108ed30d3b6742384163bd9dbb5689b9a5f48ed --- /dev/null +++ b/app.py @@ -0,0 +1,32 @@ +"""HuggingFace Spaces entry point for the Gradio demo (docs/04_ROADMAP.md, Phase 9). + +Not used for local development -- `scripts/demo.py` is the local launcher, +with a required `--checkpoint` flag pointing anywhere on disk. A Space has no +such flag: HF Spaces runs `python app.py` with no arguments, so this file +fetches the checkpoint from the separate model repo +(huggingface.co/adarshcod30/drdetect-dr-screening -- research-use-only, see +MODEL_CARD.md there) via `hf_hub_download` instead of bundling a 46MB weight +file into this Space's own repo. `hf_hub_download` caches the file locally +after the first fetch, so a Space restart does not re-download it. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) + +MODEL_REPO = "adarshcod30/drdetect-dr-screening" + +if __name__ == "__main__": + from huggingface_hub import hf_hub_download + + checkpoint = hf_hub_download(repo_id=MODEL_REPO, filename="best.ckpt") + + from drdetect.serve.demo import build_interface + + # regression loss, not ce: it won the locked external evaluation decisively + # (docs/22_PHASE8_VALIDATION_RESULTS.md, referable AUC 0.924 vs 0.888, DeLong p=6.1e-10) + demo = build_interface(checkpoint, backbone="efficientnet_b0", loss_name="regression", size=512) + demo.launch() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..867fab5f70d3459807af50813b749a5bd452b3e2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,33 @@ +# Generated from pyproject.toml β€” prefer `pip install -e ".[dev]"` +torch>=2.2 +torchvision>=0.17 +lightning>=2.2 +timm>=1.0 +monai>=1.3 +segmentation-models-pytorch>=0.3 +opencv-python-headless>=4.9 +scikit-image>=0.22 +albumentations>=1.4 +Pillow>=10.2 +numpy>=1.26 +pandas>=2.2 +scipy>=1.12 +scikit-learn>=1.4 +statsmodels>=0.14 +torchmetrics>=1.3 +hydra-core>=1.3 +omegaconf>=2.3 +wandb>=0.16 +grad-cam>=1.5 +simpy>=4.1 +reportlab>=4.1 +fastapi>=0.110 +uvicorn>=0.29 +onnx>=1.16 +onnxruntime>=1.17 +gradio>=4.25 +matplotlib>=3.8 +seaborn>=0.13 +tqdm>=4.66 +kaggle>=1.6 +python-dotenv>=1.0 diff --git a/src/drdetect.egg-info/PKG-INFO b/src/drdetect.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..5920f786e018804c8750e7278e8db181983bb94d --- /dev/null +++ b/src/drdetect.egg-info/PKG-INFO @@ -0,0 +1,643 @@ +Metadata-Version: 2.4 +Name: drdetect +Version: 0.1.0 +Summary: Explainable, quality-aware AI for diabetic retinopathy screening +Author: Adarsh Dwivedi +License: MIT License + + Copyright (c) 2026 Adarsh Dwivedi + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. + + --- + + NOTE ON SCOPE: This licence covers the SOURCE CODE in this repository only. + + Datasets referenced by this project remain under their own licences and are NOT + redistributed here (IDRiD: CC BY 4.0; APTOS / EyePACS: Kaggle competition rules; + Messidor-2: ADCIS terms; DRIVE: research use with registration). + + Trained model weights derived from those datasets are released for RESEARCH USE + ONLY. + + This software is a research prototype. It is NOT a medical device and must not + be used for clinical diagnosis. + +Project-URL: Homepage, https://github.com/adarshcod30/Diabetic-Retinopathy-Detection +Project-URL: Issues, https://github.com/adarshcod30/Diabetic-Retinopathy-Detection/issues +Keywords: diabetic-retinopathy,medical-imaging,explainable-ai,deep-learning,fundus,computer-vision,healthcare-ai +Classifier: Development Status :: 2 - Pre-Alpha +Classifier: Intended Audience :: Science/Research +Classifier: License :: OSI Approved :: MIT License +Classifier: Programming Language :: Python :: 3.11 +Classifier: Topic :: Scientific/Engineering :: Medical Science Apps. +Requires-Python: <3.13,>=3.10 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: torch>=2.2 +Requires-Dist: torchvision>=0.17 +Requires-Dist: lightning>=2.2 +Requires-Dist: timm>=1.0 +Requires-Dist: monai>=1.3 +Requires-Dist: segmentation-models-pytorch>=0.3 +Requires-Dist: opencv-python-headless>=4.9 +Requires-Dist: scikit-image>=0.22 +Requires-Dist: albumentations>=1.4 +Requires-Dist: Pillow>=10.2 +Requires-Dist: numpy>=1.26 +Requires-Dist: pandas>=2.2 +Requires-Dist: scipy>=1.12 +Requires-Dist: scikit-learn>=1.4 +Requires-Dist: statsmodels>=0.14 +Requires-Dist: torchmetrics>=1.3 +Requires-Dist: hydra-core>=1.3 +Requires-Dist: omegaconf>=2.3 +Requires-Dist: wandb>=0.16 +Requires-Dist: grad-cam>=1.5 +Requires-Dist: simpy>=4.1 +Requires-Dist: reportlab>=4.1 +Requires-Dist: fastapi>=0.110 +Requires-Dist: uvicorn>=0.29 +Requires-Dist: onnx>=1.16 +Requires-Dist: onnxruntime>=1.17 +Requires-Dist: gradio>=4.25 +Requires-Dist: matplotlib>=3.8 +Requires-Dist: seaborn>=0.13 +Requires-Dist: tqdm>=4.66 +Requires-Dist: kaggle>=1.6 +Requires-Dist: python-dotenv>=1.0 +Provides-Extra: dev +Requires-Dist: pytest>=8.0; extra == "dev" +Requires-Dist: pytest-cov>=5.0; extra == "dev" +Requires-Dist: ruff>=0.4; extra == "dev" +Requires-Dist: pre-commit>=3.7; extra == "dev" +Requires-Dist: ipykernel>=6.29; extra == "dev" +Requires-Dist: jupyterlab>=4.1; extra == "dev" +Dynamic: license-file + +
+ +# Diabetic Retinopathy Detection + +### Explainable, quality-aware AI for diabetic retinopathy screening in rural India β€” built to be audited, not just admired + +[![Status](https://img.shields.io/badge/status-in%20development-orange)](docs/04_ROADMAP.md) +[![License](https://img.shields.io/badge/license-MIT%20(code)-blue)](LICENSE) +[![Python](https://img.shields.io/badge/python-3.11-3776AB?logo=python&logoColor=white)](#getting-started) +[![PyTorch](https://img.shields.io/badge/PyTorch-2.x-EE4C2C?logo=pytorch&logoColor=white)](#tech-stack) +[![Not a medical device](https://img.shields.io/badge/⚠️-not%20a%20medical%20device-red)](#ethical-boundary--intended-use) + +[**Analysis**](docs/01_PROJECT_ANALYSIS.md) Β· [**Literature**](docs/02_LITERATURE_REVIEW.md) Β· [**Tech Stack**](docs/03_TECH_STACK.md) Β· [**Roadmap**](docs/04_ROADMAP.md) Β· [**Prototype Scope**](docs/05_PROTOTYPE_SCOPE.md) Β· [**Report Bug**](https://github.com/adarshcod30/Diabetic-Retinopathy-Detection/issues) + +
+ +> **⚠️ Development status.** This repository is in **Phase 0** of a 20-week plan. The architecture, +> evidence base, and validation protocol are complete and documented; models are not yet trained. +> Every performance figure below is labelled **target** until it is measured on a locked test set. +> No unmeasured number is presented as a result. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Key Features](#key-features) +- [Tech Stack](#tech-stack) +- [System Architecture](#system-architecture) +- [Application Flow](#application-flow) +- [Data & ML Pipeline](#data--ml-pipeline) +- [Results & Model Performance](#results--model-performance) +- [Deployment & Infrastructure](#deployment--infrastructure) +- [Prototype Scope](#prototype-scope-tier-p) +- [Project Structure](#project-structure) +- [Getting Started](#getting-started) +- [Usage](#usage) +- [Testing](#testing) +- [Roadmap](#roadmap) +- [Ethical Boundary & Intended Use](#ethical-boundary--intended-use) +- [Contributing](#contributing) +- [License](#license) +- [Contact](#contact) + +--- + +## Overview + +**Problem.** India has roughly **101 million adults with diabetes** (ICMR-INDIAB, *Lancet Diabetes & +Endocrinology* 2023 β€” 11.4 % weighted prevalence across 113,043 participants). Diabetic retinopathy +affects **14–17 %** of them, and there is approximately **one retina specialist per 1.26 million +people**. Manual screening at that scale is arithmetically impossible. + +**Solution.** A seven-stage screening pipeline that assesses image quality before it grades, grades +using the International Clinical DR Severity Scale, and explains itself with lesion-level evidence a +clinician can verify in seconds β€” plus a discrete-event simulation of the district screening programme +the model would run inside. + +**Why it is different.** Automated DR grading is a solved-enough problem: [Gulshan 2016](https://jamanetwork.com/journals/jama/fullarticle/2588763) +reached 97.5 % sensitivity, and [IDx-DR](https://www.nature.com/articles/s41746-018-0040-6) is +FDA-authorised. The unsolved problems are the ones this project targets: + +1. **Explainability that is measured, not asserted.** Nearly every DR project ships a Grad-CAM and + stops. This one runs [Adebayo's sanity checks](https://arxiv.org/abs/1810.03292) on the saliency + method itself, then *quantifies* whether the heatmap lands on real lesions using IDRiD's + pixel-level masks (pointing game, CAM–lesion IoU). +2. **Graceful failure on field images.** [Beede et al. (CHI 2020)](https://dl.acm.org/doi/10.1145/3313831.3376718) + found that in 11 Thai clinics the dominant failure was ungradable images being silently rejected, + wasting patient trips. Here, rejection returns an actionable reason: defocus, illumination, or + field of view. +3. **Clinically grounded evidence.** Lesion segmentation feeds *both* the grade and the explanation, + mapped onto the ICDR criteria β€” so the rationale is the same evidence the model actually used. + +**Keywords:** `diabetic-retinopathy` `medical-imaging` `explainable-ai` `grad-cam` `deep-learning` +`pytorch` `image-segmentation` `fundus-photography` `computer-vision` `healthcare-ai` `ordinal-regression` +`model-calibration` `simpy` `rural-health` `screening` + +--- + +## Key Features + +| Feature | Description | +|---|---| +| **Quality gate with recapture guidance** | 3-class gradability model (EyeQ) fused with focus, illumination and FOV features; a rejected image returns *why*, so the operator can retake it while the patient is still seated | +| **Adaptive enhancement** | Ben Graham local-colour normalisation, CLAHE, illumination correction β€” every step ablated, none assumed | +| **Retinal structure segmentation** | Vessels (DRIVE), optic disc & fovea landmarks, and four lesion classes (MA, haemorrhage, hard/soft exudate) at pixel level | +| **Sub-pixel microaneurysm detection** | Candidate generation via morphological top-hat + matched filtering, then CNN classification β€” because MAs vanish under naΓ―ve downscaling | +| **Ordinal DR grading** | ICDR 0–4 with rank-consistent ordinal heads, not plain cross-entropy: grading 0 as 4 is not the same mistake as 0 as 1 | +| **Lesion-aware fusion** | CNN embedding concatenated with clinically meaningful lesion features (counts, areas, quadrant distribution) | +| **Calibrated confidence** | Temperature scaling with reliability diagrams and ECE; low-confidence cases escalate to a human grader | +| **Audited explainability** | Grad-CAM variants scored against ground-truth lesion masks, with sanity checks reported pass/fail | +| **Auto-generated clinical report** | Annotated PDF with lesion overlays, ICDR evidence table, and calibrated confidence β€” designed for a 30-second review | +| **Screening-programme simulation** | SimPy discrete-event model of a 100,000-patient/year district programme: bandwidth, throughput, grader capacity, ophthalmologist-hours freed | + +--- + +## Tech Stack + +| Layer | Technology | +|---|---| +| DL framework | PyTorch 2.x Β· PyTorch Lightning Β· timm | +| Medical imaging | MONAI Β· segmentation-models-pytorch | +| Classical CV | OpenCV Β· scikit-image Β· Albumentations | +| Explainability | pytorch-grad-cam (Grad-CAM / ++ / Score-CAM / Eigen-CAM) | +| Config & tracking | Hydra + OmegaConf Β· Weights & Biases | +| Metrics & statistics | torchmetrics Β· scikit-learn Β· scipy Β· statsmodels (bootstrap CI, DeLong, McNemar) | +| Simulation | SimPy *(optional Simulink mirror)* | +| Serving | FastAPI Β· ONNX Runtime Β· Docker | +| Demo | Gradio β†’ HuggingFace Spaces | +| Reporting | ReportLab | +| Quality | ruff Β· pytest Β· pre-commit Β· GitHub Actions | + +Rationale for every choice, including the MATLABβ†’open-source mapping: [`docs/03_TECH_STACK.md`](docs/03_TECH_STACK.md). + +--- + +## System Architecture + +```mermaid +flowchart TD + A[Fundus image
portable camera] --> B[Stage 1
Quality Assessment] + B -->|Reject| B1[Recapture guidance:
defocus Β· illumination Β· FOV] + B1 -.retake.-> A + B -->|Good / Usable| C[Stage 2
Adaptive Enhancement] + C --> D[Stage 3
Segmentation] + D --> D1[Vessels] + D --> D2[Optic disc / fovea
β†’ quadrant map] + D --> D3[Lesions: MA Β· HE Β· SE Β· EX] + D3 --> E[Lesion features
counts Β· areas Β· quadrants] + C --> F[Stage 4
Grading backbone] + E --> G[Fusion head] + F --> G + G --> H[Stage 5
Ordinal head + calibration] + H --> I{Confident?} + I -->|Yes| J[Stage 6
Explainability] + I -->|No| K[Escalate to human grader] + D3 --> J + J --> L[Annotated PDF report] + L --> M[Ophthalmologist review] + K --> M + M --> N[(Audit log)] + N -.feedback.-> F +``` + +**In plain language.** An image arrives from a portable camera in a primary health centre. Before +anything else, the system decides whether it is *gradable* β€” and if not, tells the operator exactly +what to fix. A gradable image is normalised for the wide variation in camera and lighting, then passed +down two parallel paths: one segments the retina's structures and lesions, the other classifies overall +severity. The two paths meet at a fusion head, because the lesions a clinician would look for are +exactly the features that should drive the grade. The result is calibrated into an honest probability; +uncertain cases go straight to a human. Confident cases produce a report whose explanation is grounded +in the same lesion evidence the model used β€” not a decorative heatmap. Every decision is logged, and +graders' corrections feed the next training round. + +--- + +## Application Flow + +```mermaid +sequenceDiagram + participant P as Patient + participant H as Health worker + participant E as Edge app + participant S as District server + participant O as Ophthalmologist + + P->>H: Arrives at screening camp + H->>E: Capture fundus image + E->>E: On-device quality check + alt Ungradable + E-->>H: "Image defocused β€” retake" + H->>E: Recapture + end + E->>S: Upload (store-and-forward if offline) + S->>S: Enhance β†’ segment β†’ grade β†’ calibrate + alt High confidence, no DR + S-->>H: Auto-cleared Β· routine recall + else Referable or uncertain + S->>O: Queue annotated report + O->>O: Review (~30 s target) + O-->>H: Confirmed grade + referral + end + H-->>P: Result and next step + O->>S: Correction logged for retraining +``` + +The **on-device quality check is deliberate**: a round trip to the cloud to learn the image was blurry +is precisely the failure mode documented in Thailand. The check must complete while the patient is +still in the chair. + +--- + +## Data & ML Pipeline + +### Data sources + +| Dataset | Size | Labels | Population | Role | Licence | +|---|---|---|---|---|---| +| [APTOS 2019](https://www.kaggle.com/c/aptos2019-blindness-detection) | 3,662 | ICDR 0–4 | **India** (Aravind, Madurai) | Primary training | Competition rules | +| [EyePACS 2015](https://www.kaggle.com/c/diabetic-retinopathy-detection) | 88,702 | ICDR 0–4 (noisy) | US | Pretraining *(cloud-side only)* | Competition rules | +| [IDRiD](https://ieee-dataport.org/open-access/indian-diabetic-retinopathy-image-dataset-idrid) | 516 | Grades + **pixel masks** + OD/fovea | **India** (Nanded) | Lesion segmentation, XAI ground truth | CC BY 4.0 | +| [Messidor-2](https://www.adcis.net/en/third-party/messidor2/) | 1,748 | Adjudicated grades | France | **Locked external test** | ADCIS terms | +| [DRIVE](https://drive.grand-challenge.org/) | 40 | Vessel masks | Netherlands | Vessel segmentation | Research use | +| EyeQ | 28,792 | Good / Usable / Reject | derived from EyePACS | Quality model | Research use | + +> Raw images are **never committed**. `scripts/download_data.sh` fetches them; `data/manifests/` commits +> sha256 hashes so anyone can verify byte-identical inputs. + +### Pipeline + +```mermaid +flowchart LR + R[Raw fundus] --> Q{Gradable?} + Q -->|No| RC[Recapture] + Q -->|Yes| P1[crop_from_gray] + P1 --> P2[circle_crop] + P2 --> P3[Ben Graham
local colour norm] + P3 --> P4[CLAHE + illumination norm] + P4 --> P5[Resize 512–768 px
cache to disk] + P5 --> AUG[Augment:
flip Β· rotate Β· colour jitter] + AUG --> TR[Train] + TR --> CAL[Calibrate on val] + CAL --> TEST[Locked external test
evaluated ONCE] +``` + +**Cleaning & preparation.** Circle-crop removes the black surround and standardises the field of view; +`crop_from_gray` handles off-centre captures. **Ben Graham preprocessing** β€” subtracting a heavily +blurred copy to remove local average colour β€” is the single highest-value step, as established by the +Kaggle DR 2015 winner and used by most APTOS top solutions. Preprocessed images are cached once at +512 px, which shrinks APTOS from ~10 GB to ~200 MB and makes every subsequent epoch cheaper. + +**Feature engineering.** Beyond learned features, the lesion branch yields *clinically named* features: +microaneurysm count, haemorrhage and exudate area, per-quadrant distribution (using the optic +disc–fovea axis), and distance-to-fovea. These are the ICDR criteria expressed numerically β€” they +improve the grade and simultaneously become the explanation. + +**Training approach.** +- Two-stage: pretrain on EyePACS (or initialise from [RETFound](https://www.nature.com/articles/s41586-023-06555-x), + a masked-autoencoder foundation model trained on 1.6 M retinal images), then fine-tune on APTOS. +- **Ordinal loss**, not cross-entropy β€” DR grades are ordered, and recent SOTA + ([Dual-SwinOrd](https://www.mdpi.com/2306-5354/13/4/374), AOR-DR) confirms this matters. +- **Resolution is the dominant hyperparameter.** A microaneurysm is ~10 px on a 4288 px image; at + 224 px it is sub-pixel and physically destroyed. Sweep 384/512/768 early. +- **Patient-level splits** β€” both eyes of one patient must never straddle train and test. + +**Evaluation metrics.** Quadratic Weighted Kappa (grading); sensitivity and specificity at the +referable-DR operating point (grade β‰₯ 2); AUROC; **AUPRC for lesion segmentation** (lesion pixels are +<0.1 % of an image, which makes AUROC flattering and near-meaningless); Expected Calibration Error; +all with bootstrap 95 % CIs, DeLong for AUC comparisons, and McNemar for paired sensitivity/specificity. + +--- + +## Results & Model Performance + +> **No results yet β€” models are not trained.** The table below states **targets** and the **published +> work they are benchmarked against**. It will be replaced with measured values, confidence intervals, +> and the full ablation once Phase 8 completes. + +### Targets + +| Metric | Target | Benchmark it is measured against | +|---|---|---| +| Referable DR sensitivity | **β‰₯ 90 %** | Ting 2017: 90.5 % Β· IDx-DR: 87.2 % Β· Google/Aravind: 88.9 % | +| Referable DR specificity | **β‰₯ 85 %** | Ting 2017: 91.6 % Β· IDx-DR: 90.7 % Β· Google/Aravind: 92.2 % | +| QWK (APTOS, 5-class) | **β‰₯ 0.90** | Dual-SwinOrd SOTA: 0.9370 | +| Quality classification (EyeQ) | **β‰ˆ 0.90 acc** | VISTA: 0.9066 acc, 0.8868 F1 | +| Lesion AUPRC (IDRiD) | per-class, cross-validated | IDRiD ISBI-2018 challenge leaderboard | +| Calibration (ECE) | **< 0.05** after temperature scaling | β€” | +| Grad-CAM localisation | pointing-game accuracy vs IDRiD masks | *rarely reported β€” this is the contribution* | + +### Published comparators + +| Study | Task | Sensitivity | Specificity | AUC | +|---|---|---|---|---| +| [Gulshan 2016 (JAMA)](https://jamanetwork.com/journals/jama/fullarticle/2588763) | Referable DR, EyePACS-1 | 97.5 % | 93.4 % | 0.991 | +| [Ting 2017 (JAMA)](https://pubmed.ncbi.nlm.nih.gov/29234807/) | Referable DR, multiethnic | 90.5 % | 91.6 % | 0.936 | +| [AbrΓ moff 2018 (npj Digit Med)](https://www.nature.com/articles/s41746-018-0040-6) | mtmDR, **prospective primary care** | 87.2 % | 90.7 % | β€” | +| [Gulshan 2019 (JAMA Ophthalmol)](https://research.google/pubs/performance-of-a-deep-learning-algorithm-vs-manual-grading-for-detecting-diabetic-retinopathy-in-india/) | Referable DR, **Aravind, India** | 88.9 % | 92.2 % | 0.963 | +| same | Referable DR, **Sankara Nethralaya** | 92.1 % | 95.2 % | 0.980 | + +> **A note on honesty:** a faithful reproduction of Gulshan 2016 scored AUC 0.951 on EyePACS but only +> **0.853 on Messidor-2**. External validation drops are the norm, not a failure. This project plans +> for that drop, reports it, and analyses the domain shift rather than re-tuning until it disappears. + +The planned **ablation** β€” eleven configurations, one factor added per row, evaluated on the locked +test set with significance tests β€” is specified in +[Β§7 of the analysis](docs/01_PROJECT_ANALYSIS.md#7-the-ablation-that-answers-integrated--single-technique). + +--- + +## Deployment & Infrastructure + +| Concern | Approach | +|---|---| +| **Training** | Kaggle Notebooks (T4/P100, 30 GPU-h/week, EyePACS pre-mounted) β€” chosen because the dev machine has 28 GB free disk and EyePACS is ~90 GB | +| **Local development** | Apple M4 via PyTorch MPS β€” preprocessing, IDRiD segmentation, inference, explainability, demo | +| **Model export** | PyTorch β†’ ONNX β†’ ONNX Runtime (server), CoreML (Apple), TFLite (Android capture app) | +| **Serving** | FastAPI + ONNX Runtime in Docker (`linux/amd64` + `linux/arm64`, so Jetson/Raspberry Pi remain viable for a PHC) | +| **Edge** | On-device quality check so blur is caught before upload, not after | +| **Connectivity** | Store-and-forward queue; the simulation explicitly models 1–10 Mbps rural links and outages | +| **Public demo** | Gradio on HuggingFace Spaces (free CPU tier) | +| **CI/CD** | GitHub Actions β€” ruff, pytest, and an end-to-end single-image smoke test on CPU | +| **Monitoring** | Audit log of every prediction + grader override; drift review before any retraining | +| **Reproducibility** | Hydra configs, fixed seeds, sha256 data manifests; `make setup && make evaluate` reproduces the headline table | + +--- + +## Project Structure + +``` +Diabetic-Retinopathy-Detection/ +β”œβ”€β”€ configs/ # Hydra YAMLs β€” one per experiment; this IS the ablation grid +β”œβ”€β”€ data/ # gitignored (manifests are committed) +β”‚ β”œβ”€β”€ raw/ interim/ processed/ external/ +β”‚ └── manifests/ # sha256 + labels + patient_id +β”œβ”€β”€ docs/ +β”‚ β”œβ”€β”€ 01_PROJECT_ANALYSIS.md # what this is, why it is hard, what "done" means +β”‚ β”œβ”€β”€ 02_LITERATURE_REVIEW.md # annotated evidence base +β”‚ β”œβ”€β”€ 03_TECH_STACK.md # tooling decisions and rationale +β”‚ β”œβ”€β”€ 04_ROADMAP.md # 20-week phased plan +β”‚ └── 05_PROTOTYPE_SCOPE.md # Tier-P: scaling down without breaking the science +β”œβ”€β”€ notebooks/ # exploration only β€” logic lives in src/ +β”œβ”€β”€ src/drdetect/ +β”‚ β”œβ”€β”€ data/ # datasets, patient-level splits, manifests +β”‚ β”œβ”€β”€ quality/ # Stage 1 β€” gradability + recapture guidance +β”‚ β”œβ”€β”€ enhance/ # Stage 2 β€” Ben Graham, CLAHE, illumination +β”‚ β”œβ”€β”€ segmentation/ # Stage 3 β€” vessels, OD/fovea, lesions +β”‚ β”œβ”€β”€ grading/ # Stage 4 β€” backbones, ordinal heads, fusion +β”‚ β”œβ”€β”€ calibration/ # Stage 5 β€” temperature scaling, thresholds, uncertainty +β”‚ β”œβ”€β”€ explain/ # Stage 6 β€” CAMs, sanity checks, localisation metrics, reports +β”‚ β”œβ”€β”€ eval/ # metrics, bootstrap CI, DeLong, McNemar +β”‚ β”œβ”€β”€ serve/ # FastAPI + ONNX +β”‚ └── utils/ # seeding, logging, io +β”œβ”€β”€ simulation/ +β”‚ β”œβ”€β”€ simpy/ # Stage 7 β€” district screening programme model +β”‚ └── simulink/ # optional .slx mirror +β”œβ”€β”€ scripts/ # benchmark_device.py Β· download_data.sh Β· preprocess.py Β· train.py Β· evaluate.py +β”œβ”€β”€ tests/ +β”œβ”€β”€ models/ # gitignored; released via GitHub Releases / HF Hub +└── .github/workflows/ +``` + +--- + +## Getting Started + +### Prerequisites + +- Python **3.11** (not 3.13 β€” several CV/DL wheels still lag) +- ~20 GB free disk for local datasets +- A Kaggle account (for datasets and free GPU) +- Optional: CUDA GPU. Apple Silicon works via MPS for everything except large-scale training. + +### Installation + +```bash +git clone https://github.com/adarshcod30/Diabetic-Retinopathy-Detection.git +cd Diabetic-Retinopathy-Detection +``` + +```bash +conda create -n dr python=3.11 -y && conda activate dr +``` + +```bash +pip install -e ".[dev]" +``` + +```bash +python -c "import torch; print('MPS:', torch.backends.mps.is_available(), '| CUDA:', torch.cuda.is_available())" +``` + +### Get the data + +Place your Kaggle token at `~/.config/kaggle/kaggle.json`, then: + +```bash +bash scripts/download_data.sh --datasets aptos,idrid,drive +``` + +> Messidor-2 requires accepting [ADCIS terms](https://www.adcis.net/en/third-party/messidor2/) and is +> downloaded manually into `data/external/messidor2/`. **Do not** download EyePACS locally β€” it is +> ~90 GB; train against it on Kaggle instead. + +### Preprocess + +```bash +python scripts/preprocess.py --dataset aptos --size 512 --pipeline bengraham +``` + +--- + +## Usage + +Grade a single image and produce an annotated PDF report: + +```bash +python scripts/predict.py --image path/to/fundus.jpg --report out/report.pdf +``` + +Launch the interactive demo: + +```bash +python -m drdetect.serve.demo +``` + +Run a training experiment (Hydra β€” every ablation row is a config override): + +```bash +python scripts/train.py experiment=grading_effnetv2 data.image_size=512 loss=ordinal +``` + +Evaluate on the locked external test set: + +```bash +python scripts/evaluate.py --checkpoint models/checkpoints/best.ckpt --split external_test --bootstrap 2000 +``` + +Run the screening-programme simulation: + +```bash +python -m simulation.simpy.district --patients-per-year 100000 --graders 4 --bandwidth-mbps 5 +``` + +> Commands reflect the target interface; scripts land progressively through Phases 1–7. See +> [`docs/04_ROADMAP.md`](docs/04_ROADMAP.md) for what exists today. + +--- + +## Testing + +```bash +pytest -q +``` + +```bash +ruff check . && ruff format --check . +``` + +Test strategy: +- **Unit** β€” preprocessing determinism, circle-crop geometry, metric correctness against hand-computed + cases, patient-level split integrity (*asserts no patient appears in two splits* β€” the leakage bug + that silently inflates every published number) +- **Integration** β€” one image end-to-end on CPU, asserting a valid grade and a well-formed PDF +- **CI** β€” ruff + pytest + the smoke test on every push + +--- + +## Roadmap + +| Phase | Weeks | Milestone | +|---|---|---| +| 0 Β· Foundation | 1 | Repo, environment, data audit, manifests | +| 1 Β· Baseline | 2–3 | Baseline QWK on APTOS β€” a number to beat | +| **2 Β· Vertical slice** ⭐ | 4–5 | **JPEG in β†’ annotated PDF out** | +| 3 Β· Grading depth | 6–8 | Resolution sweep, ordinal loss, pretraining β†’ QWK β‰₯ 0.90 | +| 4 Β· Segmentation | 9–11 | Vessels, OD/fovea, four lesion classes, MA detector | +| 5 Β· Integration | 12–13 | Lesion fusion + calibration + uncertainty escalation | +| **6 Β· Rigorous XAI** ⭐ | 14–15 | **Sanity checks + measured CAM localisation** | +| 7 Β· Simulation | 16–17 | District programme model, ophthalmologist-hours freed | +| 8 Β· Validation | 18–19 | Full ablation on locked test set with significance tests | +| 9 Β· Release | 20 | ONNX, Docker, HF Spaces demo, model card | + +Detail, exit criteria, and an explicit scope-cut order: [`docs/04_ROADMAP.md`](docs/04_ROADMAP.md). + +**Known scope limit:** neovascularisation segmentation is *out of scope* β€” no public dataset provides +NV pixel masks. PDR is detected at image level and the gap is documented rather than papered over. + +--- + +## Prototype Scope (Tier-P) + +Development runs on a 16 GB Apple M4 with ~28 GB free disk. That constrains **compute**, not +data β€” after caching at 512 px, APTOS, IDRiD, Messidor-2 and DRIVE together occupy **0.4 GB**. + +So the scale-down cuts compute and keeps every image: + +| | Full (Tier-F) | **Prototype (Tier-P)** | +|---|---|---| +| APTOS / IDRiD / DRIVE / Messidor-2 | full | **full β€” unchanged** | +| EyePACS pretraining | 88,702 | 12–15k stratified, cached on Kaggle | +| Backbone | EfficientNetV2-S | EfficientNet-B0 | +| Resolution | 768 px | 512 px *(floor β€” below this microaneurysms vanish)* | +| Cross-validation | 5-fold + TTA | single split + hflip TTA | + +**β‰ˆ34Γ— cheaper, zero task data discarded.** Every cut is a config override, reversible on +Kaggle without re-splitting or re-collecting anything. + +Two things are never scaled down: the **test set** (a 50-case test set gives a Β±8.5 pp +confidence interval, which makes a ">90 % sensitivity" claim unmakeable) and the **validation +protocol**. Grade 3 has only 193 images in all of APTOS β€” subsampling collapses it to 15–38 +training examples long before the dataset merely looks small. + +Full arithmetic and the claims that do and don't survive: [`docs/05_PROTOTYPE_SCOPE.md`](docs/05_PROTOTYPE_SCOPE.md). + +## Ethical Boundary & Intended Use + +**This is a research prototype. It is not a medical device and must not be used for clinical +diagnosis.** + +- Software intended for diagnosis falls under **CDSCO** medical-device rules in India; the comparable + US system (IDx-DR) required an FDA De Novo authorisation. This project has neither and seeks neither. +- Outputs are **decision support with a human in the loop**, never a diagnosis. Every generated report + carries the disclaimer. +- **Training populations are documented and limited:** APTOS and IDRiD are Indian, EyePACS is US, + Messidor-2 is French. Performance on any other population is **unverified**. +- No patient-identifying data is stored in this repository; demo images are EXIF-stripped. +- Released weights are **research use only**, consistent with the licences of the data they derive from. + +--- + +## Contributing + +Contributions are welcome β€” particularly clinical review of the explanation format, additional external +validation datasets, and reproductions of the ablation. + +1. Fork and branch (`git checkout -b feature/your-feature`) +2. `pre-commit install` +3. Add tests for anything you change +4. Ensure `pytest` and `ruff check .` pass +5. Open a pull request describing *what you measured*, not only what you changed + +Findings that contradict results here are especially welcome. Please open an issue with the +configuration and seed. + +--- + +## License + +Code is released under the **MIT License** β€” see [LICENSE](LICENSE). + +**Data and weights are not.** Datasets remain under their original licences (IDRiD CC BY 4.0; +APTOS/EyePACS competition rules; Messidor-2 ADCIS terms; DRIVE research use). Raw images are never +redistributed here. Trained weights derived from research-use data are released for **research use +only**. + +--- + +## Contact + +**Adarsh Dwivedi** Β· [@adarshcod30](https://github.com/adarshcod30) +Project: + +## Acknowledgments + +The datasets that make this possible: [APTOS](https://www.kaggle.com/c/aptos2019-blindness-detection) +(Aravind Eye Hospital), [IDRiD](https://idrid.grand-challenge.org/) (Porwal et al.), +[Messidor-2](https://www.adcis.net/en/third-party/messidor2/) (ADCIS), and +[DRIVE](https://drive.grand-challenge.org/). Full evidence base: +[`docs/02_LITERATURE_REVIEW.md`](docs/02_LITERATURE_REVIEW.md). diff --git a/src/drdetect.egg-info/SOURCES.txt b/src/drdetect.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..b9e8db745bb25f99fe5bfa08c414f9350fc0a6d6 --- /dev/null +++ b/src/drdetect.egg-info/SOURCES.txt @@ -0,0 +1,20 @@ +LICENSE +README.md +pyproject.toml +src/drdetect/__init__.py +src/drdetect.egg-info/PKG-INFO +src/drdetect.egg-info/SOURCES.txt +src/drdetect.egg-info/dependency_links.txt +src/drdetect.egg-info/requires.txt +src/drdetect.egg-info/top_level.txt +src/drdetect/calibration/__init__.py +src/drdetect/data/__init__.py +src/drdetect/enhance/__init__.py +src/drdetect/enhance/preprocessing.py +src/drdetect/eval/__init__.py +src/drdetect/explain/__init__.py +src/drdetect/grading/__init__.py +src/drdetect/quality/__init__.py +src/drdetect/segmentation/__init__.py +src/drdetect/serve/__init__.py +src/drdetect/utils/__init__.py \ No newline at end of file diff --git a/src/drdetect.egg-info/dependency_links.txt b/src/drdetect.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/drdetect.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/drdetect.egg-info/requires.txt b/src/drdetect.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..55f42e688caefa47e1fbc86023798fb9afa1d632 --- /dev/null +++ b/src/drdetect.egg-info/requires.txt @@ -0,0 +1,40 @@ +torch>=2.2 +torchvision>=0.17 +lightning>=2.2 +timm>=1.0 +monai>=1.3 +segmentation-models-pytorch>=0.3 +opencv-python-headless>=4.9 +scikit-image>=0.22 +albumentations>=1.4 +Pillow>=10.2 +numpy>=1.26 +pandas>=2.2 +scipy>=1.12 +scikit-learn>=1.4 +statsmodels>=0.14 +torchmetrics>=1.3 +hydra-core>=1.3 +omegaconf>=2.3 +wandb>=0.16 +grad-cam>=1.5 +simpy>=4.1 +reportlab>=4.1 +fastapi>=0.110 +uvicorn>=0.29 +onnx>=1.16 +onnxruntime>=1.17 +gradio>=4.25 +matplotlib>=3.8 +seaborn>=0.13 +tqdm>=4.66 +kaggle>=1.6 +python-dotenv>=1.0 + +[dev] +pytest>=8.0 +pytest-cov>=5.0 +ruff>=0.4 +pre-commit>=3.7 +ipykernel>=6.29 +jupyterlab>=4.1 diff --git a/src/drdetect.egg-info/top_level.txt b/src/drdetect.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..4a34d31eb50000d2a11b31d4acccdd1016b2c8e3 --- /dev/null +++ b/src/drdetect.egg-info/top_level.txt @@ -0,0 +1 @@ +drdetect diff --git a/src/drdetect/__init__.py b/src/drdetect/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7f4b54e24741de07bef2ca41819cf9fc4e892204 --- /dev/null +++ b/src/drdetect/__init__.py @@ -0,0 +1,6 @@ +"""drdetect β€” explainable, quality-aware diabetic retinopathy screening. + +Research prototype. NOT a medical device; not for clinical diagnosis. +""" + +__version__ = "0.1.0" diff --git a/src/drdetect/calibration/__init__.py b/src/drdetect/calibration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/calibration/mc_dropout.py b/src/drdetect/calibration/mc_dropout.py new file mode 100644 index 0000000000000000000000000000000000000000..840f89ee8c029fa962239b909b0a32f2ec974d70 --- /dev/null +++ b/src/drdetect/calibration/mc_dropout.py @@ -0,0 +1,33 @@ +"""MC-dropout uncertainty estimation for the grading model. + +EfficientNet's classifier dropout is *functional*, not a persistent +`nn.Dropout` submodule -- timm's own `EfficientNet.forward_head` calls +`F.dropout(x, p=self.drop_rate, training=self.training)` directly. There is +no Dropout module to selectively re-enable via `.apply()`, so reactivating +it at inference means putting the WHOLE model back into `.train()` mode +(which sets the `self.training` flag `F.dropout` checks), then immediately +re-freezing BatchNorm with this project's own `freeze_batchnorm` so its +running statistics stay fixed -- leaving the dropout call as the only thing +still affected by the reactivated training flag. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from drdetect.grading.model import freeze_batchnorm + +__all__ = ["mc_dropout_samples"] + + +def mc_dropout_samples(model: nn.Module, x: torch.Tensor, n_samples: int = 20) -> torch.Tensor: + """Returns logits of shape (n_samples, batch, num_outputs), each from an + independent stochastic forward pass with dropout active and BatchNorm + still frozen on its pretrained running statistics.""" + model.train() + freeze_batchnorm(model) + with torch.no_grad(): + samples = torch.stack([model(x) for _ in range(n_samples)]) + model.eval() + return samples diff --git a/src/drdetect/calibration/temperature.py b/src/drdetect/calibration/temperature.py new file mode 100644 index 0000000000000000000000000000000000000000..703e4d014fc2980eee3be4b488ea83289de23f71 --- /dev/null +++ b/src/drdetect/calibration/temperature.py @@ -0,0 +1,101 @@ +"""Temperature scaling (Guo et al. 2017) for post-hoc confidence calibration. + +Modern networks are systematically overconfident: raw softmax output is not a +probability, even when the model's argmax is highly accurate. This project's +own Phase 2 report has said so explicitly on every single output -- +"Model confidence (uncalibrated)" -- since there was nothing calibrated to +report. This module is what closes that disclaimer, for checkpoints it has +actually been run against. + +Temperature scaling fixes overconfidence with one learned scalar T, applied +AFTER training with every other weight frozen: + + calibrated_probs = softmax(logits / T) + +T > 1 softens (reduces) confidence; T < 1 sharpens it. It cannot change which +class wins (dividing every logit by the same positive number does not +reorder them), only how confident the model is entitled to sound about it. + +Only meaningful for a softmax head (`ce` / `distance_ce`). CORN's logits are +per-task sigmoids, not a single softmax, and regression has no probability +distribution to calibrate at all -- both are out of scope here, not silently +mishandled. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = ["TemperatureScaler", "fit_temperature", "save_temperature", "load_temperature"] + + +class TemperatureScaler(nn.Module): + """Wraps a frozen classifier's logits; only `log_temperature` is ever trained.""" + + def __init__(self): + super().__init__() + self.log_temperature = nn.Parameter(torch.zeros(1)) + + @property + def temperature(self) -> torch.Tensor: + # exp(), not a raw Parameter: keeps T > 0 for every value the + # optimiser can reach. logits / T is undefined at T=0 and inverted + # (sharpens instead of softens) for T<0 -- a plain Parameter could + # wander into either during optimisation; exp() cannot reach them. + return self.log_temperature.exp() + + def forward(self, logits: torch.Tensor) -> torch.Tensor: + return logits / self.temperature + + +def fit_temperature(logits: torch.Tensor, targets: torch.Tensor, *, max_iter: int = 50) -> float: + """Fit T on held-out (logits, targets) by minimising NLL. Returns the fitted T. + + `logits` must come from a split the model was not trained OR early-stopped + on. This project has only a train/val split per fold (no separate + calibration split was reserved at Phase 0), so in practice this is fit on + the same fold-0 validation set every other Phase 1/3 number came from -- + the identical "optimistic, not yet the Phase 8 number" caveat already + disclosed for the operating threshold applies here too, not a new one. + + LBFGS, not Adam/SGD: the standard choice for this exact problem in the + original paper's own reference implementation -- one parameter, a + well-behaved objective, converges in a handful of iterations. + """ + scaler = TemperatureScaler() + optimizer = torch.optim.LBFGS([scaler.log_temperature], lr=0.01, max_iter=max_iter) + + def closure(): + optimizer.zero_grad() + loss = F.cross_entropy(scaler(logits), targets) + loss.backward() + return loss + + optimizer.step(closure) + return float(scaler.temperature.item()) + + +def save_temperature(checkpoint_path: str | Path, temperature: float, *, fitted_on: str) -> Path: + """Write a sidecar JSON next to `checkpoint_path`, not inside the .ckpt. + + Keeping it separate means load_temperature() can no-op cleanly for every + checkpoint that has never been calibrated, instead of every existing + checkpoint needing to carry a temperature=1.0 placeholder it was never + actually fit for. + """ + out = Path(checkpoint_path).with_name("temperature.json") + out.write_text(json.dumps({"temperature": temperature, "fitted_on": fitted_on}, indent=2)) + return out + + +def load_temperature(checkpoint_path: str | Path) -> float: + """1.0 (uncalibrated, a no-op divisor) if this checkpoint has no fitted temperature.""" + sidecar = Path(checkpoint_path).with_name("temperature.json") + if not sidecar.exists(): + return 1.0 + return float(json.loads(sidecar.read_text())["temperature"]) diff --git a/src/drdetect/data/__init__.py b/src/drdetect/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/data/dataset.py b/src/drdetect/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..87c2d244c135e53896619fad6d9c3bc8e949a685 --- /dev/null +++ b/src/drdetect/data/dataset.py @@ -0,0 +1,134 @@ +"""Torch dataset over the preprocessed cache. + +Reads the 512 px cache written by scripts/preprocess.py, not the raw images. +The expensive work (crop, Ben Graham, resize) happened once; the training loop +only decodes a small JPEG and augments. +""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np +from torch.utils.data import Dataset + +from drdetect.data.manifest import ImageRecord, read_manifest + +__all__ = ["FundusDataset", "build_transforms", "load_split"] + + +class FundusDataset(Dataset): + """Fundus images with ICDR grades. + + Args: + records: manifest rows for this split. + root: directory the manifest's `path` fields are relative to. + transform: Albumentations transform. Must include normalisation and + ToTensorV2 -- this class does no implicit conversion, so what the + model receives is always explicit in the config. + """ + + def __init__(self, records: list[ImageRecord], root: str | Path, transform=None): + self.records = records + self.root = Path(root) + self.transform = transform + + def __len__(self) -> int: + return len(self.records) + + def __getitem__(self, idx: int): + rec = self.records[idx] + path = self.root / rec.path + image = cv2.imread(str(path), cv2.IMREAD_COLOR) + if image is None: + raise FileNotFoundError( + f"Could not read {path}. Re-run scripts/preprocess.py, or verify the " + f"cache against the manifest with drdetect.data.manifest.verify_manifest." + ) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + if self.transform is not None: + image = self.transform(image=image)["image"] + return image, rec.label + + @property + def labels(self) -> np.ndarray: + return np.array([r.label for r in self.records]) + + @property + def groups(self) -> list[str]: + return [r.group_id for r in self.records] + + +def build_transforms(size: int, train: bool): + """Augmentations. + + Train-time choices and why: + * flips and rotation -- a fundus has no canonical orientation, and left vs + right eyes are mirror images, so these are label-preserving; + * mild brightness/contrast -- models the illumination variation that + portable cameras genuinely produce in the field; + * NO heavy blur or noise -- a microaneurysm is 1-3 px at this resolution, + and those augmentations would erase the very feature that defines + grade 1. + """ + import albumentations as A + from albumentations.pytorch import ToTensorV2 + + # ImageNet statistics, because the backbone is ImageNet-pretrained. + normalise = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) + + if not train: + return A.Compose([A.Resize(size, size), normalise, ToTensorV2()]) + + return A.Compose( + [ + A.Resize(size, size), + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.5), + A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1, rotate_limit=180, p=0.7), + A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5), + normalise, + ToTensorV2(), + ] + ) + + +def load_split( + manifest_path: str | Path, + *, + fold: int = 0, + n_splits: int = 5, + seed: int = 42, + allow_ungrouped: bool = False, +) -> tuple[list[ImageRecord], list[ImageRecord], str]: + """Split a manifest into (train, val, strategy) for one fold. + + Grouping comes from the manifest's `group_id`, which scripts/preprocess.py + populates by perceptual hashing. Leakage is asserted, not assumed. + """ + from drdetect.data.splits import ( + assert_no_group_leakage, + stratified_group_split, + validate_grouping, + ) + + records = read_manifest(manifest_path) + labels = [r.label for r in records] + groups = [r.group_id for r in records] + + # Reject a collapsed grouping BEFORE splitting. A single giant group passes + # the leakage assertion trivially and produces a wildly lopsided split. + if any(groups): + validate_grouping(groups) + + folds, strategy = stratified_group_split( + labels, groups, n_splits=n_splits, seed=seed, allow_ungrouped=allow_ungrouped + ) + assert_no_group_leakage(folds, groups) + + val_idx = set(folds[fold].tolist()) + train_records = [r for i, r in enumerate(records) if i not in val_idx] + val_records = [r for i, r in enumerate(records) if i in val_idx] + return train_records, val_records, strategy diff --git a/src/drdetect/data/manifest.py b/src/drdetect/data/manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..6caef3cee41ae10566b79a924907ca4c7fe5b016 --- /dev/null +++ b/src/drdetect/data/manifest.py @@ -0,0 +1,107 @@ +"""Dataset manifests: the committed record of exactly which bytes were used. + +Raw images are never committed (licence + size). A manifest is committed instead: +one row per image with its sha256, so anyone can verify they hold byte-identical +inputs before comparing results. This is DVC's job, done with a CSV. + +Schema +------ +image_id stable identifier (dataset-unique) +path path relative to the dataset root +sha256 content hash of the raw file +dataset aptos | idrid | messidor2 | drive | eyepacs +label ICDR grade 0-4, or -1 when unlabelled +group_id patient/eye grouping key for leak-free splits, or "" if unknown +phash structural perceptual hash (decimal string), or "" -- cached so + regrouping does not require re-decoding every source image +width pixels +height pixels +""" + +from __future__ import annotations + +import csv +import hashlib +from dataclasses import asdict, dataclass, fields +from pathlib import Path + +__all__ = ["ImageRecord", "sha256_file", "write_manifest", "read_manifest", "verify_manifest"] + +CHUNK = 1 << 20 # 1 MiB + + +@dataclass(frozen=True) +class ImageRecord: + image_id: str + path: str + sha256: str + dataset: str + label: int = -1 + group_id: str = "" + phash: str = "" + width: int = 0 + height: int = 0 + + +def sha256_file(path: str | Path) -> str: + """Content hash, streamed so a 90 GB dataset does not need 90 GB of RAM.""" + h = hashlib.sha256() + with open(path, "rb") as fh: + while chunk := fh.read(CHUNK): + h.update(chunk) + return h.hexdigest() + + +def write_manifest(records: list[ImageRecord], out_path: str | Path) -> Path: + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + cols = [f.name for f in fields(ImageRecord)] + with open(out_path, "w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=cols) + writer.writeheader() + for rec in sorted(records, key=lambda r: r.image_id): # stable ordering + writer.writerow(asdict(rec)) + return out_path + + +def read_manifest(path: str | Path) -> list[ImageRecord]: + out: list[ImageRecord] = [] + with open(path, newline="") as fh: + for row in csv.DictReader(fh): + out.append( + ImageRecord( + image_id=row["image_id"], + path=row["path"], + sha256=row["sha256"], + dataset=row["dataset"], + label=int(row["label"]), + group_id=row["group_id"], + phash=row.get("phash", ""), + width=int(row["width"]), + height=int(row["height"]), + ) + ) + return out + + +def verify_manifest( + manifest_path: str | Path, root: str | Path, *, sample: int | None = None +) -> tuple[list[str], list[str]]: + """Re-hash files on disk and compare against the manifest. + + Returns (missing, mismatched) image_ids. `sample` limits the check to the + first N records -- full verification of a large dataset is I/O bound. + """ + root = Path(root) + records = read_manifest(manifest_path) + if sample is not None: + records = records[:sample] + + missing, mismatched = [], [] + for rec in records: + fpath = root / rec.path + if not fpath.exists(): + missing.append(rec.image_id) + elif sha256_file(fpath) != rec.sha256: + mismatched.append(rec.image_id) + return missing, mismatched diff --git a/src/drdetect/data/splits.py b/src/drdetect/data/splits.py new file mode 100644 index 0000000000000000000000000000000000000000..fe3d6a6cac138ffb6a560aa524ebe5e15164f11d --- /dev/null +++ b/src/drdetect/data/splits.py @@ -0,0 +1,349 @@ +"""Leak-free cross-validation splits. + +The failure this module exists to prevent +----------------------------------------- +Fundus datasets contain both eyes of the same patient, and the two eyes of one +diabetic are strongly correlated in DR severity. If one eye lands in train and +the other in test, the model has effectively seen the test set. Every reported +metric inflates, the loss curves look perfect, and nothing warns you. + +The complication specific to APTOS +---------------------------------- +APTOS-2019's `train.csv` carries only `id_code` and `diagnosis`. There is **no +patient or eye identifier**, so a true patient-level split is impossible on that +dataset as distributed. Silently falling back to an image-level split would be +exactly the failure described above, so this module: + + 1. refuses to guess -- `stratified_group_split` requires an explicit decision; + 2. offers `perceptual_hash_groups` as a documented mitigation: images that are + near-identical (plausibly the same eye or patient) are grouped so they + cannot straddle folds; + 3. records which strategy was used, so the caveat travels with the results. + +Messidor-2 *does* encode exam pairing in its filenames, so real grouping is +available there and should be used. +""" + +from __future__ import annotations + +import warnings +from collections import defaultdict + +import numpy as np + +__all__ = [ + "SplitStrategy", + "dhash", + "structural_hash", + "perceptual_hash_groups", + "validate_grouping", + "stratified_group_split", + "assert_no_group_leakage", +] + + +class SplitStrategy: + """How grouping was determined. Travels with results as a provenance record.""" + + TRUE_GROUPS = "true_groups" # real patient/eye ids -- trustworthy + PERCEPTUAL = "perceptual_hash" # near-duplicate grouping -- mitigation only + IMAGE_LEVEL = "image_level" # no grouping -- leakage possible, must be declared + + +def dhash(image: np.ndarray, hash_size: int = 8) -> int: + """Difference hash: a 64-bit perceptual fingerprint. + + Compares each pixel with its right-hand neighbour on a downscaled grayscale + image, so it is robust to resolution, JPEG artefacts and mild brightness + shifts, but sensitive to actual content. Two photographs of the same eye + typically differ by only a few bits. + + Implemented here rather than pulled from `imagehash` to avoid a dependency + for ~10 lines, and so the exact behaviour is pinned by our own tests. + """ + import cv2 + + if image.ndim == 3: + image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + resized = cv2.resize(image, (hash_size + 1, hash_size), interpolation=cv2.INTER_AREA) + diff = resized[:, 1:] > resized[:, :-1] + bits = 0 + for bit in diff.flatten(): + bits = (bits << 1) | int(bit) + return bits + + +def structural_hash(image: np.ndarray, size: int = 256, hash_size: int = 16) -> int: + """Perceptual hash of an image's *coarse structure*, for grouping. + + Deliberately computed on a plain circle-crop + resize, NOT on the + Ben Graham-preprocessed image used for training. + + Ben Graham preprocessing subtracts the local average colour, which removes + exactly the low-frequency content dHash depends on. Measured on a fundus and + a degraded copy of itself: hashing the Ben Graham output put them 8 bits + apart (above the grouping threshold, so the duplicate was missed), while + hashing the structural image put them 0 bits apart -- with unrelated images + still 14 bits away. + + The identity of an eye lives in its vessel and optic-disc layout. That is + what this hashes. + + Resolution matters more than it looks. At 8x8 (64 bits) on 128 px, measured + across APTOS the median inter-image distance was 18 bits with ~670 pairs + falling within 5 bits -- most of them false matches, because 64 bits cannot + separate two fundus photographs that are both "a bright disc on black". + At 16x16 (256 bits) on 256 px the distribution is cleanly bimodal: 148 pairs + at 0-4 bits and then NOTHING until 39 bits. Those false matches are what + percolated the grouping into a single giant component. + """ + import cv2 + + from drdetect.enhance.preprocessing import circle_crop + + cropped = circle_crop(image) + resized = cv2.resize(cropped, (size, size), interpolation=cv2.INTER_AREA) + return dhash(resized, hash_size=hash_size) + + +def _hamming(a: int, b: int) -> int: + return bin(a ^ b).count("1") + + +def perceptual_hash_groups( + image_ids: list[str], + hashes: list[int], + max_distance: int = 6, + *, + max_group_size: int | None = None, +) -> dict[str, str]: + """Group near-duplicate images via union-find over Hamming distance. + + Returns {image_id: group_id}. Images within `max_distance` bits are treated + as the same group (transitively). + + Percolation, and why `max_group_size` exists + -------------------------------------------- + Union-find takes the TRANSITIVE closure: a~b and b~c puts a, b and c in one + group even if a and c are far apart. Over a large set, a small number of + spurious edges is enough to chain almost everything together -- a random + graph needs only ~N/2 edges to form a giant connected component. + + This is not hypothetical. On APTOS (3,662 images) a 64-bit hash produced + ~670 near-matches, and the transitive closure merged 3,138 images into a + single "group" which then became an entire validation fold. The split was + train=524 / val=3138 and silently reported as `true_groups`. + + Two defences: a hash resolution that does not manufacture false edges (see + `structural_hash`), and this hard cap on group size. `max_group_size` + defaults to max(25, 1% of the dataset) -- a "patient" larger than that is a + hash failure, not a patient, so the merge is refused. + + Choosing `max_distance` + ----------------------- + Measured on APTOS with `structural_hash` over all 6.7M pairs, the cumulative + count is flat across thresholds 2-6 (148 pairs at every value) and stays + near-flat to 8 (150). It then climbs -- 207 at 10, 391 at 12 -- and chaining + appears: threshold 10 produced a 31-image group whose members were up to 33 + bits apart. The default of 6 sits mid-plateau, so the result is insensitive + to the exact value. + + Note the threshold must be measured for the hash ACTUALLY used. An earlier + value of 10 was taken from a distribution measured on Ben Graham-preprocessed + images and then applied to `structural_hash`, which runs on raw + circle-cropped images. Different transform, different distances. + + This remains a *mitigation*, not a substitute for patient IDs: it catches + repeated or near-identical captures, but two genuinely different photographs + of one patient's two eyes will not be grouped. Declare the limitation when + reporting. + """ + if len(image_ids) != len(hashes): + raise ValueError("image_ids and hashes must be the same length") + + n = len(image_ids) + if max_group_size is None: + max_group_size = max(25, n // 100) + + parent = list(range(n)) + size = [1] * n + + def find(i: int) -> int: + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + refused = 0 + + def union(i: int, j: int) -> None: + nonlocal refused + ri, rj = find(i), find(j) + if ri == rj: + return + if size[ri] + size[rj] > max_group_size: + # Refusing the merge is the safe failure: two smaller groups may + # leak, but one giant group destroys the split entirely. + refused += 1 + return + lo, hi = (ri, rj) if ri < rj else (rj, ri) + parent[hi] = lo + size[lo] += size[hi] + + # O(n^2). Fine for APTOS (3,662 -> ~6.7M comparisons, seconds). For EyePACS + # scale this would need LSH bucketing on hash prefixes. + for i in range(n): + hi_hash = hashes[i] + for j in range(i + 1, n): + if (hi_hash ^ hashes[j]).bit_count() <= max_distance: + union(i, j) + + if refused: + warnings.warn( + f"{refused} merge(s) refused because they would exceed max_group_size=" + f"{max_group_size}. This usually means the hash is producing false " + f"matches; check structural_hash resolution and max_distance.", + UserWarning, + stacklevel=2, + ) + + roots = defaultdict(list) + for idx, image_id in enumerate(image_ids): + roots[find(idx)].append(image_id) + + # Name each group after its lexicographically smallest member, NOT the + # union-find root index. The root index depends on the order the inputs + # arrived in, so identical data grouped twice produced identical partitions + # under different names -- and StratifiedGroupKFold assigns folds by group + # VALUE, so the two runs got different splits. Measured: the 512 px and + # 768 px manifests had the same 3,523 groups but shared only 139 of 733 + # validation images, making the two runs non-comparable. + # + # Naming by member makes the mapping a pure function of the partition. + mapping: dict[str, str] = {} + for members in roots.values(): + gid = f"g_{min(members)}" + for m in members: + mapping[m] = gid + return mapping + + +def validate_grouping( + groups: list[str], + *, + max_group_fraction: float = 0.05, + min_group_fraction: float = 0.5, +) -> dict: + """Reject degenerate groupings before they silently destroy a split. + + Raises if any single group holds more than `max_group_fraction` of the data, + or if the number of distinct groups falls below `min_group_fraction` of the + number of images. + + This exists because the failure it catches is invisible downstream: a + grouping that collapsed 86% of APTOS into one group still reported strategy + `true_groups`, still passed the no-leakage assertion (trivially -- one group + cannot straddle folds), and produced a train=524 / val=3138 split that would + have trained for hours and produced a meaningless number. + """ + from collections import Counter + + n = len(groups) + counts = Counter(groups) + largest_id, largest = counts.most_common(1)[0] + stats = { + "n_images": n, + "n_groups": len(counts), + "largest_group": largest, + "largest_group_id": largest_id, + "merged": n - len(counts), + } + + if largest > max_group_fraction * n: + raise ValueError( + f"Degenerate grouping: group {largest_id!r} holds {largest}/{n} images " + f"({100 * largest / n:.1f}%), above the {100 * max_group_fraction:.0f}% limit.\n" + f"A 'patient' that large is a hash collapse, not a patient. Re-run " + f"grouping with a higher-resolution hash or a tighter max_distance." + ) + if len(counts) < min_group_fraction * n: + raise ValueError( + f"Degenerate grouping: {len(counts)} groups for {n} images " + f"(< {100 * min_group_fraction:.0f}%). The hash is over-merging." + ) + return stats + + +def stratified_group_split( + labels: list[int], + groups: list[str] | None, + *, + n_splits: int = 5, + seed: int = 42, + allow_ungrouped: bool = False, +) -> tuple[list[np.ndarray], str]: + """Stratified k-fold that keeps each group whole. + + Args: + labels: class label per sample, used for stratification. + groups: grouping key per sample. `None` or all-empty means no grouping + information exists. + allow_ungrouped: must be set explicitly to proceed without groups. This + is a guard rail: falling back silently is the bug this module exists + to prevent. + + Returns: + (fold_indices, strategy) where fold_indices[k] holds the *test* indices + of fold k, and strategy is a `SplitStrategy` constant to record. + + Raises: + ValueError: if grouping is absent and `allow_ungrouped` is False. + """ + from sklearn.model_selection import StratifiedGroupKFold, StratifiedKFold + + labels_arr = np.asarray(labels) + has_groups = groups is not None and any(g for g in groups) + + if not has_groups: + if not allow_ungrouped: + raise ValueError( + "No grouping information supplied. A plain image-level split can leak " + "the same patient across folds and silently inflate every metric.\n" + "Either provide `groups` (from patient ids, or from " + "perceptual_hash_groups()), or pass allow_ungrouped=True to accept " + "the risk explicitly and declare it when reporting results." + ) + warnings.warn( + "Splitting at IMAGE level with no grouping. If this dataset contains " + "multiple images per patient, results will be optimistically biased. " + "Record SplitStrategy.IMAGE_LEVEL alongside any metric derived from it.", + UserWarning, + stacklevel=2, + ) + splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed) + folds = [test for _, test in splitter.split(np.zeros(len(labels_arr)), labels_arr)] + return folds, SplitStrategy.IMAGE_LEVEL + + groups_arr = np.asarray(groups) + splitter = StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=seed) + folds = [test for _, test in splitter.split(np.zeros(len(labels_arr)), labels_arr, groups_arr)] + return folds, SplitStrategy.TRUE_GROUPS + + +def assert_no_group_leakage(folds: list[np.ndarray], groups: list[str]) -> None: + """Raise if any group appears in more than one fold. + + Call this after splitting, every time. It is cheap, and it converts the most + expensive silent bug in this literature into a loud, immediate failure. + """ + seen: dict[str, int] = {} + for fold_idx, indices in enumerate(folds): + for i in indices: + gid = groups[i] + if not gid: + continue + prev = seen.setdefault(gid, fold_idx) + if prev != fold_idx: + raise AssertionError( + f"Group leakage: group {gid!r} appears in folds {prev} and {fold_idx}." + ) diff --git a/src/drdetect/enhance/__init__.py b/src/drdetect/enhance/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/enhance/preprocessing.py b/src/drdetect/enhance/preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..9576d3ecb6fab272a73670ee581dd7e0887e8783 --- /dev/null +++ b/src/drdetect/enhance/preprocessing.py @@ -0,0 +1,131 @@ +"""Fundus preprocessing primitives. + +The pipeline order matters and is not arbitrary: + + crop_from_gray -> circle_crop -> ben_graham -> (CLAHE) -> resize + +`crop_from_gray` first, because the black surround otherwise drags down every +subsequent statistic (the blur in `ben_graham`, the CLAHE histogram, the +illumination estimate). `ben_graham` before `resize`, because it is defined +relative to image radius and is cheapest to reason about at native scale. + +Reference: Ben Graham's winning solution, Kaggle Diabetic Retinopathy 2015. +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +__all__ = [ + "crop_from_gray", + "circle_crop", + "ben_graham", + "apply_clahe", + "preprocess", +] + + +def crop_from_gray(image: np.ndarray, tol: int = 7) -> np.ndarray: + """Trim the uninformative dark border around a fundus photograph. + + Fundus cameras produce a bright circular retina on a black rectangle. The + black region carries no signal but skews global statistics, so it goes first. + + Args: + image: HxWx3 uint8 RGB, or HxW uint8 grayscale. + tol: Intensity at or below which a pixel counts as background. + + Returns: + The tightest crop containing above-threshold content. Returns the input + unchanged if the mask is empty (a fully dark frame), rather than raising + -- an all-dark image is a *quality* failure and is Stage 1's business, + not this function's. + """ + if image.ndim == 2: + mask = image > tol + if not mask.any(): + return image + return image[np.ix_(mask.any(1), mask.any(0))] + + gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + mask = gray > tol + if not mask.any(): + return image + + rows, cols = mask.any(1), mask.any(0) + return image[np.ix_(rows, cols)] + + +def circle_crop(image: np.ndarray) -> np.ndarray: + """Mask everything outside the inscribed circular field of view. + + Standardises the FOV across cameras so the model cannot learn to exploit + corner artefacts that happen to correlate with a particular device. + """ + image = crop_from_gray(image) + h, w = image.shape[:2] + cx, cy = w // 2, h // 2 + r = min(cx, cy) + + mask = np.zeros((h, w), dtype=np.uint8) + cv2.circle(mask, (cx, cy), r, color=1, thickness=-1) + + out = image * mask[..., None] if image.ndim == 3 else image * mask + return crop_from_gray(out) + + +def ben_graham(image: np.ndarray, sigma_scale: float = 30.0, weight: float = 4.0) -> np.ndarray: + """Subtract local average colour (Ben Graham, Kaggle DR 2015). + + out = weight*img - weight*GaussianBlur(img, sigma) + 128 + + The blur sigma is tied to *image radius*, not to a fixed pixel count, which + makes the normalisation scale-invariant -- important because fundus cameras + output wildly different resolutions. + + This removes the between-camera and between-patient variation in overall + illumination and pigmentation while preserving the high-frequency detail + that lesions live in. + """ + radius = min(image.shape[:2]) / 2.0 + sigma = max(radius / sigma_scale, 1.0) + blurred = cv2.GaussianBlur(image, (0, 0), sigmaX=sigma) + out = cv2.addWeighted(image, weight, blurred, -weight, 128) + return np.clip(out, 0, 255).astype(np.uint8) + + +def apply_clahe(image: np.ndarray, clip_limit: float = 2.0, tile_grid: int = 8) -> np.ndarray: + """Contrast-limited adaptive histogram equalisation on the LAB L channel. + + Operating on L rather than per-RGB-channel avoids the colour casts that + channel-wise CLAHE introduces. + + Note: `clip_limit` is a genuine risk knob. Too high and sensor noise is + amplified into structures that look like microaneurysms. Ablate it; do not + assume 2.0 is right for your camera mix. + """ + lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB) + clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(tile_grid, tile_grid)) + lab[..., 0] = clahe.apply(lab[..., 0]) + return cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) + + +def preprocess( + image: np.ndarray, + size: int = 512, + use_ben_graham: bool = True, + use_clahe: bool = False, +) -> np.ndarray: + """Full preprocessing chain, resized to `size` x `size`. + + `use_clahe` defaults to False deliberately: it is a hypothesis to be tested + in the ablation, not a default to be inherited. See docs/01_PROJECT_ANALYSIS.md + section 6, Stage 2. + """ + out = circle_crop(image) + if use_ben_graham: + out = ben_graham(out) + if use_clahe: + out = apply_clahe(out) + return cv2.resize(out, (size, size), interpolation=cv2.INTER_AREA) diff --git a/src/drdetect/eval/__init__.py b/src/drdetect/eval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/eval/metrics.py b/src/drdetect/eval/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..341367417b9c1d29a9fb62b110a983b6b2f0c293 --- /dev/null +++ b/src/drdetect/eval/metrics.py @@ -0,0 +1,340 @@ +"""Evaluation metrics for DR grading. + +Two things here are easy to get wrong and are therefore made explicit: + +1. **The referable-DR threshold must be chosen on validation and then frozen.** + Picking it on the test set is the most common silent cheat in this + literature -- it converts "we achieved 90% sensitivity" into "there exists a + threshold at which we would have". `choose_threshold_for_sensitivity` and + `evaluate_at_threshold` are separated so the two steps cannot be conflated. + +2. **A point estimate without an interval is not a result.** On a few hundred + test images, a reported "91.3% sensitivity" can easily be consistent with + 82%. `bootstrap_ci` exists so no headline number ships bare. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +__all__ = [ + "REFERABLE_THRESHOLD", + "BinaryScores", + "quadratic_weighted_kappa", + "referable_labels", + "binary_scores", + "choose_threshold_for_sensitivity", + "sensitivity_at_specificity_floor", + "youden_j", + "evaluate_at_threshold", + "bootstrap_ci", + "expected_calibration_error", + "DeLongResult", + "delong_roc_variance", + "delong_test", +] + +# ICDR grade >= 2 (Moderate NPDR) is the clinical referral boundary. +REFERABLE_THRESHOLD = 2 + + +@dataclass(frozen=True) +class BinaryScores: + sensitivity: float + specificity: float + ppv: float + npv: float + tp: int + fp: int + tn: int + fn: int + + @property + def n_positive(self) -> int: + return self.tp + self.fn + + @property + def n_negative(self) -> int: + return self.tn + self.fp + + +def quadratic_weighted_kappa(y_true, y_pred, *, n_classes: int = 5) -> float: + """Cohen's kappa with quadratic weights -- the standard DR grading metric. + + Quadratic weighting is what makes this the right metric for an *ordinal* + task: confusing grade 0 with grade 4 is penalised 16x more than confusing 0 + with 1, matching the clinical cost structure. Plain accuracy treats both as + a single error. + """ + y_true = np.asarray(y_true, dtype=int) + y_pred = np.asarray(y_pred, dtype=int) + + obs = np.zeros((n_classes, n_classes), dtype=float) + for t, p in zip(y_true, y_pred, strict=True): + obs[t, p] += 1 + + hist_true = obs.sum(axis=1) + hist_pred = obs.sum(axis=0) + expected = np.outer(hist_true, hist_pred) / max(obs.sum(), 1) + + i, j = np.meshgrid(np.arange(n_classes), np.arange(n_classes), indexing="ij") + weights = ((i - j) ** 2) / ((n_classes - 1) ** 2) + + denom = (weights * expected).sum() + if denom == 0: + return 0.0 + return float(1.0 - (weights * obs).sum() / denom) + + +def referable_labels(grades) -> np.ndarray: + """Binarise ICDR grades into referable (>=2) vs not.""" + return (np.asarray(grades) >= REFERABLE_THRESHOLD).astype(int) + + +def binary_scores(y_true_binary, y_pred_binary) -> BinaryScores: + y_true = np.asarray(y_true_binary, dtype=int) + y_pred = np.asarray(y_pred_binary, dtype=int) + + tp = int(((y_true == 1) & (y_pred == 1)).sum()) + fp = int(((y_true == 0) & (y_pred == 1)).sum()) + tn = int(((y_true == 0) & (y_pred == 0)).sum()) + fn = int(((y_true == 1) & (y_pred == 0)).sum()) + + def safe(num: int, den: int) -> float: + return float(num / den) if den else float("nan") + + return BinaryScores( + sensitivity=safe(tp, tp + fn), + specificity=safe(tn, tn + fp), + ppv=safe(tp, tp + fp), + npv=safe(tn, tn + fn), + tp=tp, + fp=fp, + tn=tn, + fn=fn, + ) + + +def choose_threshold_for_sensitivity( + y_true_binary, y_score, *, target_sensitivity: float = 0.90 +) -> float: + """Lowest-specificity-cost threshold meeting a sensitivity floor. + + **Call this on the VALIDATION set only.** The returned threshold is then + frozen and passed to `evaluate_at_threshold` on the test set. + + Screening is deliberately asymmetric: a missed referable case may cost + sight, a false positive costs one ophthalmologist review. So we fix + sensitivity and accept whatever specificity follows, rather than optimising + a symmetric metric like accuracy or Youden's J. + """ + y_true = np.asarray(y_true_binary, dtype=int) + y_score = np.asarray(y_score, dtype=float) + + best_threshold, best_specificity = 0.0, -1.0 + for threshold in np.unique(np.concatenate([y_score, [0.0, 1.0]])): + scores = binary_scores(y_true, (y_score >= threshold).astype(int)) + if scores.sensitivity >= target_sensitivity and scores.specificity > best_specificity: + best_threshold, best_specificity = float(threshold), scores.specificity + return best_threshold + + +def sensitivity_at_specificity_floor(y_true_binary, y_score, *, spec_floor: float = 0.85) -> float: + """Best sensitivity achievable while keeping specificity >= spec_floor. + + Why this exists + ---------------- + Plain sensitivity, monitored alone with mode="max", is gameable by a + single degenerate trick: predict referable for everyone. Checked directly + against this project's own training logs, `val/sensitivity_referable` + peaks at the exact epoch every other analysis in Phase 3 identified as the + grade-2 collapse (3 of 4 CE runs checked; the fourth misses by one epoch). + At that epoch sensitivity hits ~1.00 because the model refers almost every + image, not because it discriminates well -- specificity there is ~0.80, + barely above chance for the majority class. See + docs/07_PHASE3_RESULTS.md for the measurement. + + Result 4 already showed a *different* naive fix (macro-recall selection) + fails for the same underlying reason: a single-sided metric cannot + distinguish genuine improvement from a degenerate shortcut. + + This function instead operationalises the project's own stated target + (docs/01_PROJECT_ANALYSIS.md: sensitivity >90%, specificity >85%) directly: + among all thresholds achieving specificity >= spec_floor on this data, + return the highest sensitivity. If no threshold clears the floor, return + `best_specificity - 1.0` -- a value guaranteed to be negative, so any + epoch that DOES clear the floor always outranks one that does not, and + epochs that all fail the floor are still ordered by how close they came. + + **Selection use only.** As with every other `val/*` metric monitored + during training, this sweeps thresholds on the validation split purely to + choose which epoch's weights to keep -- the same role val/qwk already + plays. It is not the reported clinical operating point; that threshold is + chosen once, on validation, and frozen before scoring the locked test set + in `scripts/evaluate.py`, per the module-level docstring above. + """ + y_true = np.asarray(y_true_binary, dtype=int) + y_score = np.asarray(y_score, dtype=float) + + best_sensitivity = -1.0 + best_specificity_overall = -1.0 + for threshold in np.unique(np.concatenate([y_score, [0.0, 1.0]])): + scores = binary_scores(y_true, (y_score >= threshold).astype(int)) + if scores.specificity > best_specificity_overall: + best_specificity_overall = scores.specificity + if scores.specificity >= spec_floor and scores.sensitivity > best_sensitivity: + best_sensitivity = scores.sensitivity + + if best_sensitivity >= 0.0: + return float(best_sensitivity) + return float(best_specificity_overall - 1.0) + + +def youden_j(y_true_binary, y_score, threshold: float = 0.5) -> float: + """Youden's J = sensitivity + specificity - 1, at a fixed threshold. + + Logged alongside sensitivity_at_specificity_floor purely for comparison in + the metrics CSV -- it is the standard symmetric alternative, included so a + later reader can see both without re-deriving them. Not used for selection + here: unlike the floor-constrained metric, J does not target this + project's stated specificity requirement and can settle on a point that + satisfies neither the sensitivity nor the specificity bar. + """ + scores = binary_scores(y_true_binary, (np.asarray(y_score) >= threshold).astype(int)) + return float(scores.sensitivity + scores.specificity - 1.0) + + +def evaluate_at_threshold(y_true_binary, y_score, threshold: float) -> BinaryScores: + """Apply a threshold frozen from validation. No tuning happens here.""" + return binary_scores(y_true_binary, (np.asarray(y_score) >= threshold).astype(int)) + + +def bootstrap_ci( + metric_fn, + *arrays, + n_resamples: int = 2000, + confidence: float = 0.95, + seed: int = 42, +) -> tuple[float, float, float]: + """Percentile bootstrap CI for any metric over paired arrays. + + Returns (point_estimate, lower, upper). Resamples cases (not predictions), + which is the correct unit -- the uncertainty being quantified is "which + patients happened to be in this test set". + """ + arrays = [np.asarray(a) for a in arrays] + n = len(arrays[0]) + if any(len(a) != n for a in arrays): + raise ValueError("all arrays must have the same length") + + rng = np.random.default_rng(seed) + point = float(metric_fn(*arrays)) + + stats = np.empty(n_resamples, dtype=float) + for b in range(n_resamples): + idx = rng.integers(0, n, size=n) + stats[b] = metric_fn(*[a[idx] for a in arrays]) + + stats = stats[np.isfinite(stats)] + if stats.size == 0: + return point, float("nan"), float("nan") + + alpha = (1.0 - confidence) / 2.0 + return point, float(np.quantile(stats, alpha)), float(np.quantile(stats, 1 - alpha)) + + +@dataclass(frozen=True) +class DeLongResult: + auc_a: float + auc_b: float + z: float + p: float + + +def _delong_structural_components( + pos_scores: np.ndarray, neg_scores: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Placement values per DeLong et al. (1988) / Sun & Xu (2014): psi(x, y) = + 1 if x>y, 0.5 if x==y, 0 if x 0, 1.0, np.where(diff == 0, 0.5, 0.0)) + return psi.mean(axis=1), psi.mean(axis=0) + + +def delong_roc_variance(y_true_binary, y_score) -> tuple[float, float]: + """AUC and its DeLong variance for one classifier. Mainly a building block + for `delong_test`; exposed because reporting AUC's own CI (not just a + difference) is sometimes wanted on its own.""" + y_true = np.asarray(y_true_binary, dtype=int) + y_score = np.asarray(y_score, dtype=float) + v10, v01 = _delong_structural_components(y_score[y_true == 1], y_score[y_true == 0]) + auc = float(v10.mean()) + var = float(v10.var(ddof=1) / len(v10) + v01.var(ddof=1) / len(v01)) + return auc, var + + +def delong_test(y_true_binary, score_a, score_b) -> DeLongResult: + """Paired DeLong test: are two classifiers' AUCs different, on the SAME + sample? (DeLong, DeLong & Clarke-Pearson, 1988; fast form: Sun & Xu, 2014.) + + Why paired, not two independent bootstrap CIs + ---------------------------------------------- + `score_a` and `score_b` are two models' scores for the *same* patients. + Comparing two independent-looking AUC confidence intervals for overlap is + not a significance test and is typically conservative: it ignores that + both models get the same cases right or wrong for the same reasons (the + same hard patients), which is exactly the correlation McNemar's test + exploits for paired sensitivity/specificity elsewhere in this project + (`scripts/compare.py`). DeLong's test is the AUC-scale analogue: it uses + the covariance between the two classifiers' per-patient placement values, + not just their marginal variances. + """ + y_true = np.asarray(y_true_binary, dtype=int) + score_a = np.asarray(score_a, dtype=float) + score_b = np.asarray(score_b, dtype=float) + pos, neg = y_true == 1, y_true == 0 + + v10_a, v01_a = _delong_structural_components(score_a[pos], score_a[neg]) + v10_b, v01_b = _delong_structural_components(score_b[pos], score_b[neg]) + auc_a, auc_b = float(v10_a.mean()), float(v10_b.mean()) + + s10 = np.cov(np.vstack([v10_a, v10_b])) + s01 = np.cov(np.vstack([v01_a, v01_b])) + s = s10 / len(v10_a) + s01 / len(v01_a) + + var_d = float(s[0, 0] + s[1, 1] - 2 * s[0, 1]) + if var_d <= 0: + # Degenerate only when both classifiers agree on every placement + # (identical scores, or one/zero cases in a class) -- there is no + # detectable difference to test. + return DeLongResult(auc_a, auc_b, 0.0, 1.0) + + from scipy import stats as scipy_stats + + z = (auc_a - auc_b) / np.sqrt(var_d) + p = float(2 * (1 - scipy_stats.norm.cdf(abs(z)))) + return DeLongResult(auc_a, auc_b, float(z), p) + + +def expected_calibration_error(y_true_binary, y_prob, *, n_bins: int = 15) -> float: + """ECE: mean |confidence - accuracy| over equal-width probability bins. + + Modern networks are systematically overconfident (Guo et al., 2017), so raw + softmax output is not a probability. This is the number that must fall after + temperature scaling in Phase 5. + """ + y_true = np.asarray(y_true_binary, dtype=int) + y_prob = np.asarray(y_prob, dtype=float) + + edges = np.linspace(0.0, 1.0, n_bins + 1) + ece = 0.0 + for lo, hi in zip(edges[:-1], edges[1:], strict=True): + mask = (y_prob > lo) & (y_prob <= hi) if lo > 0 else (y_prob >= lo) & (y_prob <= hi) + if not mask.any(): + continue + ece += mask.mean() * abs(y_prob[mask].mean() - y_true[mask].mean()) + return float(ece) diff --git a/src/drdetect/explain/__init__.py b/src/drdetect/explain/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/explain/cam_variants.py b/src/drdetect/explain/cam_variants.py new file mode 100644 index 0000000000000000000000000000000000000000..0085fe93e21c986f5dd4cf7e011b08e44c53d601 --- /dev/null +++ b/src/drdetect/explain/cam_variants.py @@ -0,0 +1,69 @@ +"""Four CAM methods, one interface, so they can be compared rather than assumed equivalent. + +docs/04_ROADMAP.md Phase 6 calls for Grad-CAM, Grad-CAM++, Score-CAM and +Eigen-CAM "side by side" -- comparing them requires calling them the same +way. Phase 2 shipped only Grad-CAM (drdetect.explain.gradcam); this module +adds the other three without duplicating the eval-mode/target-layer handling +already correct there. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from drdetect.explain.gradcam import default_target_layer + +__all__ = ["CAM_METHODS", "generate_cam_variant"] + +CAM_METHODS = ("gradcam", "gradcam++", "scorecam", "eigencam") + + +def _cam_class(name: str): + from pytorch_grad_cam import EigenCAM, GradCAM, GradCAMPlusPlus, ScoreCAM + + return { + "gradcam": GradCAM, + "gradcam++": GradCAMPlusPlus, + "scorecam": ScoreCAM, + "eigencam": EigenCAM, + }[name] + + +def generate_cam_variant( + method: str, + model: nn.Module, + input_tensor: torch.Tensor, + target_class: int, + target_layer: nn.Module | None = None, +): + """Grayscale heatmap in [0, 1] from the named method. + + Args: + method: one of CAM_METHODS. + target_class: explain the PREDICTED class -- see drdetect.explain.gradcam.generate_cam. + + EigenCAM and Score-CAM do not use gradients (Eigen-CAM is a pure + activation decomposition; Score-CAM perturbs the input and reads off + softmax scores), so they are computed under `torch.no_grad()`, unlike + Grad-CAM/Grad-CAM++ which require gradients to reach the target layer. + """ + from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget + + if method not in CAM_METHODS: + raise ValueError(f"unknown CAM method {method!r}; expected one of {CAM_METHODS}") + + layer = target_layer or default_target_layer(model) + cam_cls = _cam_class(method) + needs_grad = method in ("gradcam", "gradcam++") + + was_training = model.training + model.eval() + try: + with torch.set_grad_enabled(needs_grad), cam_cls(model=model, target_layers=[layer]) as cam: + grayscale_cam = cam( + input_tensor=input_tensor, targets=[ClassifierOutputTarget(target_class)] + ) + return grayscale_cam[0] + finally: + model.train(was_training) diff --git a/src/drdetect/explain/evidence.py b/src/drdetect/explain/evidence.py new file mode 100644 index 0000000000000000000000000000000000000000..adcdb4e935f12ff0ba714a50798b60f1b00ab2ee --- /dev/null +++ b/src/drdetect/explain/evidence.py @@ -0,0 +1,175 @@ +"""ICDR evidence table + templated natural-language rationale, and outline-based +lesion-overlay rendering for the final report (docs/04_ROADMAP.md, Phase 6). + +Reuses `drdetect.fusion.features`'s model-running internals (the same tiled +segmentation inference, OD/fovea localisation, and microaneurysm candidate +scoring already built and evaluated for Phase 5's fusion head) rather than +re-implementing lesion detection a second time -- this module only adds the +per-instance/per-quadrant structure a human-readable report needs, on top of +the aggregate scalar features Phase 5's fusion head consumes. +""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np + +from drdetect.fusion.features import ( + LOCALIZATION_SIZE, + LesionModels, + _predict_od_fovea, + _tile_predict, + lesion_centroids, +) +from drdetect.segmentation.microaneurysms import score_candidates +from drdetect.segmentation.quadrants import assign_quadrant + +__all__ = [ + "extract_lesion_evidence", + "render_lesion_overlay", + "build_icdr_rationale", + "LESION_COLORS_RGB", +] + +_SEGMENTED_LESIONS = ["hard_exudates", "soft_exudates", "haemorrhages"] +# RGB, since render_lesion_overlay draws on an RGB array throughout this module. +LESION_COLORS_RGB = { + "hard_exudates": (255, 215, 0), # amber + "soft_exudates": (0, 191, 255), # light blue + "haemorrhages": (255, 0, 0), # red + "microaneurysms": (255, 0, 255), # magenta +} + + +def extract_lesion_evidence(image_path: str | Path, models: LesionModels) -> dict: + """Per-lesion-type instance list (quadrant + size) and binary masks for + rendering, from the same models `drdetect.fusion.features` uses for the + fusion head's aggregate scalars.""" + image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR) + image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) + + od_xy, fovea_xy, _diameter = _predict_od_fovea(image_rgb, models) + scale_x = LOCALIZATION_SIZE[1] / image_rgb.shape[1] + scale_y = LOCALIZATION_SIZE[0] / image_rgb.shape[0] + + lesions: dict[str, dict] = {} + masks: dict[str, np.ndarray] = {} + for lesion_type in _SEGMENTED_LESIONS: + module = getattr(models, lesion_type) + prob_map = _tile_predict(module, image_rgb, models.device) + binary_mask = prob_map > 0.5 + masks[lesion_type] = binary_mask + centroids = lesion_centroids(prob_map, threshold=0.5) + instances = [ + {"quadrant": assign_quadrant((cx * scale_x, cy * scale_y), od_xy, fovea_xy)} + for cx, cy in centroids + ] + lesions[lesion_type] = { + "count": len(instances), + "area_pct": float(binary_mask.mean() * 100.0), + "instances": instances, + } + + 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"] > 0.95 + ma_instances = [ + { + "quadrant": assign_quadrant((c.cx * scale_x, c.cy * scale_y), od_xy, fovea_xy), + "xy": (float(c.cx), float(c.cy)), + } + for c, keep in zip(ma_result["candidates"], accepted, strict=True) + if keep + ] + lesions["microaneurysms"] = {"count": len(ma_instances), "instances": ma_instances} + + return {"od_xy": od_xy, "fovea_xy": fovea_xy, "lesions": lesions, "masks": masks} + + +def render_lesion_overlay(image_rgb: np.ndarray, evidence: dict) -> np.ndarray: + """Outline contours, not filled blobs -- the roadmap's own stated reason: + outlines beat blobs for clinical legibility (a filled overlay hides the + tissue underneath it; an outline points at it without obscuring it). + + Microaneurysms are marked with small circles rather than outlines -- + they're candidate *points* from the classifier (drdetect.segmentation. + microaneurysms), not a segmented region, so there's no contour to draw. + """ + overlay = image_rgb.copy() + for lesion_type, mask in evidence["masks"].items(): + color = LESION_COLORS_RGB[lesion_type] + contours, _ = cv2.findContours( + mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE + ) + cv2.drawContours(overlay, contours, -1, color, thickness=3) + + ma_color = LESION_COLORS_RGB["microaneurysms"] + for inst in evidence["lesions"].get("microaneurysms", {}).get("instances", []): + x, y = inst["xy"] + cv2.circle(overlay, (int(round(x)), int(round(y))), radius=12, color=ma_color, thickness=2) + + return overlay + + +_ICDR_GRADE_NAMES = { + 0: "No apparent DR", + 1: "Mild NPDR", + 2: "Moderate NPDR", + 3: "Severe NPDR", + 4: "Proliferative DR", +} + + +def build_icdr_rationale(grade: int, evidence: dict) -> str: + """A templated sentence grounding the predicted ICDR grade in the actual + detected evidence -- not a restatement of the grade, an explanation of + what supports it (or, honestly, what doesn't fully support it). + + Neovascularisation (the feature that actually defines grade 4, PDR) has + no segmentation model in this project (docs/01's own scope-honesty + note: "neovascularisation has no public pixel masks... document as + future work, do not fabricate it") -- a grade-4 rationale says so + explicitly rather than inventing NV evidence this pipeline cannot see. + """ + lesions = evidence["lesions"] + ma_n = lesions.get("microaneurysms", {}).get("count", 0) + he_n = lesions.get("haemorrhages", {}).get("count", 0) + ex_pct = lesions.get("hard_exudates", {}).get("area_pct", 0.0) + se_pct = lesions.get("soft_exudates", {}).get("area_pct", 0.0) + + quadrants_with_he = { + inst["quadrant"] for inst in lesions.get("haemorrhages", {}).get("instances", []) + } + + evidence_parts = [] + if ma_n: + evidence_parts.append(f"{ma_n} microaneurysm candidate(s)") + if he_n: + evidence_parts.append( + f"{he_n} haemorrhage region(s) across {len(quadrants_with_he)} quadrant(s)" + ) + if ex_pct > 0.01: + evidence_parts.append(f"hard exudates covering {ex_pct:.2f}% of the image") + if se_pct > 0.01: + evidence_parts.append( + f"soft exudates (cotton-wool spots) covering {se_pct:.2f}% of the image" + ) + + grade_name = _ICDR_GRADE_NAMES.get(grade, f"grade {grade}") + if not evidence_parts: + return f"ICDR grade {grade} ({grade_name}): no significant lesion evidence detected." + + body = "; ".join(evidence_parts) + note = "" + if grade == 4: + note = ( + " Note: proliferative DR (grade 4) is defined by neovascularisation, which this " + "pipeline does not detect (no public pixel masks exist to train against) -- this " + "grade is the classifier's own image-level prediction, not evidence-backed for the " + "specific feature that defines it." + ) + return f"ICDR grade {grade} ({grade_name}): {body}.{note}" diff --git a/src/drdetect/explain/gradcam.py b/src/drdetect/explain/gradcam.py new file mode 100644 index 0000000000000000000000000000000000000000..c6ac904b791c1dea60678d2b61aa116868441bc8 --- /dev/null +++ b/src/drdetect/explain/gradcam.py @@ -0,0 +1,97 @@ +"""Grad-CAM evidence overlay for the DR grader. + +Thin wrapper around `pytorch_grad_cam` rather than a hand-rolled hook, so the +CAM implementation matches a maintained, widely-checked reference rather than +a bespoke one. Phase 6 (docs/04_ROADMAP.md) is where CAM variants get compared +side by side and run through Adebayo's sanity checks -- this module is only +the single Grad-CAM used for Phase 2's end-to-end slice, not that comparison. +""" + +from __future__ import annotations + +import cv2 +import numpy as np +import torch +import torch.nn as nn + +__all__ = ["default_target_layer", "generate_cam", "overlay_cam"] + + +def default_target_layer(model: nn.Module) -> nn.Module: + """The last activation before global pooling, for a timm EfficientNet. + + `bn2`, not `conv_head`. timm's EfficientNet fuses BatchNorm+SiLU into a + single `BatchNormAct2d` module at `bn2`, so `conv_head`'s own output is + the bare pre-activation conv -- unbounded, and roughly half-negative in + practice. Grad-CAM/Grad-CAM++/Eigen-CAM mostly tolerate that (their + weights can themselves be negative and partially cancel it out), but + Score-CAM's weights are a softmax over class scores -- always positive -- + so weighting an already-negative activation map and then ReLU-ing the + result (`base_cam.py`'s `compute_cam_per_layer`) can zero it out + entirely. Measured directly: `conv_head` produced an all-zero Score-CAM + heatmap on a real image; `bn2`, whose output has already passed through + SiLU, did not. `bn2` is the correct choice for every method, not a + Score-CAM-specific workaround -- it's simply the last point where the + activation reflects what the network actually forward-propagates. + """ + backbone = getattr(model, "backbone", model) + if not hasattr(backbone, "bn2"): + raise AttributeError( + f"{type(backbone).__name__} has no `bn2`; pass target_layer explicitly." + ) + return backbone.bn2 + + +def generate_cam( + model: nn.Module, + input_tensor: torch.Tensor, + target_class: int, + target_layer: nn.Module | None = None, +) -> np.ndarray: + """Grad-CAM heatmap for one image, as a single HxW array in [0, 1]. + + Args: + model: the grader, in eval mode. Frozen BatchNorm does not interfere + -- Grad-CAM only needs gradients w.r.t. activations, never an + optimiser step. + input_tensor: 1xCxHxW, already normalised exactly as at inference. + target_class: which class's score to explain. For a grading report + this should be the PREDICTED class, not the true label -- the + report is explaining what the model said, not grading the CAM. + target_layer: defaults to `default_target_layer(model)`. + """ + from pytorch_grad_cam import GradCAM + from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget + + layer = target_layer or default_target_layer(model) + was_training = model.training + model.eval() + try: + with GradCAM(model=model, target_layers=[layer]) as cam: + grayscale_cam = cam( + input_tensor=input_tensor, targets=[ClassifierOutputTarget(target_class)] + ) + return grayscale_cam[0] + finally: + model.train(was_training) + + +def overlay_cam( + rgb_image: np.ndarray, grayscale_cam: np.ndarray, alpha: float = 0.45 +) -> np.ndarray: + """Blend a Grad-CAM heatmap onto its source image. + + Args: + rgb_image: HxWx3 uint8, the exact image the CAM was computed on + (post-preprocessing) -- overlaying on a differently-cropped raw + image would misalign the heatmap with the lesions it is marking. + grayscale_cam: output of `generate_cam`, same H and W as `rgb_image`. + alpha: heatmap opacity. + """ + from pytorch_grad_cam.utils.image import show_cam_on_image + + if grayscale_cam.shape[:2] != rgb_image.shape[:2]: + grayscale_cam = cv2.resize(grayscale_cam, (rgb_image.shape[1], rgb_image.shape[0])) + float_rgb = rgb_image.astype(np.float32) / 255.0 + overlay = show_cam_on_image(float_rgb, grayscale_cam, use_rgb=True, image_weight=1 - alpha) + return overlay diff --git a/src/drdetect/explain/localization_metrics.py b/src/drdetect/explain/localization_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..ff5a7103dd26ce577c71a55986e19025b2f499f7 --- /dev/null +++ b/src/drdetect/explain/localization_metrics.py @@ -0,0 +1,55 @@ +"""CAM-vs-ground-truth localisation quality: pointing game + IoU. + +Phase 6's own exit criterion (docs/04_ROADMAP.md): "a table of saliency +methods x (sanity-check pass/fail, pointing-game accuracy, IoU)". This +module computes the pointing-game and IoU halves of that table, against +IDRiD's real per-lesion-type pixel masks -- the sanity-check half already +exists in drdetect.explain.sanity_checks (docs/08_PHASE6_RESULTS.md). +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +__all__ = ["upsample_cam", "pointing_game_hit", "cam_mask_iou"] + + +def upsample_cam(cam: np.ndarray, target_shape: tuple[int, int]) -> np.ndarray: + """Resize a CAM (native resolution: the model's input size, e.g. 512x512) + up to `target_shape` (height, width) -- typically a lesion mask's own + full resolution. A CAM is already a smooth, interpolated heatmap, so + upsampling it preserves the mask's own pixel precision, rather than + downsampling the mask and losing the smallest lesions -- microaneurysms + are already known, from this project's own Phase 3 findings, to vanish + below ~1024px.""" + h, w = target_shape + return cv2.resize(cam, (w, h), interpolation=cv2.INTER_LINEAR) + + +def pointing_game_hit(cam: np.ndarray, mask: np.ndarray) -> bool: + """True if the CAM's single point of maximum attention falls on a + positive mask pixel -- the standard pointing-game definition (Zhang et + al. 2018), threshold-free by construction.""" + y, x = np.unravel_index(np.argmax(cam), cam.shape) + return bool(mask[y, x] > 0) + + +def cam_mask_iou(cam: np.ndarray, mask: np.ndarray, threshold: float = 0.5) -> float: + """IoU between the CAM thresholded at a fixed value and the binary + lesion mask. + + Expect this to read low for genuinely sparse lesion types + (haemorrhages and microaneurysms are well under 0.1% of image pixels, + see drdetect.segmentation.metrics) even for a CAM correctly pointing at + the right general region -- Grad-CAM produces a smooth, diffuse + attention region by construction, not a pixel-precise segmentation, so + IoU against a sparse mask is a much harder bar than the pointing game. + NaN, not 0, when the union is empty (both CAM-positive and mask are + empty) -- 0 would misrepresent "nothing to compare" as "total miss". + """ + binary_cam = cam > threshold + binary_mask = mask.astype(bool) + intersection = np.logical_and(binary_cam, binary_mask).sum() + union = np.logical_or(binary_cam, binary_mask).sum() + return float(intersection / union) if union > 0 else float("nan") diff --git a/src/drdetect/explain/sanity_checks.py b/src/drdetect/explain/sanity_checks.py new file mode 100644 index 0000000000000000000000000000000000000000..2619e8ffbac5b6a93cb6749b662c981c22692479 --- /dev/null +++ b/src/drdetect/explain/sanity_checks.py @@ -0,0 +1,157 @@ +"""Adebayo-style sanity checks for saliency methods. + +Adebayo et al. 2018 (NeurIPS), "Sanity Checks for Saliency Maps": several +widely-used saliency methods are empirically insensitive to the model's +learned parameters -- they produce nearly the same map whether the model is +trained or has its weights replaced with noise. A method that fails this is +not explaining the MODEL; it behaves closer to an edge detector applied to +the input, dressed up as an explanation. Nobody knew whether that applied +here until it was checked -- this project's own culture treats an unchecked +assumption as no different from a wrong one. + +Two checks, both from the paper: + + Model parameter randomisation (`model_randomization_test`) -- cascading: + progressively randomise the model from the OUTPUT layer backward + toward the input, and measure how much each CAM method's heatmap + changes at each step, against its own original heatmap. A method that + PASSES should diverge quickly; a method whose heatmap barely moves as + the whole model turns to noise is not using what the model learned. + + Data randomisation (`compare_against_shuffled_labels`) -- compare CAMs + from the real trained model against CAMs from an otherwise-identical + model trained on the SAME images with shuffled labels. A method that + PASSES should look different between the two, since the two models + learned genuinely different things; a method that looks similar + regardless of what the underlying model actually learned is not + sensitive to that. + +Similarity is reported as both Spearman rank correlation and SSIM, matching +the two metrics the original paper used -- rank correlation alone can be +fooled by a saliency map that reorders pixels similarly but at a different +spatial scale, which SSIM is sensitive to and rank correlation is not. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn as nn + +from drdetect.explain.cam_variants import CAM_METHODS, generate_cam_variant + +__all__ = [ + "CascadeStep", + "cam_similarity", + "cascading_randomization_layers", + "model_randomization_test", + "compare_against_shuffled_labels", +] + + +@dataclass(frozen=True) +class CascadeStep: + layer_randomized: str + spearman: dict[str, float] + ssim: dict[str, float] + + +def cam_similarity(cam_a: np.ndarray, cam_b: np.ndarray) -> tuple[float, float]: + """(Spearman rank correlation, SSIM) between two same-shape grayscale CAMs.""" + from scipy.stats import spearmanr + from skimage.metrics import structural_similarity + + rho = float(spearmanr(cam_a.ravel(), cam_b.ravel()).statistic) + ssim = float(structural_similarity(cam_a, cam_b, data_range=1.0)) + return rho, ssim + + +def cascading_randomization_layers(backbone: nn.Module) -> list[tuple[str, list[nn.Module]]]: + """Named layer groups of a timm EfficientNet backbone, OUTPUT-to-INPUT order. + + This order is the point: Adebayo's cascading test randomises from the + output backward, one layer group at a time, so that at step k the model + is "the real model, except everything from the output back through + group k is now noise." Reversed (input-to-output) would be a different, + weaker test -- early layers dominate less of what a CAM at `bn2` + (drdetect.explain.gradcam.default_target_layer) actually reflects. + """ + layers = [ + ("classifier", [backbone.classifier]), + ("conv_head+bn2", [backbone.conv_head, backbone.bn2]), + ] + for i in reversed(range(len(backbone.blocks))): + layers.append((f"blocks[{i}]", [backbone.blocks[i]])) + layers.append(("conv_stem+bn1", [backbone.conv_stem, backbone.bn1])) + return layers + + +def _reset_in_place(modules: list[nn.Module]) -> None: + for module in modules: + for child in module.modules(): + if hasattr(child, "reset_parameters"): + child.reset_parameters() + + +def model_randomization_test( + model: nn.Module, + input_tensor: torch.Tensor, + target_class: int, + methods: tuple[str, ...] = CAM_METHODS, +) -> list[CascadeStep]: + """Cascading model-randomisation sanity check across every CAM method at once. + + Works on a deep copy -- the caller's trained model is never mutated. + Every method's CAM is computed against the SAME progressively-randomised + model at each step, so the comparison at a given step is apples to apples + across methods, not just against each method's own baseline. + """ + randomized = copy.deepcopy(model) + backbone = getattr(randomized, "backbone", randomized) + layer_groups = cascading_randomization_layers(backbone) + + originals = { + m: generate_cam_variant(m, randomized, input_tensor, target_class) for m in methods + } + + steps = [] + for name, modules in layer_groups: + _reset_in_place(modules) + spearman, ssim = {}, {} + for method in methods: + cam = generate_cam_variant(method, randomized, input_tensor, target_class) + rho, s = cam_similarity(originals[method], cam) + spearman[method] = rho + ssim[method] = s + steps.append(CascadeStep(layer_randomized=name, spearman=spearman, ssim=ssim)) + return steps + + +def compare_against_shuffled_labels( + real_model: nn.Module, + shuffled_model: nn.Module, + input_tensor: torch.Tensor, + target_class_real: int, + target_class_shuffled: int, + methods: tuple[str, ...] = CAM_METHODS, +) -> dict[str, tuple[float, float]]: + """Data-randomisation check: CAM(real model) vs CAM(shuffled-label model), same image. + + Args: + target_class_real / target_class_shuffled: each model's OWN predicted + class for this image -- the two models were trained on different + label associations and will usually disagree, and each CAM must + explain what its own model actually predicted. + + Returns: {method: (spearman, ssim)}. + """ + return { + method: cam_similarity( + generate_cam_variant(method, real_model, input_tensor, target_class_real), + generate_cam_variant(method, shuffled_model, input_tensor, target_class_shuffled), + ) + for method in methods + } diff --git a/src/drdetect/fusion/__init__.py b/src/drdetect/fusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd38b51a663fedb8d875871989d527e55ad9084d --- /dev/null +++ b/src/drdetect/fusion/__init__.py @@ -0,0 +1 @@ +"""Phase 5: fusing the grading model's CNN embedding with Phase 4's lesion evidence.""" diff --git a/src/drdetect/fusion/embedding.py b/src/drdetect/fusion/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..4e674e63381e4b58793cd66ca6def840aff8bcc2 --- /dev/null +++ b/src/drdetect/fusion/embedding.py @@ -0,0 +1,75 @@ +"""CNN embedding extraction from the canonical grading checkpoint, for +Phase 5's fusion head. + +Uses `baseline_effb0_512_fold0` (QWK 0.8930), not the 1024px checkpoint -- +Phase 3's own conclusion was "the baseline stands" (docs/07_PHASE3_RESULTS.md), +and Result 14 retracted the 1024px checkpoint's causal story (its edge over +the baseline is now suspected to be a backend artifact, not resolution). +Building the fusion head on a result the project already walked back would +undercut that finding. + +Reuses the checkpoint-loading convention scripts/evaluate.py and +scripts/calibrate.py already established (strip the Lightning `model.` +prefix, load into a bare `FundusGrader`) rather than inventing a new one. + +**Images come from the Ben-Graham-preprocessed cache +(`data/processed/aptos_512/`), not the raw APTOS files** -- unlike +`drdetect.fusion.features`, which deliberately uses raw images because the +Phase 4 lesion models were trained on IDRiD's un-enhanced originals. The +grading model was trained on the preprocessed cache, so its embedding is +only meaningful computed on the same distribution it actually learned from. +Two different "the right image for this model" answers for two different +models feeding the same fusion vector -- not an inconsistency. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from drdetect.data.dataset import build_transforms +from drdetect.grading.model import FundusGrader, build_model + +__all__ = ["load_grading_backbone", "extract_cnn_embedding", "EMBEDDING_DIM"] + +EMBEDDING_DIM = 1280 # efficientnet_b0's pre-classifier feature width + + +def load_grading_backbone( + checkpoint: str, *, backbone: str = "efficientnet_b0", device: torch.device +) -> FundusGrader: + model = build_model(backbone, num_outputs=5, pretrained=False, freeze_bn=True) + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) + state = ckpt.get("state_dict", ckpt) + state = {k.removeprefix("model."): v for k, v in state.items() if k.startswith("model.")} + model.load_state_dict(state) + return model.to(device).eval() + + +def extract_cnn_embedding( + image_path: str | Path, model: FundusGrader, device: torch.device, *, size: int = 512 +) -> torch.Tensor: + """The 1280-d pre-classifier pooled feature (`forward_head(..., + pre_logits=True)`), not the 5-way logits `model(x)` would give.""" + embedding, _logits = extract_cnn_embedding_and_logits(image_path, model, device, size=size) + return embedding + + +def extract_cnn_embedding_and_logits( + image_path: str | Path, model: FundusGrader, device: torch.device, *, size: int = 512 +) -> tuple[torch.Tensor, torch.Tensor]: + """Both the fusion head's input (the embedding) and the baseline + grading-only prediction (the 5-way logits) from one forward pass -- + used when comparing fusion against grading-alone on the same images, so + the baseline's own prediction doesn't need a second pass over the model.""" + import cv2 + + transform = build_transforms(size, train=False) + image = cv2.cvtColor(cv2.imread(str(image_path), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) + tensor = transform(image=image)["image"].unsqueeze(0).to(device) + with torch.no_grad(): + features = model.backbone.forward_features(tensor) + embedding = model.backbone.forward_head(features, pre_logits=True) + logits = model.backbone.forward_head(features) + return embedding.squeeze(0).cpu(), logits.squeeze(0).cpu() diff --git a/src/drdetect/fusion/features.py b/src/drdetect/fusion/features.py new file mode 100644 index 0000000000000000000000000000000000000000..e25acba326b81c6bacfeb6523c8cccba0962a70b --- /dev/null +++ b/src/drdetect/fusion/features.py @@ -0,0 +1,260 @@ +"""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}, + } diff --git a/src/drdetect/fusion/model.py b/src/drdetect/fusion/model.py new file mode 100644 index 0000000000000000000000000000000000000000..e92cc4da4f743ac5ae101f363689d19b149c2c20 --- /dev/null +++ b/src/drdetect/fusion/model.py @@ -0,0 +1,35 @@ +"""The fusion head itself: concat(CNN embedding, lesion features) -> ordinal head. + +A small MLP, not another CNN -- the input is already a fixed-length vector +(the grading backbone's pooled embedding plus a handful of scalar lesion +features), not an image needing spatial processing. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +__all__ = ["FusionHead"] + + +class FusionHead(nn.Module): + def __init__( + self, + *, + embedding_dim: int, + lesion_dim: int, + hidden_dim: int = 128, + num_outputs: int = 4, + dropout: float = 0.3, + ): + super().__init__() + self.net = nn.Sequential( + nn.Linear(embedding_dim + lesion_dim, hidden_dim), + nn.ReLU(inplace=True), + nn.Dropout(dropout), + nn.Linear(hidden_dim, num_outputs), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) diff --git a/src/drdetect/grading/__init__.py b/src/drdetect/grading/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/grading/losses.py b/src/drdetect/grading/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..3c72649193f5faa749d3fda4f4ded0edb252abf3 --- /dev/null +++ b/src/drdetect/grading/losses.py @@ -0,0 +1,330 @@ +"""Ordinal losses for DR grading. + +Why cross-entropy is the wrong loss here +---------------------------------------- +DR grades are ordered: 0 < 1 < 2 < 3 < 4. Cross-entropy treats them as +unordered labels, so predicting 4 for a true grade 1 costs exactly as much as +predicting 2. The loss-minimising strategy under that objective is to hedge +toward the majority middle class. + +The Phase 1 baseline did precisely that. From its confusion matrix, errors +collapsed onto grade 2: 29 of 74 grade-1 cases, 21 of 39 grade-3, and 22 of 59 +grade-4 were all predicted Moderate. See docs/06_PHASE1_RESULTS.md. + +Three alternatives are provided. + +`CornLoss` -- rank-consistent ordinal regression (Shi, Cao & Raschka, 2023). + The model emits K-1 logits; logit j models P(y > j | y > j-1). Because the + probabilities chain multiplicatively, predictions are rank-monotonic *by + construction*: the model cannot output "probably worse than grade 2 but not + worse than grade 1". Each task trains only on the conditional subset, which + is what distinguishes CORN from the earlier CORAL and removes its shared-bias + restriction. + +`OrdinalRegressionLoss` -- single continuous output, thresholds fitted after + training. This is the approach behind most winning Kaggle DR solutions. It + optimises QWK almost directly, at the cost of a post-hoc threshold search. + +`DistanceWeightedCE` -- ordinary cross-entropy with soft targets whose mass + decays with ordinal distance. The least invasive option: it keeps the K-way + head and only changes the targets. +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = [ + "CornLoss", + "decode_output", + "naive_referable_cut", + "OrdinalRegressionLoss", + "DistanceWeightedCE", + "corn_probabilities", + "corn_predict", + "corn_task_pos_weights", + "regression_predict", + "fit_thresholds", + "build_loss", + "outputs_for_loss", +] + + +class CornLoss(nn.Module): + """Conditional ordinal regression, rank-consistent by construction. + + Expects `num_classes - 1` logits. Task j is trained only on samples with + y >= j, so each binary problem is genuinely conditional. + """ + + def __init__( + self, + num_classes: int = 5, + class_weights: torch.Tensor | None = None, + task_pos_weights: torch.Tensor | None = None, + ): + super().__init__() + self.num_classes = num_classes + self.register_buffer( + "class_weights", + class_weights if class_weights is not None else torch.ones(num_classes), + ) + self.register_buffer( + "task_pos_weights", + task_pos_weights if task_pos_weights is not None else torch.ones(num_classes - 1), + ) + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + if logits.shape[1] != self.num_classes - 1: + raise ValueError( + f"CornLoss expects {self.num_classes - 1} logits, got {logits.shape[1]}. " + f"Build the model with num_outputs={self.num_classes - 1}." + ) + + losses, total = [], 0.0 + for j in range(self.num_classes - 1): + subset = targets >= j # conditional on having exceeded j-1 + if not subset.any(): + continue + binary = (targets[subset] > j).float() + w = self.class_weights[targets[subset]] + bce = F.binary_cross_entropy_with_logits( + logits[subset, j], + binary, + weight=w, + pos_weight=self.task_pos_weights[j], + reduction="sum", + ) + losses.append(bce) + total += float(w.sum()) + + if not losses: + return logits.sum() * 0.0 # keeps the graph connected + return torch.stack(losses).sum() / max(total, 1.0) + + +def corn_task_pos_weights(labels, num_classes: int = 5) -> list[float]: + """Per-task positive weights that rebalance CORN's conditional subsets. + + CORN trains task j only on samples with y >= j, asking whether y > j. Those + subsets inherit the label distribution and are usually NOT balanced. Measured + on the APTOS training split: + + task j=0 2929 samples, 50.7% positive balanced + task j=1 1485 samples, 80.1% positive badly skewed + task j=2 1189 samples, 32.8% positive skewed + task j=3 390 samples, 60.5% positive balanced + + Task j=1 asks "given grade >= 1, is it worse than 1?" and is 80% positive, + because grade 1 is only 296 of 1485 samples with grade >= 1. Unweighted, the + loss-minimising answer is almost always yes, which pushes grade-1 cases into + grade 2. Measured consequence: unweighted CORN sent 49 of 74 grade-1 + validation cases to grade 2, against 29 for plain cross-entropy, and grade-1 + recall fell from 0.541 to 0.297. + + Returns n_negative / n_positive per task, the standard BCE `pos_weight`, + which makes each conditional task cost-balanced. + """ + y = np.asarray(labels) + weights = [] + for j in range(num_classes - 1): + subset = y[y >= j] + pos = int((subset > j).sum()) + neg = int(len(subset) - pos) + weights.append(float(neg / pos) if pos > 0 else 1.0) + return weights + + +def corn_probabilities(logits: torch.Tensor) -> torch.Tensor: + """Cumulative probabilities P(y > j) from CORN logits, via the chain rule. + + The cumulative product is what guarantees monotonicity: P(y>j) can never + exceed P(y>j-1), because it is that value multiplied by a sigmoid. + """ + return torch.cumprod(torch.sigmoid(logits), dim=1) + + +def corn_predict(logits: torch.Tensor) -> torch.Tensor: + """Predicted grade = number of thresholds exceeded.""" + return (corn_probabilities(logits) > 0.5).sum(dim=1) + + +def corn_class_probabilities(logits: torch.Tensor) -> torch.Tensor: + """Per-class probabilities from cumulative ones, for referable scoring. + + P(y=k) = P(y>k-1) - P(y>k), with P(y>-1)=1 and P(y>K-1)=0. + """ + cum = corn_probabilities(logits) + ones = torch.ones_like(cum[:, :1]) + zeros = torch.zeros_like(cum[:, :1]) + upper = torch.cat([ones, cum], dim=1) + lower = torch.cat([cum, zeros], dim=1) + return (upper - lower).clamp(min=0.0) + + +class OrdinalRegressionLoss(nn.Module): + """Single-output regression on the grade, with SmoothL1. + + SmoothL1 rather than MSE: MSE lets a handful of grade-0-predicted-as-4 + outliers dominate the gradient, which at batch 4 is destabilising. + """ + + def __init__(self, beta: float = 1.0): + super().__init__() + self.beta = beta + + def forward(self, output: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + return F.smooth_l1_loss(output.squeeze(-1), targets.float(), beta=self.beta) + + +def regression_predict(output: torch.Tensor, thresholds: list[float] | None = None) -> torch.Tensor: + """Discretise a continuous prediction using cut-points.""" + if thresholds is None: + thresholds = [0.5, 1.5, 2.5, 3.5] + x = output.squeeze(-1) + pred = torch.zeros_like(x, dtype=torch.long) + for t in thresholds: + pred += (x > t).long() + return pred + + +def fit_thresholds(predictions, targets, *, n_classes: int = 5, n_rounds: int = 8) -> list[float]: + """Coordinate-ascent search for cut-points maximising QWK. + + **Fit on validation, never on test.** The returned thresholds are part of the + model and must be frozen alongside the weights. + """ + from drdetect.eval.metrics import quadratic_weighted_kappa + + preds = np.asarray(predictions, dtype=float).ravel() + targs = np.asarray(targets, dtype=int).ravel() + thresholds = [i + 0.5 for i in range(n_classes - 1)] + + def score(th: list[float]) -> float: + digitised = np.zeros_like(preds, dtype=int) + for t in th: + digitised += (preds > t).astype(int) + return quadratic_weighted_kappa(targs, digitised, n_classes=n_classes) + + best = score(thresholds) + for _ in range(n_rounds): + improved = False + for i in range(len(thresholds)): + lo = thresholds[i - 1] + 0.05 if i > 0 else preds.min() - 0.5 + hi = thresholds[i + 1] - 0.05 if i < len(thresholds) - 1 else preds.max() + 0.5 + if hi <= lo: + continue + for candidate in np.linspace(lo, hi, 40): + trial = list(thresholds) + trial[i] = float(candidate) + s = score(trial) + if s > best: + best, thresholds, improved = s, trial, True + if not improved: + break + return thresholds + + +class DistanceWeightedCE(nn.Module): + """Cross-entropy against soft targets that decay with ordinal distance. + + Target mass for class k given true grade y is proportional to + exp(-|k - y| / temperature), so predicting an adjacent grade is penalised + far less than predicting a distant one -- the property plain CE lacks. + """ + + def __init__( + self, + num_classes: int = 5, + temperature: float = 0.7, + class_weights: torch.Tensor | None = None, + ): + super().__init__() + self.num_classes = num_classes + self.temperature = temperature + grid = torch.arange(num_classes).float() + soft = torch.exp(-(grid[None, :] - grid[:, None]).abs() / temperature) + self.register_buffer("soft_targets", soft / soft.sum(dim=1, keepdim=True)) + self.register_buffer( + "class_weights", + class_weights if class_weights is not None else torch.ones(num_classes), + ) + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + soft = self.soft_targets[targets] + per_sample = -(soft * F.log_softmax(logits, dim=1)).sum(dim=1) + w = self.class_weights[targets] + return (per_sample * w).sum() / w.sum().clamp(min=1e-8) + + +def outputs_for_loss(loss_name: str, num_classes: int = 5) -> int: + """How many head outputs a given loss requires.""" + return { + "ce": num_classes, + "distance_ce": num_classes, + "corn": num_classes - 1, + "regression": 1, + }[loss_name] + + +def build_loss( + name: str, + *, + num_classes: int = 5, + class_weights: list[float] | None = None, + task_pos_weights: list[float] | None = None, +): + w = torch.tensor(class_weights, dtype=torch.float) if class_weights else None + if name == "ce": + return nn.CrossEntropyLoss(weight=w) + if name == "corn": + tw = ( + torch.tensor(task_pos_weights, dtype=torch.float) + if task_pos_weights is not None + else None + ) + return CornLoss(num_classes, class_weights=w, task_pos_weights=tw) + if name == "regression": + return OrdinalRegressionLoss() + if name == "distance_ce": + return DistanceWeightedCE(num_classes, class_weights=w) + raise ValueError(f"unknown loss {name!r}; expected ce, corn, regression or distance_ce") + + +def decode_output(output: torch.Tensor, loss_name: str) -> tuple[np.ndarray, np.ndarray]: + """Map a raw head output to (predicted grade, referable score). + + Single source of truth for decoding, shared by the training module and the + evaluation script. Duplicating this logic once caused an evaluation to load a + 5-output head against a 4-output CORN checkpoint; keeping one implementation + makes that class of mismatch impossible. + + Each loss parameterises the head differently: + ce / distance_ce -- K-way softmax; referable = sum of P(class >= 2) + corn -- K-1 chained sigmoids; referable = P(y > 1) directly + regression -- one continuous value, which itself ranks severity + + The referable score need only be monotone in severity; threshold selection + handles the scale. + """ + if loss_name == "corn": + cum = corn_probabilities(output) # P(y > j), monotone by construction + return (cum > 0.5).sum(dim=1).numpy(), cum[:, 1].numpy() + + if loss_name == "regression": + return regression_predict(output).numpy(), output.squeeze(-1).numpy() + + probs = torch.softmax(output, dim=1).numpy() + return probs.argmax(axis=1), probs[:, 2:].sum(axis=1) + + +def naive_referable_cut(loss_name: str) -> float: + """The 'no calibration' cut point, whose scale depends on the head. + + Only a progress signal during training -- the reported operating point is + chosen for target sensitivity in scripts/evaluate.py. + """ + return 1.5 if loss_name == "regression" else 0.5 diff --git a/src/drdetect/grading/model.py b/src/drdetect/grading/model.py new file mode 100644 index 0000000000000000000000000000000000000000..c6667bf911482ebf8ddf08d00da5b357db49ec8f --- /dev/null +++ b/src/drdetect/grading/model.py @@ -0,0 +1,92 @@ +"""DR grading backbones. + +Two decisions here come from measurements rather than convention. + +Frozen BatchNorm +---------------- +Benchmarking on this project's target machine showed batch size is capped at 4 +by available RAM (docs/05_PROTOTYPE_SCOPE.md section 6). EfficientNet uses +BatchNorm, whose running statistics are estimated per micro-batch -- at batch 4 +they are noisy, and gradient accumulation does NOT help, because accumulation +batches the optimiser step, not the normalisation. Freezing BN to use the +pretrained ImageNet statistics is the standard fix for small-batch fine-tuning +and is the default here. + +Ordinal-ready head +------------------ +Phase 1 is a plain cross-entropy baseline, deliberately. But DR grades are +ordinal, and Phase 3 replaces the loss. `num_outputs` is therefore separated +from `num_classes` so the same backbone serves a 5-way softmax, a 1-output +regression head, or a 4-output CORAL head without rewriting the model. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +__all__ = ["build_model", "freeze_batchnorm", "count_parameters"] + + +def freeze_batchnorm(model: nn.Module) -> int: + """Put every BatchNorm into eval mode and stop its affine parameters training. + + Returns the number of layers frozen, for logging -- a silent no-op here + would be indistinguishable from a working call. + """ + frozen = 0 + for module in model.modules(): + if isinstance(module, nn.modules.batchnorm._BatchNorm): + module.eval() + module.weight.requires_grad_(False) + module.bias.requires_grad_(False) + frozen += 1 + return frozen + + +class FundusGrader(nn.Module): + def __init__(self, backbone: nn.Module, freeze_bn: bool = True): + super().__init__() + self.backbone = backbone + self.freeze_bn = freeze_bn + if freeze_bn: + self.n_frozen_bn = freeze_batchnorm(self.backbone) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.backbone(x) + + def train(self, mode: bool = True): + """Keep frozen BN in eval mode even when the module is set to train. + + Without this override, `model.train()` silently re-enables running- + statistic updates and undoes the freeze -- a bug that produces no error + and only shows up as unstable validation metrics. + """ + super().train(mode) + if self.freeze_bn: + for module in self.backbone.modules(): + if isinstance(module, nn.modules.batchnorm._BatchNorm): + module.eval() + return self + + +def build_model( + name: str = "efficientnet_b0", + *, + num_outputs: int = 5, + pretrained: bool = True, + freeze_bn: bool = True, + drop_rate: float = 0.2, +) -> FundusGrader: + import timm + + backbone = timm.create_model( + name, pretrained=pretrained, num_classes=num_outputs, drop_rate=drop_rate + ) + return FundusGrader(backbone, freeze_bn=freeze_bn) + + +def count_parameters(model: nn.Module) -> tuple[int, int]: + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return total, trainable diff --git a/src/drdetect/grading/module.py b/src/drdetect/grading/module.py new file mode 100644 index 0000000000000000000000000000000000000000..f991048b735c8c3de0d2467071d8ac1558956460 --- /dev/null +++ b/src/drdetect/grading/module.py @@ -0,0 +1,256 @@ +"""Lightning module for DR grading. + +Phase 1 is a plain cross-entropy baseline. It is deliberately un-tuned: its job +is to produce a number to beat, and every later improvement is measured as a +delta against it. + +Per-class recall is logged every epoch, not only aggregate QWK. On APTOS, grade +3 is 5.3% of the data and grade 1 is 10.1%; a model can post a respectable QWK +while completely failing those classes, and aggregate metrics will not say so. +""" + +from __future__ import annotations + +import lightning as L +import numpy as np +import torch +import torch.nn as nn + +from drdetect.eval.metrics import ( + binary_scores, + quadratic_weighted_kappa, + referable_labels, + sensitivity_at_specificity_floor, + youden_j, +) +from drdetect.grading.losses import build_loss, decode_output, naive_referable_cut + +__all__ = ["GradingModule"] + +CLASS_NAMES = ["No DR", "Mild", "Moderate", "Severe", "PDR"] + + +class GradingModule(L.LightningModule): + def __init__( + self, + model: nn.Module, + *, + lr: float = 1e-4, + weight_decay: float = 1e-4, + num_classes: int = 5, + loss_name: str = "ce", + class_weights: list[float] | None = None, + task_pos_weights: list[float] | None = None, + warmup_epochs: int = 3, + max_epochs: int = 40, + spec_floor: float = 0.85, + ): + super().__init__() + # `model` is a live module; saving it into the checkpoint hyperparameters + # would serialise the whole network twice. + self.save_hyperparameters(ignore=["model"]) + self.model = model + self.num_classes = num_classes + self.loss_name = loss_name + self.criterion = build_loss( + loss_name, + num_classes=num_classes, + class_weights=class_weights, + task_pos_weights=task_pos_weights, + ) + + self._val_logits: list[torch.Tensor] = [] + self._val_targets: list[torch.Tensor] = [] + self._first_epoch_loss: float | None = None + self._grad_norms: list[float] = [] + + @property + def _epoch(self) -> int: + """current_epoch without requiring an attached trainer. + + Used only in error messages. A guard that raises AttributeError while + constructing its own diagnostic is worse than no guard. + """ + try: + return int(self.current_epoch) + except (AttributeError, RuntimeError): + return -1 + + def forward(self, x): + return self.model(x) + + def training_step(self, batch, batch_idx): + x, y = batch + loss = self.criterion(self(x), y) + if not torch.isfinite(loss): + raise RuntimeError( + f"Training loss is {loss.item()} at epoch {self._epoch}, step " + f"{batch_idx}. Diverged -- lower --lr or check the data pipeline." + ) + self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def on_train_epoch_end(self): + """Abort a diverged run immediately rather than after hours. + + Small-batch training (batch 4, forced here by available RAM) produces + noisy gradients; with too high a learning rate the loss can explode by + two orders of magnitude in a single epoch. Measured on this project: + train loss went 0.786 -> 239.35 when warmup ended at lr 3e-4, and the + model collapsed to predicting grade 0 for everything. Detecting that in + one epoch instead of forty is the difference between five minutes and + two and a half hours. + """ + if self._grad_norms: + import numpy as _np + + arr = _np.asarray(self._grad_norms) + clip = getattr(self.trainer, "gradient_clip_val", None) or 0.0 + self.log("train/grad_norm", float(arr.mean())) + self.log("train/grad_norm_median", float(_np.median(arr))) + if clip: + self.log("train/grad_clipped_frac", float((arr > clip).mean())) + self._grad_norms.clear() + + loss = self.trainer.callback_metrics.get("train/loss") + if loss is None: + return + loss = float(loss) + if self._first_epoch_loss is None: + self._first_epoch_loss = loss + return + if loss > max(10.0 * self._first_epoch_loss, 5.0): + raise RuntimeError( + f"Training diverged: loss {loss:.2f} at epoch {self._epoch} vs " + f"{self._first_epoch_loss:.3f} at epoch 0.\n" + f"With batch size 4 this usually means the learning rate is too high. " + f"Try --lr 5e-5, and confirm gradient clipping is enabled." + ) + + def on_before_optimizer_step(self, optimizer): + """Collect the PRE-clip gradient norm; logging happens at epoch end. + + This hook fires before configure_gradient_clipping, so it sees the + unclipped gradient. It is the direct instrument for the epoch-3 question: + clipping acts on the accumulation MEAN, whose norm is smaller than a + 4-sample gradient's, so clip 1.0 binds less often at higher accumulation. + Without it, "the clip bound less" and "the noise was lower" are + indistinguishable in the logs. + + The values are BUFFERED rather than logged here. Logging new metric keys + mid-epoch makes CSVLogger flush rows, discover keys it has not seen, and + then fail in _rewrite_with_new_header with "dict contains fields not in + fieldnames" -- which killed two runs after one epoch. Emitting them from + on_train_epoch_end keeps every train/* key on the same cadence. + """ + from lightning.pytorch.utilities import grad_norm + + total = grad_norm(self, norm_type=2).get("grad_2.0_norm_total") + if total is not None: + self._grad_norms.append(float(total)) + + def validation_step(self, batch, batch_idx): + x, y = batch + logits = self(x) + loss = self.criterion(logits, y) + self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + self._val_logits.append(logits.detach().float().cpu()) + self._val_targets.append(y.detach().cpu()) + + def decode(self, output: torch.Tensor) -> tuple[np.ndarray, np.ndarray]: + """Delegate to the shared decoder so training and evaluation cannot drift.""" + return decode_output(output, self.loss_name) + + def on_validation_epoch_end(self): + if not self._val_logits: + return + + logits = torch.cat(self._val_logits) + targets = torch.cat(self._val_targets).numpy() + preds, p_ref = self.decode(logits) + + self.log("val/qwk", quadratic_weighted_kappa(targets, preds), prog_bar=True) + self.log("val/acc", float((preds == targets).mean())) + + # Referable DR (grade >= 2) at the default 0.5 operating point. The + # calibrated, sensitivity-targeted threshold is chosen in Phase 5 -- this + # is only a progress signal, not the reported operating point. + y_ref = referable_labels(targets) + # 0.5 is a naive cut and is only a progress signal; the reported + # operating point is chosen for target sensitivity in scripts/evaluate.py. + cut = naive_referable_cut(self.loss_name) + scores = binary_scores(y_ref, (p_ref >= cut).astype(int)) + self.log("val/sensitivity_referable", scores.sensitivity, prog_bar=True) + self.log("val/specificity_referable", scores.specificity) + + # Sensitivity-at-a-specificity-floor, and Youden's J for comparison. + # Plain sensitivity is not safe to select on: checked against this + # project's own logs, val/sensitivity_referable peaks at the exact + # epoch that collapses to predicting referable for nearly everyone + # (grade-2 recall ~1.0, specificity ~0.80) in 3 of 4 CE runs. The + # floor-constrained version operationalises this project's own stated + # target (sens >90%, spec >85%) instead of a symmetric combination, so + # it cannot be satisfied by that shortcut. See + # drdetect.eval.metrics.sensitivity_at_specificity_floor. + # + # The floor is configurable (train.py --spec-floor) because 0.85 + # saturates on this task: it compressed 16 epochs spanning 0.106 QWK + # into a 0.013 band, and ModelCheckpoint then selected on noise + # (docs/07_PHASE3_RESULTS.md, Result 9). The logged key stays + # "val/sens_at_spec85" regardless of the floor actually used, for + # continuity with --monitor's fixed choices and every existing CSV -- + # the TRUE floor for a given run is always in that run's + # summary.json config, not in this key's name. + self.log( + "val/sens_at_spec85", + sensitivity_at_specificity_floor(y_ref, p_ref, spec_floor=self.hparams.spec_floor), + prog_bar=True, + ) + self.log("val/youden_j", youden_j(y_ref, p_ref, threshold=cut)) + + # Per-class recall: aggregate QWK can look fine while a rare class is + # never predicted at all. + for cls in range(self.num_classes): + mask = targets == cls + recall = float((preds[mask] == cls).mean()) if mask.any() else float("nan") + self.log(f"val/recall_{cls}_{CLASS_NAMES[cls].replace(' ', '')}", recall) + + # Macro-averaged recall: the unweighted mean over classes, so a rare + # class counts as much as a common one. Logged as a selection + # alternative to QWK, which on this data tracks grade-2 recall + # (r=+0.51 over 70 epochs) and, for balanced CORN, was strongly + # ANTI-correlated with grade-1 recall (r=-0.70). Selecting that run on + # QWK cost 0.100 macro-recall to gain 0.0097 QWK. + recalls = [ + float((preds[targets == c] == c).mean()) + for c in range(self.num_classes) + if (targets == c).any() + ] + self.log("val/macro_recall", float(np.mean(recalls)) if recalls else 0.0) + + self._val_logits.clear() + self._val_targets.clear() + + def configure_optimizers(self): + # Only parameters that actually require grad -- frozen BN affine weights + # must not be handed to the optimiser, or AdamW allocates moment buffers + # for tensors it will never update. + params = [p for p in self.model.parameters() if p.requires_grad] + optimiser = torch.optim.AdamW( + params, lr=self.hparams.lr, weight_decay=self.hparams.weight_decay + ) + + warmup = max(self.hparams.warmup_epochs, 0) + total = max(self.hparams.max_epochs, warmup + 1) + + def lr_lambda(epoch: int) -> float: + if epoch < warmup: + return (epoch + 1) / (warmup + 1) + progress = (epoch - warmup) / max(total - warmup, 1) + return float(0.5 * (1.0 + np.cos(np.pi * min(progress, 1.0)))) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimiser, lr_lambda) + return { + "optimizer": optimiser, + "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, + } diff --git a/src/drdetect/quality/__init__.py b/src/drdetect/quality/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/quality/assessment.py b/src/drdetect/quality/assessment.py new file mode 100644 index 0000000000000000000000000000000000000000..c4410ee8c80a8e1c04549a83837b7db8d5940b4d --- /dev/null +++ b/src/drdetect/quality/assessment.py @@ -0,0 +1,127 @@ +"""Handcrafted fundus image quality gate. + +This is the "handcrafted quality features" half of Phase 2's quality stage +(docs/04_ROADMAP.md), not the learned half. The roadmap also calls for an +EfficientNet-B0 classifier trained on EyeQ; that needs a new dataset download +and a training run and is deliberately deferred. Everything here is an +uncalibrated heuristic, not a validated clinical cutoff -- the thresholds are +defensible starting points from classical blur/exposure-detection practice, +not numbers measured against this project's own data. Replace them with the +EyeQ classifier before trusting this gate on real field images, the same way +`apply_clahe`'s clip_limit is flagged as a knob to ablate, not a fact. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import cv2 +import numpy as np + +from drdetect.enhance.preprocessing import circle_crop, crop_from_gray + +__all__ = ["QualityResult", "assess_quality"] + +# Laplacian variance below this, on the FOV resized to a fixed working size, +# reads as blurry. Variance of the Laplacian scales with image size and FOV +# fraction, which is why the FOV is isolated and resized before measuring it +# -- otherwise the same physical sharpness would score differently depending +# on the camera's raw resolution. +# +# 80.0 (a generic blur-detection tutorial default) was the first value tried +# here, and it rejected 28 of 30 real APTOS training images -- images that +# were sharp enough for an ophthalmologist to grade. That default was tuned +# for a different normalisation scale, not this one. 15.0 is recalibrated +# against this project's own pipeline on a 60-image APTOS sample (measured +# range 7-108, with a visible cluster of genuinely-degraded outliers at 7-13 +# and a continuous spread from ~17 up); it is still an eyeballed cut on a +# small sample, not the learned EyeQ classifier the roadmap calls for. +_BLUR_WORKING_SIZE = 512 +_MIN_SHARPNESS = 15.0 + +# Mean intensity of the FOV, on a 0-255 scale. Outside this band the image is +# globally under- or over-exposed rather than locally noisy. Lower bound has +# a margin below the measured real-data minimum (36 on the same 60-image +# sample); upper bound is untested against real overexposed captures. +_MIN_MEAN_BRIGHTNESS = 30.0 +_MAX_MEAN_BRIGHTNESS = 225.0 + +# Fraction of FOV pixels that are crushed (near-black) or blown out +# (near-white). High values mean local exposure failure even when the mean +# brightness looks fine. +_MAX_DARK_FRACTION = 0.45 +_MAX_BRIGHT_FRACTION = 0.15 + +# Fraction of the full frame the detected FOV must occupy. Too small usually +# means the camera missed the retina (misalignment, eyelash occlusion) rather +# than a genuinely small pupil. +_MIN_FOV_FRACTION = 0.25 + + +@dataclass(frozen=True) +class QualityResult: + usable: bool + reasons: list[str] = field(default_factory=list) + sharpness: float = 0.0 + mean_brightness: float = 0.0 + dark_fraction: float = 0.0 + bright_fraction: float = 0.0 + fov_fraction: float = 0.0 + + +def assess_quality(image: np.ndarray) -> QualityResult: + """Gate a raw (unpreprocessed) fundus photo before grading. + + Args: + image: HxWx3 uint8 RGB, as captured -- not yet Ben Graham-normalised. + Judging exposure and blur on the enhanced image would be + circular: `ben_graham` actively rewrites local contrast. + + Returns: + QualityResult with `usable=False` and `reasons` naming every check + that failed, not just the first -- a rejected image should tell the + operator what to fix on recapture, not just that something is wrong. + """ + fov = crop_from_gray(image) + if fov.size == 0 or min(fov.shape[:2]) < 8: + return QualityResult(usable=False, reasons=["no retinal field of view detected"]) + + gray = cv2.cvtColor(fov, cv2.COLOR_RGB2GRAY) + + working = cv2.resize( + gray, (_BLUR_WORKING_SIZE, _BLUR_WORKING_SIZE), interpolation=cv2.INTER_AREA + ) + sharpness = float(cv2.Laplacian(working, cv2.CV_64F).var()) + + mean_brightness = float(gray.mean()) + dark_fraction = float((gray < 20).mean()) + bright_fraction = float((gray > 245).mean()) + + circle = circle_crop(image) + fov_fraction = float((cv2.cvtColor(circle, cv2.COLOR_RGB2GRAY) > 0).sum()) / ( + image.shape[0] * image.shape[1] + ) + + reasons = [] + if sharpness < _MIN_SHARPNESS: + reasons.append(f"image too blurry (sharpness {sharpness:.0f} < {_MIN_SHARPNESS:.0f})") + if mean_brightness < _MIN_MEAN_BRIGHTNESS: + reasons.append(f"underexposed (mean brightness {mean_brightness:.0f})") + if mean_brightness > _MAX_MEAN_BRIGHTNESS: + reasons.append(f"overexposed (mean brightness {mean_brightness:.0f})") + if dark_fraction > _MAX_DARK_FRACTION: + reasons.append(f"too much of the field is near-black ({dark_fraction:.0%})") + if bright_fraction > _MAX_BRIGHT_FRACTION: + reasons.append(f"too much of the field is blown out ({bright_fraction:.0%})") + if fov_fraction < _MIN_FOV_FRACTION: + reasons.append(f"retinal field of view too small ({fov_fraction:.0%} of frame)") + + return QualityResult( + usable=not reasons, + reasons=reasons, + sharpness=sharpness, + mean_brightness=mean_brightness, + dark_fraction=dark_fraction, + bright_fraction=bright_fraction, + fov_fraction=fov_fraction, + ) diff --git a/src/drdetect/segmentation/__init__.py b/src/drdetect/segmentation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/segmentation/dataset.py b/src/drdetect/segmentation/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..2658a13579de31843f851e208f69dc799739aa85 --- /dev/null +++ b/src/drdetect/segmentation/dataset.py @@ -0,0 +1,174 @@ +"""IDRiD lesion segmentation dataset: on-the-fly patch sampling from full-resolution images. + +IDRiD images are 2848x4288 -- about 17x the pixel count of the 1024px grading +images this project trains at elsewhere, and holding even one at full +resolution through a segmentation decoder is not something this project's +measured memory budget (docs/05_PROTOTYPE_SCOPE.md) supports. The roadmap's +own plan is "patch sampling at full resolution" (docs/04_ROADMAP.md, Phase 4), +not downscaling -- a hard exudate here is 0.07% of the image by pixel count +(measured directly, IDRiD_41_EX.tif), so shrinking the image shrinks the +lesion out of existence before the model ever sees it. + +Patches are drawn fresh every `__getitem__` call, not pre-extracted once, +biased toward lesion-containing regions (`lesion_patch_prob`) -- a uniformly +random 512x512 crop of a 2848x4288 image with <0.1% positive pixels would +return an all-background patch almost every time, and a model trained on +"predict background" batches learns exactly that. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import albumentations as A +import cv2 +import numpy as np +from albumentations.pytorch import ToTensorV2 +from torch.utils.data import Dataset + +__all__ = [ + "LesionImagePair", + "find_idrid_lesion_pairs", + "IDRiDLesionDataset", + "build_patch_transforms", +] + +LESION_FOLDERS = { + "microaneurysms": "1. Microaneurysms", + "haemorrhages": "2. Haemorrhages", + "hard_exudates": "3. Hard Exudates", + "soft_exudates": "4. Soft Exudates", + "optic_disc": "5. Optic Disc", +} +# IDRiD's own mask filename suffixes, per lesion type -- not derivable from +# the folder name alone (e.g. "Hard Exudates" -> "_EX", not "_HardExudates"). +LESION_SUFFIXES = { + "microaneurysms": "_MA", + "haemorrhages": "_HE", + "hard_exudates": "_EX", + "soft_exudates": "_SE", + "optic_disc": "_OD", +} + + +@dataclass(frozen=True) +class LesionImagePair: + image_id: str + image_path: Path + mask_path: Path + + +def find_idrid_lesion_pairs( + idrid_root: str | Path, lesion: str, split: str +) -> list[LesionImagePair]: + """Match IDRiD images to their mask for one lesion type and split. + + Args: + idrid_root: e.g. "data/raw/idrid". + lesion: a key of LESION_FOLDERS. + split: "train" or "test". + """ + if lesion not in LESION_FOLDERS: + raise ValueError(f"unknown lesion {lesion!r}; expected one of {sorted(LESION_FOLDERS)}") + root = Path(idrid_root) + split_dir = "a. Training Set" if split == "train" else "b. Testing Set" + image_dir = root / "A. Segmentation" / "1. Original Images" / split_dir + mask_dir = ( + root + / "A. Segmentation" + / "2. All Segmentation Groundtruths" + / split_dir + / LESION_FOLDERS[lesion] + ) + + pairs = [] + for image_path in sorted(image_dir.glob("*.jpg")): + mask_path = mask_dir / f"{image_path.stem}{LESION_SUFFIXES[lesion]}.tif" + if mask_path.exists(): + pairs.append( + LesionImagePair( + image_id=image_path.stem, image_path=image_path, mask_path=mask_path + ) + ) + return pairs + + +def build_patch_transforms(patch_size: int, train: bool) -> A.Compose: + """Same reasoning as drdetect.data.dataset.build_transforms: flips/rotation + are label-preserving for a fundus with no canonical orientation, and no + blur/noise, since it would erase small lesions the same way it would + erase microaneurysms in the grading pipeline.""" + normalise = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) + if not train: + return A.Compose([normalise, ToTensorV2()]) + return A.Compose( + [ + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.5), + A.RandomRotate90(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5), + normalise, + ToTensorV2(), + ] + ) + + +class IDRiDLesionDataset(Dataset): + """One lesion type, patch-sampled on the fly. + + `patches_per_image` fixes an epoch's length to something meaningful for a + 54-image training set -- without it, `len(dataset) == len(pairs)` would + make one epoch see 54 patches total, most of a 2848x4288 image never + sampled at all. + """ + + def __init__( + self, + pairs: list[LesionImagePair], + *, + patch_size: int = 512, + train: bool = True, + patches_per_image: int = 20, + lesion_patch_prob: float = 0.8, + seed: int = 42, + ): + self.pairs = pairs + self.patch_size = patch_size + self.train = train + self.patches_per_image = patches_per_image + self.lesion_patch_prob = lesion_patch_prob + self.transform = build_patch_transforms(patch_size, train) + self._rng = np.random.default_rng(seed) + + def __len__(self) -> int: + return len(self.pairs) * self.patches_per_image + + def _load(self, pair: LesionImagePair) -> tuple[np.ndarray, np.ndarray]: + image = cv2.imread(str(pair.image_path), cv2.IMREAD_COLOR) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + mask = cv2.imread(str(pair.mask_path), cv2.IMREAD_GRAYSCALE) + mask = (mask > 0).astype(np.uint8) + return image, mask + + def _sample_patch(self, image: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + h, w = mask.shape + p = self.patch_size + ys, xs = np.where(mask > 0) + want_lesion_patch = len(ys) > 0 and self._rng.random() < self.lesion_patch_prob + if want_lesion_patch: + i = self._rng.integers(0, len(ys)) + cy, cx = int(ys[i]), int(xs[i]) + top = int(np.clip(cy - p // 2 + self._rng.integers(-p // 4, p // 4 + 1), 0, h - p)) + left = int(np.clip(cx - p // 2 + self._rng.integers(-p // 4, p // 4 + 1), 0, w - p)) + else: + top = int(self._rng.integers(0, max(h - p, 1))) + left = int(self._rng.integers(0, max(w - p, 1))) + return image[top : top + p, left : left + p], mask[top : top + p, left : left + p] + + def __getitem__(self, idx: int): + pair = self.pairs[idx % len(self.pairs)] + image, mask = self._load(pair) + patch_img, patch_mask = self._sample_patch(image, mask) + out = self.transform(image=patch_img, mask=patch_mask) + return out["image"], out["mask"].float().unsqueeze(0) diff --git a/src/drdetect/segmentation/localization.py b/src/drdetect/segmentation/localization.py new file mode 100644 index 0000000000000000000000000000000000000000..61c3b662c0c4d7275f73e3d1eee1a9e17679f351 --- /dev/null +++ b/src/drdetect/segmentation/localization.py @@ -0,0 +1,289 @@ +"""IDRiD optic disc / fovea localisation: heatmap regression, not segmentation. + +A genuinely different task from lesion or vessel segmentation -- the target +is a single (x, y) point per structure, not a per-pixel mask. Following the +roadmap's own specified approach ("heatmap regression"), each target point +becomes a 2D Gaussian blob at working resolution; the model predicts a +2-channel heatmap (OD, fovea) and inference reads the point back off as the +argmax of each channel. + +IDRiD's "C. Localization" task ships a genuine official split -- 413 +training / 103 test images, each in its own CSV with (x, y) pixel +coordinates at full (2848x4288) resolution, for both the optic disc and the +fovea. This is the largest, cleanest-labelled subset of IDRiD used in this +project so far (segmentation's hard-exudate harness has 81 images total; +this one has 516). + +OD-diameter normalisation: the roadmap's success criterion is "mean +localisation error < 0.5 x OD diameter", but per-image OD diameter isn't +available for the localisation set itself -- only the disjoint 81-image +segmentation set ships OD masks, and the two subsets use different filename +numbering (IDRiD_01..54 vs IDRiD_001..413/103), so they cannot be joined by +ID. The mean OD diameter measured directly from the 54 real segmentation +masks (527.7px at full 2848x4288 resolution) is used as a fixed proxy +constant instead of a per-image value or an invented one. +""" + +from __future__ import annotations + +import csv +import math +from dataclasses import dataclass +from pathlib import Path + +import albumentations as A +import cv2 +import lightning as L +import numpy as np +import torch +import torch.nn as nn +from albumentations.pytorch import ToTensorV2 +from torch.utils.data import Dataset + +__all__ = [ + "IDRiDLocalizationPair", + "find_idrid_localization_pairs", + "build_localization_transforms", + "make_gaussian_heatmap", + "heatmap_argmax", + "IDRiDLocalizationDataset", + "LocalizationModule", + "MEAN_OD_DIAMETER_PX_FULL_RES", + "working_res_od_diameter", + "euclidean", +] + +# Measured directly from the 54 real IDRiD segmentation-task OD masks (see +# module docstring for why this is a cross-subset proxy, not a per-image +# ground truth value). +MEAN_OD_DIAMETER_PX_FULL_RES = 527.7 +_FULL_RES_H = 2848 + + +@dataclass(frozen=True) +class IDRiDLocalizationPair: + image_id: str + image_path: Path + od_xy: tuple[float, float] + fovea_xy: tuple[float, float] + + +def _read_coordinate_csv(path: Path) -> dict[str, tuple[float, float]]: + """IDRiD's markup CSVs ship ~50 trailing empty columns per row and a + header row with inconsistent spacing ("X- Coordinate", "Y - Coordinate") + -- read positionally (columns 0-2), not by header name.""" + coords = {} + with open(path, newline="") as f: + rows = list(csv.reader(f)) + for row in rows[1:]: + if not row or not row[0].strip(): + continue + coords[row[0].strip()] = (float(row[1]), float(row[2])) + return coords + + +def find_idrid_localization_pairs( + idrid_root: str | Path, split: str +) -> list[IDRiDLocalizationPair]: + """Args: idrid_root e.g. "data/raw/idrid"; split "train" or "test".""" + root = Path(idrid_root) / "C. Localization" + split_dir = "a. Training Set" if split == "train" else "b. Testing Set" + split_label = "Training Set" if split == "train" else "Testing Set" + image_dir = root / "1. Original Images" / split_dir + + od_csv = ( + root + / "2. Groundtruths" + / "1. Optic Disc Center Location" + / f"{'a' if split == 'train' else 'b'}. IDRiD_OD_Center_{split_label}_Markups.csv" + ) + fovea_csv = ( + root + / "2. Groundtruths" + / "2. Fovea Center Location" + / f"IDRiD_Fovea_Center_{split_label}_Markups.csv" + ) + od_coords = _read_coordinate_csv(od_csv) + fovea_coords = _read_coordinate_csv(fovea_csv) + + pairs = [] + for image_path in sorted(image_dir.glob("*.jpg")): + image_id = image_path.stem + if image_id in od_coords and image_id in fovea_coords: + pairs.append( + IDRiDLocalizationPair( + image_id=image_id, + image_path=image_path, + od_xy=od_coords[image_id], + fovea_xy=fovea_coords[image_id], + ) + ) + return pairs + + +def build_localization_transforms(size: tuple[int, int], train: bool) -> A.Compose: + """`size` is (height, width). Flips are included, unlike segmentation's + IDRiD transforms sharing the same rotation-invariance reasoning does NOT + apply here -- flipping DOES change which anatomical side is which, but + that's fine: the (x, y) keypoints are flipped consistently with the + image by albumentations' own keypoint handling, and the model is trained + to predict two independent points, not a rotation-invariant blob.""" + normalise = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) + keypoint_params = A.KeypointParams(format="xy", remove_invisible=False) + if not train: + return A.Compose( + [A.Resize(*size), normalise, ToTensorV2()], keypoint_params=keypoint_params + ) + return A.Compose( + [ + A.Resize(*size), + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.5), + A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5), + normalise, + ToTensorV2(), + ], + keypoint_params=keypoint_params, + ) + + +def make_gaussian_heatmap(h: int, w: int, cx: float, cy: float, sigma: float) -> np.ndarray: + """A single-channel float32 Gaussian blob peaking at 1.0 at (cx, cy).""" + ys, xs = np.mgrid[0:h, 0:w] + heatmap = np.exp(-(((xs - cx) ** 2 + (ys - cy) ** 2) / (2.0 * sigma**2))) + return heatmap.astype(np.float32) + + +def heatmap_argmax(heatmap: np.ndarray) -> tuple[float, float]: + """(x, y) of a 2D array's peak.""" + y, x = np.unravel_index(np.argmax(heatmap), heatmap.shape) + return float(x), float(y) + + +class IDRiDLocalizationDataset(Dataset): + def __init__( + self, + pairs: list[IDRiDLocalizationPair], + *, + size: tuple[int, int] = (512, 768), + train: bool = True, + sigma: float = 12.0, + ): + self.pairs = pairs + self.size = size + self.sigma = sigma + self.transform = build_localization_transforms(size, train) + + def __len__(self) -> int: + return len(self.pairs) + + def __getitem__(self, idx: int): + pair = self.pairs[idx] + image = cv2.cvtColor(cv2.imread(str(pair.image_path), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) + keypoints = [pair.od_xy, pair.fovea_xy] + + out = self.transform(image=image, keypoints=keypoints) + h, w = self.size + (od_x, od_y), (fov_x, fov_y) = out["keypoints"] + + od_heatmap = make_gaussian_heatmap(h, w, od_x, od_y, self.sigma) + fovea_heatmap = make_gaussian_heatmap(h, w, fov_x, fov_y, self.sigma) + heatmaps = torch.from_numpy(np.stack([od_heatmap, fovea_heatmap])) + + targets_xy = torch.tensor([[od_x, od_y], [fov_x, fov_y]], dtype=torch.float32) + return out["image"], heatmaps, targets_xy + + +def working_res_od_diameter(size: tuple[int, int]) -> float: + """MEAN_OD_DIAMETER_PX_FULL_RES rescaled to a `size=(h, w)` working + resolution, assuming the uniform-scale resize `build_localization_transforms` + actually performs (height ratio == width ratio, since the working + aspect ratio is chosen to match IDRiD's real one).""" + scale = size[0] / _FULL_RES_H + return MEAN_OD_DIAMETER_PX_FULL_RES * scale + + +def euclidean(a: tuple[float, float], b: tuple[float, float]) -> float: + return math.hypot(a[0] - b[0], a[1] - b[1]) + + +class LocalizationModule(L.LightningModule): + """MSE between predicted (sigmoid) heatmaps and target Gaussians -- + standard for heatmap-regression landmark detection. Validation reads the + point back off each predicted channel's argmax and reports Euclidean + error in OD-diameter units, the roadmap's own stated success unit. + """ + + def __init__( + self, + model: nn.Module, + *, + lr: float = 1e-4, + weight_decay: float = 1e-4, + working_od_diameter: float = 1.0, + max_epochs: int = 40, + ): + super().__init__() + self.save_hyperparameters(ignore=["model"]) + self.model = model + self._val_pred_heatmaps: list[torch.Tensor] = [] + self._val_target_xy: list[torch.Tensor] = [] + + def forward(self, x): + return self.model(x) + + def _loss(self, logits: torch.Tensor, target_heatmaps: torch.Tensor) -> torch.Tensor: + probs = torch.sigmoid(logits) + return nn.functional.mse_loss(probs, target_heatmaps) + + def training_step(self, batch, batch_idx): + x, heatmaps, _ = batch + logits = self(x) + loss = self._loss(logits, heatmaps) + if not torch.isfinite(loss): + raise RuntimeError(f"Localisation loss is {loss.item()} at step {batch_idx}. Diverged.") + self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def validation_step(self, batch, batch_idx): + x, heatmaps, targets_xy = batch + logits = self(x) + loss = self._loss(logits, heatmaps) + self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + self._val_pred_heatmaps.append(torch.sigmoid(logits).detach().float().cpu()) + self._val_target_xy.append(targets_xy.detach().cpu()) + + def on_validation_epoch_end(self): + if not self._val_pred_heatmaps: + return + preds = torch.cat(self._val_pred_heatmaps).numpy() + targets = torch.cat(self._val_target_xy).numpy() + + od_errors, fovea_errors = [], [] + diameter = self.hparams.working_od_diameter + for i in range(preds.shape[0]): + od_pred = heatmap_argmax(preds[i, 0]) + fovea_pred = heatmap_argmax(preds[i, 1]) + od_errors.append(euclidean(od_pred, tuple(targets[i, 0])) / diameter) + fovea_errors.append(euclidean(fovea_pred, tuple(targets[i, 1])) / diameter) + + od_mean = float(np.mean(od_errors)) + fovea_mean = float(np.mean(fovea_errors)) + self.log("val/od_error_diameters", od_mean, prog_bar=True) + self.log("val/fovea_error_diameters", fovea_mean, prog_bar=True) + self.log("val/mean_error_diameters", (od_mean + fovea_mean) / 2.0, prog_bar=True) + + self._val_pred_heatmaps.clear() + self._val_target_xy.clear() + + def configure_optimizers(self): + optimiser = torch.optim.AdamW( + self.model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimiser, T_max=self.hparams.max_epochs + ) + return { + "optimizer": optimiser, + "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, + } diff --git a/src/drdetect/segmentation/metrics.py b/src/drdetect/segmentation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..a1624aa37e013adbda4a4bcc857e3bcbc31d191f --- /dev/null +++ b/src/drdetect/segmentation/metrics.py @@ -0,0 +1,131 @@ +"""Pixel-level segmentation metrics. + +AUPRC, not AUROC, is the metric of record here -- the roadmap says why +(docs/04_ROADMAP.md, Phase 4): "AUROC is meaningless at <0.1% positive +pixels." Measured directly on this project's own data (IDRiD_41_EX.tif): +0.068% positive. AUROC's false-positive-rate denominator is dominated by the +~99.9% true negatives regardless of model quality, so a model that is barely +better than chance can still post a high AUROC; AUPRC's precision term is not +protected by that denominator and actually penalises a flood of false +positives. +""" + +from __future__ import annotations + +import numpy as np + +__all__ = [ + "pixel_auprc", + "pixel_auroc", + "dice_coefficient", + "best_dice_threshold", + "best_f1_threshold", +] + + +def pixel_auprc(y_true: np.ndarray, y_score: np.ndarray) -> float: + """Average precision over every pixel passed in, pooled -- not averaged + per image first. Per-image AUPRC is unstable (often undefined) for an + image with zero or a handful of positive pixels; pooling across images + before scoring is what IDRiD-scale lesion evaluation actually needs. + + Args: + y_true: any shape, values in {0, 1} (or bool). + y_score: same shape, predicted probability in [0, 1]. + """ + from sklearn.metrics import average_precision_score + + y_true = np.asarray(y_true).ravel().astype(int) + y_score = np.asarray(y_score).ravel().astype(float) + if y_true.sum() == 0: + raise ValueError("no positive pixels in y_true -- AUPRC is undefined without any") + return float(average_precision_score(y_true, y_score)) + + +def pixel_auroc(y_true: np.ndarray, y_score: np.ndarray) -> float: + """AUROC, pooled the same way as pixel_auprc. Meaningless at IDRiD-lesion-scale + imbalance (<0.1% positive, see pixel_auprc's own docstring) but the roadmap + names it explicitly as the vessel-segmentation target metric -- vessels are + a measured ~7.5% of DRIVE pixels, not extreme enough for AUROC's + false-positive-rate denominator to be swamped by trivial true negatives the + way it is for lesions.""" + from sklearn.metrics import roc_auc_score + + y_true = np.asarray(y_true).ravel().astype(int) + y_score = np.asarray(y_score).ravel().astype(float) + if y_true.sum() == 0 or y_true.sum() == len(y_true): + raise ValueError("AUROC is undefined with only one class present in y_true") + return float(roc_auc_score(y_true, y_score)) + + +def dice_coefficient(y_true: np.ndarray, y_pred_binary: np.ndarray, *, eps: float = 1e-7) -> float: + """2*|A∩B| / (|A|+|B|) at a fixed threshold -- reported alongside AUPRC, + not instead of it, since Dice needs a threshold AUPRC does not.""" + y_true = np.asarray(y_true).ravel().astype(bool) + y_pred = np.asarray(y_pred_binary).ravel().astype(bool) + intersection = np.logical_and(y_true, y_pred).sum() + return float((2.0 * intersection + eps) / (y_true.sum() + y_pred.sum() + eps)) + + +def best_dice_threshold( + y_true: np.ndarray, y_score: np.ndarray, *, thresholds: np.ndarray | None = None +) -> tuple[float, float]: + """Threshold that maximises Dice, swept on a fixed grid (not every unique + score value, unlike drdetect.eval.metrics.choose_threshold_for_sensitivity + -- pixel-level score arrays here run into the hundreds of millions of + values, and a per-unique-value sweep at that scale does not finish). + + **Call this on the VALIDATION set only** -- e.g. the internal fold's held- + out images, never the 27 official IDRiD test images. Freeze the returned + threshold and apply it to the test set with dice_coefficient(), the same + selection/evaluation separation drdetect.eval.metrics already documents + for the grading model's operating point. + + Fixed 0.5 is not a safe default here: an undertrained model's sigmoid + outputs can sit almost entirely below 0.5 even where its pixel ranking is + good, so Dice@0.5 can read close to zero while AUPRC (threshold-free) + reads far higher -- measured directly on a 2-epoch checkpoint (AUPRC + 0.409, Dice@0.5 0.037) versus the fully-trained one (AUPRC 0.830, Dice@0.5 + 0.703) on the identical 27 test images. + + Returns: + (best_threshold, best_dice) from the grid. + """ + y_true = np.asarray(y_true).ravel() + y_score = np.asarray(y_score).ravel() + if thresholds is None: + thresholds = np.linspace(0.025, 0.975, 39) + + best_threshold, best_dice = 0.5, -1.0 + for t in thresholds: + dice = dice_coefficient(y_true, y_score > t) + if dice > best_dice: + best_threshold, best_dice = float(t), dice + return best_threshold, best_dice + + +def best_f1_threshold( + y_true: np.ndarray, y_score: np.ndarray, *, thresholds: np.ndarray | None = None +) -> tuple[float, float]: + """Same grid-sweep pattern as `best_dice_threshold`, scored with F1 + instead of Dice -- used where examples are discrete items (microaneurysm + *candidates*, one score each) rather than pixels, so a pixel-overlap + ratio doesn't apply but the precision/recall trade-off it's standing in + for still does. Tune on validation only, freeze, apply once to test -- + same separation `best_dice_threshold` already documents. + """ + y_true = np.asarray(y_true).ravel().astype(bool) + y_score = np.asarray(y_score).ravel() + if thresholds is None: + thresholds = np.linspace(0.025, 0.975, 39) + + best_threshold, best_f1 = 0.5, -1.0 + for t in thresholds: + pred = y_score > t + tp = int(np.logical_and(y_true, pred).sum()) + fp = int(np.logical_and(~y_true, pred).sum()) + fn = int(np.logical_and(y_true, ~pred).sum()) + f1 = (2 * tp) / (2 * tp + fp + fn) if (2 * tp + fp + fn) > 0 else 0.0 + if f1 > best_f1: + best_threshold, best_f1 = float(t), f1 + return best_threshold, best_f1 diff --git a/src/drdetect/segmentation/microaneurysms.py b/src/drdetect/segmentation/microaneurysms.py new file mode 100644 index 0000000000000000000000000000000000000000..42bae392ab539e5828585542a135c34408b44a3d --- /dev/null +++ b/src/drdetect/segmentation/microaneurysms.py @@ -0,0 +1,495 @@ +"""Microaneurysm detection: candidate generation + small-CNN classification. + +The roadmap (docs/04_ROADMAP.md, Phase 4) is explicit that this lesion type +does *not* use the same DeepLabV3+ pixel-segmentation harness as the other +four: "morphological top-hat + matched filter candidate generation -> small +CNN classifier. Do not expect plain segmentation to work here." Measured +directly from all 81 IDRiD microaneurysm masks (train+test): 3,497 instances, +equivalent-diameter median 17.3px, p5 10.0px, p95 29.9px at full (2848x4288) +resolution -- resolvable at full res, but this project's own Phase 3 findings +(docs/07_PHASE3_RESULTS.md) already established that lesions this small are +destroyed by the working resolutions (384-1024px) used elsewhere in this +project, and a patch-sampled decoder trained at 512px (the other lesions' +harness) would be learning from images where a microaneurysm is ~2px wide. + +Candidate generation runs a black top-hat (highlights small dark blobs +against a brighter surround -- microaneurysms are dark red dots) on the green +channel at full resolution, where the classes of interest are still several +pixels wide. This is a classical, deterministic, untrained step -- it proposes +candidates, it does not decide; the small CNN afterwards is what actually +distinguishes a true microaneurysm from a vessel cross-section, noise speckle, +or the countless other things a black top-hat also lights up. +""" + +from __future__ import annotations + +import gc +from dataclasses import dataclass + +import cv2 +import lightning as L +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset + +__all__ = [ + "Candidate", + "tophat_response", + "find_candidates", + "label_candidates", + "generation_recall", + "extract_patch", + "build_candidate_examples", + "PatchDataset", + "PatchClassifier", + "score_candidates", + "ClassifierModule", +] + +DEFAULT_PATCH_SIZE = 33 + + +@dataclass(frozen=True) +class Candidate: + cx: float + cy: float + area: int + + +def tophat_response(image_bgr: np.ndarray, *, disk_radius: int = 16) -> np.ndarray: + """Black top-hat (closing - image) on the median-denoised green channel. + + `disk_radius=16` (33px-diameter structuring element): the disk must be + larger than a microaneurysm for the closing step to erase it and leave a + peak in the top-hat response -- 33px comfortably exceeds the measured p95 + equivalent diameter (29.9px) without being so large it also erases small + real vessel gaps, which would flood the candidate set with vessel debris. + """ + green = image_bgr[:, :, 1] + green = cv2.medianBlur(green, 3) + ksize = 2 * disk_radius + 1 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize)) + return cv2.morphologyEx(green, cv2.MORPH_BLACKHAT, kernel) + + +def find_candidates( + image_bgr: np.ndarray, + *, + disk_radius: int = 16, + response_percentile: float = 95.0, + min_area: int = 3, + max_area: int = 1500, + min_gray: float = 15.0, +) -> list[Candidate]: + """Percentile-threshold the top-hat response, then connected-component it + into candidate blobs. + + **Not Otsu, measured, not assumed**: Otsu was the first thing tried here, + and it fails badly on this response. Checked directly on IDRiD_01: Otsu + picks a threshold of 4 (out of a max response of 71), marking 33.8% of + the *entire image* as foreground -- because the top-hat response over a + 12.2-megapixel fundus image is a skewed, long-tailed distribution + dominated by low-level texture and JPEG noise, not the clean bimodal + signal Otsu assumes. The 33.8% "foreground" merges into huge connected + blobs (one component covered 1.76M px), which then fail `max_area` and + get discarded entirely -- taking any real microaneurysm sitting inside + them down with it. Measured recall on IDRiD_01 with Otsu: 3/18 true + instances. A fixed percentile of the response's own distribution sidesteps + the bimodality assumption entirely and adapts per-image; swept on 5 real + training images (IDRiD_01/02/10/20/33), the 95th percentile recovered + 91-100% of true instances per image (100%, 97.6%, 92.1%, 91.3%, 93.7%) + vs. Otsu's 16.7% on IDRiD_01 (3/18) -- at a candidate count Otsu actually + made *worse*, not better, despite being far more permissive. + + `min_area=3, max_area=1500`: measured true-instance area at full res + ranges 13-2231px (median 234, p95 704) -- the lower bound is set below + the smallest real instance to avoid discarding true microaneurysms before + the classifier even sees them (a generation-stage miss can never be + recovered later; a false positive can still be rejected), the upper + bound well above p95 to keep merged/clustered annotations reachable while + still rejecting obviously-too-large blobs (a stray haemorrhage or vessel + segment surviving the top-hat). + + `min_gray=15.0`: IDRiD images carry a near-black margin outside the + circular fundus; without this check, JPEG ringing at that boundary + produces top-hat "candidates" that are pure imaging artifact, not tissue. + """ + response = tophat_response(image_bgr, disk_radius=disk_radius) + threshold = float(np.percentile(response, response_percentile)) + binary = (response > threshold).astype(np.uint8) * 255 + n, _labels, stats, centroids = cv2.connectedComponentsWithStats(binary, connectivity=8) + gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY) + + candidates = [] + for i in range(1, n): # 0 is the background component + area = int(stats[i, cv2.CC_STAT_AREA]) + if area < min_area or area > max_area: + continue + cx, cy = centroids[i] + if gray[int(round(cy)), int(round(cx))] < min_gray: + continue + candidates.append(Candidate(cx=float(cx), cy=float(cy), area=area)) + return candidates + + +def label_candidates(candidates: list[Candidate], mask: np.ndarray) -> np.ndarray: + """True where a candidate's centroid lands on a positive mask pixel.""" + labels = np.zeros(len(candidates), dtype=bool) + h, w = mask.shape[:2] + for i, c in enumerate(candidates): + x, y = int(round(c.cx)), int(round(c.cy)) + if 0 <= y < h and 0 <= x < w: + labels[i] = mask[y, x] > 0 + return labels + + +def generation_recall(candidates: list[Candidate], mask: np.ndarray) -> tuple[int, int]: + """(instances with >=1 candidate landing inside them, total true instances). + + This is the ceiling end-to-end recall can ever reach: the classifier only + accepts or rejects candidates the generator already proposed, so a true + microaneurysm the generator never proposed is unrecoverable downstream. + Reported separately from the classifier's own recall so a low end-to-end + number can be attributed to the right stage. + """ + binary_mask = (mask > 0).astype(np.uint8) + n_components, comp_labels = cv2.connectedComponents(binary_mask, connectivity=8) + n_true = n_components - 1 + if n_true == 0: + return 0, 0 + h, w = mask.shape[:2] + recovered_ids: set[int] = set() + for c in candidates: + x, y = int(round(c.cx)), int(round(c.cy)) + if 0 <= y < h and 0 <= x < w: + comp_id = int(comp_labels[y, x]) + if comp_id > 0: + recovered_ids.add(comp_id) + return len(recovered_ids), n_true + + +def extract_patch(image_bgr: np.ndarray, cx: float, cy: float, size: int) -> np.ndarray: + """A `size`x`size` crop centred on (cx, cy), reflect-padded at borders + (a candidate near the image edge is common -- fundus lesions don't avoid + the periphery -- and reflect padding keeps the patch's local texture + statistics sane instead of introducing a hard black edge a real + microaneurysm patch would never have). + + Pads only the small overhang beyond the image edge, not the whole image. + An earlier version called `cv2.copyMakeBorder` on the full `image_bgr` + first and cropped afterwards -- measured directly, that allocates a full + ~36MB padded copy of the *entire* 2848x4288 image on every single call, + and this is called once per kept candidate (positives plus up to + hundreds of negatives, per image). Across even a handful of images that + drove this project's development machine to a 39GB peak memory + footprint and an aborted process; cropping first and padding only the + (at most `size`-pixel) overhang produces bit-identical output at a + fraction of a percent of the cost. + """ + h, w = image_bgr.shape[:2] + half = size // 2 + x, y = int(round(cx)), int(round(cy)) + top, left = y - half, x - half + bottom, right = top + size, left + size + + top_clip, left_clip = max(top, 0), max(left, 0) + bottom_clip, right_clip = min(bottom, h), min(right, w) + crop = image_bgr[top_clip:bottom_clip, left_clip:right_clip] + + pad_top, pad_left = top_clip - top, left_clip - left + pad_bottom, pad_right = bottom - bottom_clip, right - right_clip + if pad_top or pad_left or pad_bottom or pad_right: + crop = cv2.copyMakeBorder( + crop, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_REFLECT_101 + ) + return crop + + +def build_candidate_examples( + pairs, + *, + disk_radius: int = 16, + response_percentile: float = 95.0, + patch_size: int = DEFAULT_PATCH_SIZE, + negative_ratio: float = 10.0, + max_negatives_per_image: int = 300, + seed: int = 42, +) -> tuple[np.ndarray, np.ndarray, dict]: + """Run candidate generation + ground-truth matching across a list of + `LesionImagePair`s (see drdetect.segmentation.dataset), returning + `(patches[N,size,size,3] uint8 BGR, labels[N] int64, stats)`. + + Negatives are heavily subsampled per image: a full-resolution top-hat + routinely proposes far more spurious blobs than there are true + microaneurysms in an image, and keeping every one would both blow up + memory and swamp the true positives the way naive whole-image BCE + already does for pixel segmentation (drdetect.segmentation.module). Kept + at `negative_ratio` times that image's positive count, floored at 50 so + an image with zero recovered positives still contributes hard-negative + signal rather than none, capped at `max_negatives_per_image` so no single + image dominates the pool. + """ + rng = np.random.default_rng(seed) + all_patches: list[np.ndarray] = [] + all_labels: list[bool] = [] + total_recovered, total_true = 0, 0 + per_image = [] + + for pair in pairs: + image = cv2.imread(str(pair.image_path)) + mask = cv2.imread(str(pair.mask_path), cv2.IMREAD_GRAYSCALE) + candidates = find_candidates( + image, disk_radius=disk_radius, response_percentile=response_percentile + ) + labels = label_candidates(candidates, mask) + recovered, n_true = generation_recall(candidates, mask) + total_recovered += recovered + total_true += n_true + + pos_idx = np.flatnonzero(labels) + neg_idx = np.flatnonzero(~labels) + n_neg_keep = min(len(neg_idx), max(50, int(len(pos_idx) * negative_ratio))) + n_neg_keep = min(n_neg_keep, max_negatives_per_image) + neg_keep = ( + rng.choice(neg_idx, size=n_neg_keep, replace=False) if n_neg_keep > 0 else neg_idx[:0] + ) + keep_idx = np.concatenate([pos_idx, neg_keep]) + + for idx in keep_idx: + c = candidates[idx] + all_patches.append(extract_patch(image, c.cx, c.cy, patch_size)) + all_labels.append(bool(labels[idx])) + + per_image.append( + { + "image_id": pair.image_id, + "n_candidates": len(candidates), + "n_positive_candidates": int(labels.sum()), + "n_recovered_instances": recovered, + "n_true_instances": n_true, + "n_examples_kept": int(len(keep_idx)), + } + ) + # Measured directly: candidate generation on one full-resolution + # (2848x4288) IDRiD image peaks around 640MB (connectedComponentsWithStats' + # internal bookkeeping over tens of thousands of components, on top of + # the image/response/binary/labels/gray arrays it holds at once). That + # is fine for one image; left to Python's own GC timing across a full + # train+val pair list it measurably compounded into severe system-wide + # memory pressure on this project's 16GB development machine. An + # explicit del + collect between images is cheap insurance against that. + del image, mask, candidates, labels + gc.collect() + + if all_patches: + patches = np.stack(all_patches).astype(np.uint8) + else: + patches = np.zeros((0, patch_size, patch_size, 3), dtype=np.uint8) + labels_arr = np.array(all_labels, dtype=np.int64) + + stats = { + "n_images": len(pairs), + "n_true_instances": total_true, + "n_recovered_instances": total_recovered, + "generation_recall": (total_recovered / total_true) if total_true else float("nan"), + "n_candidates_total": int(sum(p["n_candidates"] for p in per_image)), + "n_positive_candidates": int(labels_arr.sum()), + "n_examples_kept": int(len(labels_arr)), + "per_image": per_image, + } + return patches, labels_arr, stats + + +_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) + + +class PatchDataset(Dataset): + """Pre-extracted candidate patches, normalised on the fly. + + Patches are tiny (33x33x3 uint8 by default, ~3KB each) and already + extracted by `build_candidate_examples` -- unlike `IDRiDLesionDataset`, + there's no full-resolution image I/O left to amortise per `__getitem__`, + so a plain in-memory array is simpler than another on-the-fly sampler. + """ + + def __init__(self, patches: np.ndarray, labels: np.ndarray, *, train: bool = True): + self.patches = patches + self.labels = labels + self.train = train + + def __len__(self) -> int: + return len(self.labels) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + patch = self.patches[idx] + if self.train: + if np.random.rand() < 0.5: + patch = patch[:, ::-1] + if np.random.rand() < 0.5: + patch = patch[::-1, :] + patch = np.rot90(patch, k=int(np.random.randint(4))) + rgb = np.ascontiguousarray(patch[:, :, ::-1]).astype(np.float32) / 255.0 + rgb = (rgb - _IMAGENET_MEAN) / _IMAGENET_STD + tensor = torch.from_numpy(rgb.transpose(2, 0, 1).copy()).float() + label = torch.tensor(float(self.labels[idx]), dtype=torch.float32) + return tensor, label + + +class PatchClassifier(nn.Module): + """Small CNN: a candidate patch in, one accept/reject logit out. + + Deliberately tiny (~26K params) next to the 22.4M-param DeepLabV3+ used + for pixel segmentation elsewhere in this project (drdetect.segmentation. + model) -- the roadmap specifies a "small CNN classifier" because the + input here is a fixed, already-localised, tiny patch needing only a yes/no + answer, not a full image needing spatial decoding back to pixel masks. + """ + + def __init__(self, in_channels: int = 3): + super().__init__() + self.features = nn.Sequential( + nn.Conv2d(in_channels, 16, 3, padding=1), + nn.BatchNorm2d(16), + nn.ReLU(inplace=True), + nn.MaxPool2d(2), + nn.Conv2d(16, 32, 3, padding=1), + nn.BatchNorm2d(32), + nn.ReLU(inplace=True), + nn.MaxPool2d(2), + nn.Conv2d(32, 64, 3, padding=1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d(1), + ) + self.classifier = nn.Linear(64, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + features = self.features(x).flatten(1) + return self.classifier(features).squeeze(-1) + + +def score_candidates( + image_bgr: np.ndarray, + mask: np.ndarray, + model: nn.Module, + *, + device: str = "cpu", + disk_radius: int = 16, + response_percentile: float = 95.0, + patch_size: int = DEFAULT_PATCH_SIZE, + batch_size: int = 256, +) -> dict: + """Run generation + classification on one image end to end, in batches -- + used by scripts/evaluate_microaneurysms.py once per official test image. + + Unlike `build_candidate_examples` (which subsamples negatives for + training), this keeps *every* candidate: evaluation needs the true, + unsampled precision/recall/false-positive-rate the deployed pipeline + would actually see, not a class-balanced training view of it. Patches are + extracted and scored one batch at a time rather than materialising every + candidate patch for the image up front -- a single full-resolution IDRiD + image can propose tens of thousands of candidates. + """ + candidates = find_candidates( + image_bgr, disk_radius=disk_radius, response_percentile=response_percentile + ) + labels = label_candidates(candidates, mask) + recovered, n_true = generation_recall(candidates, mask) + + scores = np.zeros(len(candidates), dtype=np.float32) + model.eval() + with torch.no_grad(): + for start in range(0, len(candidates), batch_size): + batch = candidates[start : start + batch_size] + patches = np.stack([extract_patch(image_bgr, c.cx, c.cy, patch_size) for c in batch]) + rgb = patches[:, :, :, ::-1].astype(np.float32) / 255.0 + rgb = (rgb - _IMAGENET_MEAN) / _IMAGENET_STD + tensor = torch.from_numpy(rgb.transpose(0, 3, 1, 2).copy()).float().to(device) + logits = model(tensor) + scores[start : start + len(batch)] = torch.sigmoid(logits).cpu().numpy() + + return { + "candidates": candidates, + "scores": scores, + "labels": labels, + "n_recovered": recovered, + "n_true": n_true, + } + + +class ClassifierModule(L.LightningModule): + """BCE-only (no Dice term): unlike pixel segmentation, there is no + spatial overlap to score here -- each example is one candidate, one + label, so plain weighted BCE is the whole loss, with `pos_weight` set + from the same-style empirical measurement used throughout this project + (the true ratio in the actually-sampled, negative-subsampled candidate + pool, not a theoretical one).""" + + def __init__( + self, + model: nn.Module, + *, + lr: float = 1e-3, + weight_decay: float = 1e-4, + pos_weight: float = 1.0, + max_epochs: int = 30, + ): + super().__init__() + self.save_hyperparameters(ignore=["model"]) + self.model = model + self.register_buffer("_pos_weight", torch.tensor(pos_weight)) + self._val_logits: list[torch.Tensor] = [] + self._val_targets: list[torch.Tensor] = [] + + def forward(self, x): + return self.model(x) + + def training_step(self, batch, batch_idx): + x, y = batch + logits = self(x) + loss = F.binary_cross_entropy_with_logits(logits, y, pos_weight=self._pos_weight) + if not torch.isfinite(loss): + raise RuntimeError(f"Classifier loss is {loss.item()} at step {batch_idx}. Diverged.") + self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def validation_step(self, batch, batch_idx): + x, y = batch + logits = self(x) + loss = F.binary_cross_entropy_with_logits(logits, y, pos_weight=self._pos_weight) + self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + self._val_logits.append(logits.detach().float().cpu()) + self._val_targets.append(y.detach().cpu()) + + def on_validation_epoch_end(self): + if not self._val_logits: + return + logits = torch.cat(self._val_logits) + targets = torch.cat(self._val_targets) + probs = torch.sigmoid(logits).numpy() + targets_np = targets.numpy() + + if targets_np.sum() > 0: + # pixel_auprc is a generic ravel + average_precision_score despite + # its name -- reused as-is rather than duplicated for candidates. + from drdetect.segmentation.metrics import pixel_auprc + + self.log("val/auprc", pixel_auprc(targets_np, probs), prog_bar=True) + self.log("val/accuracy", float(((probs > 0.5) == (targets_np > 0.5)).mean())) + self._val_logits.clear() + self._val_targets.clear() + + def configure_optimizers(self): + optimiser = torch.optim.AdamW( + self.model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimiser, T_max=self.hparams.max_epochs + ) + return { + "optimizer": optimiser, + "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, + } diff --git a/src/drdetect/segmentation/model.py b/src/drdetect/segmentation/model.py new file mode 100644 index 0000000000000000000000000000000000000000..8afaae5a766d104281b858b57d5b9d4126835f29 --- /dev/null +++ b/src/drdetect/segmentation/model.py @@ -0,0 +1,47 @@ +"""Lesion segmentation backbone. + +DeepLabV3+ (Chen et al. 2018), not plain U-Net: the roadmap names both as +options, and DeepLabV3+'s atrous spatial pyramid pooling gives it multi-scale +context that plain U-Net's decoder lacks -- useful here because IDRiD's +lesion classes span very different physical sizes (a single-pixel +microaneurysm vs. a hard-exudate cluster), even though this module only +trains hard exudates so far. + +Uses segmentation_models_pytorch (already a declared but, until now, unused +dependency -- same situation grad-cam and reportlab were in before Phase 2), +not a hand-rolled decoder, for the same reason drdetect.grading.model uses +timm instead of a hand-rolled EfficientNet: a maintained, widely-checked +reference implementation over a bespoke one. +""" + +from __future__ import annotations + +import torch.nn as nn + +__all__ = ["build_segmentation_model"] + + +def build_segmentation_model( + encoder_name: str = "resnet34", + *, + pretrained: bool = True, + classes: int = 1, +) -> nn.Module: + """Binary (or multi-class) lesion segmentation model. + + `classes=1`: one lesion type per model, matching how this project's IDRiD + masks ship -- a separate binary mask file per lesion type, not one + multi-class mask. Training one model per lesion type (rather than a + shared multi-head model) also means a lesion with very little training + signal (soft exudates, per the roadmap's own scope-cut list) cannot drag + down a better-supported class's gradient. + """ + import segmentation_models_pytorch as smp + + return smp.DeepLabV3Plus( + encoder_name=encoder_name, + encoder_weights="imagenet" if pretrained else None, + in_channels=3, + classes=classes, + activation=None, # raw logits out; loss/inference apply sigmoid explicitly + ) diff --git a/src/drdetect/segmentation/module.py b/src/drdetect/segmentation/module.py new file mode 100644 index 0000000000000000000000000000000000000000..4be9b780d18814721f17afca098ea891040b570e --- /dev/null +++ b/src/drdetect/segmentation/module.py @@ -0,0 +1,120 @@ +"""Lightning module for lesion segmentation. + +BCE + Dice, not plain BCE: at <0.1% positive pixels (measured, see +dataset.py), unweighted BCE's gradient is dominated by the trivial +"predict background everywhere" solution long before it ever learns the +lesion class -- the same class-imbalance failure mode this project already +diagnosed once for CORN's conditional subsets (docs/07_PHASE3_RESULTS.md, +Result 2), showing up again in a different part of the pipeline. Dice loss +is imbalance-robust by construction (it scores overlap ratio, not per-pixel +classification), so combining it with BCE keeps BCE's well-behaved gradients +while Dice supplies the signal BCE alone would drown out. +""" + +from __future__ import annotations + +import lightning as L +import torch +import torch.nn as nn +import torch.nn.functional as F + +from drdetect.segmentation.metrics import dice_coefficient, pixel_auprc + +__all__ = ["SegmentationModule", "soft_dice_loss"] + + +def soft_dice_loss( + logits: torch.Tensor, targets: torch.Tensor, *, eps: float = 1.0 +) -> torch.Tensor: + """Differentiable Dice loss on probabilities, not thresholded predictions. + + `eps=1.0` (not 1e-7): with this few positive pixels per patch, a tiny eps + lets a patch with zero true positives and zero predicted positives produce + a huge, meaningless gradient from floating-point noise in the ratio. A + Dice-smoothing epsilon of 1.0 is standard practice for exactly this + regime, not an arbitrary choice. + """ + probs = torch.sigmoid(logits) + intersection = (probs * targets).sum(dim=(1, 2, 3)) + union = probs.sum(dim=(1, 2, 3)) + targets.sum(dim=(1, 2, 3)) + dice = (2.0 * intersection + eps) / (union + eps) + return 1.0 - dice.mean() + + +class SegmentationModule(L.LightningModule): + def __init__( + self, + model: nn.Module, + *, + lr: float = 1e-4, + weight_decay: float = 1e-4, + pos_weight: float = 1.0, + dice_weight: float = 1.0, + max_epochs: int = 40, + ): + super().__init__() + self.save_hyperparameters(ignore=["model"]) + self.model = model + self.register_buffer("_pos_weight", torch.tensor(pos_weight)) + + self._val_logits: list[torch.Tensor] = [] + self._val_targets: list[torch.Tensor] = [] + + def forward(self, x): + return self.model(x) + + def _loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + bce = F.binary_cross_entropy_with_logits(logits, targets, pos_weight=self._pos_weight) + dice = soft_dice_loss(logits, targets) + return bce + self.hparams.dice_weight * dice + + def training_step(self, batch, batch_idx): + x, y = batch + logits = self(x) + loss = self._loss(logits, y) + if not torch.isfinite(loss): + raise RuntimeError(f"Segmentation loss is {loss.item()} at step {batch_idx}. Diverged.") + self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def validation_step(self, batch, batch_idx): + x, y = batch + logits = self(x) + loss = self._loss(logits, y) + self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + self._val_logits.append(logits.detach().float().cpu()) + self._val_targets.append(y.detach().cpu()) + + def on_validation_epoch_end(self): + if not self._val_logits: + return + logits = torch.cat(self._val_logits) + targets = torch.cat(self._val_targets) + probs = torch.sigmoid(logits).numpy() + targets_np = targets.numpy() + + if targets_np.sum() > 0: + auprc = pixel_auprc(targets_np, probs) + self.log("val/auprc", auprc, prog_bar=True) + dice = dice_coefficient(targets_np, probs > 0.5) + self.log("val/dice", dice, prog_bar=True) + else: + # A validation split with zero positive patches this epoch -- + # log nothing rather than a fabricated 0.0 that would look like a + # real (bad) score instead of an undefined one. + pass + + self._val_logits.clear() + self._val_targets.clear() + + def configure_optimizers(self): + optimiser = torch.optim.AdamW( + self.model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimiser, T_max=self.hparams.max_epochs + ) + return { + "optimizer": optimiser, + "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, + } diff --git a/src/drdetect/segmentation/quadrants.py b/src/drdetect/segmentation/quadrants.py new file mode 100644 index 0000000000000000000000000000000000000000..097616b9f908a44ec4b3acfbeb5fd7ff38f6ffbb --- /dev/null +++ b/src/drdetect/segmentation/quadrants.py @@ -0,0 +1,69 @@ +"""Quadrant mapping from the OD-fovea axis. + +Needed for ICDR's 4-2-1 rule (severe NPDR: haemorrhages in 4 quadrants, or +venous beading in 2, or IRMA in 1), which counts lesions per retinal +quadrant. This module defines the four quadrants geometrically, as the +roadmap specifies: one dividing line through the optic disc along the +OD-fovea axis, one perpendicular to it through the same point. + +**Labelling deliberately avoids true anatomical nasal/temporal names.** +Whether the "away from fovea" side of the OD is nasal or temporal depends on +which eye (OD/OS) the image is of, and neither IDRiD nor this project's +other datasets carry reliable per-image eye-laterality labels to resolve +that. The labels used here -- "foveal-side" / "disc-side" for the axis- +parallel direction, "superior" / "inferior" for the perpendicular one -- +describe the geometry precisely without asserting an anatomical mapping +this project cannot verify. A deployment with real laterality metadata +could remap these to true nasal/temporal in one place (`QUADRANT_LABELS`) +without touching the geometry. +""" + +from __future__ import annotations + +import numpy as np + +__all__ = ["quadrant_axes", "assign_quadrant", "QUADRANT_LABELS"] + +QUADRANT_LABELS = { + (1, 1): "superior-foveal-side", + (1, -1): "inferior-foveal-side", + (-1, 1): "superior-disc-side", + (-1, -1): "inferior-disc-side", +} + + +def quadrant_axes( + od_xy: tuple[float, float], fovea_xy: tuple[float, float] +) -> tuple[np.ndarray, np.ndarray]: + """Unit vectors (along_axis, perpendicular_axis) through the OD, where + along_axis points from OD toward the fovea. Raises if OD and fovea + coincide (undefined axis) -- that should never happen for real data and + signals a genuine upstream bug (e.g. a broken localisation checkpoint) + if it does. + """ + od = np.asarray(od_xy, dtype=float) + fovea = np.asarray(fovea_xy, dtype=float) + axis = fovea - od + norm = np.linalg.norm(axis) + if norm < 1e-6: + raise ValueError(f"OD and fovea coincide ({od_xy} == {fovea_xy}) -- cannot define an axis.") + along = axis / norm + # Perpendicular, image coordinates (y grows downward): rotate -90 degrees + # so positive `perp` points toward smaller y (up the image, superior). + perp = np.array([along[1], -along[0]]) + return along, perp + + +def assign_quadrant( + point_xy: tuple[float, float], od_xy: tuple[float, float], fovea_xy: tuple[float, float] +) -> str: + """Which of the 4 quadrants `point_xy` falls in, relative to the OD as + origin and the OD-fovea line as one dividing axis. See module docstring + for why the labels describe geometry, not asserted anatomy.""" + along, perp = quadrant_axes(od_xy, fovea_xy) + od = np.asarray(od_xy, dtype=float) + v = np.asarray(point_xy, dtype=float) - od + + along_sign = 1 if np.dot(v, along) >= 0 else -1 + perp_sign = 1 if np.dot(v, perp) >= 0 else -1 + return QUADRANT_LABELS[(along_sign, perp_sign)] diff --git a/src/drdetect/segmentation/vessels.py b/src/drdetect/segmentation/vessels.py new file mode 100644 index 0000000000000000000000000000000000000000..140075e00f720ebd48809cc3f4da6f8bff5e71e6 --- /dev/null +++ b/src/drdetect/segmentation/vessels.py @@ -0,0 +1,230 @@ +"""DRIVE vessel segmentation: dataset, model, and FOV-aware Lightning module. + +Deliberately separate from dataset.py/module.py (the IDRiD lesion harness), +not sharing a base class with them, for reasons that are architectural, not +just organisational: + + * DRIVE images are small and uniform (584x565) -- full images are trained + on directly, no patch sampling. + * Vessels are a measured ~7.5% of pixels, not IDRiD's <0.1% -- AUROC is a + meaningful metric here (see metrics.pixel_auroc), where it is not for + lesions. + * Every DRIVE image has a field-of-view mask, and pixels outside it (the + black surround around the circular retinal image) are not tissue at + all -- they must be excluded from both loss and metrics, or a model that + trivially gets the black border right looks better than it is. IDRiD's + lesion harness has no equivalent concept. + +DRIVE's own public distribution ships vessel ground truth (`1st_manual/`) +only for the 20 "training" images -- the 20 "test" images have no publicly +released annotations (confirmed directly: `data/raw/drive/test/` contains +only `images/` and `mask/`, no manual-annotation directory). All 20 labelled +images are therefore used for k-fold cross-validation; there is no separate +locked test split for this lesion type the way IDRiD has one for hard +exudates. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import albumentations as A +import cv2 +import lightning as L +import numpy as np +import torch +import torch.nn.functional as F +from albumentations.pytorch import ToTensorV2 +from torch.utils.data import Dataset + +from drdetect.segmentation.metrics import dice_coefficient, pixel_auroc + +__all__ = [ + "DriveVesselPair", + "find_drive_vessel_pairs", + "DriveVesselDataset", + "VesselSegmentationModule", +] + + +@dataclass(frozen=True) +class DriveVesselPair: + image_id: str + image_path: Path + mask_path: Path + fov_path: Path + + +def find_drive_vessel_pairs(drive_root: str | Path) -> list[DriveVesselPair]: + """The 20 DRIVE "training" images -- the only ones with public vessel + ground truth. Numbered 21-40 in DRIVE's own convention.""" + root = Path(drive_root) / "training" + image_dir, mask_dir, fov_dir = root / "images", root / "1st_manual", root / "mask" + pairs = [] + for image_path in sorted(image_dir.glob("*.tif")): + image_id = image_path.stem.split("_")[0] + mask_path = mask_dir / f"{image_id}_manual1.gif" + fov_path = fov_dir / f"{image_id}_training_mask.gif" + if mask_path.exists() and fov_path.exists(): + pairs.append( + DriveVesselPair( + image_id=image_id, image_path=image_path, mask_path=mask_path, fov_path=fov_path + ) + ) + return pairs + + +def _pad_to_multiple(arr: np.ndarray, multiple: int = 32, *, reflect: bool) -> np.ndarray: + """Pad to a SQUARE size (both axes padded to the same target, from the + larger of the two), not just each axis to its own nearest multiple. + + DRIVE images are rectangular (584x565) and `RandomRotate90` is one of + this dataset's augmentations -- a 90 degree turn swaps height and width, + so a rectangular pad target produces a *different* output shape for a + rotated image than an unrotated one in the same batch, and + `torch.stack` cannot collate a batch of mismatched shapes. Padding both + axes to one shared square target makes every rotation shape-invariant. + """ + h, w = arr.shape[:2] + target = -(-max(h, w) // multiple) * multiple # round up to the next multiple + ph, pw = target - h, target - w + border = cv2.BORDER_REFLECT_101 if reflect else cv2.BORDER_CONSTANT + return cv2.copyMakeBorder(arr, 0, ph, 0, pw, border, value=0) + + +def build_vessel_transforms(train: bool) -> A.Compose: + """Same reasoning as the grading and lesion pipelines: flips/rotation are + label-preserving for a fundus with no canonical orientation; no blur, + since vessels include capillary-scale structures a blur would erase.""" + normalise = A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) + if not train: + return A.Compose([normalise, ToTensorV2()], additional_targets={"fov": "mask"}) + return A.Compose( + [ + A.HorizontalFlip(p=0.5), + A.VerticalFlip(p=0.5), + A.RandomRotate90(p=0.5), + normalise, + ToTensorV2(), + ], + additional_targets={"fov": "mask"}, + ) + + +class DriveVesselDataset(Dataset): + """One full (padded) DRIVE image per item -- no patch sampling, unlike + the IDRiD lesion dataset, since these images are already small enough + to train on whole.""" + + def __init__(self, pairs: list[DriveVesselPair], *, train: bool = True): + self.pairs = pairs + self.transform = build_vessel_transforms(train) + + def __len__(self) -> int: + return len(self.pairs) + + def __getitem__(self, idx: int): + pair = self.pairs[idx] + image = cv2.cvtColor(cv2.imread(str(pair.image_path), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) + mask = (cv2.imread(str(pair.mask_path), cv2.IMREAD_GRAYSCALE) > 0).astype(np.uint8) + fov = (cv2.imread(str(pair.fov_path), cv2.IMREAD_GRAYSCALE) > 0).astype(np.uint8) + + # Reflect-pad the image (real content), zero-pad mask/FOV (synthetic + # border pixels are never vessel and never inside the field of view, + # so this makes the padding automatically excluded everywhere the + # FOV mask is already used to exclude the real black surround). + image = _pad_to_multiple(image, reflect=True) + mask = _pad_to_multiple(mask, reflect=False) + fov = _pad_to_multiple(fov, reflect=False) + + out = self.transform(image=image, mask=mask, fov=fov) + return out["image"], out["mask"].float().unsqueeze(0), out["fov"].float().unsqueeze(0) + + +class VesselSegmentationModule(L.LightningModule): + def __init__( + self, + model, + *, + lr: float = 1e-4, + weight_decay: float = 1e-4, + pos_weight: float = 1.0, + dice_weight: float = 1.0, + max_epochs: int = 40, + ): + super().__init__() + self.save_hyperparameters(ignore=["model"]) + self.model = model + self.register_buffer("_pos_weight", torch.tensor(pos_weight)) + self._val_probs: list[torch.Tensor] = [] + self._val_targets: list[torch.Tensor] = [] + self._val_fov: list[torch.Tensor] = [] + + def forward(self, x): + return self.model(x) + + def _loss(self, logits: torch.Tensor, targets: torch.Tensor, fov: torch.Tensor) -> torch.Tensor: + bce_per_pixel = F.binary_cross_entropy_with_logits( + logits, targets, pos_weight=self._pos_weight, reduction="none" + ) + bce = (bce_per_pixel * fov).sum() / fov.sum().clamp(min=1.0) + + probs = torch.sigmoid(logits) * fov + targets_in_fov = targets * fov + intersection = (probs * targets_in_fov).sum(dim=(1, 2, 3)) + union = probs.sum(dim=(1, 2, 3)) + targets_in_fov.sum(dim=(1, 2, 3)) + dice_loss = 1.0 - ((2.0 * intersection + 1.0) / (union + 1.0)).mean() + + return bce + self.hparams.dice_weight * dice_loss + + def training_step(self, batch, batch_idx): + x, y, fov = batch + logits = self(x) + loss = self._loss(logits, y, fov) + if not torch.isfinite(loss): + raise RuntimeError( + f"Vessel segmentation loss is {loss.item()} at step {batch_idx}. Diverged." + ) + self.log("train/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def validation_step(self, batch, batch_idx): + x, y, fov = batch + logits = self(x) + loss = self._loss(logits, y, fov) + self.log("val/loss", loss, prog_bar=True, on_step=False, on_epoch=True) + self._val_probs.append(torch.sigmoid(logits).detach().float().cpu()) + self._val_targets.append(y.detach().cpu()) + self._val_fov.append(fov.detach().cpu()) + + def on_validation_epoch_end(self): + if not self._val_probs: + return + probs = torch.cat(self._val_probs).numpy() + targets = torch.cat(self._val_targets).numpy() + fov = torch.cat(self._val_fov).numpy().astype(bool) + + in_fov_probs = probs[fov] + in_fov_targets = targets[fov] + if in_fov_targets.sum() > 0: + auroc = pixel_auroc(in_fov_targets, in_fov_probs) + self.log("val/auroc", auroc, prog_bar=True) + dice = dice_coefficient(in_fov_targets, in_fov_probs > 0.5) + self.log("val/dice", dice, prog_bar=True) + + self._val_probs.clear() + self._val_targets.clear() + self._val_fov.clear() + + def configure_optimizers(self): + optimiser = torch.optim.AdamW( + self.model.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimiser, T_max=self.hparams.max_epochs + ) + return { + "optimizer": optimiser, + "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"}, + } diff --git a/src/drdetect/serve/__init__.py b/src/drdetect/serve/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/serve/api.py b/src/drdetect/serve/api.py new file mode 100644 index 0000000000000000000000000000000000000000..3d78914e64069cadb6d482446fdf3ce938e42bd9 --- /dev/null +++ b/src/drdetect/serve/api.py @@ -0,0 +1,123 @@ +"""FastAPI service wrapping the Phase 2 pipeline (docs/04_ROADMAP.md, Phase 9). + +CPU-only, single global model instance loaded once at startup -- the same +design constraint `scripts/predict.py` and `load_grader` already state +explicitly: a district server is not assumed to have a GPU. This is a +research prototype, not a medical device (docs/01_PROJECT_ANALYSIS.md Β§11); +every response carries that disclaimer rather than only the PDF report doing so. + +Run locally: + uvicorn drdetect.serve.api:app --host 0.0.0.0 --port 8000 + +Configuration is via environment variables (not CLI flags) because this is +meant to run inside a container (Dockerfile) where env vars are the natural +knob, not argv: + DRDETECT_CHECKPOINT path to the grading checkpoint (required) + DRDETECT_BACKBONE default: efficientnet_b0 + DRDETECT_LOSS default: ce + DRDETECT_SIZE default: 512 +""" + +from __future__ import annotations + +import os +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +import cv2 +import numpy as np +from fastapi import FastAPI, File, HTTPException, UploadFile +from pydantic import BaseModel + +from drdetect.calibration.temperature import load_temperature +from drdetect.serve.pipeline import load_grader, run_pipeline + +NOT_A_MEDICAL_DEVICE = ( + "Research prototype only. NOT a medical device and NOT a substitute for " + "clinical judgement. Every result requires confirmation by a qualified " + "ophthalmologist before any care decision is made." +) + +_state: dict[str, Any] = {} + + +@asynccontextmanager +async def lifespan(app: FastAPI): + checkpoint = os.environ.get("DRDETECT_CHECKPOINT") + if not checkpoint or not Path(checkpoint).exists(): + raise RuntimeError( + "DRDETECT_CHECKPOINT env var must point to an existing checkpoint file. " + f"Got: {checkpoint!r}" + ) + backbone = os.environ.get("DRDETECT_BACKBONE", "efficientnet_b0") + loss_name = os.environ.get("DRDETECT_LOSS", "ce") + _state["size"] = int(os.environ.get("DRDETECT_SIZE", "512")) + _state["loss_name"] = loss_name + _state["checkpoint"] = checkpoint + _state["model"] = load_grader(checkpoint, backbone=backbone, loss_name=loss_name, device="cpu") + _state["temperature"] = load_temperature(checkpoint) + yield + _state.clear() + + +app = FastAPI( + title="drdetect: DR screening API", + description=NOT_A_MEDICAL_DEVICE, + version="0.1.0", + lifespan=lifespan, +) + + +class GradeResponse(BaseModel): + usable: bool + quality_reasons: list[str] + grade: int | None = None + grade_name: str | None = None + referable: bool | None = None + confidence: float | None = None + calibrated: bool = False + class_probs: list[float] | None = None + disclaimer: str = NOT_A_MEDICAL_DEVICE + + +@app.get("/health") +def health() -> dict: + return { + "status": "ok", + "checkpoint": _state.get("checkpoint"), + "disclaimer": NOT_A_MEDICAL_DEVICE, + } + + +@app.post("/grade", response_model=GradeResponse) +async def grade(image: UploadFile = File(...)) -> GradeResponse: # noqa: B008 -- FastAPI's own required-file-upload idiom + if "model" not in _state: + raise HTTPException(status_code=503, detail="Model not loaded") + + raw_bytes = await image.read() + arr = np.frombuffer(raw_bytes, dtype=np.uint8) + bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR) + if bgr is None: + raise HTTPException(status_code=400, detail="Could not decode image") + rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + + result = run_pipeline( + rgb, + _state["model"], + loss_name=_state["loss_name"], + size=_state["size"], + device="cpu", + temperature=_state["temperature"], + ) + + return GradeResponse( + usable=result.quality.usable, + quality_reasons=result.quality.reasons, + grade=result.grade, + grade_name=result.grade_name, + referable=result.referable, + confidence=result.confidence, + calibrated=result.calibrated, + class_probs=result.class_probs, + ) diff --git a/src/drdetect/serve/demo.py b/src/drdetect/serve/demo.py new file mode 100644 index 0000000000000000000000000000000000000000..8b8318b7f512a25300f11f4c6aa2dd116672d288 --- /dev/null +++ b/src/drdetect/serve/demo.py @@ -0,0 +1,104 @@ +"""Local Gradio demo: upload a fundus photo, see the graded result live. + +Unlike `scripts/predict.py` this runs on whatever accelerator is available +(MPS/CUDA/CPU) -- the CPU-only constraint is specifically about the +deployable CLI path, not an interactive local demo (docs/04_ROADMAP.md). +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from drdetect.calibration.temperature import load_temperature +from drdetect.serve.pipeline import PredictionResult, load_grader, run_pipeline + +__all__ = ["build_interface"] + + +def _format_summary(result: PredictionResult) -> str: + q = result.quality + lines = [ + f"**Quality:** sharpness {q.sharpness:.0f} | brightness {q.mean_brightness:.0f} | " + f"field-of-view {q.fov_fraction:.0%}" + ] + if not q.usable: + lines.append("\n**REJECTED -- recapture requested.** Reasons:") + lines.extend(f"- {r}" for r in q.reasons) + return "\n".join(lines) + + calib_label = "temperature-scaled" if result.calibrated else "uncalibrated" + lines.append(f"\n### Grade {result.grade}: {result.grade_name}") + lines.append(f"**Referable (grade >= 2):** {'YES' if result.referable else 'no'}") + if result.confidence is not None: + lines.append(f"**Confidence ({calib_label}):** {result.confidence:.2%}") + if result.class_probs: + from drdetect.grading.module import CLASS_NAMES + + lines.append("\n| Grade | Probability |\n|---|---|") + for name, p in zip(CLASS_NAMES, result.class_probs, strict=True): + lines.append(f"| {name} | {p:.2%} |") + lines.append( + "\n*Research prototype only. Not a medical device. " + "Every case must be reviewed by a qualified clinician.*" + ) + return "\n".join(lines) + + +def build_interface( + checkpoint: str | Path, + *, + backbone: str = "efficientnet_b0", + loss_name: str = "ce", + size: int = 512, +): + import gradio as gr + import torch + + device = ( + "mps" + if torch.backends.mps.is_available() + else ("cuda" if torch.cuda.is_available() else "cpu") + ) + model = load_grader(checkpoint, backbone=backbone, loss_name=loss_name, device=device) + temperature = load_temperature(checkpoint) + + def infer(image: np.ndarray, force: bool): + if image is None: + return None, "Upload a fundus photo to begin." + result = run_pipeline( + image, + model, + loss_name=loss_name, + size=size, + device=device, + skip_quality_gate=force, + temperature=temperature, + ) + overlay = result.cam_overlay if result.cam_overlay is not None else image + return overlay, _format_summary(result) + + calib_note = ( + f"Confidence is temperature-scaled (T={temperature:.3f})." + if temperature != 1.0 + else "This checkpoint has not been calibrated -- confidence is raw softmax output." + ) + with gr.Blocks(title="DR Screening Demo") as demo: + gr.Markdown( + "# Diabetic Retinopathy Screening -- Local Demo\n" + "Upload a fundus photo. This is a research prototype, not a diagnosis. " + f"{calib_note}" + ) + with gr.Row(): + with gr.Column(): + image_in = gr.Image(type="numpy", label="Fundus photo") + force = gr.Checkbox(label="Grade anyway if quality gate rejects", value=False) + run_btn = gr.Button("Grade", variant="primary") + with gr.Column(): + image_out = gr.Image(label="Grad-CAM evidence") + summary_out = gr.Markdown() + + run_btn.click(infer, inputs=[image_in, force], outputs=[image_out, summary_out]) + + return demo diff --git a/src/drdetect/serve/pipeline.py b/src/drdetect/serve/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..8421c2c9b66e60521345742fd817bee4d5f5ccc4 --- /dev/null +++ b/src/drdetect/serve/pipeline.py @@ -0,0 +1,189 @@ +"""End-to-end single-image inference: quality gate -> grade -> Grad-CAM. + +This is the Phase 2 vertical slice (docs/04_ROADMAP.md): image in, graded and +explained result out. `scripts/predict.py` and the Gradio demo are both thin +callers of `run_pipeline` so the two entry points cannot drift the way +training and evaluation once did over `decode_output` (docs/07_PHASE3_RESULTS.md). + +Confidence is calibrated ONLY when the caller passes a `temperature` fitted +by `scripts/calibrate.py` (see `drdetect.calibration.temperature`) -- the +default, 1.0, is a no-op divisor that reports raw softmax output. Check +`PredictionResult.calibrated` before presenting a confidence number as a +validated probability rather than a raw model output; the two look +identical in this dataclass except for that flag. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from drdetect.data.dataset import build_transforms +from drdetect.enhance.preprocessing import preprocess +from drdetect.eval.metrics import REFERABLE_THRESHOLD +from drdetect.explain.gradcam import generate_cam, overlay_cam +from drdetect.grading.losses import decode_output, outputs_for_loss +from drdetect.grading.model import build_model +from drdetect.grading.module import CLASS_NAMES +from drdetect.quality.assessment import QualityResult, assess_quality + +if TYPE_CHECKING: + # Not a real import: drdetect.fusion.features pulls in lightning and + # segmentation_models_pytorch, which the Phase 2 vertical slice + # (scripts/predict.py, the demo) has no reason to load unless a caller + # actually asks for lesion evidence -- see add_lesion_evidence below, + # which imports it lazily at call time instead. + from drdetect.fusion.features import LesionModels + +__all__ = ["PredictionResult", "load_grader", "run_pipeline", "add_lesion_evidence"] + + +@dataclass(frozen=True) +class PredictionResult: + quality: QualityResult + preprocessed: np.ndarray | None = None + grade: int | None = None + grade_name: str | None = None + referable: bool | None = None + confidence: float | None = None + class_probs: list[float] | None = None + cam_overlay: np.ndarray | None = None + calibrated: bool = False + # Phase 6: populated only by `add_lesion_evidence`, since it needs 5 extra + # loaded models `run_pipeline` itself has no reason to require -- None + # here means "not computed", not "no lesions found". + lesion_overlay: np.ndarray | None = None + lesion_evidence: dict | None = None + icdr_rationale: str | None = None + + +def load_grader( + checkpoint: str | Path, + *, + backbone: str = "efficientnet_b0", + loss_name: str = "ce", + device: str = "cpu", +) -> torch.nn.Module: + """Load a trained grader for inference, on CPU by default. + + CPU-only is deliberate, not a fallback: the roadmap's exit criterion for + this script is "single image, CPU, no GPU required" -- a district + screening kiosk is not assumed to have one. + """ + n_outputs = outputs_for_loss(loss_name) + model = build_model(backbone, num_outputs=n_outputs, pretrained=False, freeze_bn=True) + ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) + state = ckpt.get("state_dict", ckpt) + state = {k.removeprefix("model."): v for k, v in state.items() if k.startswith("model.")} + model.load_state_dict(state) + model.to(device).eval() + return model + + +def run_pipeline( + raw_image_rgb: np.ndarray, + model: torch.nn.Module, + *, + loss_name: str = "ce", + size: int = 512, + device: str = "cpu", + skip_quality_gate: bool = False, + temperature: float = 1.0, +) -> PredictionResult: + """Run the full quality -> grade -> explain chain on one raw image. + + Args: + raw_image_rgb: HxWx3 uint8 RGB, as captured -- not preprocessed. + model: from `load_grader`. + skip_quality_gate: force grading through even on a rejected image. + Used by the demo so a user can see *why* an image was flagged + instead of only being refused a result. + temperature: from `drdetect.calibration.temperature.load_temperature` + for this checkpoint, or the default 1.0 (a no-op divisor) if it + has never been calibrated. Dividing logits by T before softmax + changes reported confidence, never the predicted grade (argmax + is invariant to a positive rescaling) -- see + scripts/calibrate.py for what temperature scaling is and why. + """ + quality = assess_quality(raw_image_rgb) + if not quality.usable and not skip_quality_gate: + return PredictionResult(quality=quality) + + pre = preprocess(raw_image_rgb, size=size, use_ben_graham=True) + tensor = build_transforms(size, train=False)(image=pre)["image"].unsqueeze(0).to(device) + + with torch.no_grad(): + output = model(tensor) + preds, _p_ref = decode_output(output.cpu(), loss_name) + grade = int(preds[0]) + + class_probs = None + confidence = None + if loss_name in ("ce", "distance_ce"): + probs = torch.softmax(output / temperature, dim=1)[0].detach().cpu().tolist() + class_probs = probs + confidence = float(probs[grade]) + + cam = generate_cam(model, tensor, target_class=grade) + overlay = overlay_cam(pre, cam) + + return PredictionResult( + quality=quality, + preprocessed=pre, + grade=grade, + grade_name=CLASS_NAMES[grade], + referable=grade >= REFERABLE_THRESHOLD, + confidence=confidence, + class_probs=class_probs, + cam_overlay=overlay, + calibrated=temperature != 1.0, + ) + + +def add_lesion_evidence( + result: PredictionResult, image_path: str, lesion_models: LesionModels +) -> PredictionResult: + """Phase 6: attach a lesion-overlay render, an ICDR evidence table, and a + templated rationale to an already-graded `PredictionResult`. + + A separate step from `run_pipeline`, not folded into it, because it needs + 5 extra loaded models (`drdetect.fusion.features.load_lesion_models`) that + the Phase 2 vertical slice (`scripts/predict.py`, the Gradio demo) has no + reason to always pay the cost of loading. Takes a file path, not the + in-memory array `run_pipeline` works from, because the lesion models + apply their own preprocessing convention (raw images, full resolution, + tiled inference) that differs from the grading model's Ben-Graham cache + -- see `drdetect.fusion.embedding`'s own docstring for why these two + models deliberately see different versions of the same image. + + No-ops (returns `result` unchanged) if the image was rejected by the + quality gate -- there's no grade to explain evidence for. + """ + if result.grade is None: + return result + + from dataclasses import replace + + from drdetect.explain.evidence import ( + build_icdr_rationale, + extract_lesion_evidence, + render_lesion_overlay, + ) + + evidence = extract_lesion_evidence(image_path, lesion_models) + overlay = render_lesion_overlay(_load_rgb(image_path), evidence) + rationale = build_icdr_rationale(result.grade, evidence) + + return replace( + result, lesion_overlay=overlay, lesion_evidence=evidence, icdr_rationale=rationale + ) + + +def _load_rgb(image_path: str): + import cv2 + + return cv2.cvtColor(cv2.imread(str(image_path), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) diff --git a/src/drdetect/serve/report.py b/src/drdetect/serve/report.py new file mode 100644 index 0000000000000000000000000000000000000000..0e9ec65ecde6f57ab47beb382752fc11ef7caee2 --- /dev/null +++ b/src/drdetect/serve/report.py @@ -0,0 +1,209 @@ +"""One-page PDF report: image, Grad-CAM evidence, grade, and a plain disclaimer. + +The 30-second review target (docs/04_ROADMAP.md, Phase 6) is a property of +this layout that has not been measured yet -- Phase 6 is where it gets timed +against real reviewers. This module only produces the artefact to be timed. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from PIL import Image +from reportlab.lib import colors +from reportlab.lib.pagesizes import A4 +from reportlab.lib.units import mm +from reportlab.pdfgen import canvas + +from drdetect.serve.pipeline import PredictionResult + +__all__ = ["build_report_pdf"] + +_DISCLAIMER = ( + "Research prototype output only. NOT a medical device and NOT a diagnosis. " + "Every flagged case must be reviewed by a qualified clinician before any " + "action is taken." +) + + +def _draw_lesion_legend(c: canvas.Canvas, x: float, y: float) -> float: + """A small colour key for the lesion-overlay panel -- an outline is only + legible if the viewer knows what each colour means.""" + from drdetect.explain.evidence import LESION_COLORS_RGB + + labels = { + "hard_exudates": "Hard exudates", + "soft_exudates": "Soft exudates", + "haemorrhages": "Haemorrhages", + "microaneurysms": "Microaneurysms", + } + c.setFont("Helvetica", 8) + for lesion_type, label in labels.items(): + r, g, b = LESION_COLORS_RGB[lesion_type] + c.setFillColorRGB(r / 255, g / 255, b / 255) + c.rect(x, y - 7, 8, 8, fill=1, stroke=0) + c.setFillColor(colors.black) + c.drawString(x + 11, y - 6, label) + x += 11 + len(label) * 4.3 + 14 + return y - 14 + + +def _draw_wrapped( + c: canvas.Canvas, text: str, x: float, y: float, width: float, font_size: int = 9 +) -> float: + """Draw `text` wrapped to `width` points, returning the y below the last line.""" + from reportlab.pdfbase.pdfmetrics import stringWidth + + c.setFont("Helvetica", font_size) + words = text.split() + line = "" + for word in words: + trial = f"{line} {word}".strip() + if stringWidth(trial, "Helvetica", font_size) > width and line: + c.drawString(x, y, line) + y -= font_size + 2 + line = word + else: + line = trial + if line: + c.drawString(x, y, line) + y -= font_size + 2 + return y + + +def build_report_pdf( + result: PredictionResult, + image_name: str, + out_path: str | Path, + *, + checkpoint_name: str = "", +) -> Path: + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + + page_w, page_h = A4 + margin = 18 * mm + c = canvas.Canvas(str(out_path), pagesize=A4) + + c.setFont("Helvetica-Bold", 16) + c.drawString(margin, page_h - margin, "Diabetic Retinopathy Screening Report") + c.setFont("Helvetica", 9) + c.setFillColor(colors.grey) + c.drawString(margin, page_h - margin - 14, f"Image: {image_name}") + if checkpoint_name: + c.drawString(margin, page_h - margin - 26, f"Model checkpoint: {checkpoint_name}") + c.setFillColor(colors.black) + + y = page_h - margin - 45 + + q = result.quality + if not q.usable: + c.setFillColor(colors.red) + c.setFont("Helvetica-Bold", 13) + c.drawString(margin, y, "IMAGE REJECTED -- RECAPTURE REQUIRED") + c.setFillColor(colors.black) + y -= 22 + c.setFont("Helvetica-Bold", 10) + c.drawString(margin, y, "Reasons:") + y -= 14 + for reason in q.reasons: + y = _draw_wrapped(c, f"- {reason}", margin, y, page_w - 2 * margin) + y -= 10 + c.setFont("Helvetica-Bold", 10) + c.drawString(margin, y, "Quality metrics (for reference):") + y -= 14 + else: + panels = [ + (Image.fromarray(result.preprocessed.astype(np.uint8)), "Preprocessed image"), + ( + Image.fromarray(result.cam_overlay.astype(np.uint8)), + "Grad-CAM evidence (predicted grade)", + ), + ] + if result.lesion_overlay is not None: + panels.append( + ( + Image.fromarray(result.lesion_overlay.astype(np.uint8)), + "Lesion evidence (outlines)", + ) + ) + + gap = 10 + img_w = (page_w - 2 * margin - gap * (len(panels) - 1)) / len(panels) + img_h = img_w # preprocessed images are square (circle_crop -> resize) + + for i, (panel_img, label) in enumerate(panels): + x = margin + i * (img_w + gap) + c.drawImage(_pil_reader(panel_img), x, y - img_h, width=img_w, height=img_h) + c.setFont("Helvetica", 8) + c.setFillColor(colors.grey) + c.drawCentredString(x + img_w / 2, y - img_h - 12, label) + c.setFillColor(colors.black) + y = y - img_h - 26 + + if result.lesion_overlay is not None: + y = _draw_lesion_legend(c, margin, y) + y -= 6 + + c.setFont("Helvetica-Bold", 14) + referable_tag = "REFERABLE" if result.referable else "not referable" + c.drawString( + margin, y, f"ICDR grade {result.grade}: {result.grade_name} ({referable_tag})" + ) + y -= 20 + + if result.icdr_rationale: + c.setFont("Helvetica-Bold", 10) + c.drawString(margin, y, "Evidence:") + y -= 14 + c.setFont("Helvetica", 9) + y = _draw_wrapped(c, result.icdr_rationale, margin + 8, y, page_w - 2 * margin - 8) + y -= 6 + + calib_label = "temperature-scaled" if result.calibrated else "uncalibrated" + c.setFont("Helvetica", 10) + if result.confidence is not None: + c.drawString(margin, y, f"Model confidence ({calib_label}): {result.confidence:.2%}") + y -= 16 + + if result.class_probs: + c.setFont("Helvetica-Bold", 10) + c.drawString(margin, y, f"Per-grade probability ({calib_label}):") + y -= 14 + c.setFont("Helvetica", 9) + from drdetect.grading.module import CLASS_NAMES + + for name, p in zip(CLASS_NAMES, result.class_probs, strict=True): + c.drawString(margin + 8, y, f"{name:<10} {p:.2%}") + y -= 12 + y -= 6 + + c.setFont("Helvetica-Bold", 10) + c.drawString(margin, y, "Image quality (passed):") + y -= 14 + + c.setFont("Helvetica", 9) + c.drawString( + margin + 8, + y, + f"sharpness {q.sharpness:.0f} | mean brightness {q.mean_brightness:.0f} | " + f"field-of-view {q.fov_fraction:.0%} of frame", + ) + y -= 24 + + c.setStrokeColor(colors.grey) + c.line(margin, y, page_w - margin, y) + y -= 16 + c.setFillColor(colors.grey) + y = _draw_wrapped(c, _DISCLAIMER, margin, y, page_w - 2 * margin, font_size=8) + + c.showPage() + c.save() + return out_path + + +def _pil_reader(image: Image.Image): + from reportlab.lib.utils import ImageReader + + return ImageReader(image) diff --git a/src/drdetect/utils/__init__.py b/src/drdetect/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/drdetect/utils/seed.py b/src/drdetect/utils/seed.py new file mode 100644 index 0000000000000000000000000000000000000000..ae39b98194dc0a4464ad1482e22250480a8b60a5 --- /dev/null +++ b/src/drdetect/utils/seed.py @@ -0,0 +1,58 @@ +"""Determinism helpers. + +Reproducibility is a stated requirement of this project (a clean clone must +reproduce the headline table), so seeding is not optional plumbing. +""" + +from __future__ import annotations + +import os +import random + +__all__ = ["seed_everything", "worker_init_fn"] + + +def seed_everything(seed: int = 42, deterministic: bool = True) -> int: + """Seed Python, NumPy and torch. Returns the seed for logging. + + `deterministic=True` also disables cuDNN autotuning. On MPS most kernels are + deterministic already, but the flag is honoured where it applies. + """ + os.environ["PYTHONHASHSEED"] = str(seed) + random.seed(seed) + + try: + import numpy as np + + np.random.seed(seed) + except ImportError: + pass + + try: + import torch + + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + if deterministic: + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + except ImportError: + pass + + return seed + + +def worker_init_fn(worker_id: int) -> None: + """Give each DataLoader worker a distinct but reproducible seed. + + Without this, forked workers share the parent's NumPy state and every worker + generates the *same* augmentation stream -- a subtle bug that silently + reduces effective augmentation diversity. + """ + import numpy as np + import torch + + seed = (torch.initial_seed() + worker_id) % 2**32 + np.random.seed(seed) + random.seed(seed)