HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix Pose export from YAML (#3312)

Browse source
This commit is contained in:
Glenn Jocher 2023-06-21 18:04:22 +02:00 • committed by GitHub
parent 9d1e5567de
commit 35c02c8319
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/yolo/engine/exporter.py
View file

@ -224,7 +224,7 @@ class Exporter:
'imgsz': self.imgsz, 'imgsz': self.imgsz,
'names': model.names} # model metadata 'names': model.names} # model metadata
if model.task == 'pose': if model.task == 'pose':
self.metadata['kpt_shape'] = model.kpt_shape self.metadata['kpt_shape'] = model.model[-1].kpt_shape
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with input shape {tuple(im.shape)} BCHW and " LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with input shape {tuple(im.shape)} BCHW and "
f'output shape(s) {self.output_shape} ({file_size(file):.1f} MB)') f'output shape(s) {self.output_shape} ({file_size(file):.1f} MB)')