Update app.py

app.py

@@ -1,3 +1,36 @@
+import gradio as gr
+from PIL import Image, ImageDraw, ImageFont
+from ultralytics import YOLO
+import numpy as np
+import os
+
+# === Load Model ===
+try:
+    model_path = "best.pt"
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"❌ File {model_path} tidak ditemukan. Upload 'best.pt' ke root.")
+    model = YOLO(model_path)
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print("❌ Gagal load model:", e)
+    model = None
+
+# === Peta Label dan Warna ===
+label_map = {
+    0: "coral or rock",
+    1: "pipeline",
+    2: "ripple marks",
+    3: "shipwreck"
+}
+
+color_map = {
+    0: (0, 255, 0),      # Hijau
+    1: (255, 0, 0),      # Merah
+    2: (255, 165, 0),    # Oranye
+    3: (0, 0, 255)       # Biru
+}
+
+# === Fungsi Prediksi dan Visualisasi ===
 def predict_segmentation_with_legend(image):
     try:
         if model is None:
@@ -16,7 +49,6 @@ def predict_segmentation_with_legend(image):
         class_ids = result.boxes.cls.cpu().numpy().astype(int) if result.boxes else []
 
         image_np = np.array(image).copy()
-        draw = ImageDraw.Draw(image)
 
         try:
             font = ImageFont.truetype("DejaVuSans.ttf", 18)
@@ -26,16 +58,16 @@ def predict_segmentation_with_legend(image):
         CONFIDENCE_THRESHOLD = 0.5
 
         if len(boxes) == 0:
+            draw = ImageDraw.Draw(image)
             draw.text((10, 10), "No detections", fill="red", font=font)
             return image
 
+        # Ambil index dengan skor tertinggi atau di atas threshold
         indices_to_use = [i for i, s in enumerate(scores) if s >= CONFIDENCE_THRESHOLD]
-
         if len(indices_to_use) == 0:
-            top_idx = int(np.argmax(scores))
-            indices_to_use = [top_idx]
+            indices_to_use = [int(np.argmax(scores))]
 
-        # Tampilkan mask
+        # === MASKING ===
         for i in indices_to_use:
             mask = masks[i]
             class_id = class_ids[i]
@@ -52,12 +84,10 @@ def predict_segmentation_with_legend(image):
                                 image_np * 0.5 + color_mask * 0.5, image_np)
 
         final_image = Image.fromarray(image_np.astype(np.uint8)).convert("RGBA")
-
-        # Layer teks transparan
         overlay = Image.new("RGBA", final_image.size, (255, 255, 255, 0))
         draw_overlay = ImageDraw.Draw(overlay)
 
-        # Tampilkan box & label
+        # === Gambar box dan label ===
         for i in indices_to_use:
             box = boxes[i].astype(int).tolist()
             score = scores[i]
@@ -66,10 +96,9 @@ def predict_segmentation_with_legend(image):
             color = color_map.get(class_id, (255, 255, 0))
             text_color = "white"
 
-            # Gambar kotak
             draw_overlay.rectangle(box, outline=color + (255,), width=2)
 
-            # Teks dan ukurannya
+            # Teks dan background transparan
             label_text = label
             score_text = f"Conf: {score:.2f}"
             label_size = draw_overlay.textbbox((0, 0), label_text, font=font)
@@ -78,7 +107,6 @@ def predict_segmentation_with_legend(image):
             text_x = max(box[0], 0)
             text_y = max(box[1] - (label_size[3] + score_size[3] + 8), 0)
 
-            # Background teks semi-transparan
             draw_overlay.rectangle(
                 [text_x - 2, text_y - 2,
                  text_x + max(label_size[2], score_size[2]) + 4,
@@ -86,14 +114,12 @@ def predict_segmentation_with_legend(image):
                 fill=(0, 0, 0, 160)
             )
 
-            # Tulis teks
             draw_overlay.text((text_x, text_y), label_text, fill=text_color, font=font)
             draw_overlay.text((text_x, text_y + label_size[3] + 2), score_text, fill=text_color, font=font)
 
-        # Gabungkan overlay ke gambar
         final_image = Image.alpha_composite(final_image, overlay).convert("RGB")
 
-        # Buat legenda
+        # === Buat legenda ===
         legend = Image.new("RGB", (500, 50), (255, 255, 255))
         draw_legend = ImageDraw.Draw(legend)
         x = 10
@@ -102,7 +128,7 @@ def predict_segmentation_with_legend(image):
             draw_legend.text((x + 25, 10), label, fill="black", font=font)
             x += 130
 
-        # Gabungkan dengan gambar
+        # === Gabungkan gambar & legenda ===
         combined = Image.new("RGB", (final_image.width, final_image.height + 50), (255, 255, 255))
         combined.paste(final_image, (0, 0))
         combined.paste(legend, (10, final_image.height))
@@ -112,3 +138,15 @@ def predict_segmentation_with_legend(image):
     except Exception as e:
         print("❌ Error saat segmentasi:", e)
         return image
+
+iface = gr.Interface(
+    fn=predict_segmentation_with_legend,
+    inputs=gr.Image(type="pil", label="Upload Citra Side Scan Sonar"),
+    outputs=gr.Image(type="pil", label="Hasil Segmentasi"),
+    title="YOLOv8/YOLOv11 Segmentasi Citra Sonar",
+    description="Upload citra sonar dan dapatkan hasil segmentasi lengkap dengan bounding box, mask, label, confidence, dan legenda warna.",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    iface.launch(share=True)