Blame for src/model.py

import segmentation_models_pytorch as smp


def build_model(config):
    model_name = config["model"]["model_name"]
    encoder_name = config["model"]["encoder_name"]
    encoder_weights = config["model"]["encoder_weights"]
    num_classes = config["segmentation"]["num_classes"]

    if model_name == "Unet":
        model = smp.Unet(
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            classes=num_classes,
            activation=None,
        )

    elif model_name == "DeepLabV3":
        model = smp.DeepLabV3(
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            classes=num_classes,
            activation=None,
        )

    elif model_name == "DeepLabV3Plus":
        model = smp.DeepLabV3Plus(
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            classes=num_classes,
            activation=None,
        )

    else:
        raise ValueError(f"Unsupported model: {model_name}")

    return model