From 2dac2a6a79eb4a035435b9552561d5b6cbafb2d9 Mon Sep 17 00:00:00 2001
From: jiakunhao <jiakunhao@star-rising.com>
Date: Thu, 17 Nov 2022 09:52:26 +0800
Subject: [PATCH] =?UTF-8?q?=E8=BE=93=E5=87=BA=E6=A8=A1=E5=9E=8Bimgsz?=
 =?UTF-8?q?=E5=8F=82=E6=95=B0=E6=B7=BB=E5=8A=A0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/controller/AlgorithmController.py | 9 ++++++---
 app/yolov5/export.py                  | 9 +++++----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/app/controller/AlgorithmController.py b/app/controller/AlgorithmController.py
index 70eca5f..647bc68 100644
--- a/app/controller/AlgorithmController.py
+++ b/app/controller/AlgorithmController.py
@@ -360,7 +360,8 @@ def Export_model_RODY(params_str):
     exp_inputPath = params.get('exp_inputPath').value  # 模型路径
     print('输入模型：', exp_inputPath)
     exp_device = params.get('device').value
-    modellist = Start_Model_Export(exp_inputPath, exp_device)
+    imgsz = params.get('imgsz').value
+    modellist = Start_Model_Export(exp_inputPath, exp_device, imgsz)
     exp_outputPath = exp_inputPath.replace('pt', 'zip')  # 压缩文件
     print('模型路径：',exp_outputPath)
     zipf = zipfile.ZipFile(exp_outputPath, 'w')
@@ -437,7 +438,7 @@ def returnDetectParams():
         {"index": 1, "name": "outputPath", "value": 'E:/aicheck/data_set/11442136178662604800/val_results/',
          "description": '输出结果路径',
          "default": './app/maskrcnn/datasets/M006B_waibi/res', "type": "S", 'show': False},
-        {"index": 0, "name": "modPath", "value": "E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt",
+        {"index": 2, "name": "modPath", "value": "E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt",
          "description": '模型路径', "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},
         {"index": 3, "name": "device", "value": "0", "description": '推理核', "default": "cpu", "type": "S",
          'show': False},
@@ -455,7 +456,9 @@ def returnDownloadParams():
          "default": 'E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt/',
          "type": "S", 'show': False},
         {"index": 1, "name": "device", "value": 'gpu', "description": 'CPU或GPU', "default": 'gpu', "type": "S",
-         'show': False}
+         'show': False},
+        {"index": 2, "name": "imgsz", "value": 640, "description": '图像大小', "default": 640, "type": "I",
+         'show': True}
     ]
     params_str = json.dumps(params_list)
     return params_str
diff --git a/app/yolov5/export.py b/app/yolov5/export.py
index ce12ab7..89ba977 100644
--- a/app/yolov5/export.py
+++ b/app/yolov5/export.py
@@ -566,11 +566,12 @@ def run(
     return f  # return list of exported files/dirs
 
 
-def parse_opt(weights,device):
+def parse_opt(weights,device,imgsz):
+    imgsz = [imgsz,imgsz]
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
     parser.add_argument('--weights', nargs='+', type=str, default=weights, help='model.pt path(s)')
-    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=imgsz, help='image (h, w)')  #default=[640, 640]
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--device', default=device, help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
@@ -604,13 +605,13 @@ def main(opt):
         f = run(**vars(opt))
     return f
 
-def Start_Model_Export(weights,device):
+def Start_Model_Export(weights,device,imgsz):
     # 判断cpu or gpu
     if device == 'gpu':
         device = '0'
     else:
         device = 'cpu'
-    opt = parse_opt(weights,device)
+    opt = parse_opt(weights,device,imgsz)
     f = main(opt)
     return f