app.py

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

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

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

    # วาด bounding boxes และ labels บนภาพ
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()
        labels = result.names
        confidences = result.boxes.conf.cpu().numpy()
        
        for box, confidence in zip(boxes, confidences):
            x1, y1, x2, y2 = map(int, box)
            label = labels[box[5]]  # Assuming the label is stored in the last column (index 5)

            # วาด 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.Image(type="numpy"), outputs="image")
demo.launch()