From 5e290e0d282cfb8f1301703e9fe79e264541ea82 Mon Sep 17 00:00:00 2001
From: Laughing <61612323+Laughing-q@users.noreply.github.com>
Date: Tue, 10 Jan 2023 00:25:20 +0800
Subject: [PATCH] fix dataset (#167)

---
 ultralytics/yolo/data/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ultralytics/yolo/data/base.py b/ultralytics/yolo/data/base.py
index cff9cfa..0474053 100644
--- 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