add examples

app.py

@@ -7,6 +7,8 @@
 import argparse
 import csv
 import sys
+from pathlib import Path
+
 
 import gradio as gr
 import torch
@@ -25,6 +27,9 @@ model_name_tmp = ""
 # 设备临时变量
 device_tmp = ""
 
+# 文件后缀
+suffix_list = [".csv", ".yaml"]
+
 
 def parse_args(known=False):
     parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det v0.1")
@@ -144,7 +149,25 @@ def yaml_parse(file_path):
     return yaml.safe_load(open(file_path, "r", encoding="utf-8").read())
 
 
+# yaml csv 文件解析
+def yaml_csv(file_path, file_tag):
+    file_suffix = Path(file_path).suffix
+    if file_suffix == suffix_list[0]:
+        # 模型名称
+        file_names = [i[0] for i in list(csv.reader(open(file_path)))]  # csv版
+    elif file_suffix == suffix_list[1]:
+        # 模型名称
+        file_names = yaml_parse(file_path).get(file_tag)  # yaml版
+    else:
+        print(f"{file_path}格式不正确！程序退出！")
+        sys.exit()
+
+    return file_names
+
+
 def main(args):
+    gr.close_all()
+
     global model
 
     slider_step = 0.05  # 滑动步长
@@ -157,15 +180,21 @@ def main(args):
     cls_name = args.cls_name
     device = args.device
 
+    # # 模型加载
+    # model = model_loading(model_name, device)
+    # # 模型名称
+    # # model_names = [i[0] for i in list(csv.reader(open(model_cfg)))] # csv版
+    # model_names = yaml_parse(model_cfg).get("model_names")  # yaml版
+
+    # # 类别名称
+    # # model_cls_name = [i[0] for i in list(csv.reader(open(cls_name)))] # csv版
+    # model_cls_name = yaml_parse(cls_name).get("model_cls_name")  # yaml版
+
     # 模型加载
     model = model_loading(model_name, device)
-    # 模型名称
-    # model_names = [i[0] for i in list(csv.reader(open(model_cfg)))] # csv版
-    model_names = yaml_parse(model_cfg).get("model_names")  # yaml版
 
-    # 类别名称
-    # model_cls_name = [i[0] for i in list(csv.reader(open(cls_name)))] # csv版
-    model_cls_name = yaml_parse(cls_name).get("model_cls_name")  # yaml版
+    model_names = yaml_csv(model_cfg, "model_names")
+    model_cls_name = yaml_csv(cls_name, "model_cls_name")
 
     # -------------------输入组件-------------------
     inputs_img = gr.inputs.Image(type="pil", label="原始图片")
@@ -205,7 +234,36 @@ def main(args):
     # 描述
     description = "<div align='center'>可自定义目标检测模型、安装简单、使用方便</div>"
 
-    gr.close_all()
+    # 示例图片
+    examples = [
+        [
+            "./img_example/bus.jpg",
+            "cpu",
+            "yolov5s",
+            0.6,
+            0.5,
+            True,
+            ["人", "公交车"],
+        ],
+        [
+            "./img_example/Millenial-at-work.jpg",
+            "0",
+            "yolov5l",
+            0.5,
+            0.45,
+            True,
+            ["人", "椅子", "杯子", "笔记本电脑"],
+        ],
+        [
+            "./img_example/zidane.jpg",
+            "0",
+            "yolov5m",
+            0.25,
+            0.5,
+            False,
+            ["人", "领带"],
+        ],
+    ]
 
     # 接口
     gr.Interface(
@@ -214,6 +272,7 @@ def main(args):
         outputs=[outputs, outputs02],
         title=title,
         description=description,
+        examples=examples,
         theme="seafoam",
         # live=True, # 实时变更输出
         flagging_dir="run"  # 输出目录