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
Download Patch File Download Diff File Expand all files Collapse all files

9
app/yolov5/export.py
View File

@ -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