|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import argparse |
|
import csv |
|
import sys |
|
|
|
csv.field_size_limit(sys.maxsize) |
|
|
|
import gc |
|
import json |
|
import random |
|
from collections import Counter |
|
from pathlib import Path |
|
|
|
import cv2 |
|
import gradio as gr |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
import pandas as pd |
|
import plotly.express as px |
|
from matplotlib import font_manager |
|
|
|
ROOT_PATH = sys.path[0] |
|
|
|
|
|
SimSun_path = f"{ROOT_PATH}/fonts/SimSun.ttf" |
|
TimesNesRoman_path = f"{ROOT_PATH}/fonts/TimesNewRoman.ttf" |
|
|
|
SimSun = font_manager.FontProperties(fname=SimSun_path, size=12) |
|
|
|
TimesNesRoman = font_manager.FontProperties(fname=TimesNesRoman_path, size=12) |
|
|
|
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] |
|
|
|
|
|
yolov5_path = "ultralytics/yolov5" |
|
|
|
|
|
local_model_path = f"{ROOT_PATH}/models" |
|
|
|
|
|
GYD_VERSION = "Gradio YOLOv5 Det v0.5" |
|
|
|
|
|
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 v0.5") |
|
parser.add_argument("--source", "-src", default="upload", type=str, help="image input source") |
|
parser.add_argument("--source_video", "-src_v", default="upload", type=str, help="video input source") |
|
parser.add_argument("--img_tool", "-it", default="editor", type=str, help="input image tool") |
|
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_p6_all.yaml", |
|
type=str, |
|
help="model config", |
|
) |
|
parser.add_argument( |
|
"--cls_name", |
|
"-cls", |
|
default="./cls_name/cls_name_zh.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( |
|
"--device", |
|
"-dev", |
|
default="cuda:0", |
|
type=str, |
|
help="cuda or cpu", |
|
) |
|
parser.add_argument("--inference_size", "-isz", default=640, type=int, help="model inference size") |
|
parser.add_argument("--max_detnum", "-mdn", default=50, type=float, help="model max det num") |
|
parser.add_argument("--slider_step", "-ss", default=0.05, type=float, help="slider step") |
|
parser.add_argument( |
|
"--is_login", |
|
"-isl", |
|
action="store_true", |
|
default=False, |
|
help="is login", |
|
) |
|
parser.add_argument('--usr_pwd', |
|
"-up", |
|
nargs='+', |
|
type=str, |
|
default=["admin", "admin"], |
|
help="user & password for login") |
|
parser.add_argument( |
|
"--is_share", |
|
"-is", |
|
action="store_true", |
|
default=False, |
|
help="is login", |
|
) |
|
|
|
args = parser.parse_known_args()[0] if known else parser.parse_args() |
|
return args |
|
|
|
|
|
|
|
def yaml_parse(file_path): |
|
return yaml.safe_load(open(file_path, encoding="utf-8").read()) |
|
|
|
|
|
|
|
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)))] |
|
elif file_suffix == suffix_list[1]: |
|
|
|
file_names = yaml_parse(file_path).get(file_tag) |
|
else: |
|
print(f"{file_path}格式不正确!程序退出!") |
|
sys.exit() |
|
|
|
return file_names |
|
|
|
|
|
|
|
def check_online(): |
|
|
|
|
|
import socket |
|
try: |
|
socket.create_connection(("1.1.1.1", 443), 5) |
|
return True |
|
except OSError: |
|
return False |
|
|
|
|
|
|
|
def model_loading(model_name, device, opt=[]): |
|
|
|
|
|
try: |
|
model = torch.hub.load( |
|
yolov5_path, |
|
model_name, |
|
device=device, |
|
force_reload=[True if "refresh_yolov5" in opt and check_online() else False][0], |
|
_verbose=True, |
|
) |
|
except Exception as e: |
|
print("模型加载失败!") |
|
print(e) |
|
return False |
|
else: |
|
print(f"🚀 欢迎使用{GYD_VERSION},{model_name}加载成功!") |
|
|
|
return model |
|
|
|
|
|
|
|
def export_json(results, img_size): |
|
|
|
return [[{ |
|
"ID": i, |
|
"CLASS": int(result[i][5]), |
|
"CLASS_NAME": model_cls_name_cp[int(result[i][5])], |
|
"BOUNDING_BOX": { |
|
"XMIN": round(result[i][:4].tolist()[0], 6), |
|
"YMIN": round(result[i][:4].tolist()[1], 6), |
|
"XMAX": round(result[i][:4].tolist()[2], 6), |
|
"YMAX": round(result[i][:4].tolist()[3], 6),}, |
|
"CONF": round(float(result[i][4]), 2), |
|
"FPS": round(1000 / float(results.t[1]), 2), |
|
"IMG_WIDTH": img_size[0], |
|
"IMG_HEIGHT": img_size[1],} for i in range(len(result))] for result in results.xyxyn] |
|
|
|
|
|
|
|
def color_set(cls_num): |
|
color_list = [] |
|
for i in range(cls_num): |
|
color = tuple(np.random.choice(range(256), size=3)) |
|
|
|
color_list.append(color) |
|
|
|
return color_list |
|
|
|
|
|
|
|
def random_color(cls_num, is_light=True): |
|
color_list = [] |
|
for i in range(cls_num): |
|
color = ( |
|
random.randint(0, 127) + int(is_light) * 128, |
|
random.randint(0, 127) + int(is_light) * 128, |
|
random.randint(0, 127) + int(is_light) * 128, |
|
) |
|
color_list.append(color) |
|
|
|
return color_list |
|
|
|
|
|
|
|
def pil_draw(img, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list, opt): |
|
|
|
img_pil = ImageDraw.Draw(img) |
|
id = 0 |
|
|
|
for score, (xmin, ymin, xmax, ymax), label, cls_index in zip(score_l, bbox_l, cls_l, cls_index_l): |
|
|
|
img_pil.rectangle([xmin, ymin, xmax, ymax], fill=None, outline=color_list[cls_index], width=2) |
|
countdown_msg = f"{id}-{label} {score:.2f}" |
|
text_w, text_h = textFont.getsize(countdown_msg) |
|
if "label" in opt: |
|
|
|
img_pil.rectangle( |
|
(xmin, ymin, xmin + text_w, ymin + text_h), |
|
fill=color_list[cls_index], |
|
outline=color_list[cls_index], |
|
) |
|
|
|
|
|
img_pil.multiline_text( |
|
(xmin, ymin), |
|
countdown_msg, |
|
fill=(0, 0, 0), |
|
font=textFont, |
|
align="center", |
|
) |
|
|
|
id += 1 |
|
|
|
return img |
|
|
|
|
|
|
|
def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_cls, opt): |
|
|
|
global model, model_name_tmp, device_tmp |
|
|
|
if img is None or img == "": |
|
|
|
print("图片不存在!") |
|
return None, None, None, None, None, None, None |
|
|
|
det_img = img.copy() |
|
|
|
s_obj, m_obj, l_obj = 0, 0, 0 |
|
|
|
area_obj_all = [] |
|
|
|
score_det_stat = [] |
|
bbox_det_stat = [] |
|
cls_det_stat = [] |
|
cls_index_det_stat = [] |
|
|
|
pdf_csv_xlsx = [] |
|
|
|
if model_name_tmp != model_name: |
|
|
|
model_name_tmp = model_name |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
elif device_tmp != device: |
|
|
|
device_tmp = device |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
else: |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
|
|
|
|
model.conf = conf |
|
model.iou = iou |
|
model.max_det = int(max_num) |
|
model.classes = model_cls |
|
|
|
color_list = random_color(len(model_cls_name_cp), True) |
|
|
|
img_size = img.size |
|
|
|
results = model(img, size=infer_size) |
|
|
|
|
|
is_results_null = results.xyxyn[0].shape == torch.Size([0, 6]) |
|
|
|
if not is_results_null: |
|
|
|
|
|
crops = results.crop(save=False) |
|
img_crops = [] |
|
for i in range(len(crops)): |
|
img_crops.append(crops[i]["im"][..., ::-1]) |
|
|
|
|
|
dataframe = results.pandas().xyxy[0].round(2) |
|
|
|
report = "./Det_Report.pdf" |
|
det_csv = "./Det_Report.csv" |
|
det_excel = "./Det_Report.xlsx" |
|
|
|
if "csv" in opt: |
|
dataframe.to_csv(det_csv, index=False) |
|
pdf_csv_xlsx.append(det_csv) |
|
else: |
|
det_csv = None |
|
|
|
if "excel" in opt: |
|
dataframe.to_excel(det_excel, sheet_name='sheet1', index=False) |
|
pdf_csv_xlsx.append(det_excel) |
|
else: |
|
det_excel = None |
|
|
|
|
|
yaml_index = cls_name.index(".yaml") |
|
cls_name_lang = cls_name[yaml_index - 2:yaml_index] |
|
|
|
if cls_name_lang == "zh": |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE) |
|
elif cls_name_lang in ["en", "ru", "es", "ar"]: |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/TimesNewRoman.ttf"), size=FONTSIZE) |
|
elif cls_name_lang == "ko": |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE) |
|
|
|
for result in results.xyxyn: |
|
for i in range(len(result)): |
|
|
|
obj_cls_index = int(result[i][5]) |
|
cls_index_det_stat.append(obj_cls_index) |
|
|
|
obj_cls = model_cls_name_cp[obj_cls_index] |
|
cls_det_stat.append(obj_cls) |
|
|
|
|
|
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) |
|
bbox_det_stat.append((x0, y0, x1, y1)) |
|
|
|
conf = float(result[i][4]) |
|
score_det_stat.append(conf) |
|
|
|
|
|
|
|
|
|
w_obj = x1 - x0 |
|
h_obj = y1 - y0 |
|
area_obj = w_obj * h_obj |
|
area_obj_all.append(area_obj) |
|
|
|
det_img = pil_draw(img, score_det_stat, bbox_det_stat, cls_det_stat, cls_index_det_stat, textFont, color_list, |
|
opt) |
|
|
|
det_json = export_json(results, img.size)[0] |
|
det_json_format = json.dumps(det_json, sort_keys=False, indent=4, separators=(",", ":"), |
|
ensure_ascii=False) |
|
if "json" not in opt: |
|
det_json = None |
|
|
|
|
|
if "pdf" in opt: |
|
pdf_generate(f"{det_json_format}", report, GYD_VERSION) |
|
pdf_csv_xlsx.append(report) |
|
else: |
|
report = None |
|
|
|
|
|
for i in range(len(area_obj_all)): |
|
if (0 < area_obj_all[i] <= 32 ** 2): |
|
s_obj = s_obj + 1 |
|
elif (32 ** 2 < area_obj_all[i] <= 96 ** 2): |
|
m_obj = m_obj + 1 |
|
elif (area_obj_all[i] > 96 ** 2): |
|
l_obj = l_obj + 1 |
|
|
|
sml_obj_total = s_obj + m_obj + l_obj |
|
objSize_dict = {} |
|
objSize_dict = {obj_style[i]: [s_obj, m_obj, l_obj][i] / sml_obj_total for i in range(3)} |
|
|
|
|
|
clsRatio_dict = {} |
|
clsDet_dict = Counter(cls_det_stat) |
|
clsDet_dict_sum = sum(clsDet_dict.values()) |
|
for k, v in clsDet_dict.items(): |
|
clsRatio_dict[k] = v / clsDet_dict_sum |
|
|
|
return det_img, img_crops, objSize_dict, clsRatio_dict, dataframe, det_json, pdf_csv_xlsx |
|
else: |
|
print("图片目标不存在!") |
|
return None, None, None, None, None, None, None |
|
|
|
|
|
|
|
def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, model_cls, opt, draw_style): |
|
|
|
global model, model_name_tmp, device_tmp |
|
|
|
if video is None or video == "": |
|
|
|
print("视频不存在!") |
|
return None, None, None |
|
|
|
|
|
s_obj, m_obj, l_obj = 0, 0, 0 |
|
|
|
area_obj_all = [] |
|
s_list, m_list, l_list = [], [], [] |
|
|
|
score_det_stat = [] |
|
bbox_det_stat = [] |
|
cls_det_stat = [] |
|
cls_index_det_stat = [] |
|
|
|
fps_list = [] |
|
|
|
frame_count = 0 |
|
fps = 0 |
|
|
|
os.system(""" |
|
if [ -e './output.mp4' ]; then |
|
rm ./output.mp4 |
|
fi |
|
""") |
|
|
|
if model_name_tmp != model_name: |
|
|
|
model_name_tmp = model_name |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
elif device_tmp != device: |
|
|
|
device_tmp = device |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
else: |
|
print(f"正在加载模型{model_name_tmp}......") |
|
model = model_loading(model_name_tmp, device, opt) |
|
|
|
|
|
model.conf = conf |
|
model.iou = iou |
|
model.max_det = int(max_num) |
|
model.classes = model_cls |
|
|
|
color_list = random_color(len(model_cls_name_cp), True) |
|
|
|
|
|
yaml_index = cls_name.index(".yaml") |
|
cls_name_lang = cls_name[yaml_index - 2:yaml_index] |
|
|
|
if cls_name_lang == "zh": |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE) |
|
elif cls_name_lang in ["en", "ru", "es", "ar"]: |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/TimesNewRoman.ttf"), size=FONTSIZE) |
|
elif cls_name_lang == "ko": |
|
|
|
textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE) |
|
|
|
|
|
gc.collect() |
|
output_video_path = "./output.avi" |
|
cap = cv2.VideoCapture(video) |
|
fourcc = cv2.VideoWriter_fourcc(*"I420") |
|
|
|
out = cv2.VideoWriter(output_video_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4)))) |
|
if cap.isOpened(): |
|
while cap.isOpened(): |
|
ret, frame = cap.read() |
|
|
|
if not ret: |
|
break |
|
|
|
frame_count += 1 |
|
|
|
results = model(frame, size=infer_size) |
|
|
|
h, w, _ = frame.shape |
|
img_size = (w, h) |
|
|
|
for result in results.xyxyn: |
|
for i in range(len(result)): |
|
|
|
obj_cls_index = int(result[i][5]) |
|
cls_index_det_stat.append(obj_cls_index) |
|
|
|
obj_cls = model_cls_name_cp[obj_cls_index] |
|
cls_det_stat.append(obj_cls) |
|
|
|
|
|
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) |
|
bbox_det_stat.append((x0, y0, x1, y1)) |
|
|
|
conf = float(result[i][4]) |
|
score_det_stat.append(conf) |
|
|
|
fps = f"{(1000 / float(results.t[1])):.2f}" |
|
|
|
|
|
w_obj = x1 - x0 |
|
h_obj = y1 - y0 |
|
area_obj = w_obj * h_obj |
|
area_obj_all.append(area_obj) |
|
|
|
|
|
|
|
is_results_null = results.xyxyn[0].shape == torch.Size([0, 6]) |
|
if not is_results_null: |
|
fps_list.append(float(fps)) |
|
else: |
|
fps_list.append(0.0) |
|
|
|
frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) |
|
frame = pil_draw(frame, score_det_stat, bbox_det_stat, cls_det_stat, cls_index_det_stat, textFont, |
|
color_list, opt) |
|
frame = cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR) |
|
|
|
|
|
out.write(frame) |
|
|
|
|
|
score_det_stat = [] |
|
bbox_det_stat = [] |
|
cls_det_stat = [] |
|
cls_index_det_stat = [] |
|
|
|
|
|
for i in range(len(area_obj_all)): |
|
if (0 < area_obj_all[i] <= 32 ** 2): |
|
s_obj = s_obj + 1 |
|
elif (32 ** 2 < area_obj_all[i] <= 96 ** 2): |
|
m_obj = m_obj + 1 |
|
elif (area_obj_all[i] > 96 ** 2): |
|
l_obj = l_obj + 1 |
|
|
|
s_list.append(s_obj) |
|
m_list.append(m_obj) |
|
l_list.append(l_obj) |
|
|
|
|
|
s_obj, m_obj, l_obj = 0, 0, 0 |
|
|
|
area_obj_all = [] |
|
|
|
out.release() |
|
cap.release() |
|
|
|
|
|
df_objSize = pd.DataFrame({"fID": list(range(frame_count))}) |
|
df_objSize[obj_style[0]] = tuple(s_list) |
|
df_objSize[obj_style[1]] = tuple(m_list) |
|
df_objSize[obj_style[2]] = tuple(l_list) |
|
print(df_objSize) |
|
|
|
if draw_style == "Plotly": |
|
|
|
fig_objSize = px.scatter(df_objSize, x="fID", y=obj_style) |
|
|
|
fig_objSize.update_layout(title="帧数-目标尺寸数", xaxis_title="帧数", yaxis_title="目标尺寸数") |
|
|
|
|
|
fig_fps = px.scatter(df_objSize, x="fID", y=fps_list) |
|
|
|
fig_fps.update_layout(title="帧数-FPS", xaxis_title="帧数", yaxis_title="FPS") |
|
|
|
elif draw_style == "Matplotlib": |
|
|
|
fig_objSize = plt.figure() |
|
|
|
|
|
plt.scatter(df_objSize['fID'], df_objSize[obj_style[0]]) |
|
plt.scatter(df_objSize['fID'], df_objSize[obj_style[1]]) |
|
plt.scatter(df_objSize['fID'], df_objSize[obj_style[2]]) |
|
|
|
plt.title("帧数-目标尺寸数图", fontsize=12, fontproperties=SimSun) |
|
plt.xlabel("帧数", fontsize=12, fontproperties=SimSun) |
|
plt.ylabel("目标尺寸数", fontsize=12, fontproperties=SimSun) |
|
plt.legend(obj_style, prop=SimSun, fontsize=12, loc="best") |
|
|
|
|
|
fig_fps = plt.figure() |
|
|
|
plt.scatter(df_objSize['fID'], fps_list) |
|
|
|
plt.title("帧数-FPS", fontsize=12, fontproperties=SimSun) |
|
plt.xlabel("帧数", fontsize=12, fontproperties=SimSun) |
|
plt.ylabel("FPS", fontsize=12, fontproperties=SimSun) |
|
|
|
return output_video_path, fig_objSize, fig_fps |
|
|
|
else: |
|
print("视频加载失败!") |
|
return None, None, None |
|
|
|
|
|
def main(args): |
|
gr.close_all() |
|
|
|
global model, model_cls_name_cp, cls_name |
|
|
|
source = args.source |
|
source_video = args.source_video |
|
img_tool = args.img_tool |
|
nms_conf = args.nms_conf |
|
nms_iou = args.nms_iou |
|
model_name = args.model_name |
|
model_cfg = args.model_cfg |
|
cls_name = args.cls_name |
|
device = args.device |
|
inference_size = args.inference_size |
|
max_detnum = args.max_detnum |
|
slider_step = args.slider_step |
|
is_login = args.is_login |
|
usr_pwd = args.usr_pwd |
|
is_share = args.is_share |
|
|
|
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.Image(image_mode="RGB", source=source, tool=img_tool, type="pil", label="原始图片") |
|
inputs_device01 = gr.Radio(choices=["cuda:0", "cpu"], value=device, label="设备") |
|
inputs_model01 = gr.Dropdown(choices=model_names, value=model_name, type="value", label="模型") |
|
inputs_size01 = gr.Slider(384, 1536, step=128, value=inference_size, label="推理尺寸") |
|
input_conf01 = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="置信度阈值") |
|
inputs_iou01 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU 阈值") |
|
inputs_maxnum01 = gr.Number(value=max_detnum, label="最大检测数") |
|
inputs_clsName01 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="类别") |
|
inputs_opt01 = gr.CheckboxGroup(choices=["refresh_yolov5", "label", "pdf", "json", "csv", "excel"], |
|
value=["label", "pdf"], |
|
type="value", |
|
label="操作") |
|
|
|
|
|
inputs_video = gr.Video(format="mp4", source=source_video, mirror_webcam=False, label="原始视频") |
|
inputs_device02 = gr.Radio(choices=["cuda:0", "cpu"], value=device, label="设备") |
|
inputs_model02 = gr.Dropdown(choices=model_names, value=model_name, type="value", label="模型") |
|
inputs_size02 = gr.Slider(384, 1536, step=128, value=inference_size, label="推理尺寸") |
|
input_conf02 = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="置信度阈值") |
|
inputs_iou02 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU 阈值") |
|
inputs_maxnum02 = gr.Number(value=max_detnum, label="最大检测数") |
|
inputs_clsName02 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="类别") |
|
inputs_opt02 = gr.CheckboxGroup(choices=["refresh_yolov5", "label"], value=["label"], type="value", label="操作") |
|
inputs_draw02 = gr.Radio(choices=["Matplotlib", "Plotly"], value="Matplotlib", label="绘图") |
|
|
|
|
|
inputs_img_list = [ |
|
inputs_img, |
|
inputs_device01, |
|
inputs_model01, |
|
inputs_size01, |
|
input_conf01, |
|
inputs_iou01, |
|
inputs_maxnum01, |
|
inputs_clsName01, |
|
inputs_opt01, |
|
] |
|
|
|
|
|
inputs_video_list = [ |
|
inputs_video, |
|
inputs_device02, |
|
inputs_model02, |
|
inputs_size02, |
|
input_conf02, |
|
inputs_iou02, |
|
inputs_maxnum02, |
|
inputs_clsName02, |
|
inputs_opt02, |
|
inputs_draw02, |
|
] |
|
|
|
|
|
outputs_img = gr.Image(type="pil", label="检测图片") |
|
outputs_df = gr.Dataframe(max_rows=5, overflow_row_behaviour="paginate", type="pandas", label="检测信息列表") |
|
outputs_crops = gr.Gallery(label="目标裁剪") |
|
outputs_objSize = gr.Label(label="目标尺寸占比统计") |
|
outputs_clsSize = gr.Label(label="类别检测占比统计") |
|
outputs_json = gr.JSON(label="检测信息") |
|
outputs_pdf = gr.File(label="检测报告") |
|
|
|
|
|
outputs_video = gr.Video(format='mp4', label="检测视频") |
|
outputs_frame_objSize_plot = gr.Plot(label="帧数-目标尺寸数") |
|
outputs_frame_fps_plot = gr.Plot(label="帧数-FPS") |
|
|
|
|
|
outputs_img_list = [ |
|
outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf] |
|
|
|
|
|
outputs_video_list = [outputs_video, outputs_frame_objSize_plot, outputs_frame_fps_plot] |
|
|
|
|
|
title = "Gradio YOLOv5 Det v0.5" |
|
|
|
|
|
description = "<div align='center'>可自定义目标检测模型、安装简单、使用方便</div>" |
|
|
|
|
|
|
|
examples_img = [ |
|
[ |
|
"./img_examples/bus.jpg", |
|
"cpu", |
|
"yolov5s", |
|
640, |
|
0.6, |
|
0.5, |
|
10, |
|
["人", "公交车"], |
|
["label", "pdf"],], |
|
[ |
|
"./img_examples/giraffe.jpg", |
|
"cpu", |
|
"yolov5l", |
|
320, |
|
0.5, |
|
0.45, |
|
12, |
|
["长颈鹿"], |
|
["label", "pdf"],], |
|
[ |
|
"./img_examples/zidane.jpg", |
|
"cpu", |
|
"yolov5m", |
|
640, |
|
0.6, |
|
0.5, |
|
15, |
|
["人", "领带"], |
|
["pdf", "json"],], |
|
[ |
|
"./img_examples/Millenial-at-work.jpg", |
|
"cpu", |
|
"yolov5s6", |
|
1280, |
|
0.5, |
|
0.5, |
|
20, |
|
["人", "椅子", "杯子", "笔记本电脑"], |
|
["label", "pdf", "csv", "excel"],],] |
|
|
|
examples_video = [ |
|
[ |
|
"./video_examples/test01.mp4", |
|
"cpu", |
|
"yolov5s", |
|
640, |
|
0.5, |
|
0.45, |
|
12, |
|
["鸟"], |
|
["label"], |
|
"Matplotlib",], |
|
[ |
|
"./video_examples/test02.mp4", |
|
"cpu", |
|
"yolov5m", |
|
640, |
|
0.6, |
|
0.5, |
|
15, |
|
["马"], |
|
["label"], |
|
"Matplotlib",], |
|
[ |
|
"./video_examples/test03.mp4", |
|
"cpu", |
|
"yolov5s6", |
|
1280, |
|
0.5, |
|
0.5, |
|
20, |
|
["人", "风筝"], |
|
["label"], |
|
"Plotly",],] |
|
|
|
|
|
gyd_img = gr.Interface( |
|
fn=yolo_det_img, |
|
inputs=inputs_img_list, |
|
outputs=outputs_img_list, |
|
title=title, |
|
description=description, |
|
|
|
examples=examples_img, |
|
cache_examples=False, |
|
|
|
|
|
flagging_dir="run", |
|
allow_flagging="manual", |
|
flagging_options=["good", "generally", "bad"], |
|
) |
|
|
|
gyd_video = gr.Interface( |
|
fn=yolo_det_video, |
|
inputs=inputs_video_list, |
|
outputs=outputs_video_list, |
|
title=title, |
|
description=description, |
|
|
|
examples=examples_video, |
|
cache_examples=False, |
|
|
|
|
|
flagging_dir="run", |
|
allow_flagging="manual", |
|
flagging_options=["good", "generally", "bad"], |
|
) |
|
|
|
gyd = gr.TabbedInterface(interface_list=[gyd_img, gyd_video], tab_names=["图片模式", "视频模式"]) |
|
|
|
if not is_login: |
|
gyd.launch( |
|
inbrowser=True, |
|
show_tips=True, |
|
share=is_share, |
|
favicon_path="./icon/logo.ico", |
|
show_error=True, |
|
quiet=True, |
|
) |
|
else: |
|
gyd.launch( |
|
inbrowser=True, |
|
show_tips=True, |
|
auth=usr_pwd, |
|
share=is_share, |
|
favicon_path="./icon/logo.ico", |
|
show_error=True, |
|
quiet=True, |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
args = parse_args() |
|
main(args) |
|
|