import gradio as gr
import torch
import cv2
from PIL import Image
import numpy as np
from ultralytics import YOLO

# โหลดโมเดล YOLOv8 ที่ฝึกมาเอง
model = YOLO('epoch80.pt')  # เปลี่ยน 'best_V5.pt' เป็นโมเดลของคุณ

def predict(image):
    # ทำการทำนาย
    results = model(image)

    detected = False  # ตัวแปรเพื่อตรวจสอบว่ามีการตรวจพบหรือไม่

    # วาด bounding boxes และ labels บนภาพ
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()
        confidences = result.boxes.conf.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy()
        
        for box, confidence, class_id in zip(boxes, confidences, class_ids):
            x1, y1, x2, y2 = map(int, box[:4])
            label = result.names[int(class_id)]  # ดึง label จาก class_id

            # วาด bounding box
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # วาด label และ confidence
            label_text = f"{label} {confidence:.2f}"
            cv2.putText(image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            detected = True  # ตั้งค่าตัวแปรเป็น True ถ้ามีการตรวจพบ

    if not detected:  # ถ้าไม่มีการตรวจพบ
        cv2.putText(image, "No detected", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    # แปลงภาพกลับเป็นรูปแบบที่ Gradio สามารถแสดงได้
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    return pil_image

demo = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="image")
demo.launch()