Spaces:
Running
Running
| from torchvision.datasets import CIFAR10, CelebA | |
| from torch.utils.data import DataLoader, Dataset | |
| from torchvision.transforms import Compose, ToTensor, Lambda, CenterCrop, Resize, RandomHorizontalFlip | |
| import os | |
| import torch | |
| import json | |
| from PIL import Image as im | |
| from helper.tokenizer import Tokenizer | |
| from transformers import AutoProcessor | |
| def center_crop_and_resize(img, crop_size, resize_size): | |
| width, height = img.size | |
| # 1. Center Crop | |
| left = (width - crop_size) / 2 | |
| top = (height - crop_size) / 2 | |
| right = (width + crop_size) / 2 | |
| bottom = (height + crop_size) / 2 | |
| img_cropped = img.crop((left, top, right, bottom)) | |
| # 2. Resize | |
| img_resized = img_cropped.resize((resize_size, resize_size), im.Resampling.BICUBIC) | |
| return img_resized | |
| class UnlabelDataset(Dataset): | |
| def __init__(self, path, transform): | |
| self.path = path | |
| self.file_list = os.listdir(path) | |
| self.transform = transform | |
| def __len__(self) : | |
| return len(self.file_list) | |
| def __getitem__(self, index): | |
| img_path = self.path + self.file_list[index] | |
| image = im.open(img_path) | |
| image = self.transform(image) | |
| return image | |
| class CompositeDataset(Dataset): | |
| def __init__(self, path, text_path, processor: AutoProcessor = None): | |
| self.path = path | |
| self.text_path = text_path | |
| self.tokenizer = Tokenizer() | |
| self.processor = processor | |
| self.file_numbers = os.listdir(path) | |
| self.file_numbers = [ os.path.splitext(filename)[0] for filename in self.file_numbers ] | |
| self.transform = Compose([ | |
| ToTensor(), | |
| CenterCrop(400), | |
| Resize(256, antialias=True), | |
| RandomHorizontalFlip(), | |
| Lambda(lambda x: (x - 0.5) * 2) | |
| ]) | |
| def __len__(self) : | |
| return len(self.file_numbers) | |
| def get_text(self, text_path): | |
| with open(text_path, encoding = 'CP949') as f: | |
| text = json.load(f)['description']['impression']['description'] | |
| return text | |
| def __getitem__(self, idx) : | |
| img_path = self.path + self.file_numbers[idx] + '.png' | |
| text_path = self.text_path + self.file_numbers[idx] + '.json' | |
| image = im.open(img_path) | |
| text = self.get_text(text_path) | |
| if self.processor is not None: | |
| image = center_crop_and_resize(image, 400, 256) | |
| inputs = self.processor( | |
| text=text, | |
| images=image, | |
| return_tensors="pt", | |
| padding='max_length', | |
| max_length=77, | |
| truncation=True, | |
| ) | |
| for j in inputs: | |
| inputs[j] = inputs[j].squeeze(0) | |
| return inputs | |
| else: | |
| image = self.transform(image) | |
| text = self.tokenizer.tokenize(text) | |
| for j in text: | |
| text[j] = text[j].squeeze(0) | |
| return image, text | |
| class DataGenerator(): | |
| def __init__(self, num_workers: int = 4, pin_memory: bool = True): | |
| self.transform = Compose([ | |
| ToTensor(), | |
| Lambda(lambda x: (x - 0.5) * 2) | |
| ]) | |
| self.num_workers = num_workers | |
| self.pin_memory = pin_memory | |
| def cifar10(self, path = './datasets', batch_size : int = 64, train : bool = True): | |
| train_data = CIFAR10(path, download = True, train = train, transform = self.transform) | |
| dl = DataLoader(train_data, batch_size, shuffle = True, num_workers=self.num_workers, pin_memory=self.pin_memory) | |
| return dl | |
| def celeba(self, path = './datasets', batch_size : int = 16): | |
| train_data = CelebA(path, transform = Compose([ | |
| ToTensor(), | |
| CenterCrop(178), | |
| Resize(128), | |
| Lambda(lambda x: (x - 0.5) * 2) | |
| ])) | |
| dl = DataLoader(train_data, batch_size, shuffle = True, num_workers=self.num_workers, pin_memory=self.pin_memory) | |
| return dl | |
| def composite(self, path, text_path, batch_size : int = 16, is_process: bool = False): | |
| processor = None | |
| if is_process: | |
| model_name = "Bingsu/clip-vit-base-patch32-ko" | |
| processor = AutoProcessor.from_pretrained(model_name, use_fast=False) | |
| dataset = CompositeDataset(path, text_path, processor) | |
| return DataLoader(dataset, batch_size=batch_size, shuffle=True, | |
| num_workers=self.num_workers, pin_memory=self.pin_memory) | |
| def random_data(self, size, batch_size : int = 4): | |
| train_data = torch.randn(size) | |
| return DataLoader(train_data, batch_size) | |