Update app.py

app.py

@@ -1,7 +1,7 @@
-# Gradio YOLOv8 Det v1.0.0
+# Gradio YOLOv8 Det v1.1.0
 # 创建人：曾逸夫
-# 创建时间：2023-10-28
-# pip install gradio>=4.0.2
+# 创建时间：2023-11-04
+# pip install gradio>=4.1.1
 
 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.0.0"
+GYD_VERSION = "Gradio YOLOv8 Det v1.1.0"
 
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
@@ -47,9 +47,7 @@ obj_style = ["小目标", "中目标", "大目标"]
 
 
 def parse_args(known=False):
-    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 = argparse.ArgumentParser(description="Gradio YOLOv8 Det v1.1.0")
     parser.add_argument("--model_name", "-mn", default="yolov8s", type=str, help="model name")
     parser.add_argument(
         "--model_cfg",
@@ -96,7 +94,7 @@ def parse_args(known=False):
         default=False,
         help="is login",
     )
-    parser.add_argument("--server_port", "-sp", default=7861, type=int, help="server port")
+    parser.add_argument("--server_port", "-sp", default=7860, type=int, help="server port")
 
     args = parser.parse_known_args()[0] if known else parser.parse_args()
     return args
@@ -228,7 +226,7 @@ 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)
 
@@ -237,6 +235,15 @@ def model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_mod
     return results
 
 
+# 图像分类模型加载
+def model_cls_loading(img_path, yolo_model="yolov8s-cls.pt"):
+    model = YOLO(yolo_model)
+
+    results = model(source=img_path)
+    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):
 
@@ -381,12 +388,35 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
         raise gr.Error("图片检测失败！")
 
 
+# YOLOv8图片分类函数
+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
+
+    clas_name_list = []
+    for i in clas_index_list:
+        clas_name_list.append(predict_results.names[i])
+
+    clsRatio_dict = {}
+    index_cls = 0
+    clsDet_dict = Counter(clas_name_list)
+    for k, v in clsDet_dict.items():
+        clsRatio_dict[k] = clas_ratio_list[index_cls]
+        index_cls+=1
+
+    return det_img, clsRatio_dict
+
+
 def main(args):
     gr.close_all()
 
     global model_cls_name_cp, cls_name
 
-    source = args.source
     nms_conf = args.nms_conf
     nms_iou = args.nms_iou
     model_name = args.model_name
@@ -406,8 +436,6 @@ def main(args):
     # ------------ 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>"
@@ -416,32 +444,51 @@ def main(args):
             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.Tabs():
+                    with gr.TabItem("目标检测与图像分割"):
+                        with gr.Row():
+                            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="设备")
+                        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.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="模型")
+                        with gr.Row():
+                            gr.ClearButton(inputs_img, value="清除")
+                            det_btn_img_cls = 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.Tabs():
+                    with gr.TabItem("目标检测与图像分割"):
+                        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.TabItem("图像分类"):
+                        with gr.Row():
+                            outputs_img_cls = gr.Image(type="pil", label="检测图片")
+                        with gr.Row():
+                            outputs_ratio_cls = gr.Label(label="图像分类结果")
+
 
         with gr.Row():
             example_list = [
@@ -463,6 +510,11 @@ def main(args):
                               obj_size],
                           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])
+
     return gyd