Upload 6 files

app.py

@@ -1,7 +1,102 @@
 import gradio as gr  
+import numpy as np
+from PIL import Image
+import random 
+import matplotlib.pyplot as plt
+import torch
+from transformers import SegformerForSemanticSegmentation
+from transformers import SegformerImageProcessor
 
+image_list = [
+    "data/1.png",
+    "data/2.png",
+    "data/3.png",
+    "data/4.png",
+]
 
+def visualize_instance_seg_mask(mask):
+    # Initialize image with zeros with the image resolution
+    # of the segmentation mask and 3 channels
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
 
-#examples=["example_Benign1.png" ,"example_Benign2.png", "example_Malign2.png"]
-gr.Interface.load("huggingface/deprem-ml/deprem_satellite_semantic_whu",
-                  title="deprem_satellite_semantic_whu",cache_examples=False).launch()
+    # Create labels
+    labels = np.unique(mask)
+    label2color = {
+        label: (
+            random.randint(0, 255),
+            random.randint(0, 255),
+            random.randint(0, 255),
+        )
+        for label in labels
+    }
+
+    for height in range(image.shape[0]):
+        for width in range(image.shape[1]):
+            image[height, width, :] = label2color[mask[height, width]]
+
+    image = image / 255
+    return image
+
+
+def Segformer_Segmentation(image_path, model_id): 
+    output_save = "output.png"
+    
+    test_image = Image.open(image_path)
+
+    model = SegformerForSemanticSegmentation.from_pretrained(model_id)
+    proccessor = SegformerImageProcessor(model_id)
+    
+    inputs = proccessor(images=test_image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        
+    result = proccessor.post_process_semantic_segmentation(outputs)[0]
+    result = np.array(result)
+    result = visualize_instance_seg_mask(result)
+    plt.figure(figsize=(10, 10))
+    for plot_index in range(2):
+        if plot_index == 0:
+            plot_image = test_image
+            title = "Original"
+        else:
+            plot_image = result
+            title = "Segmentation"
+        
+        plt.subplot(1, 2, plot_index+1)
+        plt.imshow(plot_image)
+        plt.title(title)
+        plt.axis("off")
+    plt.savefig(output_save)
+    
+    return output_save
+
+inputs = [
+    gr.inputs.Image(type="filepath", label="Input Image"),
+    gr.inputs.Dropdown(
+        choices=[
+            "deprem-ml/deprem_satellite_semantic_whu"
+        ],
+        label="Model ID",
+        default="deprem-ml/deprem_satellite_semantic_whu",
+    )
+]
+
+outputs = gr.Image(type="filepath", label="Segmentation")
+
+examples = [[image_list[0], "deprem-ml/deprem_satellite_semantic_whu"],
+            [image_list[1], "deprem-ml/deprem_satellite_semantic_whu"],
+            [image_list[2], "deprem-ml/deprem_satellite_semantic_whu"],
+            [image_list[3], "deprem-ml/deprem_satellite_semantic_whu"]]
+
+title = "Deprem ML - Segformer Semantic Segmentation"
+
+demo_app = gr.Interface(
+    Segformer_Segmentation,
+    inputs,
+    outputs,
+    examples=examples,
+    title=title,
+    cache_examples=True
+)
+
+demo_app.launch(debug=True, enable_queue=True)

requirements.txt

@@ -0,0 +1,6 @@
+gradio==3.18.0
+matplotlib==3.6.2
+numpy==1.24.2
+Pillow==9.4.0
+torch==1.12.1
+transformers==4.26.0