Spaces:
Running
Running
| import torch, os | |
| from PIL import Image | |
| import numpy as np | |
| import torchvision.transforms as transforms | |
| import torch.utils.data as data | |
| from einops import rearrange | |
| class ImageLoader: | |
| def __init__(self, root): | |
| self.img_dir = root | |
| def __call__(self, img): | |
| file = f'{self.img_dir}/{img}' | |
| img = Image.open(file).convert('RGB') | |
| return img | |
| def imagenet_transform(phase): | |
| if phase == 'train': | |
| transform = transforms.Compose([ | |
| transforms.RandomResizedCrop(224), | |
| transforms.RandomHorizontalFlip(), | |
| transforms.ToTensor() | |
| ]) | |
| elif phase == 'test': | |
| transform = transforms.Compose([ | |
| transforms.Resize([224,224]), | |
| transforms.ToTensor() | |
| ]) | |
| return transform | |
| class Dataset_embedding(data.Dataset): | |
| def __init__(self, cfg_data, phase='train'): | |
| self.transform = imagenet_transform(phase) | |
| self.type_name = cfg_data.type_name | |
| self.type2idx = {self.type_name[i]: i for i in range(len(self.type_name))} | |
| if phase == 'train': | |
| self.loader = ImageLoader(cfg_data.train_dir) | |
| name = os.listdir(f'{cfg_data.train_dir}/{self.type_name[0]}') | |
| self.data = [] | |
| for i in range(len(self.type_name)): | |
| for j in range(len(name)): | |
| self.data.append([self.type_name[i], name[j]]) | |
| elif phase == 'test': | |
| self.loader = ImageLoader(cfg_data.test_dir) | |
| name = os.listdir(f'{cfg_data.test_dir}/{self.type_name[0]}') | |
| self.data = [] | |
| for i in range(1, len(self.type_name)): | |
| for j in range(len(name)): | |
| self.data.append([self.type_name[i], name[j]]) | |
| print(f'The amount of {phase} data is {len(self.data)}') | |
| def __getitem__(self, index): | |
| type_name, image_name = self.data[index] | |
| scene = self.type2idx[type_name] | |
| image = self.transform(self.loader(f'{type_name}/{image_name}')) | |
| return (scene, image) | |
| def __len__(self): | |
| return len(self.data) | |
| def init_embedding_data(cfg_em, phase): | |
| if phase == 'train': | |
| train_dataset = Dataset_embedding(cfg_em, 'train') | |
| test_dataset = Dataset_embedding(cfg_em, 'test') | |
| train_loader = data.DataLoader(train_dataset, | |
| batch_size=cfg_em.batch, | |
| shuffle=True, | |
| num_workers=cfg_em.num_workers, | |
| pin_memory=True) | |
| test_loader = data.DataLoader(test_dataset, | |
| batch_size=cfg_em.batch, | |
| shuffle=False, | |
| num_workers=cfg_em.num_workers, | |
| pin_memory=True) | |
| print(len(train_dataset),len(test_dataset)) | |
| elif phase == 'inference': | |
| test_dataset = Dataset_embedding(cfg_em, 'test') | |
| test_loader = data.DataLoader(test_dataset, | |
| batch_size=1, | |
| shuffle=False, | |
| num_workers=cfg_em.num_workers, | |
| pin_memory=True) | |
| return train_loader, test_loader |