代码拉取完成,页面将自动刷新
同步操作将从 CV_Lab/Gradio-YOLOv5-Det-Blocks 强制同步,此操作会覆盖自 Fork 仓库以来所做的任何修改,且无法恢复!!!
确定后同步将在后台操作,完成时将刷新页面,请耐心等待。
# Gradio YOLOv5 Det Blocks 10_03
# 创建人:曾逸夫
# 创建时间:2022-08-09
# 功能描述:图片批量检测,批量下载
from util.gradio_version_opt import gr_v_opt
gr_v_opt()
import argparse
import csv
import gc
import io
import json
import os
import shutil
import sys
import zipfile
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 10_03"
# 模型名称临时变量
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 10_03")
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 dir_opt(target_dir):
if os.path.exists(target_dir):
shutil.rmtree(target_dir)
os.mkdir(target_dir)
else:
os.mkdir(target_dir)
# 文件压缩
def zipDir(origin_dir, compress_file):
# --------------- 压缩 ---------------
zip = zipfile.ZipFile(f"{compress_file}", "w", zipfile.ZIP_DEFLATED)
for path, dirnames, filenames in os.walk(f"{origin_dir}"):
fpath = path.replace(f"{origin_dir}", '')
for filename in filenames:
zip.write(os.path.join(path, filename), os.path.join(fpath, filename))
zip.close()
# 模型加载
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(imgs, 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 # 最大检测框数
det_imgs = []
img_save_dir = "./blocks_10"
dir_opt(img_save_dir)
for img in imgs:
img_ = Image.open(img.name)
results = model(img_, size=infer_size) # 检测
results.render() # 渲染
det_img = Image.fromarray(results.ims[0]) # 检测图片
det_imgs.append(det_img)
img_name_ = (img.name[:-12] + img.name[-4:]).split("/")[2]
det_img.save(f"{img_save_dir}/{img_name_}")
zipDir(img_save_dir, "det_imgs.zip") # 下载打包文件
return det_imgs, "det_imgs.zip"
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.File(file_count="multiple", type="file", 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():
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_01 = gr.Button(value='Detect 01', variant="primary")
with gr.Column():
with gr.Row():
outputs_img_p5 = gr.Gallery(label="检测图片")
with gr.Row():
outputs_file_p5 = gr.File(type="file", label="检测图片")
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, outputs_file_p5])
gyd.launch(inbrowser=True)
if __name__ == '__main__':
args = parse_args()
main(args)
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。