Create app.py

app.py

@@ -0,0 +1,102 @@
+import gradio as gr
+import tempfile
+import os
+import track
+import shutil
+from pathlib import Path
+from yolov5 import detect
+from PIL import Image
+
+# 目标检测
+def Detect(image):
+    # 创建临时文件夹
+    temp_path = tempfile.TemporaryDirectory(dir="./")
+    temp_dir = temp_path.name
+    # 临时图片的路径
+    temp_image_path = os.path.join(temp_dir, f"temp.jpg")
+    # 存储临时图片
+    img = Image.fromarray(image)
+    img.save(temp_image_path)
+    # 结果图片的存储目录
+    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)
+    # 结果图片的路径
+    temp_result_path = os.path.join(temp_result_path, os.listdir(temp_result_path)[0])
+    # 读取结果图片
+    result_image = Image.open(temp_result_path).copy()
+    # 删除临时文件夹
+    temp_path.cleanup()
+    return result_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"
+    ]
+
+# 目标追踪
+def Track(video, tracking_method):
+    # 存储临时视频的文件夹
+    temp_dir = "./temp"
+    # 先清空temp文件夹
+    shutil.rmtree("./temp")
+    os.mkdir("./temp")
+    # 获取视频的名字
+    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)
+    # 结果视频的路径
+    temp_result_path = os.path.join(f'./{temp_dir}', "tempresult", video_name)
+    # 返回结果视频的路径
+    return temp_result_path
+
+# 候选视频
+example_video= [
+    ["./video/5.mp4", None],
+    ["./video/bicyclecity.mp4", None],
+    ["./video/9.mp4", None],
+    ["./video/8.mp4", None],
+    ["./video/4.mp4", None],
+    ["./video/car.mp4", None],
+    ]
+
+iface_Image = gr.Interface(fn=Detect,
+                           inputs=gr.Image(label="上传一张红外图像,仅支持jpg格式"),
+                           outputs=gr.Image(label="检测结果"),
+                           examples=example_image)
+
+iface_video = gr.Interface(fn=Track,
+                           inputs=[gr.Video(label="上传段红外视频,仅支持mp4格式"), 
+                                   gr.Radio(["bytetrack", "strongsort"], 
+                                            label="track methond", 
+                                            info="选择追踪器",
+                                            value="bytetrack")],
+                           outputs=gr.Video(label="追踪结果"),
+                           examples=example_video)
+
+demo = gr.TabbedInterface([iface_video, iface_Image], tab_names=["目标追踪", "目标检测"], title="红外目标检测追踪")
+
+demo.launch(share=True)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+