Spaces:
Sleeping
Sleeping
| import torch | |
| import segmentation_models_pytorch as smp | |
| def get_model(weights_path=None): | |
| model = smp.Unet( | |
| encoder_name="resnet34", | |
| in_channels=11, | |
| classes=1 | |
| ) | |
| if weights_path: | |
| model.load_state_dict(torch.load(weights_path, map_location='cpu')) | |
| model.eval() | |
| return model | |