|
|
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 |
|
|
|
