Spaces:
Running on Zero
Running on Zero
File size: 1,813 Bytes
b379113 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """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
)
|