Update app.py

app.py

@@ -5,18 +5,15 @@ import gradio as gr
 import pandas as pd
 from ultralytics import YOLO
 from pathlib import Path
-from datetime import datetime
 
 # === Konfigurasi ===
-model_path = "best.pt"  # pastikan file ini ada di root repo Hugging Face
-gt_label_dir = "data/labels"  # sesuaikan dengan struktur foldermu
+model_path = "best.pt"
 class_names = ['coral or rock', 'pipeline', 'ripple marks', 'shipwreck']
 
-# Load model
 model = YOLO(model_path)
 
-# === Fungsi untuk menggambar segmentasi ===
-def draw_mask(img, segments, class_ids, color=(255, 0, 0), alpha=0.5, label_type="Pred", confs=None):
+# === Fungsi untuk menggambar mask hasil deteksi ===
+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)
@@ -24,57 +21,31 @@ def draw_mask(img, segments, class_ids, color=(255, 0, 0), alpha=0.5, label_type
         polygon[:, 1] *= img.shape[0]
         polygon = polygon.astype(np.int32)
 
-        cls_id = int(class_ids[i])  # FIX: pastikan integer
-        label = f"{label_type}: {class_names[cls_id]}"
-        if confs and i < len(confs):
-            label += f" ({confs[i]:.2f})"
-
+        label = f"{class_names[int(class_ids[i])]} ({confs[i]:.2f})"
         cv2.fillPoly(overlay, [polygon], color)
         x, y, w, h = cv2.boundingRect(polygon)
         cv2.putText(overlay, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,255,255), 1)
+        cv2.rectangle(overlay, (x, y), (x + w, y + h), color, 1)
 
     return cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
 
-# === Proses prediksi + simpan CSV ===
-def process_image_and_report(image):
-    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-    temp_path = f"/tmp/temp_{timestamp}.jpg"
-    csv_path = f"/tmp/detection_report_{timestamp}.csv"
-
+# === Fungsi utama proses dan simpan report ===
+def process_image_and_generate_report(image):
+    temp_path = "temp.jpg"
     image.save(temp_path)
     img = cv2.imread(temp_path)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
-    # === Load ground truth ===
-    label_path = os.path.join(gt_label_dir, Path(temp_path).with_suffix(".txt").name)
-    img_gt = img.copy()
-    if os.path.exists(label_path):
-        with open(label_path, "r") as f:
-            gt_segments = []
-            gt_ids = []
-            for line in f:
-                parts = line.strip().split()
-                if len(parts) < 3:
-                    continue
-                cls_id = int(parts[0])
-                coords = list(map(float, parts[1:]))
-                gt_segments.append(np.array(coords).reshape(-1, 2))
-                gt_ids.append(cls_id)
-            img_gt = draw_mask(img_gt, gt_segments, gt_ids, color=(0, 255, 0), label_type="GT")
-
-    # === Deteksi dengan YOLOv8 ===
+    # Prediksi dari 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 []
 
-    img_pred = draw_mask(img.copy(), segs, cls_ids, color=(255, 0, 0), label_type="Pred", confs=confs)
-
-    # === Gabung GT dan Prediksi
-    combined = np.concatenate((img_gt, img_pred), axis=1)
+    img_result = draw_predictions(img.copy(), segs, cls_ids, confs)
 
-    # === Simpan laporan ke CSV
+    # Simpan CSV
     rows = []
     for i in range(len(cls_ids)):
         rows.append({
@@ -88,24 +59,27 @@ def process_image_and_report(image):
         })
 
     df = pd.DataFrame(rows)
+    csv_path = "detection_result.csv"
     df.to_csv(csv_path, index=False)
 
-    return combined, csv_path
+    return img_result, csv_path
 
 # === UI Gradio ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 YOLOv8 Segmentasi Viewer + Report Generator")
-    gr.Markdown("Upload gambar, lihat segmentasi prediksi vs GT, dan unduh laporan CSV deteksi.")
+    gr.Markdown("## 🎯 YOLOv8 Segmentasi: Prediksi dan Laporan Otomatis")
+    gr.Markdown("Upload gambar untuk mendeteksi objek dan mengunduh hasil deteksi sebagai CSV")
 
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Image")
-            submit_btn = gr.Button("Generate Visual & Report")
+            run_btn = gr.Button("Deteksi dan Simpan Report")
         with gr.Column():
-            image_output = gr.Image(type="numpy", label="GT vs Prediction")
-            download_csv = gr.File(label="Download Detection CSV")
+            result_image = gr.Image(type="numpy", label="Hasil Prediksi")
+            download_btn = gr.File(label="Unduh CSV Deteksi")
 
-    submit_btn.click(fn=process_image_and_report, inputs=image_input, outputs=[image_output, download_csv])
+    run_btn.click(fn=process_image_and_generate_report,
+                  inputs=image_input,
+                  outputs=[result_image, download_btn])
 
 if __name__ == "__main__":
     demo.launch()