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:
@ -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)))
|
||||
|
Reference in New Issue
Block a user