Update UI/Main.py

UI/Main.py

@@ -2,13 +2,19 @@ import os
 import sys
 sys.path.append(os.getcwd())
 
-import numpy as np
-import gradio as gr
 import cv2
-from Lib.Core import CablePoleSegmentation
+import gradio as gr
+import numpy as np
 
+from Lib.Consts import COLOR_MAP, LABELS
+from Lib.Core import CablePoleSegmentation
 
 demoImages = [
+    "data/DJI_20240905095004_0007_W.JPG",
+    "data/DJI_20240905091530_0003_W.JPG",
+    "data/DJI_20240905094647_0003_W.JPG",
+    "data/DJI_20240905094647_0003_Z.JPG",
+    "data/DJI_20240905101846_0005_W.JPG",
     "data/16_3450.png", 
     "data/16_3735.png",
     "data/16_3900.png",
@@ -16,8 +22,7 @@ demoImages = [
     "data/25_00272.png",
     "data/67_02661.png"
 ]
-labels = {0: "Boş", 1: "Çelik Direkler", 2: "Kafes Kule", 3: "Kablo", 4: "Ahşap Kule"}
-color_map = {"Boş":"#ffffff", "Çelik Direkler":"#0000ff", "Kablo":"#00ff00", "Kafes Kule":"#ff0000", "Kablo":"#00ff00", "Ahşap Kule":"#ff0000"}
+
 MODEL = CablePoleSegmentation(retina_mask=False)
 
 def Warning():
@@ -26,46 +31,64 @@ def Warning():
 with gr.Blocks(css="footer{display:none !important}") as block:
     gr.Markdown("## Yüksek Gerilim Hattı Kablo ve Direk Tespit ve Segmentasyon Uygulaması - Demo")
     
-
     with gr.Row():
         image = gr.Image(label="Fotoğraf")
-        processedImage = gr.Image(label="Sonuçlar")
-
-    with gr.Row() as imageGalleryRow:
-        imageGallery = gr.Gallery(
-            demoImages,
-            label="Demo Görsellerden Seçebilirsiniz",
-            rows=2, 
-            columns=4, 
+        processedImageGallery = gr.Gallery(
+            label="Sonuçlar",
+            rows=1, 
+            columns=2, 
             object_fit="contain", 
             height="auto"
         )
+    
+    with gr.Row() as imageGalleryRow:
+        with gr.Accordion("Demo Görsellerden Seçebilirsiniz", open=False):
+            imageGallery = gr.Gallery(
+                demoImages,
+                label="Demo Görseller",
+                rows=3, 
+                columns=4, 
+                object_fit="contain", 
+                height="auto"
+            )
 
 
     results = gr.Textbox(label="Log")
+
+    thresholdSlider = gr.Slider(0, 1, value=0.5, label="Model Eşik Değeri", info="0 ve 1 arası seçiniz.")
+
     processButton = gr.Button("Tespit Et")
-    annotatedImage = gr.AnnotatedImage(color_map=color_map)
+    annotatedImage = gr.AnnotatedImage(color_map=COLOR_MAP)
 
-    @processButton.click(outputs=[processedImage, annotatedImage, results], inputs=image)
-    def Process(image):
-        MODEL = CablePoleSegmentation()
-        masks, boxes, classes, result_plot = MODEL.Process(image)
+    @processButton.click(outputs=[processedImageGallery, annotatedImage, results], inputs=[image, thresholdSlider])
+    def Process(image, model_threshold):
+        if image is None:
+            raise gr.Warning("Lütfen görüntü yükleyiniz veya hazır seçiniz!", duration=3)
 
-        sections = []
+        batches = MODEL.Process(image, model_threshold)
+
+        if len(batches) == 0:
+            raise gr.Error("Bir Hata ile Karşılaşıldı: Görüntüde Tespit Yapılamadı 💥!", duration=5)
+
+        scores, classes, masks, boxes = batches[0] # First
         
-        for m, b, c in zip(masks, boxes, classes):
-            sections+=[(m, labels[c])]
-            # sections+=[([*b], labels[c])]
+        canvas, mask = MODEL.DrawResults(image, scores, classes, masks, boxes, class_filter=None)
 
-        return result_plot, (image, sections), "Görüntü İşlendi!"
+        sections = []
+        for m, b, c in zip(masks.cpu().numpy(), boxes.cpu().numpy(),classes.cpu().numpy()):
+            sections+=[(m, LABELS[c])]
+            # sections+=[(list(b.astype(np.int32)), LABELS[c])]
+        
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        return [image, canvas], (image, sections), "Görüntü İşlendi!"
     
 
-    @imageGallery.select(inputs=imageGallery, outputs=[processedImage, annotatedImage, results])
+    @imageGallery.select(inputs=imageGallery, outputs=[processedImageGallery, annotatedImage, results])
     def GalleryProcess(image, evt_data: gr.EventData):
         selectedIdx = evt_data._data["index"]
         imagePath = demoImages[selectedIdx]
         image = cv2.imread(imagePath)
-        return Process(image)
+        return Process(image, thresholdSlider.value)
 
     
     block.load(Warning)