Spaces:
Running on Zero
Running on Zero
| """ models from the segmentation_models_pytorch library | |
| install the library before using these models | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import segmentation_models_pytorch as smp | |
| class SegModel(nn.Module): | |
| def __init__(self, architecture='Unet', | |
| encoder_name='mobilenet_v2', | |
| encoder_weights='imagenet', | |
| encoder_depth=4, | |
| in_channels=3, | |
| classes=4): | |
| super().__init__() | |
| """ | |
| architecture = { | |
| Unet, UnetPlusPlus, EfficientUnetPlusPlus, DeepLabV3, DeepLabV3+ | |
| } | |
| encoder_name = { | |
| resnet18 / 11M, resnet34 / 21M, resnet50 / 23M, resnet101 / 42M | |
| densenet121 / 6M, densenet169 / 12M, densenet201 / 18M | |
| efficientnet-b0 / 4M, -b1 / 6M, -b2 / 7M, -b3 / 10M, -b4 / 17M, -b5 / 28M | |
| mobilenet_v2 / 2M | |
| } | |
| encoder_weights = { | |
| "imagenet", None | |
| } | |
| in_channels = 3 | |
| classes = 4 | |
| """ | |
| if encoder_depth == 4: | |
| decoder_channels = (256, 128, 64, 32) | |
| exec("""self.model = smp.%s( | |
| encoder_name='%s', | |
| encoder_weights='%s', | |
| encoder_depth=%s, | |
| decoder_channels=%s, | |
| in_channels=%s, | |
| classes=%s, | |
| activation=None, | |
| ) | |
| """ % (architecture, encoder_name, encoder_weights, encoder_depth, decoder_channels, in_channels, classes)) | |
| def forward(self, x): | |
| return self.model(x) | |