Upload 3 files

dataloader.py

@@ -0,0 +1,55 @@
+import albumentations as albu
+import numpy as np
+import cv2
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+
+class Dataset:
+    def __init__(
+            self, 
+            image_path, 
+            augmentation=None, 
+            preprocessing=None,
+    ):
+        self.pil_image = image_path
+        self.augmentation = augmentation
+        self.preprocessing = preprocessing
+    
+    def get(self):
+        # pil image > numpy array
+        image = np.array(self.pil_image)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # apply augmentations
+        if self.augmentation:
+            sample = self.augmentation(image=image)
+            image = sample['image']
+        
+        # apply preprocessing
+        if self.preprocessing:
+            sample = self.preprocessing(image=image)
+            image = sample['image']
+            
+        return image
+
+
+def get_validation_augmentation():
+    """Add paddings to make image shape divisible by 32"""
+    test_transform = [
+        albu.PadIfNeeded(384, 480)
+    ]
+    return albu.Compose(test_transform)
+
+
+def to_tensor(x, **kwargs):
+    return x.transpose(2, 0, 1).astype('float32')
+
+
+def get_preprocessing(preprocessing_fn):
+    
+    _transform = [
+        albu.Lambda(image=preprocessing_fn),
+        albu.Lambda(image=to_tensor),
+    ]
+    return albu.Compose(_transform)    

download.py

@@ -0,0 +1,17 @@
+def attempt_download_from_hub(repo_id, hf_token=None):
+    # https://github.com/fcakyon/yolov5-pip/blob/main/yolov5/utils/downloads.py
+    from huggingface_hub import hf_hub_download, list_repo_files
+    from huggingface_hub.utils._errors import RepositoryNotFoundError
+    from huggingface_hub.utils._validators import HFValidationError
+    try:
+        repo_files = list_repo_files(repo_id=repo_id, repo_type='model', token=hf_token)
+        model_file = [f for f in repo_files if f.endswith('.pth')][0]
+        file = hf_hub_download(
+            repo_id=repo_id,
+            filename=model_file,
+            repo_type='model',
+            token=hf_token,
+        )
+        return file
+    except (RepositoryNotFoundError, HFValidationError):
+        return None

istanbul_unet.py

@@ -0,0 +1,21 @@
+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'