File size: 1,176 Bytes
b8fae22 | 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 | """segmentation_models.pytorch (SMP) backbones.
SMP gives us a large architecture x encoder zoo behind one constructor and is a
plain nn.Module (DDP- and bf16-friendly). Grayscale (in_channels=1) with ImageNet
encoder weights is handled by SMP's built-in channel adaptation.
"""
from __future__ import annotations
import segmentation_models_pytorch as smp
# our arch name -> SMP architecture key
_SMP_ARCH = {
"unet": "Unet",
"unetpp": "UnetPlusPlus",
"unetplusplus": "UnetPlusPlus",
"manet": "MAnet",
"linknet": "Linknet",
"fpn": "FPN",
"pspnet": "PSPNet",
"deeplabv3": "DeepLabV3",
"deeplabv3plus": "DeepLabV3Plus",
"pan": "PAN",
}
def is_smp_arch(arch: str) -> bool:
return arch.lower() in _SMP_ARCH
def build_smp(arch: str, in_channels: int, num_classes: int,
encoder: str = "resnet34", encoder_weights: str = "imagenet", **_):
weights = None if encoder_weights in ("", "none", None) else encoder_weights
return smp.create_model(
arch=_SMP_ARCH[arch.lower()],
encoder_name=encoder,
encoder_weights=weights,
in_channels=in_channels,
classes=num_classes,
)
|