# Streamlit YOLOv5 Model2X v0.2
# 创建人：曾逸夫
# 创建时间：2022-07-17
# 功能描述：多选，多项模型转换和打包下载

import os
import shutil
import time
import zipfile

import streamlit as st


# 目录操作
def dir_opt(target_dir):
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)
        os.mkdir(target_dir)
    else:
        os.mkdir(target_dir)


# 文件下载
def download_file(uploaded_file):
    # --------------- 下载 ---------------
    with open(f"{uploaded_file}", 'rb') as fmodel:
        # 读取转换的模型文件（pt2x）
        f_download_model = fmodel.read()
    st.download_button(label='下载转换后的模型', data=f_download_model, file_name=f"{uploaded_file}")
    fmodel.close()


# 文件压缩
def zipDir(origin_dir, compress_file):
    # --------------- 压缩 ---------------
    zip = zipfile.ZipFile(f"{compress_file}", "w", zipfile.ZIP_DEFLATED)
    for path, dirnames, filenames in os.walk(f"{origin_dir}"):
        fpath = path.replace(f"{origin_dir}", '')
        for filename in filenames:
            zip.write(os.path.join(path, filename), os.path.join(fpath, filename))
    zip.close()


# params_include_list = ["torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", "pb", "tflite", "tfjs"]
def cb_opt(device, imgSize, weight_name, btn_model_list, params_include_list, iou_conf, tflite_options, onnx_options,
           torchscript_options):

    for i in range(len(btn_model_list)):
        if btn_model_list[i]:
            st.info(f"正在转换{params_include_list[i]}......")
            s = time.time()
            if i == 0:  # torchscript
                os.system(
                    f"python export.py --device {device} --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]} "
                    + "".join([f"--{x} " for x in torchscript_options]))
            if i == 1:  # onnx
                os.system(
                    f"python export.py --device {device} --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]} "
                    + "".join([f"--{x} " for x in onnx_options]))
            if i == 3:
                # TensorRT需要在GPU模式下导出
                pass
                # os.system(
                #     f"python export.py --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]} --device 0"
                # )
            elif i == 8:  # tfjs
                os.system(
                    f"python export.py --device {device} --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]} --iou-thres {iou_conf[0]} --conf-thres {iou_conf[1]}"
                )
            elif i == 7:  # tflite
                # 参考：https://github.com/zldrobit/yolov5
                os.system(
                    f"python export.py --device {device} --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]} "
                    + "".join([f"--{x} " for x in tflite_options]))
            else:
                os.system(
                    f"python export.py --device {device} --imgsz {imgSize} --weights ./weights/{weight_name} --include {params_include_list[i]}"
                )
            e = time.time()
            st.success(f"{params_include_list[i]}转换完成，用时{round((e-s), 2)}秒")

    zipDir("./weights", "convert_weights.zip")  # 打包weights目录，包括原始权重和转换后的权重
    download_file("convert_weights.zip")  # 下载打包文件


def main():
    with st.container():
        st.title("Streamlit YOLOv5 Model2X")
        st.text("基于Streamlit的YOLOv5模型转换工具")

        st.write("-------------------------------------------------------------")

        dir_opt("./weights")

        uploaded_file = st.file_uploader("选择YOLOv5模型文件（.pt）")
        if uploaded_file is not None:

            # 读取上传的模型文件（.pt）
            weight_name = uploaded_file.name

            st.info(f"正在写入{weight_name}......")

            bytes_data = uploaded_file.getvalue()
            with open(f"./weights/{weight_name}", 'wb') as fb:
                fb.write(bytes_data)
            fb.close()
            st.success(f"{weight_name}写入成功！")

            device = st.radio("请选择设备", ('cpu', 'cuda:0'), index=0)
            imgSize = st.radio("请选择图片尺寸", (320, 640, 1280), index=1)

            st.text("请选择转换的类型：")
            cb_torchscript = st.checkbox('TorchScript')

            # ------------- torchscript -------------
            if cb_torchscript:
                torchscript_options = st.multiselect('onnx选项', ['optimize'])
            else:
                torchscript_options = []

            cb_onnx = st.checkbox('ONNX')
            # ------------- onnx -------------
            if cb_onnx:
                onnx_options = st.multiselect('onnx选项', ['dynamic', 'simplify'])
            else:
                onnx_options = []

            cb_openvino = st.checkbox('OpenVINO')
            cb_engine = st.checkbox('TensorRT')
            cb_coreml = st.checkbox('CoreML')
            cb_saved_model = st.checkbox('TensorFlow SavedModel')
            cb_pb = st.checkbox('TensorFlow GraphDef')
            cb_tflite = st.checkbox('TensorFlow Lite')

            # ------------- tflite -------------
            if cb_tflite:
                tflite_options = st.multiselect('tflite选项', ['int8', 'nms', 'agnostic-nms'])
            else:
                tflite_options = []

            # cb_edgetpu = st.checkbox('TensorFlow Edge TPU')
            cb_tfjs = st.checkbox('TensorFlow.js')

            # ------------- tfjs -------------
            if cb_tfjs:
                iou_thres = st.slider(label='NMS IoU', min_value=0.0, max_value=1.0, value=0.45, step=0.05)
                conf_thres = st.slider(label='NMS CONF', min_value=0.0, max_value=1.0, value=0.5, step=0.05)
            else:
                iou_thres, conf_thres = 0.45, 0.5

            btn_convert = st.button('转换')

            btn_model_list = [
                cb_torchscript, cb_onnx, cb_openvino, cb_engine, cb_coreml, cb_saved_model, cb_pb, cb_tflite, cb_tfjs]

            params_include_list = [
                "torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", "pb", "tflite", "tfjs"]

            if btn_convert:
                cb_opt(device, imgSize, weight_name, btn_model_list, params_include_list, [iou_thres, conf_thres],
                       tflite_options, onnx_options, torchscript_options)

    st.write("-------------------------------------------------------------")


if __name__ == "__main__":
    main()