app v02 update

app.py

@@ -1,6 +1,6 @@
-# Streamlit YOLOv5 Model2X v0.1
+# Streamlit YOLOv5 Model2X v0.2
 # 创建人：曾逸夫
-# 创建时间：2022-07-14
+# 创建时间：2022-07-17
 # 功能描述：多选，多项模型转换和打包下载
 
 import os
@@ -41,14 +41,41 @@ def zipDir(origin_dir, compress_file):
     zip.close()
 
 
-# params_include_list = ["torchscript", "onnx", "openvino", "coreml", "saved_model", "pb", "tflite", "tfjs"]
-def cb_opt(weight_name, btn_model_list, params_include_list):
+# 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()
-            os.system(f'python export.py --weights ./weights/{weight_name} --include {params_include_list[i]}')
+            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)}秒")
 
@@ -59,7 +86,6 @@ def cb_opt(weight_name, btn_model_list, params_include_list):
 def main():
     with st.container():
         st.title("Streamlit YOLOv5 Model2X")
-        st.subheader('创建人：曾逸夫（Zeng Yifu）')
         st.text("基于Streamlit的YOLOv5模型转换工具")
 
         st.write("-------------------------------------------------------------")
@@ -80,28 +106,59 @@ def main():
             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_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_coreml, cb_saved_model, cb_pb, cb_tflite, cb_tfjs]
+                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(weight_name, btn_model_list, params_include_list)
+                cb_opt(device, imgSize, weight_name, btn_model_list, params_include_list, [iou_thres, conf_thres],
+                       tflite_options, onnx_options, torchscript_options)
 
     st.write("-------------------------------------------------------------")