Model interface enhancement (#106)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia 2022-12-28 18:05:01 +05:30 committed by GitHub
parent 38d6df55cb
commit 384f0ef1c6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 39 additions and 22 deletions

View file

@ -86,10 +86,15 @@ class BasePredictor:
# data
if self.data:
if self.data.endswith(".yaml"):
self.data = check_dataset_yaml(self.data)
else:
self.data = check_dataset(self.data)
try:
if self.data.endswith(".yaml"):
self.data = check_dataset_yaml(self.data)
else:
self.data = check_dataset(self.data)
except AssertionError as e:
LOGGER.info(f"Error ocurred: {e}")
finally:
LOGGER.info("Predictor will continue without reading the dataset")
# model
device = select_device(self.args.device)