Update app.py

app.py

@@ -1,11 +1,12 @@
+from transformers import SegformerForSemanticSegmentation
+from transformers import SegformerImageProcessor
+from PIL import Image
 import gradio as gr  
 import numpy as np
-from PIL import Image
 import random 
-import matplotlib.pyplot as plt
+import cv2
 import torch
-from transformers import SegformerForSemanticSegmentation
-from transformers import SegformerImageProcessor
+
 
 image_list = [
     "data/1.png",
@@ -14,6 +15,8 @@ image_list = [
     "data/4.png",
 ]
 
+model_path = ['deprem-ml/deprem_satellite_semantic_whu']
+
 def visualize_instance_seg_mask(mask):
     # Initialize image with zeros with the image resolution
     # of the segmentation mask and 3 channels
@@ -53,35 +56,9 @@ def Segformer_Segmentation(image_path, model_id):
     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",
-    )
-]
+    cv2.imwrite(output_save, result*255)
 
-outputs = gr.Image(type="filepath", label="Segmentation")
+    return image_path, output_save
 
 examples = [[image_list[0], "deprem-ml/deprem_satellite_semantic_whu"],
             [image_list[1], "deprem-ml/deprem_satellite_semantic_whu"],
@@ -90,13 +67,24 @@ examples = [[image_list[0], "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
-)
+app = gr.Blocks()
+with app:
+    gr.HTML("<h1 style='text-align: center'>{}</h1>".format(title))    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("Video")
+            input_video = gr.Image(type='filepath')
+            model_id =  gr.Dropdown(value=model_path[0], choices=model_path)
+            input_video_button = gr.Button(value="Predict")    
+                    
+        with gr.Column():
+            output_orijinal_image = gr.Image(type='filepath')
+            
+        with gr.Column():
+            output_mask_image =  gr.Image(type='filepath')
+
+        
+    gr.Examples(examples, inputs=[input_video, model_id], outputs=[output_orijinal_image, output_mask_image], fn=Segformer_Segmentation, cache_examples=True)
+    input_video_button.click(Segformer_Segmentation, inputs=[input_video, model_id], outputs=[output_orijinal_image, output_mask_image])
 
-demo_app.launch(debug=True, enable_queue=True)
+app.launch()