# Gradio YOLOv5 Det Blocks 02 # 创建人:曾逸夫 # 创建时间:2022-06-05 # 功能描述:单图片,双端操作 from util.gradio_version_opt import gr_v_opt gr_v_opt() import argparse import csv import gc import json import sys from collections import Counter from pathlib import Path import cv2 import gradio as gr import numpy as np import pandas as pd import plotly.express as px import torch import yaml from PIL import Image, ImageDraw, ImageFont ROOT_PATH = sys.path[0] # 根目录 # yolov5路径 yolov5_path = "ultralytics/yolov5" # 本地模型路径 local_model_path = f"{ROOT_PATH}/models" # Gradio YOLOv5 Det版本 GYD_VERSION = "Gradio YOLOv5 Det block 02" # 模型名称临时变量 model_name_tmp = "" # 设备临时变量 device_tmp = "" # 文件后缀 suffix_list = [".csv", ".yaml"] # 字体大小 FONTSIZE = 25 # 目标尺寸 obj_style = ["小目标", "中目标", "大目标"] def parse_args(known=False): parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det block 02") parser.add_argument( "--model_cfg_p5", "-mc5", default="./model_config/model_name_p5_all.yaml", type=str, help="model config", ) 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("--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): # 加载本地模型 try: torch.hub._validate_not_a_forked_repo = lambda a, b, c: True model = torch.hub.load( yolov5_path, "custom", path=f"{local_model_path}/{model_name}", device="cuda:0", force_reload=False, _verbose=True, ) except Exception as e: print("模型加载失败!") print(e) return False else: print(f"🚀 欢迎使用{GYD_VERSION},{model_name}加载成功!") return model # YOLOv5图片检测函数 def yolo_det(img, model_name, infer_size, conf, iou): global model, model_name_tmp if model_name_tmp != model_name: # 模型判断,避免反复加载 model_name_tmp = model_name print(f"正在加载模型{model_name_tmp}......") model = model_loading(model_name_tmp) else: print(f"正在加载模型{model_name_tmp}......") model = model_loading(model_name_tmp) # -----------模型调参----------- model.conf = conf # NMS 置信度阈值 model.iou = iou # NMS IOU阈值 model.max_det = 1000 # 最大检测框数 results = model(img, size=infer_size) # 检测 results.render() # 渲染 det_img = Image.fromarray(results.ims[0]) # 检测图片 return det_img def main(args): gr.close_all() slider_step = 0.05 # 滑动步长 nms_conf = args.nms_conf nms_iou = args.nms_iou model_cfg_p5 = args.model_cfg_p5 inference_size = args.inference_size # 模型加载 model_names_p5 = yaml_csv(model_cfg_p5, "model_names") with gr.Blocks() as gyd: with gr.Box(): with gr.Row(): gr.Markdown("### P5检测") with gr.Row(): with gr.Column(): with gr.Row(): inputs_img_p5 = gr.Image(image_mode="RGB", source="upload", type="pil", label="原始图片") with gr.Row(): inputs_model_p5 = gr.Radio(choices=model_names_p5, value="yolov5s", label="P5模型") with gr.Row(): inputs_size_p5 = gr.Radio(choices=[320, 640, 1280], value=inference_size, label="推理尺寸") with gr.Row(): det_btn_01 = gr.Button(value='Detect 01') with gr.Column(): with gr.Row(): outputs_img_p5 = gr.Image(type="pil", label="检测图片") with gr.Row(): input_conf_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_conf, label="置信度阈值") with gr.Row(): inputs_iou_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_iou, label="IoU 阈值") with gr.Row(): det_btn_02 = gr.Button(value='Detect 02') det_btn_01.click(fn=yolo_det, inputs=[inputs_img_p5, inputs_model_p5, inputs_size_p5, input_conf_p5, inputs_iou_p5], outputs=[ outputs_img_p5,]) det_btn_02.click(fn=yolo_det, inputs=[inputs_img_p5, inputs_model_p5, inputs_size_p5, input_conf_p5, inputs_iou_p5], outputs=[ outputs_img_p5,]) gyd.launch(inbrowser=True) if __name__ == '__main__': args = parse_args() main(args)