app update

app.py

@@ -1,6 +1,6 @@
 # Gradio YOLOv5 Det v0.5
 # author: Zeng Yifu（曾逸夫）
-# creation time: 2022-07-22
+# creation time: 2022-08-05
 # email: [email protected]
 # project homepage: https://gitee.com/CV_Lab/gradio_yolov5_det
 
@@ -24,7 +24,6 @@ import numpy as np
 import pandas as pd
 import plotly.express as px
 from matplotlib import font_manager
-from matplotlib.ticker import MaxNLocator
 
 ROOT_PATH = sys.path[0]  # 项目根目录
 
@@ -49,7 +48,7 @@ ROOT_PATH = sys.path[0]  # 根目录
 yolov5_path = "ultralytics/yolov5"
 
 # 本地模型路径
-# local_model_path = f"{ROOT_PATH}/models"
+local_model_path = f"{ROOT_PATH}/models"
 
 # Gradio YOLOv5 Det版本
 GYD_VERSION = "Gradio YOLOv5 Det v0.5"
@@ -101,7 +100,7 @@ def parse_args(known=False):
     parser.add_argument(
         "--device",
         "-dev",
-        default="cpu",
+        default="cuda:0",
         type=str,
         help="cuda or cpu",
     )
@@ -154,6 +153,18 @@ def yaml_csv(file_path, file_tag):
     return file_names
 
 
+# 检查网络连接
+def check_online():
+    # 参考：https://github.com/ultralytics/yolov5/blob/master/utils/general.py
+    # Check internet connectivity
+    import socket
+    try:
+        socket.create_connection(("1.1.1.1", 443), 5)  # check host accessibility
+        return True
+    except OSError:
+        return False
+
+
 #  模型加载
 def model_loading(model_name, device, opt=[]):
 
@@ -163,7 +174,7 @@ def model_loading(model_name, device, opt=[]):
             yolov5_path,
             model_name,
             device=device,
-            force_reload=[True if "refresh_yolov5" in opt else False][0],
+            force_reload=[True if "refresh_yolov5" in opt and check_online() else False][0],
             _verbose=True,
         )
     except Exception as e:
@@ -472,7 +483,7 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
     model.max_det = int(max_num)  # 最大检测框数
     model.classes = model_cls  # 模型类别
 
-    color_list = color_set(len(model_cls_name_cp))  # 设置颜色
+    color_list = random_color(len(model_cls_name_cp), True)
 
     # ---------------- 加载字体 ----------------
     yaml_index = cls_name.index(".yaml")
@@ -746,7 +757,7 @@ def main(args):
     # article="https://gitee.com/CV_Lab/gradio_yolov5_det"
 
     # 示例图片
-    examples = [
+    examples_img = [
         [
             "./img_examples/bus.jpg",
             "cpu",
@@ -788,6 +799,41 @@ def main(args):
             ["人", "椅子", "杯子", "笔记本电脑"],
             ["label", "pdf", "csv", "excel"],],]
 
+    examples_video = [
+        [
+            "./video_examples/test01.mp4",
+            "cpu",
+            "yolov5s",
+            640,
+            0.5,
+            0.45,
+            12,
+            ["鸟"],
+            ["label"],
+            "Matplotlib",],
+        [
+            "./video_examples/test02.mp4",
+            "cpu",
+            "yolov5m",
+            640,
+            0.6,
+            0.5,
+            15,
+            ["马"],
+            ["label"],
+            "Matplotlib",],
+        [
+            "./video_examples/test03.mp4",
+            "cpu",
+            "yolov5s6",
+            1280,
+            0.5,
+            0.5,
+            20,
+            ["人", "风筝"],
+            ["label"],
+            "Plotly",],]
+
     # 接口
     gyd_img = gr.Interface(
         fn=yolo_det_img,
@@ -796,13 +842,13 @@ def main(args):
         title=title,
         description=description,
         # article=article,
-        examples=examples,
+        examples=examples_img,
         cache_examples=False,
         # theme="seafoam",
         # live=True, # 实时变更输出
         flagging_dir="run",  # 输出目录
-        # allow_flagging="manual",
-        # flagging_options=["good", "generally", "bad"],
+        allow_flagging="manual",
+        flagging_options=["good", "generally", "bad"],
     )
 
     gyd_video = gr.Interface(
@@ -812,12 +858,13 @@ def main(args):
         title=title,
         description=description,
         # article=article,
-        # examples=examples,
+        examples=examples_video,
+        cache_examples=False,
         # theme="seafoam",
         # live=True, # 实时变更输出
         flagging_dir="run",  # 输出目录
-        allow_flagging="never",
-        # flagging_options=["good", "generally", "bad"],
+        allow_flagging="manual",
+        flagging_options=["good", "generally", "bad"],
     )
 
     gyd = gr.TabbedInterface(interface_list=[gyd_img, gyd_video], tab_names=["图片模式", "视频模式"])