fix dataset (#167)

2 years ago · 5e290e0d28
parent b84ed678de
commit 5e290e0d28
1 changed files with 1 additions and 1 deletions
--- a/ultralytics/yolo/data/base.py
+++ b/ultralytics/yolo/data/base.py
@ -50,7 +50,7 @@ class BaseDataset(Dataset):
        if self.single_cls:
            self.update_labels(include_class=[])
-        self.ni = len(self.im_files)
+        self.ni = len(self.labels)
        # rect stuff
        self.rect = rect