Upload app.py

app.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import gradio as gr
+from PIL import Image, ImageDraw, ImageFont
+from ultralytics import YOLO
+import numpy as np
+import os
+
+# === 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_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
+}
+
+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]
+
+    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()
+
+    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)
+
+    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))
+        score = scores[i]
+        box = boxes[i].astype(int)
+
+        # Blend mask ke image
+        color_mask = np.zeros_like(image_np)
+        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)
+
+        # 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)
+
+    # Convert kembali ke PIL image
+    final_image = Image.fromarray(image_np.astype(np.uint8))
+
+    # Legend
+    legend_height = 50
+    legend = Image.new("RGB", (final_image.width, legend_height), (255, 255, 255))
+    draw_legend = ImageDraw.Draw(legend)
+    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)
+        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
+
+# === Gradio Interface ===
+iface = gr.Interface(
+    fn=predict_segmentation_with_legend,
+    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."
+)
+
+if __name__ == "__main__":
+    iface.launch()