Cache images only mode: iterate all data and write images to LMDB cache.
(self)
| 90 | self._run_pipeline() |
| 91 | |
| 92 | def _run_cache_images(self): |
| 93 | """Cache images only mode: iterate all data and write images to LMDB cache.""" |
| 94 | self.logger.info("Start caching images") |
| 95 | |
| 96 | total_images = 0 |
| 97 | start_time = time.time() |
| 98 | |
| 99 | for dataloader in self.dataloaders: |
| 100 | if dataloader.is_empty(): |
| 101 | self.logger.info(f"No items to process for: {dataloader.name}") |
| 102 | continue |
| 103 | |
| 104 | self.logger.info(f"Caching images for: {dataloader.name}") |
| 105 | |
| 106 | pbar = tqdm(dataloader, desc=f"Caching {dataloader.name}") |
| 107 | for batch in pbar: |
| 108 | batch_images = sum( |
| 109 | len(item.get("image")) if isinstance(item.get("image"), list) else 1 |
| 110 | for item in batch |
| 111 | if item.get("image") |
| 112 | ) |
| 113 | total_images += batch_images |
| 114 | pbar.set_description(f"Cached {total_images} images") |
| 115 | |
| 116 | # Free image_pil to prevent memory accumulation |
| 117 | for item in batch: |
| 118 | item.pop("image_pil", None) |
| 119 | |
| 120 | elapsed = time.time() - start_time |
| 121 | self.logger.info( |
| 122 | f"Image caching completed! Total: {total_images}, Time: {elapsed:.2f}s" |
| 123 | ) |
| 124 | |
| 125 | if self.model is not None: |
| 126 | self.model.shutdown() |
| 127 | |
| 128 | def _run_pipeline(self): |
| 129 | """Normal pipeline execution mode.""" |