demo.py

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)