From 53fa0f4550471c1938a9bfceedcdff724a875f5d Mon Sep 17 00:00:00 2001
From: Laughing <61612323+Laughing-q@users.noreply.github.com>
Date: Wed, 11 Jan 2023 18:25:10 +0800
Subject: [PATCH] Fix `__len__` and `single_cls` (#249)

---
 ultralytics/yolo/data/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ultralytics/yolo/data/base.py b/ultralytics/yolo/data/base.py
index 1745f8f..d702b1b 100644
--- a/ultralytics/yolo/data/base.py
+++ b/ultralytics/yolo/data/base.py
@@ -110,7 +110,7 @@ class BaseDataset(Dataset):
                 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 @@ class BaseDataset(Dataset):
         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"""