Update app.py

app.py

@@ -9,7 +9,14 @@ from yolov5 import detect
 from PIL import Image
 
 # 目标检测
-def Detect(image):
+def Detect(image, image_type):
+    if image_type == "红外图像":
+        pt = "best.pt"
+        cnf = "FLIR.yaml"
+    else:
+        pt = "yolov5s.pt"
+        cnf = "coco128.yaml"
+
     # 创建临时文件夹
     temp_path = tempfile.TemporaryDirectory(dir="./")
     temp_dir = temp_path.name
@@ -21,7 +28,7 @@ def Detect(image):
     # 结果图片的存储目录
     temp_result_path = os.path.join(temp_dir, "tempresult")
     # 对临时图片进行检测
-    detect.run(source=temp_image_path, data="test_image/FLIR.yaml", weights="weights/best.pt", project=f'./{temp_dir}',name = 'tempresult', hide_conf=False,  conf_thres=0.35)
+    detect.run(source=temp_image_path, data=f"test_image/{cnf}", weights=f"weights/{pt}", project=f'./{temp_dir}',name = 'tempresult', hide_conf=False,  conf_thres=0.35)
     # 结果图片的路径
     temp_result_path = os.path.join(temp_result_path, os.listdir(temp_result_path)[0])
     # 读取结果图片
@@ -32,25 +39,34 @@ def Detect(image):
 
 # 候选图片
 example_image= [
-    "./test_image/video-2SReBn5LtAkL5HMj2-frame-005072-MA7NCLQGoqq9aHaiL.jpg",
-    "./test_image/video-2rsjnZFyGQGeynfbv-frame-003708-6fPQbB7jtibwaYAE7.jpg",
-    "./test_image/video-2SReBn5LtAkL5HMj2-frame-000317-HTgPBFgZyPdwQnNvE.jpg",
-    "./test_image/video-jNQtRj6NGycZDEXpe-frame-002515-J3YntG8ntvZheKK3P.jpg",
-    "./test_image/video-kDDWXrnLSoSdHCZ7S-frame-003063-eaKjPvPskDPjenZ8S.jpg",
-    "./test_image/video-r68Yr9RPWEp5fW2ZF-frame-000333-X6K5iopqbmjKEsSqN.jpg"
+    ["./test_image/1.jpg", "红外图像"],
+    ["./test_image/2.jpg", "红外图像"],
+    ["./test_image/3.jpg", "红外图像"],
+    ["./test_image/8.jpg", "红外图像"],
+    ["./test_image/5.jpg", "红外图像"],
+    # ["./test_image/6.jpg]", "红外图像"],
+    ["./test_image/4.jpg", "可见光图像"],
+    ["./test_image/7.jpg", "可见光图像"]
     ]
 
 # 目标追踪
-def Track(video, tracking_method):
+def Track(video, video_type, tracking_method):
     # 存储临时视频的文件夹
     temp_dir = "./temp"
     # 先清空temp文件夹
     shutil.rmtree("./temp")
     os.mkdir("./temp")
+
+    # 获取视频的形式
+    if video_type == "红外视频":
+        pt = "best2.pt"
+    else:
+        pt = "yolov5s.pt"
+
     # 获取视频的名字
     video_name = os.path.basename(video)
     # 对视频进行检测
-    track.run(source=video, yolo_weights=Path("weights/best2.pt"),reid_weights=Path("weights/osnet_x0_25_msmt17.pt") , project=Path(f'./{temp_dir}'),name = 'tempresult', tracking_method=tracking_method)
+    track.run(source=video, yolo_weights=Path(f"weights/{pt}"),reid_weights=Path("weights/osnet_x0_25_msmt17.pt") , project=Path(f'./{temp_dir}'), name = 'tempresult', tracking_method=tracking_method)
     # 结果视频的路径
     temp_result_path = os.path.join(f'./{temp_dir}', "tempresult", video_name)
     # 返回结果视频的路径
@@ -58,24 +74,35 @@ def Track(video, tracking_method):
 
 # 候选视频
 example_video= [
-    ["./video/5.mp4", "bytetrack"],
-    ["./video/bicyclecity.mp4", "strongsort"],
-    ["./video/9.mp4", "bytetrack"],
-    ["./video/8.mp4", "strongsort"],
-    ["./video/4.mp4", "bytetrack"],
-    ["./video/car.mp4", "strongsort"],
+    ["./video/5.mp4", "红外视频", "bytetrack"],
+    ["./video/bicyclecity.mp4","红外视频", "strongsort"],
+    ["./video/9.mp4", "红外视频", "bytetrack"],
+    ["./video/8.mp4", "红外视频", "strongsort"],
+    ["./video/4.mp4", "红外视频", "bytetrack"],
+    ["./video/car.mp4", "红外视频", "strongsort"],
+    ["./video/caixukun.mp4", "可见光视频", "bytetrack"],
+    ["./video/palace.mp4", "可见光视频", "strongsort"],
     ]
 
 iface_Image = gr.Interface(fn=Detect,
-                           inputs=gr.Image(label="上传一张红外图像,仅支持jpg格式"),
+                           inputs=[gr.Image(label="上传一张图像(jpg格式)"),
+                                   gr.Radio(["红外图像", "可见光图像"],
+                                            label="image type",
+                                            info="选择图片的形式",
+                                            value="红外图像")],
                            outputs=gr.Image(label="检测结果"),
-                       examples=example_image)
+                           examples=example_image
+                           )
 
 iface_video = gr.Interface(fn=Track,
-                           inputs=[gr.Video(label="上传段红外视频,仅支持mp4格式"), 
-                                   gr.Radio(["bytetrack", "strongsort"], 
-                                            label="track methond", 
-                                            info="选择追踪器",
+                           inputs=[gr.Video(label="上传一段视频(mp4格式)"),
+                                   gr.Radio(["红外视频", "可见光视频"],
+                                            label="video type",
+                                            info="选择视频的形式",
+                                            value="bytetrack"),
+                                   gr.Radio(["bytetrack", "strongsort"],
+                                            label="track methond",
+                                            info="建议使用bytetrack, strongsort在cpu上运行很慢",
                                             value="bytetrack")],
                            outputs=gr.Video(label="追踪结果"),
                            examples=example_video
@@ -84,22 +111,4 @@ iface_video = gr.Interface(fn=Track,
 demo = gr.TabbedInterface([iface_video, iface_Image], tab_names=["目标追踪", "目标检测"], title="红外目标检测追踪")
 
 demo.launch()
-#iface_Image.launch()
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-