demo.py

import gradio.inputs
import torch
import gradio as gr
from PIL import Image
import tempfile
import detect
import os
import shutil


# result = detect.run(source=img, data="VOCdevkit/FLIR.yaml", weights="./best.pt")

def de(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="./best.pt", project=f'./{temp_dir}',name = 'tempresult', hide_conf=True)
    # 结果图片的路径
    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"]
    ]

iface = gr.Interface(fn=de, inputs=gr.Image(label="上传一张红外图像,仅支持jpg格式"), outputs="image", examples=example_image, share=True)
iface.launch(share=True)