Support both *.yml and *.yaml files (#4086)

Co-authored-by: ChristopherRogers1991 <ChristopherRogers1991@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sungjoo(Dennis) Hwang <48212469+Denny-Hwang@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-08-01 15:41:09 +02:00
committed by GitHub
parent f2ed85790f
commit b507e3a032
10 changed files with 17 additions and 17 deletions

View File

@ -192,7 +192,7 @@ class ProfileModels:
output = []
for file in files:
engine_file = file.with_suffix('.engine')
if file.suffix in ('.pt', '.yaml'):
if file.suffix in ('.pt', '.yaml', '.yml'):
model = YOLO(str(file))
model.fuse() # to report correct params and GFLOPs in model.info()
model_info = model.info()
@ -229,7 +229,7 @@ class ProfileModels:
if path.is_dir():
extensions = ['*.pt', '*.onnx', '*.yaml']
files.extend([file for ext in extensions for file in glob.glob(str(path / ext))])
elif path.suffix in {'.pt', '.yaml'}: # add non-existing
elif path.suffix in ('.pt', '.yaml', '.yml'): # add non-existing
files.append(str(path))
else:
files.extend(glob.glob(str(path)))