import torch
import gradio as gr

model = torch.hub.load('./', 'custom', path='best.pt', source='local')
#  #torch.hub 是一个PyTorch提供的工具，用于从GitHub仓库或者本地文件加载模型
# load() 函数用于从指定的地址中加载模型，支持两种来源：source='github' 表示从GitHub仓库加载，source='local' 表示从本地文件加载
# "ultralytics/yolov5" 是模型所在的GitHub仓库名
# "custom" 是指模型的名称，表示这是一个自定义训练的模型
# "path/to/weights.pt" 是模型文件的路径，它指向了已经保存的权重文件，权重文件是模型的核心，包含了模型的所有参数和权重信息
title = "安全帽、火焰识别平台"
desc = "这是一个基于yolov5的安全帽、火焰识别平台"

base_conf,base_iou = 0.25, 0.45

def det_image(img, conf, iou):
    model.conf = conf #conf名字可以改
    model.iou = iou
    return model(img).render()[0]
gr.Interface(inputs=[gr.Webcam(), gr.Slider(minimum=0, maximum=1, value=base_conf), gr.Slider(minimum=0, maximum=1, value=base_iou)],  # 启动摄像头
#gr.Interface(inputs=["image", gr.Slider(minimum=0, maximum=1, value=base_conf), gr.Slider(minimum=0, maximum=1, value=base_iou)],# slider是滑块
             outputs=["image"],
             fn=det_image,
             title=title,
             description=desc,
             live=True, # 输入端变化时输出端也变化
             examples=[["./data/images/fire2.jpeg", base_conf, base_iou], ["./data/images/helmet2.jpeg", base_conf, base_iou]]).launch(share=True)
                                                                                                                                 #share=Ture生成公网地址