| 9 | |
| 10 | |
| 11 | class MyDataset(Dataset): |
| 12 | def __init__(self): |
| 13 | self.data = [] |
| 14 | root = './data/prompt.json' |
| 15 | with open(root, 'rt') as f: |
| 16 | for line in f: |
| 17 | self.data.append(json.loads(line)) |
| 18 | |
| 19 | def __len__(self): |
| 20 | return len(self.data) |
| 21 | |
| 22 | def __getitem__(self, idx): |
| 23 | item = self.data[idx] |
| 24 | |
| 25 | source_filename = item['source'] |
| 26 | target_filename = item['target'] |
| 27 | prompt = item['prompt'] |
| 28 | |
| 29 | source = Image.open(source_filename).convert('L') |
| 30 | source_array = np.array(source) |
| 31 | threshold = 127 |
| 32 | binary_array = np.where(source_array > threshold, 255, 0).astype(np.uint8) |
| 33 | binary_image = Image.fromarray(binary_array) |
| 34 | source = binary_image.convert('RGB') |
| 35 | |
| 36 | target = Image.open(target_filename).convert('RGB') |
| 37 | |
| 38 | source = np.array(source).astype(np.uint8) |
| 39 | target = np.array(target).astype(np.uint8) |
| 40 | |
| 41 | preprocess = self.transform()(image=target, mask=source) |
| 42 | source, target = preprocess['mask'], preprocess['image'] |
| 43 | |
| 44 | ############ Mask-Image Pair ############ |
| 45 | source = source.astype(np.float32) / 255.0 |
| 46 | target = target.astype(np.float32) / 127.5 - 1.0 |
| 47 | |
| 48 | return dict(jpg=target, txt=prompt, hint=source) |
| 49 | |
| 50 | |
| 51 | def transform(self, size=384): |
| 52 | transforms = albumentations.Compose( |
| 53 | [ |
| 54 | albumentations.Resize(height=size, width=size) |
| 55 | |
| 56 | ] |
| 57 | ) |
| 58 | return transforms |
| 59 | |