xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 21:44:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix cache in new loader (#112)

Browse Source 
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
This commit is contained in:
Laughing 2022-12-29 17:36:28 +08:00 committed by GitHub
parent 100e21ee80
commit 401bc15345
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
ultralytics/yolo/data/base.py
View File

@ -65,7 +65,7 @@ class BaseDataset(Dataset):
self.ims = [None] * self.ni self.ims = [None] * self.ni
self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files] self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
if cache: if cache:
self.cache_images() self.cache_images(cache)
# transforms # transforms
self.transforms = self.build_transforms(hyp=hyp) self.transforms = self.build_transforms(hyp=hyp)
@ -127,20 +127,20 @@ class BaseDataset(Dataset):
return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
def cache_images(self): def cache_images(self, cache):
# cache images to memory or disk # cache images to memory or disk
gb = 0 # Gigabytes of cached images gb = 0 # Gigabytes of cached images
self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni
fcn = self.cache_images_to_disk if self.cache == "disk" else self.load_image fcn = self.cache_images_to_disk if cache == "disk" else self.load_image
results = ThreadPool(NUM_THREADS).imap(fcn, range(self.ni)) results = ThreadPool(NUM_THREADS).imap(fcn, range(self.ni))
pbar = tqdm(enumerate(results), total=self.ni, bar_format=TQDM_BAR_FORMAT, disable=LOCAL_RANK > 0) pbar = tqdm(enumerate(results), total=self.ni, bar_format=TQDM_BAR_FORMAT, disable=LOCAL_RANK > 0)
for i, x in pbar: for i, x in pbar:
if self.cache == "disk": if cache == "disk":
gb += self.npy_files[i].stat().st_size gb += self.npy_files[i].stat().st_size
else: # 'ram' else: # 'ram'
self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i) self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
gb += self.ims[i].nbytes gb += self.ims[i].nbytes
pbar.desc = f"{self.prefix}Caching images ({gb / 1E9:.1f}GB {self.cache})" pbar.desc = f"{self.prefix}Caching images ({gb / 1E9:.1f}GB {cache})"
pbar.close() pbar.close()
def cache_images_to_disk(self, i): def cache_images_to_disk(self, i):