Spaces:
Sleeping
Sleeping
| from download import attempt_download_from_hub | |
| import segmentation_models_pytorch as smp | |
| from dataloader import * | |
| import torch | |
| def unet_prediction(input_path, model_path): | |
| model_path = attempt_download_from_hub(model_path) | |
| best_model = torch.load(model_path) | |
| preprocessing_fn = smp.encoders.get_preprocessing_fn('efficientnet-b6', 'imagenet') | |
| test_dataset = Dataset(input_path, augmentation=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) | |
| image = test_dataset.get() | |
| x_tensor = torch.from_numpy(image).to("cuda").unsqueeze(0) | |
| pr_mask = best_model.predict(x_tensor) | |
| pr_mask = (pr_mask.squeeze().cpu().numpy().round())*255 | |
| # Save the predicted mask | |
| cv2.imwrite("output.png", pr_mask) | |
| return 'output.png' |