ultralytics 8.0.44 export and task fixes (#1088)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Mehran Ghandehari <mehran.maps@gmail.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-02-24 03:11:25 +01:00 committed by GitHub
parent fe61018975
commit 3ea659411b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
32 changed files with 439 additions and 480 deletions

View file

@ -208,12 +208,15 @@ class Exporter:
self.file = file
self.output_shape = tuple(y.shape) if isinstance(y, torch.Tensor) else tuple(tuple(x.shape) for x in y)
self.pretty_name = self.file.stem.replace('yolo', 'YOLO')
description = f'Ultralytics {self.pretty_name} model' + f'trained on {Path(self.args.data).name}' \
if self.args.data else '(untrained)'
self.metadata = {
'description': f'Ultralytics {self.pretty_name} model trained on {Path(self.args.data).name}',
'description': description,
'author': 'Ultralytics',
'license': 'GPL-3.0 https://ultralytics.com/license',
'version': __version__,
'stride': int(max(model.stride)),
'task': model.task,
'names': model.names} # model metadata
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with input shape {tuple(im.shape)} BCHW and "