| --- |
| license: cc-by-sa-4.0 |
| pipeline_tag: image-classification |
| tags: |
| - medical-imaging |
| - ophthalmology |
| - fundus |
| - retina |
| - multi-label-classification |
| - pytorch |
| - timm |
| - dinov2 |
| - lora |
| - efficientnetv2 |
| - retfound |
| --- |
| |
| # 2-Net: a two-network ensemble for multi-label fundus classification |
|
|
| Trained weights for the paper **"A Two-Network Ensemble for Multi-Label Fundus Classification"** |
| by Théodore Cousin¹, Coralie Hemptinne², Sébastien Jodogne¹ |
| (¹ ICTEAM, UCLouvain, Louvain-la-Neuve, Belgium; ² Cliniques Universitaires Saint-Luc, Brussels, Belgium). |
|
|
| Code: <https://forge.uclouvain.be/ThodoreCousin/2-net-implem> |
|
|
| The 2-Net detects retinal conditions from a colour fundus photograph as a multi-label problem. |
| It combines two networks that fail differently: a general-purpose DINOv2 ViT-B/14 adapted with |
| LoRA, and an EfficientNetV2-L that sees the image at 896 px to capture fine-grained detail. |
| Each network is calibrated per class on validation data, and the two are fused by a geometric |
| mean of their probabilities. The 2-Net performs on par with the 30-member winner of the RIADD |
| challenge, with 4.5 times fewer parameters and a CPU inference time of 4.4 s instead of 77.5 s |
| per image. |
|
|
| ## Files |
|
|
| One checkpoint per network and dataset, all trained with seed 42 (one of the five seeds of the paper). |
|
|
| | File | Network | Input | Parameters | |
| |---|---|---|---| |
| | `DINOv2-LoRA/{MuReD,RFMiD}/DINOv2-LoRA-<Dataset>.pt` | DINOv2 ViT-B/14, LoRA rank 16, token-pooling head | 518 × 518 | 87.8 M (1.2 M trained) | |
| | `EfficientNetV2-L/{MuReD,RFMiD}/EfficientNetV2-L-<Dataset>.pt` | EfficientNetV2-L, linear head | 896 × 896 | 117.3 M | |
| | `RETFound/{MuReD,RFMiD}/RETFound-<Dataset>.pt` | RETFound ViT-L/16, fully fine-tuned, token-pooling head | 384 × 384 | ≈ 304 M | |
| | `calibration/2-Net-<Dataset>.json` | 2-Net calibration: per-class Platt parameters of both members and decision thresholds, fitted on the validation split | | | |
|
|
| - **2-Net** = DINOv2-LoRA + EfficientNetV2-L (the main system of the paper). |
| - **2-Net (RETFound)** = RETFound + EfficientNetV2-L, the variant used to compare against |
| retina-specific pretraining. |
| - The ViT checkpoints are PyTorch `state_dict`s. The EfficientNetV2-L checkpoints are dicts with |
| the `state_dict` and its metadata (architecture, class names, seed, …). All of them are |
| loaded with the code of the repository above. |
|
|
| ## Outputs |
|
|
| One probability per class, in this order: |
|
|
| - **MuReD (20)**: DR, NORMAL, MH, ODC, TSLN, ARMD, DN, MYA, BRVO, ODP, CRVO, CNV, RS, ODE, LS, |
| CSR, HTR, ASR, CRS, OTHER. |
| - **RFMiD (29)**: Disease_Risk, DR, ARMD, MH, DN, MYA, BRVO, TSLN, ERM, LS, MS, CSR, ODC, CRVO, |
| TV, AH, ODP, ODE, ST, AION, PT, RT, RS, CRS, EDN, RPEC, MHL, RP, OTHER (the 29 labels of the |
| RIADD challenge; OTHER groups the remaining 18 RFMiD labels). |
| |
| ## Training data |
| |
| - **MuReD** (2,208 images, 20 labels). No official split: we use a class-stratified split of |
| 1,504 / 260 / 444 images (train / validation / test), provided in the code repository. |
| - **RFMiD** (3,200 images). Official split of 1,920 / 640 / 640 images, 29 labels. |
| |
| Images are cropped to the fundus disc and enhanced with CLAHE. Training uses |
| oversampling of rare-class images, data augmentation (flips, small rotations, photometric |
| changes, cutout, random crops) and MixUp / CutMix. DINOv2-LoRA and RETFound use a |
| distribution-balanced loss, EfficientNetV2-L an asymmetric loss. |
| |
| ## Results |
| |
| Test set, macro F1 / macro mAP / Recall@n. Our models use per-class Platt calibration and |
| decision thresholds tuned on the validation split. |
| |
| **These checkpoints (seed 42)** |
| |
| | Model | F1 MuReD | F1 RFMiD | mAP MuReD | mAP RFMiD | R@n MuReD | R@n RFMiD | |
| |---|---|---|---|---|---|---| |
| | EfficientNetV2-L | .637 | .535 | .720 | .551 | .791 | .854 | |
| | RETFound, full fine-tuning | .621 | .481 | .683 | .534 | .764 | .864 | |
| | DINOv2, LoRA | .675 | .516 | .726 | .581 | .791 | .872 | |
| | 2-Net (RETFound) | .671 | .510 | .745 | .601 | .800 | .877 | |
| | **2-Net (DINOv2)** | **.706** | **.544** | **.747** | **.607** | **.804** | **.873** | |
| |
| **Paper (mean ± 1 SD over five seeds)** |
| |
| | Model | F1 MuReD | F1 RFMiD | mAP MuReD | mAP RFMiD | R@n MuReD | R@n RFMiD | |
| |---|---|---|---|---|---|---| |
| | EfficientNetV2-L | .629 ± .009 | .511 ± .031 | .704 ± .012 | .556 ± .010 | .779 ± .008 | .858 ± .006 | |
| | RETFound, full fine-tuning | .618 ± .015 | .489 ± .012 | .677 ± .012 | .546 ± .010 | .751 ± .020 | .866 ± .003 | |
| | DINOv2, LoRA | .655 ± .018 | .506 ± .030 | .711 ± .009 | .565 ± .012 | .792 ± .011 | .869 ± .003 | |
| | 2-Net (RETFound) | .658 ± .013 | .511 ± .003 | .730 ± .011 | .602 ± .010 | .799 ± .011 | .876 ± .001 | |
| | **2-Net (DINOv2)** | **.673 ± .020** | **.527 ± .012** | **.735 ± .008** | .589 ± .011 | **.806 ± .004** | .875 ± .003 | |
| |
| Comparisons with C-Tran, RetExpert, RET-CLIP and the RIADD winner are in the paper. |
| |
| ## How to use |
| |
| ```bash |
| git clone https://forge.uclouvain.be/ThodoreCousin/2-net-implem && cd 2-net-implem |
| pip install -r requirements.txt |
| hf download TheodoreCousin/2-Net --local-dir weights |
| |
| python infer.py --dataset MURED --images path/to/images/ --out predictions.csv |
| ``` |
| |
| `--dataset` selects the label set and the networks trained on it (`MURED` or `RFMiD`). Each |
| image goes through DINOv2-LoRA and EfficientNetV2-L; their probabilities are calibrated with |
| `calibration/2-Net-<Dataset>.json`, fused by a geometric mean and thresholded per class. The |
| CSV has one row per image with the probability of every class and the predicted labels. |
| Preprocessing (fundus crop, CLAHE, resizing) is done by the script. The code repository also |
| explains how to reproduce every result of the paper. |
| |
| ## Intended use and limitations |
| |
| - **Research use only.** These models are not medical devices and have not been validated |
| in clinical practice. They must not be used for diagnosis. |
| - They were trained and evaluated on MuReD and RFMiD only; performance on other cameras, |
| populations or acquisition protocols is unknown, and no cross-dataset evaluation was done. |
| - Rare diseases remain hard: macro F1 on the rarest classes stays around 0.3. |
| - Calibration and thresholds are fitted on each dataset's validation split and do not transfer |
| as is to new data. |
| - The label sets mix pathological signs and diseases, and the models give probabilities |
| without an explanation of the evidence behind them. |
| |
| ## License |
| |
| The weights are released under [CC BY-SA 4.0](https://creativecommons.org/licenses/by-sa/4.0/): |
| you may share and adapt them, including commercially, provided you give appropriate credit and |
| distribute your contributions under the same license. |
| |
| The RETFound checkpoints are fine-tuned from RETFound, itself released under CC BY-NC 4.0, so |
| its non-commercial restriction still applies to them. DINOv2-LoRA starts from Meta's DINOv2 and |
| EfficientNetV2-L from ImageNet weights distributed with `timm`. |
| |
| ## Citation |
| |
| TO BE COMPLETED |
| |
