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

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

def predict(image):
    # ทำการทำนาย
    results = model(image)
    df = results.pandas().xyxy[0]  # ผลลัพธ์การทำนายในรูปแบบ pandas DataFrame

    # วาด bounding boxes และ labels บนภาพ
    for _, row in df.iterrows():
        label = row['name']
        confidence = row['confidence']
        x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])

        # วาด 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)

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

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