Spaces:
Runtime error
Runtime error
| # Gradio YOLOv5 Det v0.2.2 | |
| # 创建人：曾逸夫 | |
| # 创建时间：2022-05-11 | |
| # email：zyfiy1314@163.com | |
| # 项目主页：https://gitee.com/CV_Lab/gradio_yolov5_det | |
| # import os | |
| # os.system("pip install gradio==2.9b50") | |
| import argparse | |
| import csv | |
| import json | |
| import sys | |
| from pathlib import Path | |
| import gradio as gr | |
| import torch | |
| import yaml | |
| from PIL import Image, ImageDraw, ImageFont | |
| from util.fonts_opt import is_fonts | |
| from util.pdf_opt import pdf_generate | |
| ROOT_PATH = sys.path[0] # 根目录 | |
| # 模型路径 | |
| model_path = "ultralytics/yolov5" | |
| # Gradio YOLOv5 Det版本 | |
| GYD_VERSION = "Gradio YOLOv5 Det v0.2.2" | |
| # 模型名称临时变量 | |
| model_name_tmp = "" | |
| # 设备临时变量 | |
| device_tmp = "" | |
| # 文件后缀 | |
| suffix_list = [".csv", ".yaml"] | |
| # 字体大小 | |
| FONTSIZE = 25 | |
| def parse_args(known=False): | |
| parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det v0.2.2") | |
| parser.add_argument( | |
| "--model_name", "-mn", default="yolov5s", type=str, help="model name" | |
| ) | |
| parser.add_argument( | |
| "--model_cfg", | |
| "-mc", | |
| default="./model_config/model_name_p5_all.yaml", | |
| type=str, | |
| help="model config", | |
| ) | |
| parser.add_argument( | |
| "--cls_name", | |
| "-cls", | |
| default="./cls_name/cls_name.yaml", | |
| type=str, | |
| help="cls name", | |
| ) | |
| parser.add_argument( | |
| "--nms_conf", | |
| "-conf", | |
| default=0.5, | |
| type=float, | |
| help="model NMS confidence threshold", | |
| ) | |
| parser.add_argument( | |
| "--nms_iou", "-iou", default=0.45, type=float, help="model NMS IoU threshold" | |
| ) | |
| parser.add_argument( | |
| "--label_dnt_show", | |
| "-lds", | |
| action="store_true", | |
| default=False, | |
| help="label show", | |
| ) | |
| parser.add_argument( | |
| "--device", "-dev", default="cpu", type=str, help="cuda or cpu", | |
| ) | |
| parser.add_argument( | |
| "--inference_size", "-isz", default=640, type=int, help="model inference size" | |
| ) | |
| args = parser.parse_known_args()[0] if known else parser.parse_args() | |
| return args | |
| # yaml文件解析 | |
| def yaml_parse(file_path): | |
| return yaml.safe_load(open(file_path, encoding="utf-8").read()) | |
| # yaml csv 文件解析 | |
| def yaml_csv(file_path, file_tag): | |
| file_suffix = Path(file_path).suffix | |
| if file_suffix == suffix_list[0]: | |
| # 模型名称 | |
| file_names = [i[0] for i in list(csv.reader(open(file_path)))] # csv版 | |
| elif file_suffix == suffix_list[1]: | |
| # 模型名称 | |
| file_names = yaml_parse(file_path).get(file_tag) # yaml版 | |
| else: | |
| print(f"{file_path}格式不正确！程序退出！") | |
| sys.exit() | |
| return file_names | |
| # 模型加载 | |
| def model_loading(model_name, device): | |
| # 加载模型 | |
| model = torch.hub.load( | |
| model_path, model_name, force_reload=True, device=device, _verbose=False | |
| ) | |
| return model | |
| # 检测信息 | |
| def export_json(results, model, img_size): | |
| return [ | |
| [ | |
| { | |
| "id": i, | |
| "class": int(result[i][5]), | |
| # "class_name": model.model.names[int(result[i][5])], | |
| "class_name": model_cls_name_cp[int(result[i][5])], | |
| "normalized_box": { | |
| "x0": round(result[i][:4].tolist()[0], 6), | |
| "y0": round(result[i][:4].tolist()[1], 6), | |
| "x1": round(result[i][:4].tolist()[2], 6), | |
| "y1": round(result[i][:4].tolist()[3], 6), | |
| }, | |
| "confidence": round(float(result[i][4]), 2), | |
| "fps": round(1000 / float(results.t[1]), 2), | |
| "width": img_size[0], | |
| "height": img_size[1], | |
| } | |
| for i in range(len(result)) | |
| ] | |
| for result in results.xyxyn | |
| ] | |
| # 帧转换 | |
| def pil_draw(img, countdown_msg, textFont, xyxy, font_size, label_opt): | |
| img_pil = ImageDraw.Draw(img) | |
| img_pil.rectangle(xyxy, fill=None, outline="green") # 边界框 | |
| if label_opt: | |
| text_w, text_h = textFont.getsize(countdown_msg) # 标签尺寸 | |
| img_pil.rectangle( | |
| (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h), | |
| fill="green", | |
| outline="green", | |
| ) # 标签背景 | |
| img_pil.multiline_text( | |
| (xyxy[0], xyxy[1]), | |
| countdown_msg, | |
| fill=(205, 250, 255), | |
| font=textFont, | |
| align="center", | |
| ) | |
| return img | |
| # YOLOv5图片检测函数 | |
| def yolo_det( | |
| img, device, model_name, inference_size, conf, iou, label_opt, model_cls, opt | |
| ): | |
| global model, model_name_tmp, device_tmp | |
| if model_name_tmp != model_name: | |
| # 模型判断，避免反复加载 | |
| model_name_tmp = model_name | |
| model = model_loading(model_name_tmp, device) | |
| elif device_tmp != device: | |
| device_tmp = device | |
| model = model_loading(model_name_tmp, device) | |
| # -----------模型调参----------- | |
| model.conf = conf # NMS 置信度阈值 | |
| model.iou = iou # NMS IOU阈值 | |
| model.max_det = 1000 # 最大检测框数 | |
| model.classes = model_cls # 模型类别 | |
| results = model(img, size=inference_size) # 检测 | |
| img_size = img.size # 帧尺寸 | |
| # 加载字体 | |
| textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE) | |
| det_img = img.copy() | |
| for result in results.xyxyn: | |
| for i in range(len(result)): | |
| id = int(i) # 实例ID | |
| obj_cls_index = int(result[i][5]) # 类别索引 | |
| obj_cls = model_cls_name_cp[obj_cls_index] # 类别 | |
| # ------------边框坐标------------ | |
| x0 = float(result[i][:4].tolist()[0]) | |
| y0 = float(result[i][:4].tolist()[1]) | |
| x1 = float(result[i][:4].tolist()[2]) | |
| y1 = float(result[i][:4].tolist()[3]) | |
| # ------------边框实际坐标------------ | |
| x0 = int(img_size[0] * x0) | |
| y0 = int(img_size[1] * y0) | |
| x1 = int(img_size[0] * x1) | |
| y1 = int(img_size[1] * y1) | |
| conf = float(result[i][4]) # 置信度 | |
| # fps = f"{(1000 / float(results.t[1])):.2f}" # FPS | |
| det_img = pil_draw( | |
| img, | |
| f"{id}-{obj_cls}:{conf:.2f}", | |
| textFont, | |
| [x0, y0, x1, y1], | |
| FONTSIZE, | |
| label_opt, | |
| ) | |
| det_json = export_json(results, model, img.size)[0] # 检测信息 | |
| # JSON格式化 | |
| det_json_format = json.dumps( | |
| det_json, sort_keys=True, indent=4, separators=(",", ":"), ensure_ascii=False | |
| ) | |
| # -------pdf------- | |
| report = "./Det_Report.pdf" | |
| if "pdf" in opt: | |
| pdf_generate(f"{det_json_format}", report, GYD_VERSION) | |
| else: | |
| report = None | |
| if "json" not in opt: | |
| det_json = None | |
| return det_img, det_json, report | |
| def main(args): | |
| gr.close_all() | |
| global model, model_cls_name_cp | |
| slider_step = 0.05 # 滑动步长 | |
| nms_conf = args.nms_conf | |
| nms_iou = args.nms_iou | |
| label_opt = args.label_dnt_show | |
| model_name = args.model_name | |
| model_cfg = args.model_cfg | |
| cls_name = args.cls_name | |
| device = args.device | |
| inference_size = args.inference_size | |
| is_fonts(f"{ROOT_PATH}/fonts") # 检查字体文件 | |
| # 模型加载 | |
| model = model_loading(model_name, device) | |
| model_names = yaml_csv(model_cfg, "model_names") | |
| model_cls_name = yaml_csv(cls_name, "model_cls_name") | |
| model_cls_name_cp = model_cls_name.copy() # 类别名称 | |
| # -------------------输入组件------------------- | |
| inputs_img = gr.inputs.Image(type="pil", label="原始图片") | |
| inputs_device = gr.inputs.Dropdown( | |
| choices=["0", "cpu"], default=device, type="value", label="设备" | |
| ) | |
| inputs_model = gr.inputs.Dropdown( | |
| choices=model_names, default=model_name, type="value", label="模型" | |
| ) | |
| inputs_size = gr.inputs.Radio( | |
| choices=[320, 640], default=inference_size, label="推理尺寸" | |
| ) | |
| input_conf = gr.inputs.Slider( | |
| 0, 1, step=slider_step, default=nms_conf, label="置信度阈值" | |
| ) | |
| inputs_iou = gr.inputs.Slider( | |
| 0, 1, step=slider_step, default=nms_iou, label="IoU 阈值" | |
| ) | |
| inputs_label = gr.inputs.Checkbox(default=(not label_opt), label="标签显示") | |
| inputs_clsName = gr.inputs.CheckboxGroup( | |
| choices=model_cls_name, default=model_cls_name, type="index", label="类别" | |
| ) | |
| inputs_opt = gr.inputs.CheckboxGroup( | |
| choices=["pdf", "json"], default=["pdf"], type="value", label="操作" | |
| ) | |
| # 输入参数 | |
| inputs = [ | |
| inputs_img, # 输入图片 | |
| inputs_device, # 设备 | |
| inputs_model, # 模型 | |
| inputs_size, # 推理尺寸 | |
| input_conf, # 置信度阈值 | |
| inputs_iou, # IoU阈值 | |
| inputs_label, # 标签显示 | |
| inputs_clsName, # 类别 | |
| inputs_opt, # 检测操作 | |
| ] | |
| # 输出参数 | |
| outputs_img = gr.outputs.Image(type="pil", label="检测图片") | |
| outputs02_json = gr.outputs.JSON(label="检测信息") | |
| outputs03_pdf = gr.outputs.File(label="下载检测报告") | |
| outputs = [outputs_img, outputs02_json, outputs03_pdf] | |
| # 标题 | |
| title = "基于Gradio的YOLOv5通用目标检测系统v0.2.2" | |
| # 描述 | |
| description = "<div align='center'>可自定义目标检测模型、安装简单、使用方便</div><div align='center'>Customizable object detection model, easy to install and easy to use</div>" | |
| # 示例图片 | |
| examples = [ | |
| [ | |
| "./img_example/bus.jpg", | |
| "cpu", | |
| "yolov5s", | |
| 640, | |
| 0.6, | |
| 0.5, | |
| True, | |
| ["人", "公交车"], | |
| ["pdf"], | |
| ], | |
| [ | |
| "./img_example/Millenial-at-work.jpg", | |
| "cpu", | |
| "yolov5l", | |
| 320, | |
| 0.5, | |
| 0.45, | |
| True, | |
| ["人", "椅子", "杯子", "笔记本电脑"], | |
| ["json"], | |
| ], | |
| [ | |
| "./img_example/zidane.jpg", | |
| "cpu", | |
| "yolov5m", | |
| 640, | |
| 0.25, | |
| 0.5, | |
| False, | |
| ["人", "领带"], | |
| ["pdf", "json"], | |
| ], | |
| ] | |
| # 接口 | |
| gr.Interface( | |
| fn=yolo_det, | |
| inputs=inputs, | |
| outputs=outputs, | |
| title=title, | |
| description=description, | |
| # examples=examples, | |
| theme="seafoam", | |
| # live=True, # 实时变更输出 | |
| flagging_dir="run", # 输出目录 | |
| # flagging_options=["good", "generally", "bad"], | |
| # allow_flagging="auto", | |
| # ).launch(inbrowser=True, auth=['admin', 'admin']) | |
| ).launch( | |
| inbrowser=True, # 自动打开默认浏览器 | |
| show_tips=True, # 自动显示gradio最新功能 | |
| # favicon_path="./icon/logo.ico", | |
| ) | |
| if __name__ == "__main__": | |
| args = parse_args() | |
| main(args) | |