Update app.py

app.py

@@ -1,10 +1,11 @@
-# Gradio YOLOv8 Det v1.3.0
+# Gradio YOLOv8 Det v1.3.1
 # 创建人：曾逸夫
-# 创建时间：2023-12-15
-# pip install gradio>=4.9.1
+# 创建时间：2024-01-03
+# pip install gradio>=4.12.0
 # python gradio_yolov8_det_v1.py
 
 
+import __init__
 import argparse
 import csv
 import random
@@ -36,9 +37,6 @@ from PIL import Image, ImageDraw, ImageFont
 
 from util.fonts_opt import is_fonts
 
-# Gradio YOLOv8 Det版本
-GYD_VERSION = "Gradio YOLOv8 Det v1.2.1"
-
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
 
@@ -63,9 +61,9 @@ GYD_SUB_TITLE = """
 """
 
 EXAMPLES_DET = [
-    ["./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/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",
@@ -74,9 +72,10 @@ EXAMPLES_DET = [
         0.5,
         0.5,
         100,
+        [],
         "所有尺寸",
     ],
-    ["./img_examples/bus.jpg", "yolov8s-seg", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
+    ["./img_examples/bus.jpg", "yolov8s-seg", "cpu", 640, 0.6, 0.5, 100, [], "所有尺寸"],
     [
         "./img_examples/Millenial-at-work.jpg",
         "yolov8x-seg",
@@ -85,17 +84,18 @@ EXAMPLES_DET = [
         0.5,
         0.5,
         100,
+        [],
         "所有尺寸",
     ],
 ]
 
 EXAMPLES_CLAS = [
-    ["./img_examples/img_clas/ILSVRC2012_val_00000008.JPEG", "yolov8s-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000018.JPEG", "yolov8l-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000023.JPEG", "yolov8m-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000067.JPEG", "yolov8m-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000077.JPEG", "yolov8m-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000247.JPEG", "yolov8m-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000008.JPEG", "cpu", "yolov8s-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000018.JPEG", "cpu", "yolov8l-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000023.JPEG", "cpu", "yolov8m-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000067.JPEG", "cpu", "yolov8m-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000077.JPEG", "cpu", "yolov8m-cls"],
+    ["./img_examples/img_clas/ILSVRC2012_val_00000247.JPEG", "cpu", "yolov8m-cls"],
 ]
 
 GYD_CSS = """#disp_image {
@@ -103,8 +103,11 @@ GYD_CSS = """#disp_image {
     }"""
 
 
+custom_css = "./gyd_style.css"
+
+
 def parse_args(known=False):
-    parser = argparse.ArgumentParser(description=GYD_VERSION)
+    parser = argparse.ArgumentParser(description=__init__.__version__)
     parser.add_argument(
         "--model_name", "-mn", default="yolov8s", type=str, help="model name"
     )
@@ -315,9 +318,18 @@ def seg_output(img_path, seg_mask_list, color_list, cls_list):
 
 # 目标检测和图像分割模型加载
 def model_det_loading(
-    img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"
+    img_path,
+    device_opt,
+    conf,
+    iou,
+    infer_size,
+    max_det,
+    inputs_cls_name,
+    yolo_model="yolov8n.pt",
 ):
     model = YOLO(yolo_model)
+    if inputs_cls_name == []:
+        inputs_cls_name = None
 
     results = model(
         source=img_path,
@@ -325,6 +337,7 @@ def model_det_loading(
         imgsz=infer_size,
         conf=conf,
         iou=iou,
+        classes=inputs_cls_name,
         max_det=max_det,
     )
     results = list(results)[0]
@@ -332,17 +345,25 @@ def model_det_loading(
 
 
 # 图像分类模型加载
-def model_cls_loading(img_path, yolo_model="yolov8s-cls.pt"):
+def model_cls_loading(img_path, device_opt, yolo_model="yolov8s-cls.pt"):
     model = YOLO(yolo_model)
 
-    results = model(source=img_path)
+    results = model(source=img_path, device=device_opt)
     results = list(results)[0]
     return results
 
 
 # YOLOv8图片检测函数
 def yolo_det_img(
-    img_path, model_name, device_opt, infer_size, conf, iou, max_det, obj_size
+    img_path,
+    model_name,
+    device_opt,
+    infer_size,
+    conf,
+    iou,
+    max_det,
+    inputs_cls_name,
+    obj_size,
 ):
     global model, model_name_tmp, device_tmp
 
@@ -363,8 +384,10 @@ def yolo_det_img(
         iou,
         infer_size,
         max_det,
+        inputs_cls_name,
         yolo_model=f"{model_name}.pt",
     )
+
     # 检测参数
     xyxy_list = predict_results.boxes.xyxy.cpu().numpy().tolist()
     conf_list = predict_results.boxes.conf.cpu().numpy().tolist()
@@ -520,9 +543,11 @@ def yolo_det_img(
 
 
 # YOLOv8图片分类函数
-def yolo_cls_img(img_path, model_name):
+def yolo_cls_img(img_path, device_opt, model_name):
     # 模型加载
-    predict_results = model_cls_loading(img_path, yolo_model=f"{model_name}.pt")
+    predict_results = model_cls_loading(
+        img_path, device_opt, yolo_model=f"{model_name}.pt"
+    )
 
     det_img = Image.open(img_path)
     clas_ratio_list = predict_results.probs.top5conf.tolist()
@@ -572,8 +597,6 @@ def main(args):
         button_secondary_background_fill_hover="*neutral_200",
     )
 
-    custom_css = GYD_CSS
-
     # ------------ Gradio Blocks ------------
     with gr.Blocks(theme=custom_theme, css=custom_css) as gyd:
         with gr.Row():
@@ -581,110 +604,203 @@ def main(args):
         with gr.Row():
             gr.Markdown(GYD_SUB_TITLE)
         with gr.Row():
-                with gr.Tabs():
-                    with gr.TabItem("目标检测与图像分割"):
-                        with gr.Row():
+            with gr.Tabs():
+                with gr.TabItem("目标检测与图像分割"):
+                    with gr.Row():
+                        with gr.Group(elem_id="show_box"):
                             with gr.Column(scale=1):
                                 with gr.Row():
-                                    inputs_img = gr.Image(image_mode="RGB", type="filepath", label="原始图片")
+                                    inputs_img = gr.Image(
+                                        image_mode="RGB", type="filepath", label="原始图片"
+                                    )
                                 with gr.Row():
-                                    device_opt = gr.Radio(choices=["cpu", "0", "1", "2", "3"], value="cpu", label="设备")
+                                    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="模型")
+                                    inputs_model = gr.Dropdown(
+                                        choices=model_names,
+                                        value=model_name,
+                                        type="value",
+                                        label="模型",
+                                    )
                                 with gr.Accordion("高级设置", open=True):
                                     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="最大检测数")
+                                        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 阈值")
+                                        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="目标尺寸")
+                                        inputs_cls_name = gr.Dropdown(
+                                            choices=model_cls_name_cp,
+                                            value=[],
+                                            multiselect=True,
+                                            allow_custom_value=True,
+                                            type="index",
+                                            label="类别选择",
+                                        )
+                                    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")
+                                    det_btn_img = gr.Button(
+                                        value="检测", variant="primary"
+                                    )
+
+                        with gr.Group(elem_id="show_box"):
                             with gr.Column(scale=1):
                                 # with gr.Row():
                                 #     outputs_img = gr.Image(type="pil", label="检测图片")
                                 with gr.Row():
-                                    outputs_img_slider = ImageSlider(position=0.5, label="检测图片")
+                                    outputs_img_slider = ImageSlider(
+                                        position=0.5, label="检测图片"
+                                    )
                                 with gr.Row():
                                     outputs_imgfiles = gr.Files(label="图片下载")
                                 with gr.Row():
                                     outputs_objSize = gr.Label(label="目标尺寸占比统计")
                                 with gr.Row():
                                     outputs_clsSize = gr.Label(label="类别检测占比统计")
+
+                    with gr.Group(elem_id="show_box"):
                         with gr.Row():
                             gr.Examples(
                                 examples=EXAMPLES_DET,
                                 fn=yolo_det_img,
                                 inputs=[
-                                    inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det,
-                                    obj_size],
+                                    inputs_img,
+                                    inputs_model,
+                                    device_opt,
+                                    inputs_size,
+                                    input_conf,
+                                    inputs_iou,
+                                    max_det,
+                                    inputs_cls_name,
+                                    obj_size,
+                                ],
                                 # outputs=[outputs_img, outputs_objSize, outputs_clsSize],
-                                cache_examples=False)
+                                cache_examples=False,
+                            )
 
-                    with gr.TabItem("图像分类"):
-                        with gr.Row():
+                with gr.TabItem("图像分类"):
+                    with gr.Row():
+                        with gr.Group(elem_id="show_box"):
                             with gr.Column(scale=1):
                                 with gr.Row():
-                                    inputs_img_cls = gr.Image(image_mode="RGB", type="filepath", label="原始图片")
+                                    inputs_img_cls = gr.Image(
+                                        image_mode="RGB", type="filepath", label="原始图片"
+                                    )
                                 with gr.Row():
-                                    inputs_model_cls = gr.Dropdown(choices=[
-                                        "yolov8n-cls", "yolov8s-cls", "yolov8l-cls", "yolov8m-cls", "yolov8x-cls"],
-                                                                value="yolov8s-cls",
-                                                                type="value",
-                                                                label="模型")
+                                    device_opt_cls = gr.Radio(
+                                        choices=["cpu", "0", "1", "2", "3"],
+                                        value="cpu",
+                                        label="设备",
+                                    )
+                                with gr.Row():
+                                    inputs_model_cls = gr.Dropdown(
+                                        choices=[
+                                            "yolov8n-cls",
+                                            "yolov8s-cls",
+                                            "yolov8l-cls",
+                                            "yolov8m-cls",
+                                            "yolov8x-cls",
+                                        ],
+                                        value="yolov8s-cls",
+                                        type="value",
+                                        label="模型",
+                                    )
                                 with gr.Row():
                                     gr.ClearButton(inputs_img, value="清除")
-                                    det_btn_img_cls = gr.Button(value='检测', variant="primary")
+                                    det_btn_img_cls = gr.Button(
+                                        value="检测", variant="primary"
+                                    )
+
+                        with gr.Group(elem_id="show_box"):
                             with gr.Column(scale=1):
                                 with gr.Row():
                                     outputs_img_cls = gr.Image(type="pil", label="检测图片")
                                 with gr.Row():
                                     outputs_ratio_cls = gr.Label(label="图像分类结果")
+
+                    with gr.Group(elem_id="show_box"):
                         with gr.Row():
                             gr.Examples(
                                 examples=EXAMPLES_CLAS,
                                 fn=yolo_cls_img,
-                                inputs=[inputs_img_cls, inputs_model_cls],
+                                inputs=[
+                                    inputs_img_cls,
+                                    device_opt_cls,
+                                    inputs_model_cls,
+                                ],
                                 # outputs=[outputs_img_cls, outputs_ratio_cls],
-                                cache_examples=False)
-                            
+                                cache_examples=False,
+                            )
 
         with gr.Accordion("Gradio YOLOv8 Det 安装与使用教程"):
-            gr.Markdown(
-                """## Gradio YOLOv8 Det 安装与使用教程
-            ```shell
-            conda create -n yolo python==3.8
-            conda activate yolo # 进入环境
-            git clone https://gitee.com/CV_Lab/gradio-yolov8-det.git
-            cd gradio-yolov8-det
-            pip install -r ./requirements.txt -U
-            ```
-            ```shell
-            # 共享模式
-            python gradio_yolov8_det_v1.py -is # 在浏览器中以共享模式打开，https://**.gradio.app/
-            # 自定义模型配置
-            python gradio_yolov8_det_v1.py -mc ./model_config/model_name_all.yaml
-            # 自定义下拉框默认模型名称
-            python gradio_yolov8_det_v1.py -mn yolov8m
-            # 自定义类别名称
-            python gradio_yolov8_det_v1.py -cls ./cls_name/cls_name_zh.yaml  （目标检测与图像分割）
-            python gradio_yolov8_det_v1.py -cin ./cls_name/cls_imgnet_name_zh.yaml  （图像分类）
-            # 自定义NMS置信度阈值
-            python gradio_yolov8_det_v1.py -conf 0.8
-            # 自定义NMS IoU阈值
-            python gradio_yolov8_det_v1.py -iou 0.5
-            # 设置推理尺寸，默认为640
-            python gradio_yolov8_det_v1.py -isz 320
-            # 设置最大检测数，默认为50
-            python gradio_yolov8_det_v1.py -mdn 100
-            # 设置滑块步长，默认为0.05
-            python gradio_yolov8_det_v1.py -ss 0.01
-            ```
-            """
-            )
+            with gr.Group(elem_id="show_box"):
+                gr.Markdown(
+                    """## Gradio YOLOv8 Det 安装与使用教程
+                ```shell
+                conda create -n yolo python==3.8
+                conda activate yolo # 进入环境
+                git clone https://gitee.com/CV_Lab/gradio-yolov8-det.git
+                cd gradio-yolov8-det
+                pip install -r ./requirements.txt -U
+                ```
+                ```shell
+                # 共享模式
+                python gradio_yolov8_det_v1.py -is # 在浏览器中以共享模式打开，https://**.gradio.app/
+                # 自定义模型配置
+                python gradio_yolov8_det_v1.py -mc ./model_config/model_name_all.yaml
+                # 自定义下拉框默认模型名称
+                python gradio_yolov8_det_v1.py -mn yolov8m
+                # 自定义类别名称
+                python gradio_yolov8_det_v1.py -cls ./cls_name/cls_name_zh.yaml  （目标检测与图像分割）
+                python gradio_yolov8_det_v1.py -cin ./cls_name/cls_imgnet_name_zh.yaml  （图像分类）
+                # 自定义NMS置信度阈值
+                python gradio_yolov8_det_v1.py -conf 0.8
+                # 自定义NMS IoU阈值
+                python gradio_yolov8_det_v1.py -iou 0.5
+                # 设置推理尺寸，默认为640
+                python gradio_yolov8_det_v1.py -isz 320
+                # 设置最大检测数，默认为50
+                python gradio_yolov8_det_v1.py -mdn 100
+                # 设置滑块步长，默认为0.05
+                python gradio_yolov8_det_v1.py -ss 0.01
+                ```
+                """
+                )
 
         det_btn_img.click(
             fn=yolo_det_img,
@@ -696,6 +812,7 @@ def main(args):
                 input_conf,
                 inputs_iou,
                 max_det,
+                inputs_cls_name,
                 obj_size,
             ],
             outputs=[
@@ -708,7 +825,7 @@ def main(args):
 
         det_btn_img_cls.click(
             fn=yolo_cls_img,
-            inputs=[inputs_img_cls, inputs_model_cls],
+            inputs=[inputs_img_cls, device_opt_cls, inputs_model_cls],
             outputs=[outputs_img_cls, outputs_ratio_cls],
         )
 
@@ -726,4 +843,4 @@ if __name__ == "__main__":
         favicon_path="./icon/logo.ico",  # 网页图标
         show_error=True,  # 在浏览器控制台中显示错误信息
         quiet=True,  # 禁止大多数打印语句
-    )
+    )