Update app.py

app.py

@@ -1,23 +1,47 @@
 import gradio as gr
-from ultralytics import YOLO
-from onnxruntime import InferenceSession
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+import cv2
 
 # Muat model ONNX
-model_path = "best.onnx"
-session = InferenceSession(model_path)
-# # Muat model custom yang telah dilatih
-# model = YOLO("best.pt")  # Pastikan file best.pt ada di direktori yang sama
+onnx_model_path = "best.onnx"  # Pastikan file best.onnx ada di direktori yang sama
+session = ort.InferenceSession(onnx_model_path)
 
 # Fungsi untuk melakukan prediksi
 def predict_image(img, conf_threshold, iou_threshold):
-    results = model.predict(
-        source=img,
-        conf=conf_threshold,
-        iou=iou_threshold,
-        show_labels=True,
-        show_conf=True,
-    )
-    return results[0].plot() if results else None
+    # Konversi gambar ke format yang bisa diterima model ONNX
+    img = np.array(img)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # Ubah ke format BGR (OpenCV default)
+
+    # Preprocessing: Resize dan normalisasi
+    img_resized = cv2.resize(img, (640, 640))  # Sesuaikan dengan ukuran input model
+    img_normalized = img_resized / 255.0  # Normalisasi ke rentang 0-1
+    img_input = np.expand_dims(img_normalized, axis=0).astype(np.float32)
+
+    # Lakukan inferensi dengan ONNX Runtime
+    inputs = {session.get_inputs()[0].name: img_input}
+    outputs = session.run(None, inputs)
+
+    # Ambil hasil deteksi dan bounding box (misalnya, hasil berada di output[0])
+    boxes = outputs[0]  # Sesuaikan dengan output yang relevan dari model ONNX
+    confidences = outputs[1]  # Confidence score
+    labels = outputs[2]  # Label kelas (jika ada)
+
+    # Filter prediksi dengan threshold
+    valid_boxes = []
+    for i, conf in enumerate(confidences):
+        if conf > conf_threshold:
+            valid_boxes.append(boxes[i])
+
+    # Plot hasil deteksi pada gambar
+    for box in valid_boxes:
+        x1, y1, x2, y2 = box
+        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
+
+    # Kembalikan hasil sebagai gambar dengan bounding box
+    result_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    return result_img
 
 # Gradio Interface
 iface = gr.Interface(
@@ -28,8 +52,8 @@ iface = gr.Interface(
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),  # IoU threshold
     ],
     outputs=gr.Image(type="pil", label="Result"),  # Output gambar dengan bounding box
-    title="Ultralytics Gradio YOLO11 - Custom Model",  # Judul aplikasi
-    description="Upload images for custom YOLO11 object detection.",  # Deskripsi aplikasi
+    title="ONNX YOLO - Custom Model",  # Judul aplikasi
+    description="Upload images for custom YOLO object detection using ONNX model.",  # Deskripsi aplikasi
 )
 
 # Luncurkan aplikasi