Spaces:
Paused
Paused
| 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 |