import torch
import gradio as gr

model = torch.hub.load("yolov5-code-main/", "custom", path="多病变识别.pt", source="local",force_reload=True)
model_paths = ["多病变识别.pt", "胃部息肉识别.pt", "YOLOV5官方模型.pt","best.pt"]
title = "✨基于YOLOV5的消化道多病变检测平台✨"

desc = "<img src='https://img.zcool.cn/community/0142eb5992b0650000002129b61cfc.gif' />"
base_conf, base_iou = 0.15, 0.5

def load_model(model_path):
    model = torch.hub.load("yolov5-code-main/", "custom", path=model_path, source="local")
    return model



def det_image(img, model_path, conf_最小置信度, iou_最大置信度):
    model = load_model(model_path)
    model.conf = conf_最小置信度
    model.iou = iou_最大置信度

    return model(img).render()[0]
def add_text():
    main_text = "多病变模型可识别的病变类型为：barrett(巴雷特食道),esophageal-cancer(食道癌),cancer(胃癌),esophageal papilloma(食道乳头状瘤),gastric polyp(胃息肉),mycotic esophagitis(真菌性食道炎),GIM(胃肠上皮化生),esophageal varices(食道静脉曲张),Reflux esophagitis(反流性食道炎)"
    extra_text = ""
    return main_text + "<br>" + extra_text
    
def add1_text():
    main1_text = ""
    extra1_text = "该平台基于gradio免费服务器构建，其处理速度受公共服务器影响可能略慢，望理解！"
    return main1_text + "<br>" + extra1_text

iface = gr.Interface(inputs=["image", gr.inputs.Dropdown(model_paths, label="选择模型"), gr.Slider(minimum=0, maximum=1, value=base_conf), gr.Slider(minimum=0, maximum=1, value=base_iou)],
             outputs=["image"],
             fn=det_image,
             title=title,
             description="<img src='https://img.zcool.cn/community/01b04956170efa32f875313d664d44.gif' width='300' height=''/>"+ "<p style='text-align:center; font-size:16px; color:#FFA500; margin-top:10px;'>" + add1_text() + "</p>",
             article="<img src='https://gimg2.baidu.com/image_search/src=http%3A%2F%2Fb-ssl.duitang.com%2Fuploads%2Fitem%2F201505%2F24%2F20150524222415_hwTnY.thumb.700_0.gif&refer=http%3A%2F%2Fb-ssl.duitang.com&app=2002&size=f9999,10000&q=a80&n=0&g=0n&fmt=auto?sec=1697462434&t=be8ccd6bf414a078c71ef063f3da378a' width='100' height='100'/>" + "<p style='text-align:center; font-size:16px; color:#222; margin-top:10px;'>" + add_text() + "</p>",
             live=True,
             examples=[["CS/GIM321.jpg", "多病变识别.pt", base_conf, base_iou], ["CS/barrett278.png", "多病变识别.pt", 0.3, base_iou], ["CS/RERD65.jpg", "多病变识别.pt", base_conf, base_iou],["CS/400.jpg", "胃部息肉识别.pt", base_conf, base_iou],["CS/bus.jpg", "YOLOV5官方模型.pt", base_conf, base_iou]]
).launch(auth=("admin", "1234"), share=False)