| 77 | |
| 78 | |
| 79 | class AnyWordDataset(Dataset): |
| 80 | def __init__( |
| 81 | self, |
| 82 | json_path, |
| 83 | seed, |
| 84 | resolution=256, |
| 85 | ttf_size=64, |
| 86 | max_len=25, |
| 87 | language=None, |
| 88 | ): |
| 89 | assert isinstance(json_path, (str, list)) |
| 90 | if isinstance(json_path, str): |
| 91 | json_path = [json_path] |
| 92 | self.resolution = resolution |
| 93 | self.ttf_size = ttf_size |
| 94 | self.max_len = max_len |
| 95 | self.language = language |
| 96 | self.raw_data = [] |
| 97 | for jp in json_path: |
| 98 | self.raw_data += self.load_data(jp) |
| 99 | self._length = len(self.raw_data) |
| 100 | |
| 101 | self.transform = transforms.Compose([ |
| 102 | transforms.ToTensor(), |
| 103 | transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), |
| 104 | transforms.Resize((resolution, resolution)) |
| 105 | ]) |
| 106 | random.seed(seed) |
| 107 | |
| 108 | def __len__(self): |
| 109 | return self._length |
| 110 | |
| 111 | def __getitem__(self, index): |
| 112 | gt = self.raw_data[index] |
| 113 | img_path = gt['img_name'] |
| 114 | full_image = np.array(Image.open(img_path).convert('RGB')) |
| 115 | height, width = full_image.shape[:2] |
| 116 | im_shape = (width, height) |
| 117 | polygon = gt['polygon'] |
| 118 | location = calculate_square(full_image, polygon) |
| 119 | crop_image = full_image[location[1]:location[3], location[0]:location[2]] |
| 120 | trans_image = self.transform(crop_image) |
| 121 | mask, masked_image, mask_rect = generate_mask(trans_image, im_shape, self.resolution, polygon, location) |
| 122 | text = gt['text'] |
| 123 | draw_ttf = self.draw_text(text[:self.max_len]) |
| 124 | glyph = self.draw_glyph(text, mask_rect) |
| 125 | info = { |
| 126 | # "full_image": torch.tensor(full_image), |
| 127 | # "location": torch.tensor(location), |
| 128 | "image": trans_image, |
| 129 | 'mask': mask, |
| 130 | 'masked_image': masked_image, |
| 131 | 'ttf_img': draw_ttf, |
| 132 | 'glyph': glyph, |
| 133 | "text": text |
| 134 | } |
| 135 | return info |
| 136 | |