xkrs_manan/RODY
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

输出模型imgsz参数添加

Browse Source
This commit is contained in:
jiakunhao 2022-11-17 09:52:26 +08:00
parent cdb081654a
commit 2dac2a6a79
2 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
app/controller/AlgorithmController.py
View File

@ -360,7 +360,8 @@ def Export_model_RODY(params_str):
exp_inputPath = params.get('exp_inputPath').value # 模型路径 exp_inputPath = params.get('exp_inputPath').value # 模型路径
print('输入模型：', exp_inputPath) print('输入模型：', exp_inputPath)
exp_device = params.get('device').value 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') # 压缩文件 exp_outputPath = exp_inputPath.replace('pt', 'zip') # 压缩文件
print('模型路径：',exp_outputPath) print('模型路径：',exp_outputPath)
zipf = zipfile.ZipFile(exp_outputPath, 'w') zipf = zipfile.ZipFile(exp_outputPath, 'w')
@ -437,7 +438,7 @@ def returnDetectParams():
{"index": 1, "name": "outputPath", "value": 'E:/aicheck/data_set/11442136178662604800/val_results/', {"index": 1, "name": "outputPath", "value": 'E:/aicheck/data_set/11442136178662604800/val_results/',
"description": '输出结果路径', "description": '输出结果路径',
"default": './app/maskrcnn/datasets/M006B_waibi/res', "type": "S", 'show': False}, "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}, "description": '模型路径', "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},
{"index": 3, "name": "device", "value": "0", "description": '推理核', "default": "cpu", "type": "S", {"index": 3, "name": "device", "value": "0", "description": '推理核', "default": "cpu", "type": "S",
'show': False}, 'show': False},
@ -455,7 +456,9 @@ def returnDownloadParams():
"default": 'E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt/', "default": 'E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt/',
"type": "S", 'show': False}, "type": "S", 'show': False},
{"index": 1, "name": "device", "value": 'gpu', "description": 'CPU或GPU', "default": 'gpu', "type": "S", {"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) params_str = json.dumps(params_list)
return params_str return params_str

9
app/yolov5/export.py
View File

@ -566,11 +566,12 @@ def run(
return f # return list of exported files/dirs 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 = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') 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('--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('--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('--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') parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
@ -604,13 +605,13 @@ def main(opt):
f = run(**vars(opt)) f = run(**vars(opt))
return f return f
def Start_Model_Export(weights,device): def Start_Model_Export(weights,device,imgsz):
# 判断cpu or gpu # 判断cpu or gpu
if device == 'gpu': if device == 'gpu':
device = '0' device = '0'
else: else:
device = 'cpu' device = 'cpu'
opt = parse_opt(weights,device) opt = parse_opt(weights,device,imgsz)
f = main(opt) f = main(opt)
return f return f