from ultralytics import YOLO
from PIL import Image
import cv2
import gradio as gr
 
model_seg = YOLO('./runs/segment/train/weights/best.pt')# 加载训练好的模型，换成自己的即可
model_detect = YOLO('./runs/detect/train2/weights/best.pt')
 
def seg(image):
    results = model_seg([image])  # 返回 Results 对象列表
    for r in results:
        im_array = r.plot()  # 绘制包含预测结果的BGR numpy数组
        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL图像
        # im.show()  # 显示图像
        # im.save('results_seg.jpg')  # 保存图像
    return im
 
def det(image):
    results = model_detect([image])  # 返回 Results 对象列表
    for r in results:
        im_array = r.plot()  # 绘制包含预测结果的BGR numpy数组
        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL图像
        # im.show()  # 显示图像
        # im.save('results_seg.jpg')  # 保存图像
    return im
  
with gr.Blocks() as demo:
 
 
    with gr.Tab("语义分割"):
        gr.Markdown("# yolov8语义分割演示")
        # 在这个域里都按列
        with gr.Row():
            in_img = gr.Image(sources='upload', type='pil')
            out_img = gr.Image(type='pil')
 
        # 所有组建默认是垂直排列gr.Examples　gr.Button　gr.Markdown都是组件
        gr.Examples(examples=['./images/test/001.jpg'], inputs=[in_img])
        button = gr.Button("执行分割", variant="primary")
        button.click(seg,
                        inputs=in_img,
                        outputs=out_img,
                        concurrency_limit=10)
 
    with gr.Tab("目标检测"):
        gr.Markdown("# yolov8目标检测演示")
        with gr.Row():
            in_img = gr.Image(sources='upload', type='pil')
            out_img = gr.Image(type='pil')
 
        gr.Examples(examples=['./images/test/001.jpg'], inputs=[in_img])
        button = gr.Button("执行检测", variant="primary")
 
        button.click(det,
                        inputs=in_img,
                        outputs=out_img,
                        concurrency_limit=10)
 
gr.close_all()
demo.queue()
demo.launch()