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