From db26ccba94e34e2c5c0d1422b193d44ff5b13eff Mon Sep 17 00:00:00 2001
From: Laughing <61612323+Laughing-q@users.noreply.github.com>
Date: Wed, 18 Jan 2023 13:56:17 +0800
Subject: [PATCH] fix im_files (#460)

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

diff --git a/ultralytics/yolo/data/dataset.py b/ultralytics/yolo/data/dataset.py
index 387271e..218a385 100644
--- a/ultralytics/yolo/data/dataset.py
+++ b/ultralytics/yolo/data/dataset.py
@@ -88,6 +88,7 @@ class YOLODataset(BaseDataset):
         x["results"] = nf, nm, ne, nc, len(self.im_files)
         x["msgs"] = msgs  # warnings
         x["version"] = self.cache_version  # cache version
+        self.im_files = [lb["im_file"] for lb in x["labels"]]
         try:
             np.save(path, x)  # save cache for next time
             path.with_suffix(".cache.npy").rename(path)  # remove .npy suffix