v02 update

app.py

@@ -6,6 +6,7 @@
 
 import argparse
 import csv
+import json
 import sys
 from pathlib import Path
 
@@ -15,11 +16,12 @@ 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"
+# 本地模型路径
+local_model_path = f"{ROOT_PATH}/yolov5"
 
 # 模型名称临时变量
 model_name_tmp = ""
@@ -72,7 +74,11 @@ def parse_args(known=False):
         help="label show",
     )
     parser.add_argument(
-        "--device", "-dev", default="cpu", type=str, help="cuda or cpu",
+        "--device",
+        "-dev",
+        default="cpu",
+        type=str,
+        help="cuda or cpu",
     )
     parser.add_argument(
         "--inference_size", "-isz", default=640, type=int, help="model inference size"
@@ -82,12 +88,38 @@ def parse_args(known=False):
     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
+        local_model_path,
+        "custom",
+        path=f"{local_model_path}/{model_name}",
+        source="local",
+        device=device,
+        _verbose=False,
     )
 
     return model
@@ -146,7 +178,9 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, label_opt):
 
 
 # YOLOv5图片检测函数
-def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, model_cls):
+def yolo_det(
+    img, device, model_name, inference_size, conf, iou, label_opt, model_cls, opt
+):
 
     global model, model_name_tmp, device_tmp
 
@@ -169,7 +203,7 @@ def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, mode
     img_size = img.size  # 帧尺寸
 
     # 加载字体
-    textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttc"), size=FONTSIZE)
+    textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE)
 
     det_img = img.copy()
 
@@ -205,28 +239,22 @@ def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, mode
 
     det_json = export_json(results, model, img.size)[0]  # 检测信息
 
-    return det_img, det_json
-
-
-# yaml文件解析
-def yaml_parse(file_path):
-    return yaml.safe_load(open(file_path, encoding="utf-8").read())
-
+    # JSON格式化
+    det_json_format = json.dumps(
+        det_json, sort_keys=True, indent=4, separators=(",", ":"), ensure_ascii=False
+    )
 
-# 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版
+    # -------pdf-------
+    report = "./Det_Report.pdf"
+    if "pdf" in opt:
+        pdf_generate(f"{det_json_format}", report)
     else:
-        print(f"{file_path}格式不正确！程序退出！")
-        sys.exit()
+        report = None
 
-    return file_names
+    if "json" not in opt:
+        det_json = None
+
+    return det_img, det_json, report
 
 
 def main(args):
@@ -276,6 +304,9 @@ def main(args):
     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 = [
@@ -287,19 +318,35 @@ def main(args):
         inputs_iou,  # IoU阈值
         inputs_label,  # 标签显示
         inputs_clsName,  # 类别
+        inputs_opt,  # 检测操作
     ]
+
     # 输出参数
-    outputs = gr.outputs.Image(type="pil", label="检测图片")
-    outputs02 = gr.outputs.JSON(label="检测信息")
+    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通用目标检测系统v2"
+    title = "基于Gradio的YOLOv5通用目标检测系统v0.2"
     # 描述
     description = "<div align='center'>可自定义目标检测模型、���装简单、使用方便</div><div align='center'>Customizable target detection model, easy to install and easy to use</div>"
 
+
     # 示例图片
     examples = [
-        ["./img_example/bus.jpg", "cpu", "yolov5s", 640, 0.6, 0.5, True, ["人", "公交车"],],
+        [
+            "./img_example/bus.jpg",
+            "cpu",
+            "yolov5s",
+            640,
+            0.6,
+            0.5,
+            True,
+            ["人", "公交车"],
+            ["pdf"],
+        ],
         [
             "./img_example/Millenial-at-work.jpg",
             "cpu",
@@ -309,6 +356,7 @@ def main(args):
             0.45,
             True,
             ["人", "椅子", "杯子", "笔记本电脑"],
+            ["json"],
         ],
         [
             "./img_example/zidane.jpg",
@@ -319,6 +367,7 @@ def main(args):
             0.5,
             False,
             ["人", "领带"],
+            ["pdf", "json"],
         ],
     ]
 
@@ -326,13 +375,15 @@ def main(args):
     gr.Interface(
         fn=yolo_det,
         inputs=inputs,
-        outputs=[outputs, outputs02],
+        outputs=outputs,
         title=title,
         description=description,
         examples=examples,
         theme="seafoam",
         # live=True, # 实时变更输出
-        flagging_dir="run"  # 输出目录
+        flagging_dir="run",  # 输出目录
+        # flagging_options=["good", "generally", "bad"],
+        allow_flagging="auto",
         # ).launch(inbrowser=True, auth=['admin', 'admin'])
     ).launch(
         inbrowser=True,  # 自动打开默认浏览器

util/fonts_opt.py

@@ -11,10 +11,12 @@ from rich.console import Console
 
 ROOT_PATH = sys.path[0]  # 项目根目录
 
-fonts_list = ["SimSun.ttc"]  # 字体列表
+fonts_list = ["SimSun.ttf"]  # 字体列表
 fonts_suffix = ["ttc", "ttf", "otf"]  # 字体后缀
 
-data_url_dict = {"SimSun.ttc": "https://gitee.com/CV_Lab/opencv_webcam/attach_files/959173/download/SimSun.ttc"}
+data_url_dict = {
+    "SimSun.ttf": "https://gitee.com/CV_Lab/gradio_yolov5_det/attach_files/1049109/download/SimSun.ttf"
+}
 
 console = Console()
 
@@ -40,7 +42,9 @@ def add_fronts(font_diff):
                 sys.exit()
             else:
                 print()
-                console.print(f"{font_name} [bold green]字体文件下载完成！[/bold green] 已保存至：{file_path}")
+                console.print(
+                    f"{font_name} [bold green]字体文件下载完成！[/bold green] 已保存至：{file_path}"
+                )
 
 
 # 判断字体文件
@@ -61,4 +65,3 @@ def is_fonts(fonts_dir):
         # 字体库不存在，创建字体库
         console.print("[bold red]字体库不存在，正在创建。。。[/bold red]")
         add_fronts(fonts_list)  # 创建字体库
-

util/pdf_opt.py

@@ -0,0 +1,74 @@
+# 参考：https://pyfpdf.readthedocs.io/en/latest/Tutorial/index.html
+
+from fpdf import FPDF
+
+# title = 'Gradio YOLOv5 Det v0.2检测报告'
+title = "Gradio YOLOv5 Det v0.2"
+
+
+class PDF(FPDF):
+    def header(self):
+        # 设置中文字体
+        self.add_font("SimSun", "", "./fonts/SimSun.ttf", uni=True)
+        self.set_font("SimSun", "", 16)
+        # Calculate width of title and position
+        w = self.get_string_width(title) + 6
+        self.set_x((210 - w) / 2)
+        # Colors of frame, background and text
+        self.set_draw_color(255, 255, 255)
+        self.set_fill_color(255, 255, 255)
+        self.set_text_color(0, 0, 0)
+        # Thickness of frame (1 mm)
+        # self.set_line_width(1)
+        # Title
+        self.cell(w, 9, title, 1, 1, "C", 1)
+        # Line break
+        self.ln(10)
+
+    def footer(self):
+        # Position at 1.5 cm from bottom
+        self.set_y(-15)
+        # 设置中文字体
+        self.add_font("SimSun", "", "./fonts/SimSun.ttf", uni=True)
+        self.set_font("SimSun", "", 12)
+        # Text color in gray
+        self.set_text_color(128)
+        # Page number
+        self.cell(0, 10, "Page " + str(self.page_no()), 0, 0, "C")
+
+    def chapter_title(self, num, label):
+        # 设置中文字体
+        self.add_font("SimSun", "", "./fonts/SimSun.ttf", uni=True)
+        self.set_font("SimSun", "", 12)
+        # Background color
+        self.set_fill_color(200, 220, 255)
+        # Title
+        # self.cell(0, 6, 'Chapter %d : %s' % (num, label), 0, 1, 'L', 1)
+        self.cell(0, 6, f"检测结果：", 0, 1, "L", 1)
+        # Line break
+        self.ln(4)
+
+    def chapter_body(self, name):
+
+        # 设置中文字体
+        self.add_font("SimSun", "", "./fonts/SimSun.ttf", uni=True)
+        self.set_font("SimSun", "", 12)
+        # Output justified text
+        self.multi_cell(0, 5, name)
+        # Line break
+        self.ln()
+        self.cell(0, 5, "--------------------------------------")
+
+    def print_chapter(self, num, title, name):
+        self.add_page()
+        self.chapter_title(num, title)
+        self.chapter_body(name)
+
+
+# pdf生成函数
+def pdf_generate(input_file, output_file):
+    pdf = PDF()
+    pdf.set_title(title)
+    pdf.set_author("Zeng Yifu")
+    pdf.print_chapter(1, "A RUNAWAY REEF", input_file)
+    pdf.output(output_file)