Check im.format during dataset caching (#2042)

* Check im.format during dataset caching * Update datasets.py
ultralytics · Jan 26, 2021 · 8fcbe43 · 8fcbe43
1 parent a41d910
commit 8fcbe43
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -445,7 +445,8 @@ def cache_labels(self, path=Path('./labels.cache'), prefix=''):
                 im = Image.open(im_file)
                 im.verify()  # PIL verify
                 shape = exif_size(im)  # image size
-                assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels'
+                assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
+                assert im.format.lower() in img_formats, f'invalid image format {im.format}'
 
                 # verify labels
                 if os.path.isfile(lb_file):