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

# Khởi tạo model YOLOv8s
model = YOLO("yolov10n.pt")  # Đường dẫn tới model YOLOv10n đã huấn luyện sẵn

# Hàm xử lý dự đoán ảnh
def predict_image(image):
    # Chuyển đổi ảnh đầu vào thành định dạng NumPy
    np_image = np.array(image)
    # Chạy dự đoán với YOLO
    results = model.predict(source=np_image)
    # Vẽ bounding boxes lên ảnh gốc
    result_image = results[0].plot()
    return result_image

# Tạo giao diện Gradio
iface = gr.Interface(
    fn=predict_image,                    # Hàm xử lý đầu vào
    inputs=gr.Image(type="pil"),         # Tệp ảnh đầu vào
    outputs=gr.Image(type="pil"),        # Kết quả đầu ra là ảnh
    title="YOLOv10 Image Prediction",
    description="Tải ảnh lên và nhấn vào Predict để xem dự đoán với YOLOv10."
)

# Khởi chạy ứng dụng
iface.launch()