| import hashlib |
| import os |
| import pickle |
|
|
| import torch |
| from torch.utils.data import Dataset |
| from torchvision.transforms import CenterCrop, Normalize, Resize |
| from torchvision.transforms.functional import to_tensor |
| from PIL import Image |
|
|
| EXTs = ['.png', '.jpg', '.jpeg', ".JPEG"] |
|
|
|
|
| def is_image_file(filename): |
| return any(filename.endswith(ext) for ext in EXTs) |
|
|
|
|
| def _index_root(root: str, cache_path: str = None, require_text: bool = True): |
| """Walk ``root`` once and return a list of (image_path, text_path) pairs. |
| |
| The result is cached on disk so subsequent epochs / DDP workers do not |
| repeat the walk. Caching is keyed by ``root`` + mtime of ``root`` so the |
| cache invalidates when the top-level directory changes. |
| """ |
| if cache_path is None: |
| key = hashlib.md5(root.encode("utf-8")).hexdigest()[:16] |
| cache_path = os.path.join("/tmp", f"imagetext_index_{key}.pkl") |
|
|
| try: |
| st = os.stat(root) |
| mtime = st.st_mtime |
| except FileNotFoundError: |
| raise FileNotFoundError(f"dataset root does not exist: {root}") |
|
|
| if os.path.exists(cache_path): |
| try: |
| with open(cache_path, "rb") as f: |
| cached = pickle.load(f) |
| if cached.get("root") == root and cached.get("mtime") == mtime: |
| return cached["pairs"] |
| except Exception: |
| pass |
|
|
| pairs = [] |
| for dir_, _subdirs, files in os.walk(root): |
| |
| files_set = set(files) |
| for file in files: |
| if not is_image_file(file): |
| continue |
| base, _ = os.path.splitext(file) |
| txt_name = base + ".txt" |
| if require_text and txt_name not in files_set: |
| continue |
| image_path = os.path.join(dir_, file) |
| text_path = os.path.join(dir_, txt_name) if txt_name in files_set else None |
| pairs.append((image_path, text_path)) |
|
|
| try: |
| with open(cache_path, "wb") as f: |
| pickle.dump({"root": root, "mtime": mtime, "pairs": pairs}, f, |
| protocol=pickle.HIGHEST_PROTOCOL) |
| except Exception: |
| pass |
|
|
| return pairs |
|
|
|
|
| class ImageText(Dataset): |
| """Image-text dataset. |
| |
| Behavior change vs. the original implementation: the ``__init__`` no longer |
| opens and reads every ``.txt`` file -- it only indexes file *paths*. The |
| actual text is read on demand in ``__getitem__``. An index of (image, |
| text) path pairs is also cached on disk under ``/tmp`` so that subsequent |
| runs / DDP worker processes start instantly. |
| """ |
|
|
| def __init__(self, root: str, resolution: int, cache_path: str = None, |
| require_text: bool = True): |
| super().__init__() |
| self.root = root |
| self.pairs = _index_root(root, cache_path=cache_path, require_text=require_text) |
|
|
| self.resize = Resize(resolution) |
| self.center_crop = CenterCrop(resolution) |
| self.normalize = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) |
|
|
| @property |
| def image_paths(self): |
| return [p[0] for p in self.pairs] |
|
|
| def __getitem__(self, idx: int): |
| image_path, text_path = self.pairs[idx] |
| if text_path is not None: |
| try: |
| with open(text_path, "r") as f: |
| text = f.read() |
| except Exception: |
| text = "" |
| else: |
| text = "" |
|
|
| pil_image = Image.open(image_path).convert('RGB') |
| pil_image = self.resize(pil_image) |
| pil_image = self.center_crop(pil_image) |
| raw_image = to_tensor(pil_image) |
| normalized_image = self.normalize(raw_image) |
| metadata = { |
| "image_path": image_path, |
| "prompt": text, |
| "raw_image": raw_image, |
| } |
| return normalized_image, text, metadata |
|
|
| def __len__(self): |
| return len(self.pairs) |
|
|
|
|
|
|
