# Streamlit YOLOv5 Model2X v0.1
# 创建人：曾逸夫
# 创建时间：2022-07-13
# 功能描述：单选，单个模型转换和下载

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"./weights/{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"./weights/{compress_file}", "w", zipfile.ZIP_DEFLATED)
    for path, dirnames, filenames in os.walk(f"./weights/{origin_dir}"):
        fpath = path.replace(f"./weights/{origin_dir}", '')
        for filename in filenames:
            zip.write(os.path.join(path, filename), os.path.join(fpath, filename))
    zip.close()


# rd_convert = st.radio("请选择转换的类型：",
#                     ("TorchScript", "ONNX", "OpenVINO", "TensorRT", "CoreML", "TensorFlow SavedModel",
#                     "TensorFlow GraphDef", "TensorFlow Lite", "TensorFlow.js"))
# params_include_list = ["torchscript", "onnx", "openvino", "engine", "coreml", "saved_model", "pb", "tflite", "tfjs"]
def cb_opt(weight_name, btn_model, rd_convert_list, params_include_list):
    weight_name_ = weight_name.split(".")[0]

    for i in range(len(rd_convert_list)):
        if btn_model == rd_convert_list[i]:
            st.info(f"正在转换{btn_model}......")
            s = time.time()
            if btn_model == "TensorRT":
                os.system(
                    f'python export.py --weights ./weights/{weight_name} --include {params_include_list[i]} --device 0')
            else:
                os.system(f'python export.py --weights ./weights/{weight_name} --include {params_include_list[i]}')
            e = time.time()
            st.success(f"{btn_model}转换完成，用时{round((e-s), 2)}秒")

            # openvino, saved_model, tfjs
            if btn_model == "OpenVINO":
                # openvino
                zipDir(f"{weight_name_}_openvino_model", f"{weight_name_}_openvino_model.zip")
                download_file(f"{weight_name_}_openvino_model.zip")
            elif btn_model == "TensorFlow SavedModel":
                # saved_model
                zipDir(f"{weight_name_}_saved_model", f"{weight_name_}_saved_model.zip")
                download_file(f"{weight_name_}_saved_model.zip")
            elif btn_model == "TensorFlow.js":
                # tfjs
                zipDir(f"{weight_name_}_web_model", f"{weight_name_}_web_model.zip")
                download_file(f"{weight_name_}_web_model.zip")

            elif btn_model == "CoreML":
                # coreml
                download_file(f"{weight_name_}.mlmodel")
            elif btn_model == "TensorFlow Lite":
                download_file(f"{weight_name_}-fp16.{params_include_list[i]}")
            else:
                download_file(f"{weight_name_}.{params_include_list[i]}")


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}写入成功！")

            rd_convert = st.radio("请选择转换的类型：",
                                  ("TorchScript", "ONNX", "OpenVINO", "TensorRT", "CoreML", "TensorFlow SavedModel",
                                   "TensorFlow GraphDef", "TensorFlow Lite", "TensorFlow.js"))
            rd_convert_list = [
                "TorchScript", "ONNX", "OpenVINO", "TensorRT", "CoreML", "TensorFlow SavedModel", "TensorFlow GraphDef",
                "TensorFlow Lite", "TensorFlow.js"]

            btn_convert = st.button('转换')

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

            if btn_convert:
                cb_opt(weight_name, rd_convert, rd_convert_list, params_include_list)

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


if __name__ == "__main__":
    main()