Fix __len__ and single_cls (ultralytics#249)

bigcash · Jan 11, 2023 · 53fa0f4 · 53fa0f4
1 parent 88c9418
commit 53fa0f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ultralytics/yolo/data/base.py b/ultralytics/yolo/data/base.py
@@ -110,7 +110,7 @@ def update_labels(self, include_class: Optional[list]):
                 if segments:
                     self.labels[i]["segments"] = segments[j]
             if self.single_cls:
-                self.labels[i]["cls"] = 0
+                self.labels[i]["cls"][:, 0] = 0
 
     def load_image(self, i):
         # Loads 1 image from dataset index 'i', returns (im, resized hw)
@@ -191,7 +191,7 @@ def get_label_info(self, index):
         return label
 
     def __len__(self):
-        return len(self.im_files)
+        return len(self.labels)
 
     def update_labels_info(self, label):
         """custom your label format here"""