Update app.py

app.py

@@ -1,31 +1,21 @@
-# Gradio YOLOv8 Det v0.2.5
+# Gradio YOLOv8 Det v1.0.0
 # 创建人：曾逸夫
-# 创建时间：2023-10-21
+# 创建时间：2023-10-28
+# pip install gradio>=4.0.2
 
 import argparse
 import csv
-import gc
-import json
-import os
 import random
-import shutil
 import sys
 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
 from ultralytics import YOLO
 
-# pip_source = "https://pypi.tuna.tsinghua.edu.cn/simple some-package"
-# os.system(f"pip install gradio --upgrade -i {pip_source}")
-# os.system(f"pip install ultralytics --upgrade -i {pip_source}")
-
 ROOT_PATH = sys.path[0]  # 项目根目录
 
 # --------------------- 字体库 ---------------------
@@ -36,7 +26,6 @@ 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
 
@@ -45,7 +34,7 @@ from util.fonts_opt import is_fonts
 ROOT_PATH = sys.path[0]  # 根目录
 
 # Gradio YOLOv8 Det版本
-GYD_VERSION = "Gradio YOLOv8 Det v0.2.5"
+GYD_VERSION = "Gradio YOLOv8 Det v1.0.0"
 
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
@@ -58,11 +47,9 @@ obj_style = ["小目标", "中目标", "大目标"]
 
 
 def parse_args(known=False):
-    parser = argparse.ArgumentParser(description="Gradio YOLOv8 Det v0.2.5")
+    parser = argparse.ArgumentParser(description="Gradio YOLOv8 Det v1.0.0")
     parser.add_argument("--model_type", "-mt", default="online", type=str, help="model type")
     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="yolov8s", type=str, help="model name")
     parser.add_argument(
         "--model_cfg",
@@ -86,13 +73,6 @@ def parse_args(known=False):
         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")
@@ -155,21 +135,11 @@ def check_online():
         return False
 
 
-# 模型加载
-def model_loading(img_path, conf, iou, infer_size, yolo_model="yolov8n.pt"):
-    model = YOLO(yolo_model)
-
-    results = model(source=img_path, imgsz=infer_size, conf=conf, iou=iou)
-    results = list(results)[0]
-    return results
-
-
 # 标签和边界框颜色设置
 def color_set(cls_num):
     color_list = []
     for i in range(cls_num):
         color = tuple(np.random.choice(range(256), size=3))
-        # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
         color_list.append(color)
 
     return color_list
@@ -195,6 +165,7 @@ def pil_draw(img, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list):
     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)  # 标签尺寸 pillow 9.5.0
@@ -257,8 +228,17 @@ def seg_output(img_path, seg_mask_list, color_list, cls_list):
     return img_mask_merge
 
 
+# 模型加载
+def model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"):
+    model = YOLO(yolo_model)
+
+    results = model(source=img_path, device=device_opt, imgsz=infer_size, conf=conf, iou=iou, max_det=max_det)
+    results = list(results)[0]
+    return results
+
+
 # YOLOv8图片检测函数
-def yolo_det_img(img_path, model_name, infer_size, conf, iou):
+def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_det, obj_size):
 
     global model, model_name_tmp, device_tmp
 
@@ -269,10 +249,10 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
     score_det_stat = []  # 置信度统计
     bbox_det_stat = []  # 边界框统计
     cls_det_stat = []  # 类别数量统计
-    cls_index_det_stat = []  # 类别索引统计
+    cls_index_det_stat = []  # 1
 
     # 模型加载
-    predict_results = model_loading(img_path, conf, iou, infer_size, yolo_model=f"{model_name}.pt")
+    predict_results = model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model=f"{model_name}.pt")
     # 检测参数
     xyxy_list = predict_results.boxes.xyxy.cpu().numpy().tolist()
     conf_list = predict_results.boxes.conf.cpu().numpy().tolist()
@@ -308,11 +288,6 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
             textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE)
 
         for i in range(len(xyxy_list)):
-            obj_cls_index = int(cls_list[i])  # 类别索引
-            cls_index_det_stat.append(obj_cls_index)
-
-            obj_cls = model_cls_name_cp[obj_cls_index]  # 类别
-            cls_det_stat.append(obj_cls)
 
             # ------------ 边框坐标 ------------
             x0 = int(xyxy_list[i][0])
@@ -320,16 +295,63 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
             x1 = int(xyxy_list[i][2])
             y1 = int(xyxy_list[i][3])
 
-            bbox_det_stat.append((x0, y0, x1, y1))
-
-            conf = float(conf_list[i])  # 置信度
-            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)
+            area_obj = w_obj * h_obj  # 目标尺寸
+
+            if (obj_size == "小目标" and area_obj > 0 and area_obj <= 32 ** 2):
+                obj_cls_index = int(cls_list[i])  # 类别索引
+                cls_index_det_stat.append(obj_cls_index)
+
+                obj_cls = model_cls_name_cp[obj_cls_index]  # 类别
+                cls_det_stat.append(obj_cls)
+
+                bbox_det_stat.append((x0, y0, x1, y1))
+
+                conf = float(conf_list[i])  # 置信度
+                score_det_stat.append(conf)
+
+                area_obj_all.append(area_obj)
+            elif (obj_size == "中目标" and area_obj > 32 ** 2 and area_obj <= 96 ** 2):
+                obj_cls_index = int(cls_list[i])  # 类别索引
+                cls_index_det_stat.append(obj_cls_index)
+
+                obj_cls = model_cls_name_cp[obj_cls_index]  # 类别
+                cls_det_stat.append(obj_cls)
+
+                bbox_det_stat.append((x0, y0, x1, y1))
+
+                conf = float(conf_list[i])  # 置信度
+                score_det_stat.append(conf)
+
+                area_obj_all.append(area_obj)
+            elif (obj_size == "大目标" and area_obj > 96 ** 2):
+                obj_cls_index = int(cls_list[i])  # 类别索引
+                cls_index_det_stat.append(obj_cls_index)
+
+                obj_cls = model_cls_name_cp[obj_cls_index]  # 类别
+                cls_det_stat.append(obj_cls)
+
+                bbox_det_stat.append((x0, y0, x1, y1))
+
+                conf = float(conf_list[i])  # 置信度
+                score_det_stat.append(conf)
+
+                area_obj_all.append(area_obj)
+            elif (obj_size == "所有尺寸"):
+                obj_cls_index = int(cls_list[i])  # 类别索引
+                cls_index_det_stat.append(obj_cls_index)
+
+                obj_cls = model_cls_name_cp[obj_cls_index]  # 类别
+                cls_det_stat.append(obj_cls)
+
+                bbox_det_stat.append((x0, y0, x1, y1))
+
+                conf = float(conf_list[i])  # 置信度
+                score_det_stat.append(conf)
+
+                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)
 
@@ -353,10 +375,10 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
         for k, v in clsDet_dict.items():
             clsRatio_dict[k] = v / clsDet_dict_sum
 
+        gr.Info("图片检测成功！")
         return det_img, objSize_dict, clsRatio_dict
     else:
-        print("图片目标不存在！")
-        return None, None, None
+        raise gr.Error("图片检测失败！")
 
 
 def main(args):
@@ -365,13 +387,13 @@ def main(args):
     global model_cls_name_cp, cls_name
 
     source = args.source
-    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
     inference_size = args.inference_size
+    max_detnum = args.max_detnum
     slider_step = args.slider_step
 
     is_fonts(f"{ROOT_PATH}/fonts")  # 检查字体文件
@@ -381,100 +403,76 @@ def main(args):
 
     model_cls_name_cp = model_cls_name.copy()  # 类别名称
 
-    # ------------------- 图片模式输入组件 -------------------
-    inputs_img = gr.Image(image_mode="RGB", source=source, tool=img_tool, type="filepath", 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_img_list = [
-        inputs_img,  # 输入图片
-        inputs_model01,  # 模型
-        inputs_size01,  # 推理尺寸
-        input_conf01,  # 置信度阈值
-        inputs_iou01,  # IoU阈值
-    ]
-
-    # ------------------- 图片模式输出组件 -------------------
-    outputs_img = gr.Image(type="pil", label="检测图片")
-    outputs_objSize = gr.Label(label="目标尺寸占比统计")
-    outputs_clsSize = gr.Label(label="类别检测占比统计")
-
-    # ------------------- 图片模式输出参数 -------------------
-    outputs_img_list = [outputs_img, outputs_objSize, outputs_clsSize]
-
-    # 标题
-    title = "Gradio YOLOv8 Det"
-
-    # 描述
-    description = "<div align='center'>基于YOLOv8的目标检测与图像分割系统</div>"
-
-    # 示例图片
-    examples_imgs = [
-        [
-            "./img_examples/bus.jpg",
-            "yolov8s",
-            640,
-            0.6,
-            0.5,],
-        [
-            "./img_examples/giraffe.jpg",
-            "yolov8l",
-            320,
-            0.5,
-            0.45,],
-        [
-            "./img_examples/zidane.jpg",
-            "yolov8m",
-            640,
-            0.6,
-            0.5,],
-        [
-            "./img_examples/Millenial-at-work.jpg",
-            "yolov8x",
-            1280,
-            0.5,
-            0.5,],
-        [
-            "./img_examples/bus.jpg",
-            "yolov8s-seg",
-            640,
-            0.6,
-            0.5,],
-        [
-            "./img_examples/Millenial-at-work.jpg",
-            "yolov8x-seg",
-            1280,
-            0.5,
-            0.5,],]
-
-    # 接口
-    gyd_img = gr.Interface(
-        fn=yolo_det_img,
-        inputs=inputs_img_list,
-        outputs=outputs_img_list,
-        title=title,
-        description=description,
-        examples=examples_imgs,
-        cache_examples=False,
-        flagging_dir="run",  # 输出目录
-        allow_flagging="manual",
-        flagging_options=["good", "generally", "bad"],
-    )
-
-    return gyd_img
+    # ------------ Gradio Blocks ------------
+    with gr.Blocks() as gyd:
+        with gr.Row():
+            # gr.HTML(value="<div align='center' style='font-size:25px'>Gradio YOLOv8 Det<div> \
+            #         <div align='center' style='font-size:20px'>基于 YOLOv8 的目标检测与图像分割系统<div>")
+            gr.Markdown(value="<p align='center'><a href='https://gitee.com/CV_Lab/gradio-yolov8-det'>\
+                        <img src='https://pycver.gitee.io/ows-pics/imgs/gradio_yolov8_det_logo.png' alt='Simple Icons' ></a>\
+                        <p align='center'>基于 Gradio 的 YOLOv8 通用目标检测与图像分割演示系统</p><p align='center'>可自定义检测模型、演示便捷、安装简单</p>"
+                        )
+        with gr.Row():
+            gr.Markdown(value="作者：曾逸夫，Gitee：https://gitee.com/PyCVer ，Github：https://github.com/Zengyf-CVer")
+        with gr.Row():
+            with gr.Column(scale=1):
+                with gr.Row():
+                    inputs_img = gr.Image(image_mode="RGB", sources=source, type="filepath", label="原始图片")
+
+                with gr.Row():
+                    device_opt = gr.Radio(choices=["cpu", "0", "1", "2", "3"], value="cpu", label="设备")
+                with gr.Row():
+                    inputs_model = gr.Dropdown(choices=model_names, value=model_name, type="value", label="模型")
+                with gr.Row():
+                    inputs_size = gr.Slider(320, 1600, step=1, value=inference_size, label="推理尺寸")
+                    max_det = gr.Slider(1, 1000, step=1, value=max_detnum, label="最大检测数")
+                with gr.Row():
+                    input_conf = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="置信度阈值")
+                    inputs_iou = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU 阈值")
+                with gr.Row():
+                    obj_size = gr.Radio(choices=["所有尺寸", "小目标", "中目标", "大目标"], value="所有尺寸", label="目标尺寸")
+                with gr.Row():
+                    gr.ClearButton(inputs_img, value="清除")
+                    det_btn_img = gr.Button(value='检测', variant="primary")
+
+            with gr.Column(scale=1):
+                with gr.Row():
+                    outputs_img = gr.Image(type="pil", label="检测图片")
+                with gr.Row():
+                    outputs_objSize = gr.Label(label="目标尺寸占比统计")
+                with gr.Row():
+                    outputs_clsSize = gr.Label(label="类别检测占比统计")
+
+        with gr.Row():
+            example_list = [
+                ["./img_examples/bus.jpg", "yolov8s", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
+                ["./img_examples/giraffe.jpg", "yolov8l", "cpu", 320, 0.5, 0.45, 100, "所有尺寸"],
+                ["./img_examples/zidane.jpg", "yolov8m", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
+                ["./img_examples/Millenial-at-work.jpg", "yolov8x", "cpu", 1280, 0.5, 0.5, 100, "所有尺寸"],
+                ["./img_examples/bus.jpg", "yolov8s-seg", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
+                ["./img_examples/Millenial-at-work.jpg", "yolov8x-seg", "cpu", 1280, 0.5, 0.5, 100, "所有尺寸"],]
+            gr.Examples(example_list,
+                        [inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det, obj_size],
+                        [outputs_img, outputs_objSize, outputs_clsSize],
+                        yolo_det_img,
+                        cache_examples=False)
+
+        det_btn_img.click(fn=yolo_det_img,
+                          inputs=[
+                              inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det,
+                              obj_size],
+                          outputs=[outputs_img, outputs_objSize, outputs_clsSize])
+
+    return gyd
 
 
 if __name__ == "__main__":
     args = parse_args()
-    gyd_img = main(args)
+    gyd = main(args)
     is_share = args.is_share
 
-    gyd_img.launch(
+    gyd.queue().launch(
         inbrowser=True,  # 自动打开默认浏览器
-        show_tips=True,  # 自动显示gradio最新功能
         share=is_share,  # 项目共享，其他设备可以访问
         favicon_path="./icon/logo.ico",  # 网页图标
         show_error=True,  # 在浏览器控制台中显示错误信息