File size: 1,212 Bytes
be2715b |
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 |
from .unet import Unet
from .unetplusplus import UnetPlusPlus
from .manet import MAnet
from .linknet import Linknet
from .fpn import FPN
from .pspnet import PSPNet
from .deeplabv3 import DeepLabV3, DeepLabV3Plus
from .pan import PAN
from . import encoders
from . import utils
from . import losses
from .__version__ import __version__
from typing import Optional
import torch
def create_model(
arch: str,
encoder_name: str = "resnet34",
encoder_weights: Optional[str] = "imagenet",
in_channels: int = 3,
classes: int = 1,
**kwargs,
) -> torch.nn.Module:
"""Models wrapper. Allows to create any model just with parametes
"""
archs = [Unet, UnetPlusPlus, MAnet, Linknet, FPN, PSPNet, DeepLabV3, DeepLabV3Plus, PAN]
archs_dict = {a.__name__.lower(): a for a in archs}
try:
model_class = archs_dict[arch.lower()]
except KeyError:
raise KeyError("Wrong architecture type `{}`. Available options are: {}".format(
arch, list(archs_dict.keys()),
))
return model_class(
encoder_name=encoder_name,
encoder_weights=encoder_weights,
in_channels=in_channels,
classes=classes,
**kwargs,
) |