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