fix im_files (#460)

2023-01-18 13:56:17 +08:00 · 2023-01-18 13:56:17 +08:00 · db26ccba94
commit db26ccba94
parent 453b5f259a
1 changed files with 1 additions and 0 deletions
--- 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