Upload 5 files

app.py

@@ -3,13 +3,10 @@ from PIL import Image
 import gradio as gr
 import numpy
 import torch
-import os
 
-os.system('pip install git+https://github.com/fcakyon/ultralyticsplus.git')
-
-model_id_list = ['deprem-ml/Binafarktespit-yolo5x-v1-xview', 'SerdarHelli/deprem_satellite_labeled_yolov8', 'kadirnar/yolov7-v0.1']
+model_id_list = ['deprem-ml/Binafarktespit-yolo5x-v1-xview', 'SerdarHelli/deprem_satellite_labeled_yolov8', 'kadirnar/yolov7-v0.1', 'kadirnar/UNet-EfficientNet-b6-Istanbul']
 current_device = "cuda" if torch.cuda.is_available() else "cpu"
-model_types = ["YOLOv5", "YOLOv5 + SAHI", "YOLOv8", "YOLOv7"]
+model_types = ["YOLOv5", "YOLOv5 + SAHI", "YOLOv8", "YOLOv7", "Unet-Istanbul"]
 
 def sahi_yolov5_inference(
     image,
@@ -25,7 +22,7 @@ def sahi_yolov5_inference(
     postprocess_match_threshold=0.25,
     postprocess_class_agnostic=False,
 ):
-
+    
     rect_th = None or max(round(sum(image.size) / 2 * 0.0001), 1)
     text_th = None or max(rect_th - 2, 1)
     
@@ -100,6 +97,12 @@ def sahi_yolov5_inference(
         results = model([image], size=image_size)
         return results.render()[0]
 
+    elif model_type == "Unet-Istanbul":
+        from utils.istanbul_unet import unet_prediction
+        
+        output = unet_prediction(input_path=image, model_path=model_id)
+        return output
+    
 inputs = [
     gr.Image(type="pil", label="Original Image"),
     gr.Dropdown(choices=model_id_list,label="Choose Model",value=model_id_list[0]),
@@ -136,7 +139,7 @@ demo = gr.Interface(
     article=article,
     examples=examples,
     theme="huggingface",
-    cache_examples=True,
+    cache_examples=False,
 )
 
 demo.launch(debug=True, enable_queue=True)

requirements.txt

@@ -1,5 +1,7 @@
-torch==1.10.2
-torchvision==0.11.3
+torch==1.7.1
 yolov5==7.0.8
 sahi==0.11.11
-yolov7detect==1.0.1
+yolov7detect==1.0.1
+ultralyticsplus==0.0.26
+segmentation-models-pytorch==0.1.3
+albumentations==1.3.0

utils/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)    

utils/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

utils/istanbul_unet.py

@@ -0,0 +1,21 @@
+from utils.download import attempt_download_from_hub
+import segmentation_models_pytorch as smp
+from utils.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'