| import torch, os, torchvision |
| from torchvision import transforms |
| import pandas as pd |
| from PIL import Image |
|
|
|
|
|
|
| class TextImageDataset(torch.utils.data.Dataset): |
| def __init__(self, dataset_path, steps_per_epoch=10000, height=1024, width=1024, center_crop=True, random_flip=False): |
| self.steps_per_epoch = steps_per_epoch |
| metadata = pd.read_csv(os.path.join(dataset_path, "train/metadata.csv")) |
| self.path = [os.path.join(dataset_path, "train", file_name) for file_name in metadata["file_name"]] |
| self.text = metadata["text"].to_list() |
| self.height = height |
| self.width = width |
| self.image_processor = transforms.Compose( |
| [ |
| transforms.CenterCrop((height, width)) if center_crop else transforms.RandomCrop((height, width)), |
| transforms.RandomHorizontalFlip() if random_flip else transforms.Lambda(lambda x: x), |
| transforms.ToTensor(), |
| transforms.Normalize([0.5], [0.5]), |
| ] |
| ) |
|
|
|
|
| def __getitem__(self, index): |
| data_id = torch.randint(0, len(self.path), (1,))[0] |
| data_id = (data_id + index) % len(self.path) |
| text = self.text[data_id] |
| image = Image.open(self.path[data_id]).convert("RGB") |
| target_height, target_width = self.height, self.width |
| width, height = image.size |
| scale = max(target_width / width, target_height / height) |
| shape = [round(height*scale),round(width*scale)] |
| image = torchvision.transforms.functional.resize(image,shape,interpolation=transforms.InterpolationMode.BILINEAR) |
| image = self.image_processor(image) |
| return {"text": text, "image": image} |
|
|
|
|
| def __len__(self): |
| return self.steps_per_epoch |
|
|
