Created
December 21, 2023 08:46
-
-
Save e96031413/4e8a538f0c23292e9deb566658edda71 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def _cache_images(self): | |
logger.warning("\n********************************************************************************\n" | |
"You are using cached images in RAM to accelerate training.\n" | |
"This requires large system RAM.\n" | |
"Make sure you have 200G+ RAM and 136G available disk space for training COCO.\n" | |
"********************************************************************************\n") | |
max_h = self.img_size[0] | |
max_w = self.img_size[1] | |
cache_file = self.data_dir + "/img_resized_cache_" + self.name + ".array" | |
if not os.path.exists(cache_file): | |
logger.info("Caching images for the first time. This might take about 20 minutes for COCO") | |
self.imgs = np.memmap(cache_file, shape=(len(self.ids), max_h, max_w, 3), dtype=np.float32, mode="w+") | |
from tqdm import tqdm | |
from multiprocessing.pool import ThreadPool | |
NUM_THREADs = min(8, os.cpu_count()) | |
loaded_images = ThreadPool(NUM_THREADs).imap(lambda x: self.load_resized_img(x), range(len(self.annotations))) | |
pbar = tqdm(enumerate(loaded_images), total=len(self.annotations)) | |
for k, out in pbar: | |
self.imgs[k][: out.shape[0], : out.shape[1], :] = out.copy() | |
self.imgs.flush() | |
pbar.close() | |
else: | |
logger.warning("You are using cached imgs! Make sure your dataset is not changed!!\n" | |
"Everytime the self.input_size is changed in your exp file, you need to delete the cached data and re-generate them.\n") | |
logger.info("Loading cached imgs...") | |
self.imgs = np.memmap(cache_file, shape=(len(self.ids), max_h, max_w, 3), dtype=np.float32, mode="r+") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment