ulralytics 8.0.199 *.npy image loading exception handling (#5683)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: snyk-bot <snyk-bot@snyk.io>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-10-15 18:24:06 +02:00 committed by GitHub
parent 5b3c4cfc0e
commit cedce60f8c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 479 additions and 280 deletions

View file

@ -61,8 +61,8 @@ class BaseDataset(Dataset):
single_cls=False,
classes=None,
fraction=1.0):
super().__init__()
"""Initialize BaseDataset with given configuration and options."""
super().__init__()
self.img_path = img_path
self.imgsz = imgsz
self.augment = augment
@ -85,7 +85,7 @@ class BaseDataset(Dataset):
self.buffer = [] # buffer size = batch size
self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0
# Cache stuff
# Cache images
if cache == 'ram' and not self.check_cache_ram():
cache = False
self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
@ -123,7 +123,7 @@ class BaseDataset(Dataset):
return im_files
def update_labels(self, include_class: Optional[list]):
"""include_class, filter labels to include only these classes (optional)."""
"""Update labels to include only these classes (optional)."""
include_class_array = np.array(include_class).reshape(1, -1)
for i in range(len(self.labels)):
if include_class is not None:
@ -146,11 +146,17 @@ class BaseDataset(Dataset):
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
if im is None: # not cached in RAM
if fn.exists(): # load npy
im = np.load(fn)
try:
im = np.load(fn)
except Exception as e:
LOGGER.warning(f'{self.prefix}WARNING ⚠️ Removing corrupt *.npy image file {fn} due to: {e}')
Path(fn).unlink(missing_ok=True)
im = cv2.imread(f) # BGR
else: # read image
im = cv2.imread(f) # BGR
if im is None:
raise FileNotFoundError(f'Image Not Found {f}')
if im is None:
raise FileNotFoundError(f'Image Not Found {f}')
h0, w0 = im.shape[:2] # orig hw
if rect_mode: # resize long side to imgsz while maintaining aspect ratio
r = self.imgsz / max(h0, w0) # ratio