Update app.py

app.py

@@ -1,7 +1,7 @@
-# Gradio YOLOv8 Det v1.1.0
+# Gradio YOLOv8 Det v1.1.1
 # 创建人：曾逸夫
-# 创建时间：2023-11-04
-# pip install gradio>=4.1.1
+# 创建时间：2023-11-10
+# pip install gradio>=4.1.2
 
 import argparse
 import csv
@@ -34,7 +34,7 @@ from util.fonts_opt import is_fonts
 ROOT_PATH = sys.path[0]  # 根目录
 
 # Gradio YOLOv8 Det版本
-GYD_VERSION = "Gradio YOLOv8 Det v1.1.0"
+GYD_VERSION = "Gradio YOLOv8 Det v1.1.1"
 
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
@@ -45,9 +45,36 @@ FONTSIZE = 25
 # 目标尺寸
 obj_style = ["小目标", "中目标", "大目标"]
 
+GYD_TITLE = """
+<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>
+</p>
+"""
+
+GYD_SUB_TITLE = """
+作者：曾逸夫，Gitee：https://gitee.com/PyCVer ，Github：https://github.com/Zengyf-CVer
+"""
+
+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/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, "所有尺寸"],]
+
+
+EXAMPLES_CLAS = [
+            ["./img_examples/bus.jpg", "yolov8s-cls"],
+            ["./img_examples/giraffe.jpg", "yolov8l-cls"],
+            ["./img_examples/zidane.jpg", "yolov8m-cls"],
+            ["./img_examples/Millenial-at-work.jpg", "yolov8m-cls"],
+            ["./img_examples/bus.jpg", "yolov8m-cls"],
+            ["./img_examples/Millenial-at-work.jpg", "yolov8m-cls"],]
 
 def parse_args(known=False):
-    parser = argparse.ArgumentParser(description="Gradio YOLOv8 Det v1.1.0")
+    parser = argparse.ArgumentParser(description=GYD_VERSION)
     parser.add_argument("--model_name", "-mn", default="yolov8s", type=str, help="model name")
     parser.add_argument(
         "--model_cfg",
@@ -63,6 +90,13 @@ def parse_args(known=False):
         type=str,
         help="cls name",
     )
+    parser.add_argument(
+        "--cls_imgnet_name",
+        "-cin",
+        default="./cls_name/cls_imgnet_name_zh.yaml",
+        type=str,
+        help="cls ImageNet name",
+    )
     parser.add_argument(
         "--nms_conf",
         "-conf",
@@ -400,7 +434,8 @@ def yolo_cls_img(img_path, model_name):
 
     clas_name_list = []
     for i in clas_index_list:
-        clas_name_list.append(predict_results.names[i])
+        # clas_name_list.append(predict_results.names[i])
+        clas_name_list.append(model_cls_imagenet_name_cp[i])
 
     clsRatio_dict = {}
     index_cls = 0
@@ -415,13 +450,14 @@ def yolo_cls_img(img_path, model_name):
 def main(args):
     gr.close_all()
 
-    global model_cls_name_cp, cls_name
+    global model_cls_name_cp, model_cls_imagenet_name_cp, cls_name
 
     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
+    cls_imagenet_name = args.cls_imgnet_name  # ImageNet类别
     inference_size = args.inference_size
     max_detnum = args.max_detnum
     slider_step = args.slider_step
@@ -430,18 +466,24 @@ def main(args):
 
     model_names = yaml_csv(model_cfg, "model_names")  # 模型名称
     model_cls_name = yaml_csv(cls_name, "model_cls_name")  # 类别名称
+    model_cls_imagenet_name = yaml_csv(cls_imagenet_name, "model_cls_name")  # 类别名称
 
     model_cls_name_cp = model_cls_name.copy()  # 类别名称
+    model_cls_imagenet_name_cp = model_cls_imagenet_name.copy()  # 类别名称
+
+    custom_theme = gr.themes.Soft(primary_hue="blue").set(
+                button_secondary_background_fill="*neutral_100",
+                button_secondary_background_fill_hover="*neutral_200")
+    custom_css = '''#disp_image {
+        text-align: center; /* Horizontally center the content */
+    }'''
 
     # ------------ Gradio Blocks ------------
-    with gr.Blocks() as gyd:
+    with gr.Blocks(theme=custom_theme, css=custom_css) as gyd:
         with gr.Row():
-            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>"
-                        )
+            gr.Markdown(GYD_TITLE)
         with gr.Row():
-            gr.Markdown(value="作者：曾逸夫，Gitee：https://gitee.com/PyCVer ，Github：https://github.com/Zengyf-CVer")
+            gr.Markdown(GYD_SUB_TITLE)
         with gr.Row():
             with gr.Column(scale=1):
                 with gr.Tabs():
@@ -463,6 +505,12 @@ def main(args):
                         with gr.Row():
                             gr.ClearButton(inputs_img, value="清除")
                             det_btn_img = gr.Button(value='检测', variant="primary")
+                        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],
+                                        # outputs=[outputs_img, outputs_objSize, outputs_clsSize],
+                                        cache_examples=False)
                     
                     with gr.TabItem("图像分类"):
                         with gr.Row():
@@ -472,6 +520,12 @@ def main(args):
                         with gr.Row():
                             gr.ClearButton(inputs_img, value="清除")
                             det_btn_img_cls = gr.Button(value='检测', variant="primary")
+                        with gr.Row():
+                            gr.Examples(examples=EXAMPLES_CLAS,
+                                        fn=yolo_cls_img,
+                                        inputs=[inputs_img_cls, inputs_model_cls],
+                                        # outputs=[outputs_img_cls, outputs_ratio_cls],
+                                        cache_examples=False)
 
             with gr.Column(scale=1):
                 with gr.Tabs():
@@ -490,20 +544,6 @@ def main(args):
                             outputs_ratio_cls = 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,