Fix HUBDatasetStats for no-label edge cases (#4583)

This commit is contained in:
Glenn Jocher 2023-08-26 19:38:02 +02:00 committed by GitHub
parent 2db35afad5
commit f755ba88c3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 62 additions and 41 deletions

View file

@ -447,10 +447,17 @@ class HUBDatasetStats:
return [[int(c[0]), *(round(float(x), 4) for x in points)] for c, points in zipped]
for split in 'train', 'val', 'test':
if self.data.get(split) is None:
self.stats[split] = None # i.e. no test set
self.stats[split] = None # predefine
path = self.data.get(split)
# Check split
if path is None: # no split
continue
files = [f for f in Path(path).rglob('*.*') if f.suffix[1:].lower() in IMG_FORMATS] # image files in split
if not files: # no images
continue
# Get dataset statistics
dataset = YOLODataset(img_path=self.data[split],
data=self.data,
use_segments=self.task == 'segment',