Update app.py

app.py

@@ -1,91 +1,132 @@
-import os
-import cv2
-import numpy as np
 import gradio as gr
-import pandas as pd
+from PIL import Image, ImageDraw, ImageFont
 from ultralytics import YOLO
-from pathlib import Path
-
-# === Konfigurasi ===
-model_path = "best.pt"  # Ganti dengan path model kamu
-class_names = ['coral or rock', 'pipeline', 'ripple marks', 'shipwreck']
-model = YOLO(model_path)
-
-# === Fungsi untuk menggambar prediksi lengkap ===
-def draw_predictions(img, segments, class_ids, confs, color=(255, 0, 0), alpha=0.5):
-    overlay = img.copy()
-    for i, seg in enumerate(segments):
-        polygon = np.array(seg).reshape(-1, 2)
-        polygon[:, 0] *= img.shape[1]
-        polygon[:, 1] *= img.shape[0]
-        polygon = polygon.astype(np.int32)
-
-        cls_id = int(class_ids[i])  # pastikan integer
-        conf = confs[i]
-
-        # Draw mask
-        cv2.fillPoly(overlay, [polygon], color)
-
-        # Draw bounding box
-        x, y, w, h = cv2.boundingRect(polygon)
-        cv2.rectangle(overlay, (x, y), (x + w, y + h), color, 2)
-
-        # Draw label + confidence
-        label = f"{class_names[cls_id]} ({conf:.2f})"
-        cv2.putText(overlay, label, (x, y - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 2)
-
-    return cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
-
-# === Fungsi proses gambar dan hasil deteksi ===
-def process_image_and_report(image):
-    temp_path = "temp.jpg"
-    image.save(temp_path)
-    img = cv2.imread(temp_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-    # === Prediksi menggunakan YOLOv8
-    results = model(temp_path)[0]
-    segs = [seg.xy for seg in results.masks] if results.masks else []
-    cls_ids = results.boxes.cls.tolist() if results.boxes else []
-    confs = results.boxes.conf.tolist() if results.boxes else []
-    xywh = results.boxes.xywhn.tolist() if results.boxes else []
-
-    # === Visualisasi prediksi
-    img_pred = draw_predictions(img.copy(), segs, cls_ids, confs)
-
-    # === Buat laporan CSV
-    rows = []
-    for i in range(len(cls_ids)):
-        rows.append({
-            "Filename": Path(temp_path).name,
-            "Class": class_names[int(cls_ids[i])],
-            "Confidence": round(float(confs[i]), 4),
-            "X_center": round(xywh[i][0], 4),
-            "Y_center": round(xywh[i][1], 4),
-            "Width": round(xywh[i][2], 4),
-            "Height": round(xywh[i][3], 4),
-        })
-
-    df = pd.DataFrame(rows)
-    csv_path = "detection_report.csv"
-    df.to_csv(csv_path, index=False)
-
-    return img_pred, csv_path
-
-# === Gradio UI ===
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 YOLOv8 Segmentation Viewer + Detection Report")
-    gr.Markdown("Upload image → lihat hasil deteksi (mask, box, label, confidence) → unduh CSV")
-
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Image")
-            submit_btn = gr.Button("Predict and Generate Report")
-        with gr.Column():
-            image_output = gr.Image(type="numpy", label="Prediction Result")
-            download_csv = gr.File(label="Download Detection CSV")
-
-    submit_btn.click(fn=process_image_and_report, inputs=image_input, outputs=[image_output, download_csv])
+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
+
+# Label dan Warna untuk tiap class
+label_map = {
+    0: "coral or rock",
+    1: "pipeline",
+    2: "ripple marks",
+    3: "shipwreck"
+}
+
+color_map = {
+    0: (0, 255, 0),      # coral or rock - hijau
+    1: (255, 0, 0),      # pipeline - merah
+    2: (255, 165, 0),    # ripple marks - oranye
+    3: (0, 0, 255)       # shipwreck - biru
+}
+
+def predict_segmentation_with_legend(image):
+    try:
+        if model is None:
+            print("❌ Model belum dimuat.")
+            return image
+
+        image = image.convert("RGB")
+        print("📥 Gambar diterima:", image.size)
+
+        results = model.predict(image, conf=0.0, iou=0.5)
+        result = results[0]
+
+        masks = result.masks.data.cpu().numpy() if result.masks else []
+        boxes = result.boxes.xyxy.cpu().numpy() if result.boxes else []
+        scores = result.boxes.conf.cpu().numpy() if result.boxes else []
+        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)
+        except:
+            font = ImageFont.load_default()
+
+        CONFIDENCE_THRESHOLD = 0.5
+
+        if len(boxes) == 0:
+            draw.text((10, 10), "No detections", fill="red", font=font)
+            return image
+
+        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]
+
+        # Tampilkan mask
+        for i in indices_to_use:
+            mask = masks[i]
+            class_id = class_ids[i]
+            color = color_map.get(class_id, (255, 255, 0))
+            mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
+            mask_resized = mask_pil.resize((image_np.shape[1], image_np.shape[0]), resample=Image.BILINEAR)
+            mask_resized = np.array(mask_resized) / 255.0
+
+            color_mask = np.zeros_like(image_np)
+            for c in range(3):
+                color_mask[:, :, c] = mask_resized * color[c]
+
+            image_np = np.where(mask_resized[..., None] > 0.5,
+                                image_np * 0.5 + color_mask * 0.5, image_np)
+
+        final_image = Image.fromarray(image_np.astype(np.uint8))
+        draw = ImageDraw.Draw(final_image)
+
+        # Tampilkan box & label
+        for i in indices_to_use:
+            box = boxes[i].astype(int).tolist()
+            score = scores[i]
+            class_id = class_ids[i]
+            label = label_map.get(class_id, str(class_id))
+            color = color_map.get(class_id, (255, 255, 0))
+            text_color = "white" if score >= CONFIDENCE_THRESHOLD else "gray"
+
+            draw.rectangle(box, outline=color, width=2)
+            draw.text((box[0], box[1] - 30), label, fill=text_color, font=font)
+            draw.text((box[0], box[1] - 10), f"Confidence: {score:.2f}", fill=text_color, font=font)
+
+        # Buat legenda
+        legend = Image.new("RGB", (500, 50), (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 += 130
+
+        # Gabungkan dengan gambar
+        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))
+
+        return combined
+
+    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="YOLOv11 Segmentasi Citra Sonar",
+    description="Upload citra sonar dan dapatkan hasil segmentasi lengkap dengan warna mask, label, confidence, dan legenda warna.",
+    allow_flagging="never"
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch(share=True)