Update datasets.py

ultralytics · Feb 5, 2021 · fb32919 · fb32919
1 parent 9e5e9be
commit fb32919
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -372,9 +372,7 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r
 
         # Check cache
         self.label_files = img2label_paths(self.img_files)  # labels
-        cache_basename = os.path.basename(os.path.dirname(self.label_files[0]))
-        new_cache = self.label_files[0].replace(cache_basename, data_type)
-        cache_path = Path(new_cache).parent.with_suffix('.cache')  # cached labels
+        cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')  # cached labels
         if cache_path.is_file():
             cache = torch.load(cache_path)  # load
             if cache['hash'] != get_hash(self.label_files + self.img_files) or 'results' not in cache:  # changed