Update app.py

app.py

@@ -1,17 +1,17 @@
 import gradio as gr
-from PIL import Image, ImageFont, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 from ultralytics import YOLO
 import numpy as np
 import os
-import cv2
 
-# === Load model ===
+# ===== Load model safely =====
 model_path = "best.pt"
 if not os.path.exists(model_path):
-    raise FileNotFoundError("❌ File model best.pt tidak ditemukan.")
+    raise FileNotFoundError(f"❌ File '{model_path}' not found. Upload best.pt to the root directory.")
 model = YOLO(model_path)
+print("✅ Model loaded successfully!")
 
-# Label dan warna
+# ===== Label and Color Maps (sesuai urutan data.yaml) =====
 label_map = {
     0: "coral",
     1: "pipeline",
@@ -25,65 +25,87 @@ color_map = {
     3: (0, 0, 255)      # blue
 }
 
-def predict_segmentation(image_pil):
-    image_pil = image_pil.convert("RGB")
-    image_np = np.array(image_pil)
-    results = model.predict(image_pil, conf=0.25, iou=0.5)
+def predict_segmentation_with_legend(image):
+    if model is None:
+        return Image.new("RGB", image.size, color=(255, 255, 255))
+
+    image = image.convert("RGB")
+    results = model.predict(image, conf=0.25, iou=0.5)
     result = results[0]
 
-    # Jika tidak ada deteksi sama sekali
-    if result.boxes is None or len(result.boxes) == 0 or result.masks is None:
-        image_np = np.array(image_pil)
-        overlay = image_np.copy()
-        cv2.putText(overlay, "No detection found", (30, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
-        return Image.fromarray(overlay)
+    if result.masks is None or len(result.boxes) == 0:
+        # No detection
+        print("⚠️ No detection.")
+        image_np = np.array(image)
+        final_image = Image.fromarray(image_np)
+        draw = ImageDraw.Draw(final_image)
+        try:
+            font = ImageFont.truetype("arial.ttf", 24)
+        except:
+            font = ImageFont.load_default()
+        draw.text((10, 10), "No detection", fill="red", font=font)
+        return final_image
+
+    image_np = np.array(image).copy()
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.truetype("arial.ttf", 24)
+    except:
+        font = ImageFont.load_default()
 
     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, f"class_{class_id}")
         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))
 
-        # Warnai area mask (semi-transparent)
+        # Safety checks
+        if len(box) != 4 or mask.shape != image_np.shape[:2]:
+            continue
+
+        # Blend mask
+        color_mask = np.zeros_like(image_np)
         for c in range(3):
-            overlay[:, :, c] = np.where(mask > 0, overlay[:, :, c] * 0.5 + color[c] * 0.5, overlay[:, :, c])
+            color_mask[:, :, c] = mask * color[c]
+        image_np = np.where(color_mask > 0, image_np * 0.5 + color_mask * 0.5, image_np)
 
-        # Gambar bounding box dan label
-        cv2.rectangle(overlay, (box[0], box[1]), (box[2], box[3]), color, 2)
-        text = f"{label} {score:.2f}"
-        cv2.putText(overlay, text, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+        # Draw box and label
+        draw = ImageDraw.Draw(Image.fromarray(image_np))
+        draw.rectangle(box.tolist(), outline=color, width=3)
+        draw.text((box[0], box[1] - 10), f"{label}: {score:.2f}", fill="white", font=font)
 
-    # Tambahkan legend
-    legend_height = 60
-    h, w, _ = overlay.shape
-    combined = np.ones((h + legend_height, w, 3), dtype=np.uint8) * 255
-    combined[:h] = overlay
+    # Convert final image
+    final_image = Image.fromarray(image_np.astype(np.uint8))
 
-    # Legend
+    # Create legend
+    legend = Image.new("RGB", (400, 50), (255, 255, 255))
+    draw_legend = ImageDraw.Draw(legend)
     x = 10
     for cid, label in label_map.items():
-        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
+        draw_legend.rectangle([x, 10, x+20, 30], fill=color_map[cid])
+        draw_legend.text((x + 25, 10), label, fill="black", font=font)
+        x += 120
+
+    # Combine image and legend
+    combined = Image.new("RGB", (final_image.width, final_image.height + 50))
+    combined.paste(final_image, (0, 0))
+    combined.paste(legend, (10, final_image.height))
 
-    return Image.fromarray(combined)
+    return combined
 
-# Gradio app
+# === Gradio Interface ===
 iface = gr.Interface(
-    fn=predict_segmentation,
+    fn=predict_segmentation_with_legend,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
     title="YOLOv8 Segmentasi Sonar",
-    description="Upload citra Side Scan Sonar. Jika terdeteksi objek, akan ditampilkan mask, box, dan confidence. Jika tidak, akan muncul pesan."
+    description="Upload citra Side Scan Sonar. Hasil segmentasi ditampilkan dengan mask warna, bounding box, confidence score, dan legenda."
 )
 
 if __name__ == "__main__":