v2 update

.gitignore

@@ -40,4 +40,6 @@
 !requirements.txt
 !cls_name/*
 !model_config/*
-!img_example/*
+!img_example/*
+
+app copy.py

app.py

@@ -1,6 +1,6 @@
-# Gradio YOLOv5 Det v0.1
+# Gradio YOLOv5 Det v0.2
 # 创建人：曾逸夫
-# 创建时间：2022-04-03
+# 创建时间：2022-05-01
 # email：[email protected]
 # 项目主页：https://gitee.com/CV_Lab/gradio_yolov5_det
 
@@ -12,14 +12,15 @@ from pathlib import Path
 import gradio as gr
 import torch
 import yaml
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
+
+from util.fonts_opt import is_fonts
 
 ROOT_PATH = sys.path[0]  # 根目录
 
 # 模型路径
 model_path = "ultralytics/yolov5"
 
-
 # 模型名称临时变量
 model_name_tmp = ""
 
@@ -29,12 +30,13 @@ device_tmp = ""
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
 
+# 字体大小
+FONTSIZE = 25
+
 
 def parse_args(known=False):
-    parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det v0.1")
-    parser.add_argument(
-        "--model_name", "-mn", default="yolov5s", type=str, help="model name"
-    )
+    parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det v0.2")
+    parser.add_argument("--model_name", "-mn", default="yolov5s", type=str, help="model name")
     parser.add_argument(
         "--model_cfg",
         "-mc",
@@ -56,15 +58,13 @@ def parse_args(known=False):
         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("--nms_iou", "-iou", default=0.45, type=float, help="model NMS IoU threshold")
 
     parser.add_argument(
         "--label_dnt_show",
         "-lds",
-        action="store_false",
-        default=True,
+        action="store_true",
+        default=False,
         help="label show",
     )
     parser.add_argument(
@@ -72,11 +72,9 @@ def parse_args(known=False):
         "-dev",
         default="cpu",
         type=str,
-        help="cuda or cpu, hugging face only cpu",
-    )
-    parser.add_argument(
-        "--inference_size", "-isz", default=640, type=int, help="model inference size"
+        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
@@ -99,24 +97,44 @@ def export_json(results, model, img_size):
     return [
         [
             {
-                "id": int(i),
+                "id": i,
                 "class": int(result[i][5]),
-                "class_name": model.model.names[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),
-                },
+                    "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
-    ]
+                "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图片检测函数
@@ -139,9 +157,43 @@ def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, mode
     model.classes = model_cls  # 模型类别
 
     results = model(img, size=inference_size)  # 检测
-    results.render(labels=label_opt)  # 渲染
 
-    det_img = Image.fromarray(results.imgs[0])  # 检测图片
+    img_size = img.size  # 帧尺寸
+
+    # 加载字体
+    textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttc"), 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]  # 检测信息
 
@@ -150,7 +202,7 @@ def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, mode
 
 # yaml文件解析
 def yaml_parse(file_path):
-    return yaml.safe_load(open(file_path, "r", encoding="utf-8").read())
+    return yaml.safe_load(open(file_path, encoding="utf-8").read())
 
 
 # yaml csv 文件解析
@@ -172,7 +224,7 @@ def yaml_csv(file_path, file_tag):
 def main(args):
     gr.close_all()
 
-    global model
+    global model, model_cls_name_cp
 
     slider_step = 0.05  # 滑动步长
 
@@ -185,38 +237,30 @@ def main(args):
     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="原始图片")
-    device = gr.inputs.Dropdown(
-        choices=["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=label_opt, label="标签显示")
-    inputs_clsName = gr.inputs.CheckboxGroup(
-        choices=model_cls_name, default=model_cls_name, type="index", 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 = [
         inputs_img,  # 输入图片
-        device,  # 设备
+        inputs_device,  # 设备
         inputs_model,  # 模型
         inputs_size,  # 推理尺寸
         input_conf,  # 置信度阈值
@@ -243,8 +287,7 @@ def main(args):
             0.6,
             0.5,
             True,
-            ["人", "公交车"],
-        ],
+            ["人", "公交车"],],
         [
             "./img_example/Millenial-at-work.jpg",
             "cpu",
@@ -253,8 +296,7 @@ def main(args):
             0.5,
             0.45,
             True,
-            ["人", "椅子", "杯子", "笔记本电脑"],
-        ],
+            ["人", "椅子", "杯子", "笔记本电脑"],],
         [
             "./img_example/zidane.jpg",
             "cpu",
@@ -263,9 +305,7 @@ def main(args):
             0.25,
             0.5,
             False,
-            ["人", "领带"],
-        ],
-    ]
+            ["人", "领带"],],]
 
     # 接口
     gr.Interface(
@@ -282,7 +322,7 @@ def main(args):
     ).launch(
         inbrowser=True,  # 自动打开默认浏览器
         show_tips=True,  # 自动显示gradio最新功能
-        favicon_path="./icon/logo.ico",
+        # favicon_path="./icon/logo.ico",
     )
 
 

util/fonts_opt.py

@@ -0,0 +1,64 @@
+# 字体管理
+# 创建人：曾逸夫
+# 创建时间：2022-05-01
+
+import os
+import sys
+from pathlib import Path
+
+import wget
+from rich.console import Console
+
+ROOT_PATH = sys.path[0]  # 项目根目录
+
+fonts_list = ["SimSun.ttc"]  # 字体列表
+fonts_suffix = ["ttc", "ttf", "otf"]  # 字体后缀
+
+data_url_dict = {"SimSun.ttc": "https://gitee.com/CV_Lab/opencv_webcam/attach_files/959173/download/SimSun.ttc"}
+
+console = Console()
+
+
+# 创建字体库
+def add_fronts(font_diff):
+
+    global font_name
+
+    for k, v in data_url_dict.items():
+        if k in font_diff:
+            font_name = v.split("/")[-1]  # 字体名称
+            Path(f"{ROOT_PATH}/fonts").mkdir(parents=True, exist_ok=True)  # 创建目录
+
+            file_path = f"{ROOT_PATH}/fonts/{font_name}"  # 字体路径
+
+            try:
+                # 下载字体文件
+                wget.download(v, file_path)
+            except Exception as e:
+                print("路径错误！程序结束！")
+                print(e)
+                sys.exit()
+            else:
+                print()
+                console.print(f"{font_name} [bold green]字体文件下载完成！[/bold green] 已保存至：{file_path}")
+
+
+# 判断字体文件
+def is_fonts(fonts_dir):
+    if os.path.isdir(fonts_dir):
+        # 如果字体库存在
+        f_list = os.listdir(fonts_dir)  # 本地字体库
+
+        font_diff = list(set(fonts_list).difference(set(f_list)))
+
+        if font_diff != []:
+            # 字体不存在
+            console.print("[bold red]字体不存在，正在加载。。。[/bold red]")
+            add_fronts(font_diff)  # 创建字体库
+        else:
+            console.print(f"{fonts_list}[bold green]字体已存在！[/bold green]")
+    else:
+        # 字体库不存在，创建字体库
+        console.print("[bold red]字体库不存在，正在创建。。。[/bold red]")
+        add_fronts(fonts_list)  # 创建字体库
+