Update app.py

app.py

@@ -1,111 +1,89 @@
-#!/usr/bin/env python
-# coding: utf-8
-
 import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageFont, ImageDraw
 from ultralytics import YOLO
 import numpy as np
 import os
+import cv2
+
+# === Load model ===
+model_path = "best.pt"
+if not os.path.exists(model_path):
+    raise FileNotFoundError("❌ File model best.pt tidak ditemukan.")
+model = YOLO(model_path)
 
-# === Load model safely ===
-try:
-    model_path = "best.pt"
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"❌ File {model_path} not found. Upload best.pt to root directory.")
-    model = YOLO(model_path)
-    print("✅ Model loaded successfully!")
-except Exception as e:
-    print("❌ Failed to load YOLO model:", e)
-    model = None
-
-# Mapping label (4 classes)
+# Label dan warna
 label_map = {
     0: "coral",
     1: "pipeline",
     2: "ripple marks",
     3: "shipwreck"
 }
-
-# Warna RGB untuk setiap class
 color_map = {
-    0: (0, 255, 0),      # hijau - coral
-    1: (255, 0, 0),      # merah - pipeline
-    2: (255, 255, 0),    # kuning - ripple marks
-    3: (0, 0, 255)       # biru - shipwreck
+    0: (0, 255, 0),     # green
+    1: (255, 0, 0),     # red
+    2: (255, 255, 0),   # yellow
+    3: (0, 0, 255)      # blue
 }
 
-def predict_segmentation_with_legend(image):
-    if model is None:
-        return Image.new("RGB", image.size, color=(255, 255, 255))
+def predict_segmentation(image_pil):
+    image_pil = image_pil.convert("RGB")
+    image_np = np.array(image_pil)
 
-    image = image.convert("RGB")
-    results = model.predict(image, conf=0.25, iou=0.5)
+    results = model.predict(image_pil, conf=0.25, iou=0.5)
     result = results[0]
 
-    print("🔍 DETECTIONS:", result.boxes)
-    print("🎭 MASKS:", result.masks)
-
     if result.masks is None:
-        print("⚠️ No masks detected.")
-        return image
-
-    image_np = np.array(image).copy()
-    draw = ImageDraw.Draw(image)
-
-    try:
-        font = ImageFont.truetype("arial.ttf", 24)
-    except:
-        font = ImageFont.load_default()
+        return image_pil
 
     masks = result.masks.data.cpu().numpy()
     boxes = result.boxes.xyxy.cpu().numpy()
     scores = result.boxes.conf.cpu().numpy()
     class_ids = result.boxes.cls.cpu().numpy().astype(int)
 
+    overlay = image_np.copy()
+
     for i, mask in enumerate(masks):
         class_id = class_ids[i]
-        label = label_map.get(class_id, str(class_id))
-        color = color_map.get(class_id, (255, 255, 0))
+        color = color_map.get(class_id, (255, 255, 255))
         score = scores[i]
         box = boxes[i].astype(int)
+        label = label_map.get(class_id, str(class_id))
 
-        # Blend mask ke image
-        color_mask = np.zeros_like(image_np)
+        # Warnai area mask (semi-transparent)
         for c in range(3):
-            color_mask[:, :, c] = mask * color[c]
-        image_np = np.where(color_mask > 0, image_np * 0.5 + color_mask * 0.5, image_np)
+            overlay[:, :, c] = np.where(mask > 0, overlay[:, :, c] * 0.5 + color[c] * 0.5, overlay[:, :, c])
 
-        # Gambar bounding box dan label
-        draw.rectangle(box.tolist(), outline=color, width=3)
-        draw.text((box[0], box[1] - 10), f"{label}: {score:.2f}", fill="white", font=font)
+        # Gambar bounding box
+        cv2.rectangle(overlay, (box[0], box[1]), (box[2], box[3]), color, 2)
 
-    # Convert kembali ke PIL image
-    final_image = Image.fromarray(image_np.astype(np.uint8))
+        # Label + score
+        text = f"{label} {score:.2f}"
+        cv2.putText(overlay, text, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
 
-    # Legend
-    legend_height = 50
-    legend = Image.new("RGB", (final_image.width, legend_height), (255, 255, 255))
-    draw_legend = ImageDraw.Draw(legend)
+    # Tambah legend
+    legend_height = 60
+    h, w, _ = overlay.shape
+    combined = np.ones((h + legend_height, w, 3), dtype=np.uint8) * 255
+    combined[:h] = overlay
+
+    # Gambar legenda
     x = 10
     for cid, label in label_map.items():
-        draw_legend.rectangle([x, 10, x+20, 30], fill=color_map[cid])
-        draw_legend.text((x + 25, 10), label, fill="black", font=font)
+        color = color_map[cid]
+        cv2.rectangle(combined, (x, h + 10), (x + 20, h + 30), color, -1)
+        cv2.putText(combined, label, (x + 30, h + 27), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
         x += 150
 
-    combined = Image.new("RGB", (final_image.width, final_image.height + legend_height))
-    combined.paste(final_image, (0, 0))
-    combined.paste(legend, (0, final_image.height))
-
-    return combined
+    return Image.fromarray(combined)
 
-# === Gradio Interface ===
+# Gradio app
 iface = gr.Interface(
-    fn=predict_segmentation_with_legend,
+    fn=predict_segmentation,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="YOLOv8 Segmentasi (4 Kelas: Coral, Pipeline, Ripple Marks, Shipwreck)",
-    description="Upload citra sonar. Hasil segmentasi akan ditampilkan dengan mask, bounding box, skor, dan legenda warna."
+    title="YOLOv8 Segmentasi Sonar",
+    description="Upload citra Side Scan Sonar. Akan ditampilkan: mask, bounding box, confidence score, dan legenda kelas."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()