Update app.py

app.py

@@ -1,60 +1,40 @@
-import gradio as gr
-import onnxruntime as ort
-import numpy as np
+import torch
 from PIL import Image
-import cv2
-
-# Muat model ONNX
-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):
-    # 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
+import gradio as gr
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=predict_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),  # Input gambar
-        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),  # Confidence threshold
-        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="ONNX YOLO - Custom Model",  # Judul aplikasi
-    description="Upload images for custom YOLO object detection using ONNX model.",  # Deskripsi aplikasi
+# Load trained model
+model = torch.load("best.pt", map_location=torch.device("cpu"))
+model.eval()
+
+# Function to process image
+def predict(image):
+    # Convert to PIL image if needed
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Preprocess image (adjust as per your model's requirements)
+    results = model([image])  # Assuming YOLOv11 inference works like this
+    detections = results.xyxy[0].numpy()  # Extract bounding boxes, scores, etc.
+
+    # Draw boxes on image
+    for box in detections:
+        x1, y1, x2, y2, conf, cls = box
+        label = f"Class {int(cls)}: {conf:.2f}"
+        draw = ImageDraw.Draw(image)
+        draw.rectangle([(x1, y1), (x2, y2)], outline="red", width=3)
+        draw.text((x1, y1), label, fill="red")
+    
+    return image
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="image",
+    title="YOLOv11 Object Detection",
+    description="Upload an image to detect objects using YOLOv11.",
 )
 
-# Luncurkan aplikasi
-iface.launch()
+# Launch app
+if __name__ == "__main__":
+    interface.launch()