| import json | |
| import os.path | |
| from PIL import Image | |
| from torch.utils.data import DataLoader | |
| from transformers import CLIPProcessor | |
| from torchvision.transforms import transforms | |
| import pytorch_lightning as pl | |
| class WikiArtDataset(): | |
| def __init__(self, meta_file): | |
| super(WikiArtDataset, self).__init__() | |
| self.files = [] | |
| with open(meta_file, 'r') as f: | |
| js = json.load(f) | |
| for img_path in js: | |
| img_name = os.path.splitext(os.path.basename(img_path))[0] | |
| caption = img_name.split('_')[-1] | |
| caption = caption.split('-') | |
| j = len(caption) - 1 | |
| while j >= 0: | |
| if not caption[j].isdigit(): | |
| break | |
| j -= 1 | |
| if j < 0: | |
| continue | |
| sentence = ' '.join(caption[:j + 1]) | |
| self.files.append({'img_path': os.path.join('datasets/wikiart', img_path), 'sentence': sentence}) | |
| version = 'openai/clip-vit-large-patch14' | |
| self.processor = CLIPProcessor.from_pretrained(version) | |
| self.jpg_transform = transforms.Compose([ | |
| transforms.Resize(512), | |
| transforms.RandomCrop(512), | |
| transforms.ToTensor(), | |
| ]) | |
| def __getitem__(self, idx): | |
| file = self.files[idx] | |
| im = Image.open(file['img_path']) | |
| im_tensor = self.jpg_transform(im) | |
| clip_im = self.processor(images=im, return_tensors="pt")['pixel_values'][0] | |
| return {'jpg': im_tensor, 'style': clip_im, 'txt': file['sentence']} | |
| def __len__(self): | |
| return len(self.files) | |
| class WikiArtDataModule(pl.LightningDataModule): | |
| def __init__(self, meta_file, batch_size, num_workers): | |
| super(WikiArtDataModule, self).__init__() | |
| self.train_dataset = WikiArtDataset(meta_file) | |
| self.batch_size = batch_size | |
| self.num_workers = num_workers | |
| def train_dataloader(self): | |
| return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, | |
| pin_memory=True) | |