import os
os.system("pip install cython_bbox")
import gradio as gr
import tempfile
import track
import shutil
from pathlib import Path
from yolov5 import detect
from PIL import 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
    # 临时图片的路径
    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=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])
    # 读取结果图片
    result_image = Image.open(temp_result_path).copy()
    # 删除临时文件夹
    temp_path.cleanup()
    return result_image

# 候选图片
example_image= [
    ["./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, 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(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)
    # 返回结果视频的路径
    return temp_result_path

# 候选视频
example_video= [
    ["./video/5.mp4", "红外视频", "bytetrack"],
    ["./video/bicyclecity.mp4","红外视频", "bytetrack"],
    ["./video/9.mp4", "红外视频", "bytetrack"],
    ["./video/8.mp4", "红外视频", "strongsort"],
    ["./video/4.mp4", "红外视频", "bytetrack"],
    ["./video/car.mp4", "红外视频", "strongsort"],
    ["./video/caixukun.mp4", "可见光视频", "bytetrack"],
    ["./video/palace.mp4", "可见光视频", "bytetrack"],
    ]

iface_Image = gr.Interface(fn=Detect,
                           inputs=[gr.Image(label="上传一张图像(jpg格式)"),
                                   gr.Radio(["红外图像", "可见光图像"],
                                            label="image type",
                                            info="选择图片的形式",
                                            value="红外图像")],
                           outputs=gr.Image(label="检测结果"),
                           examples=example_image
                           )

iface_video = gr.Interface(fn=Track,
                           inputs=[gr.Video(label="上传一段视频(mp4格式)"),
                                   gr.Radio(["红外视频", "可见光视频"],
                                            label="video type",
                                            info="选择视频的形式",
                                            value="红外视频"),
                                   gr.Radio(["bytetrack", "strongsort"],
                                            label="track methond",
                                            info="建议使用bytetrack, strongsort在cpu上运行很慢",
                                            value="bytetrack")],
                           outputs=gr.Video(label="追踪结果"),
                           examples=example_video
                          )

demo = gr.TabbedInterface([iface_video, iface_Image], tab_names=["目标追踪", "目标检测"], title="红外目标检测追踪")


demo.launch()