import gradio as gr
import torch
from PIL import Image
from ultralytics import YOLO

model = YOLO(r'pcb-best.pt')

def predict(img, conf, iou):
    results = model.predict(img, conf=conf, iou=iou)
    for i, r in enumerate(results):
        # Plot results image
        im_bgr = r.plot()  # BGR-order numpy array
        im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image

        # Show results to screen (in supported environments)
        return im_rgb


base_conf, base_iou = 0.25, 0.45
title = "基于YOLO-V8的PCB电路板缺陷检测"
des = "鼠标点击上传图片即可检测缺陷，可通过鼠标调整预测置信度，还可点击网页最下方示例图片进行预测"
gr.Interface(inputs=['image',gr.Slider(maximum=1, minimum=0, value=base_conf), gr.Slider(maximum=1, minimum=0, value=base_iou)],
             outputs=["image"], examples='example1.jpg', fn=predict, title=title, description=des).launch()