Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
 import tempfile
 import os
+import cv2
 from sahi import AutoDetectionModel
 from sahi.predict import get_sliced_prediction
 from PIL import Image
+import numpy as np
 
 # Inisialisasi model deteksi menggunakan SAHI
 model_path = "best.pt"  # Ganti dengan path model YOLO lokal Anda
@@ -40,10 +42,32 @@ def detect_objects(image):
         class_count = {}
         total_count = 0  # Menyimpan total jumlah objek
 
+        # Menggambar bounding boxes pada gambar
+        output_image = np.array(image)  # Convert PIL Image to numpy array for OpenCV processing
+
         for prediction in results.object_prediction_list:
+            bbox = prediction.bbox
             class_name = prediction.category.name  # Nama kelas objek
-            class_count[class_name] = class_count.get(class_name, 0) + 1
-            total_count += 1  # Menambah jumlah objek
+            confidence = prediction.score.value  # Skor prediksi
+            
+            # Hanya gambar bounding box jika skor kepercayaan lebih besar dari threshold
+            if confidence >= confidence_threshold:
+                # Gambar bounding box
+                cv2.rectangle(output_image, 
+                              (int(bbox.minx), int(bbox.miny)), 
+                              (int(bbox.maxx), int(bbox.maxy)), 
+                              (0, 255, 0), 2)  # Gambar kotak hijau
+
+                # Gambar label dan skor
+                cv2.putText(output_image, 
+                            f"{class_name} {confidence:.2f}", 
+                            (int(bbox.minx), int(bbox.miny) - 10),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, 
+                            (0, 255, 0), 2)
+
+                # Hitung jumlah objek per kelas
+                class_count[class_name] = class_count.get(class_name, 0) + 1
+                total_count += 1  # Menambah jumlah objek
 
         # Menyusun output berupa string hasil perhitungan
         result_text = "Detected Objects:\n\n"
@@ -51,9 +75,10 @@ def detect_objects(image):
             result_text += f"{class_name}: {count}\n"
         result_text += f"\nTotal Objects: {total_count}"
 
-        # Menyimpan gambar dengan prediksi
+        # Convert output_image (numpy array) back to PIL Image to save
+        output_image_pil = Image.fromarray(output_image)
         output_image_path = "/tmp/prediction.jpg"
-        results.save(output_image_path)  # Menyimpan gambar dengan prediksi
+        output_image_pil.save(output_image_path)  # Menyimpan gambar dengan prediksi
 
     except Exception as err:
         # Menangani kesalahan lain