Update app.py

app.py

@@ -1,7 +1,7 @@
-# Gradio YOLOv8 Det v1.1.1
+# Gradio YOLOv8 Det v1.2.0
 # 创建人：曾逸夫
 # 创建时间：2023-11-10
-# pip install gradio>=4.1.2
+# pip install gradio>=4.3.0
 
 import argparse
 import csv
@@ -31,10 +31,8 @@ from PIL import Image, ImageDraw, ImageFont
 
 from util.fonts_opt import is_fonts
 
-ROOT_PATH = sys.path[0]  # 根目录
-
 # Gradio YOLOv8 Det版本
-GYD_VERSION = "Gradio YOLOv8 Det v1.1.1"
+GYD_VERSION = "Gradio YOLOv8 Det v1.2.0"
 
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
@@ -50,6 +48,9 @@ GYD_TITLE = """
 <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>
+<p align='center'>
+<a href='https://gitee.com/CV_Lab/gradio-yolov8-det'><img src='https://gitee.com/CV_Lab/gradio-yolov8-det/widgets/widget_6.svg' alt='Fork me on Gitee'></img></a>
+</p>
 """
 
 GYD_SUB_TITLE = """
@@ -57,21 +58,25 @@ 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/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, "所有尺寸"],]
-
+    ["./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"],]
+    ["./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"],]
+
+GYD_CSS = """#disp_image {
+        text-align: center; /* Horizontally center the content */
+    }"""
+
 
 def parse_args(known=False):
     parser = argparse.ArgumentParser(description=GYD_VERSION)
@@ -93,7 +98,7 @@ def parse_args(known=False):
     parser.add_argument(
         "--cls_imgnet_name",
         "-cin",
-        default="./cls_name/cls_imgnet_name_zh.yaml",
+        default="./cls_name/cls_imagenet_name_zh.yaml",
         type=str,
         help="cls ImageNet name",
     )
@@ -261,7 +266,7 @@ def seg_output(img_path, seg_mask_list, color_list, cls_list):
 
 
 # 目标检测和图像分割模型加载
-def model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"):
+def model_det_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)
@@ -293,7 +298,13 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
     cls_index_det_stat = []  # 1
 
     # 模型加载
-    predict_results = model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model=f"{model_name}.pt")
+    predict_results = model_det_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()
@@ -341,7 +352,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
             h_obj = y1 - y0
             area_obj = w_obj * h_obj  # 目标尺寸
 
-            if (obj_size == "小目标" and area_obj > 0 and area_obj <= 32 ** 2):
+            if (obj_size == obj_style[0] and area_obj > 0 and area_obj <= 32 ** 2):
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -354,7 +365,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
                 score_det_stat.append(conf)
 
                 area_obj_all.append(area_obj)
-            elif (obj_size == "中目标" and area_obj > 32 ** 2 and area_obj <= 96 ** 2):
+            elif (obj_size == obj_style[1] 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)
 
@@ -367,7 +378,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
                 score_det_stat.append(conf)
 
                 area_obj_all.append(area_obj)
-            elif (obj_size == "大目标" and area_obj > 96 ** 2):
+            elif (obj_size == obj_style[2] and area_obj > 96 ** 2):
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -427,7 +438,7 @@ def yolo_cls_img(img_path, model_name):
 
     # 模型加载
     predict_results = model_cls_loading(img_path, yolo_model=f"{model_name}.pt")
-    
+
     det_img = Image.open(img_path)
     clas_ratio_list = predict_results.probs.top5conf.tolist()
     clas_index_list = predict_results.probs.top5
@@ -442,7 +453,7 @@ def yolo_cls_img(img_path, model_name):
     clsDet_dict = Counter(clas_name_list)
     for k, v in clsDet_dict.items():
         clsRatio_dict[k] = clas_ratio_list[index_cls]
-        index_cls+=1
+        index_cls += 1
 
     return det_img, clsRatio_dict
 
@@ -471,12 +482,10 @@ def main(args):
     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 */
-    }'''
+    custom_theme = gr.themes.Soft(primary_hue="blue").set(button_secondary_background_fill="*neutral_100",
+                                                          button_secondary_background_fill_hover="*neutral_200")
+
+    custom_css = GYD_CSS
 
     # ------------ Gradio Blocks ------------
     with gr.Blocks(theme=custom_theme, css=custom_css) as gyd:
@@ -506,26 +515,34 @@ def main(args):
                             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)
-                    
+                            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():
                             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="模型")
+                            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")
                         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)
+                            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():
@@ -543,7 +560,6 @@ def main(args):
                         with gr.Row():
                             outputs_ratio_cls = gr.Label(label="图像分类结果")
 
-
         det_btn_img.click(fn=yolo_det_img,
                           inputs=[
                               inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det,
@@ -551,9 +567,8 @@ def main(args):
                           outputs=[outputs_img, outputs_objSize, outputs_clsSize])
 
         det_btn_img_cls.click(fn=yolo_cls_img,
-                          inputs=[
-                              inputs_img_cls, inputs_model_cls],
-                          outputs=[outputs_img_cls, outputs_ratio_cls])
+                              inputs=[inputs_img_cls, inputs_model_cls],
+                              outputs=[outputs_img_cls, outputs_ratio_cls])
 
     return gyd