| from __future__ import annotations | |
| import json | |
| import math | |
| from pathlib import Path | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| from einops import rearrange | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| class EditDataset(Dataset): | |
| def __init__( | |
| self, | |
| path: str, | |
| split: str = "train", | |
| splits: tuple[float, float, float] = (0.9, 0.05, 0.05), | |
| min_resize_res: int = 256, | |
| max_resize_res: int = 256, | |
| crop_res: int = 256, | |
| flip_prob: float = 0.0, | |
| ): | |
| assert split in ("train", "val", "test") | |
| assert sum(splits) == 1 | |
| self.path = path | |
| self.min_resize_res = min_resize_res | |
| self.max_resize_res = max_resize_res | |
| self.crop_res = crop_res | |
| self.flip_prob = flip_prob | |
| with open(Path(self.path, "seeds.json")) as f: | |
| self.seeds = json.load(f) | |
| split_0, split_1 = { | |
| "train": (0.0, splits[0]), | |
| "val": (splits[0], splits[0] + splits[1]), | |
| "test": (splits[0] + splits[1], 1.0), | |
| }[split] | |
| idx_0 = math.floor(split_0 * len(self.seeds)) | |
| idx_1 = math.floor(split_1 * len(self.seeds)) | |
| self.seeds = self.seeds[idx_0:idx_1] | |
| def __len__(self) -> int: | |
| return len(self.seeds) | |
| def __getitem__(self, i: int) -> dict[str, Any]: | |
| name, seeds = self.seeds[i] | |
| propt_dir = Path(self.path, name) | |
| seed = seeds[torch.randint(0, len(seeds), ()).item()] | |
| with open(propt_dir.joinpath("prompt.json")) as fp: | |
| prompt = json.load(fp)["edit"] | |
| image_0 = Image.open(propt_dir.joinpath(f"{seed}_0.jpg")) | |
| image_1 = Image.open(propt_dir.joinpath(f"{seed}_1.jpg")) | |
| reize_res = torch.randint(self.min_resize_res, self.max_resize_res + 1, ()).item() | |
| image_0 = image_0.resize((reize_res, reize_res), Image.Resampling.LANCZOS) | |
| image_1 = image_1.resize((reize_res, reize_res), Image.Resampling.LANCZOS) | |
| image_0 = rearrange(2 * torch.tensor(np.array(image_0)).float() / 255 - 1, "h w c -> c h w") | |
| image_1 = rearrange(2 * torch.tensor(np.array(image_1)).float() / 255 - 1, "h w c -> c h w") | |
| crop = torchvision.transforms.RandomCrop(self.crop_res) | |
| flip = torchvision.transforms.RandomHorizontalFlip(float(self.flip_prob)) | |
| image_0, image_1 = flip(crop(torch.cat((image_0, image_1)))).chunk(2) | |
| return dict(edited=image_1, edit=dict(c_concat=image_0, c_crossattn=prompt)) | |
| class EditDatasetEval(Dataset): | |
| def __init__( | |
| self, | |
| path: str, | |
| split: str = "train", | |
| splits: tuple[float, float, float] = (0.9, 0.05, 0.05), | |
| res: int = 256, | |
| ): | |
| assert split in ("train", "val", "test") | |
| assert sum(splits) == 1 | |
| self.path = path | |
| self.res = res | |
| with open(Path(self.path, "seeds.json")) as f: | |
| self.seeds = json.load(f) | |
| split_0, split_1 = { | |
| "train": (0.0, splits[0]), | |
| "val": (splits[0], splits[0] + splits[1]), | |
| "test": (splits[0] + splits[1], 1.0), | |
| }[split] | |
| idx_0 = math.floor(split_0 * len(self.seeds)) | |
| idx_1 = math.floor(split_1 * len(self.seeds)) | |
| self.seeds = self.seeds[idx_0:idx_1] | |
| def __len__(self) -> int: | |
| return len(self.seeds) | |
| def __getitem__(self, i: int) -> dict[str, Any]: | |
| name, seeds = self.seeds[i] | |
| propt_dir = Path(self.path, name) | |
| seed = seeds[torch.randint(0, len(seeds), ()).item()] | |
| with open(propt_dir.joinpath("prompt.json")) as fp: | |
| prompt = json.load(fp) | |
| edit = prompt["edit"] | |
| input_prompt = prompt["input"] | |
| output_prompt = prompt["output"] | |
| image_0 = Image.open(propt_dir.joinpath(f"{seed}_0.jpg")) | |
| reize_res = torch.randint(self.res, self.res + 1, ()).item() | |
| image_0 = image_0.resize((reize_res, reize_res), Image.Resampling.LANCZOS) | |
| image_0 = rearrange(2 * torch.tensor(np.array(image_0)).float() / 255 - 1, "h w c -> c h w") | |
| return dict(image_0=image_0, input_prompt=input_prompt, edit=edit, output_prompt=output_prompt) | |