| import os |
| import glob |
| import torch |
| import random |
| from PIL import Image |
| from torch.utils import data |
| from torchvision import transforms as T |
|
|
| class data_prefetcher(): |
| def __init__(self, loader): |
| self.loader = loader |
| self.dataiter = iter(loader) |
| self.stream = torch.cuda.Stream() |
| self.mean = torch.tensor([0.485, 0.456, 0.406]).cuda().view(1,3,1,1) |
| self.std = torch.tensor([0.229, 0.224, 0.225]).cuda().view(1,3,1,1) |
| |
| |
| |
| |
| self.num_images = len(loader) |
| self.preload() |
|
|
| def preload(self): |
| try: |
| self.src_image1, self.src_image2 = next(self.dataiter) |
| except StopIteration: |
| self.dataiter = iter(self.loader) |
| self.src_image1, self.src_image2 = next(self.dataiter) |
| |
| with torch.cuda.stream(self.stream): |
| self.src_image1 = self.src_image1.cuda(non_blocking=True) |
| self.src_image1 = self.src_image1.sub_(self.mean).div_(self.std) |
| self.src_image2 = self.src_image2.cuda(non_blocking=True) |
| self.src_image2 = self.src_image2.sub_(self.mean).div_(self.std) |
|
|
| def next(self): |
| torch.cuda.current_stream().wait_stream(self.stream) |
| src_image1 = self.src_image1 |
| src_image2 = self.src_image2 |
| self.preload() |
| return src_image1, src_image2 |
| |
| def __len__(self): |
| """Return the number of images.""" |
| return self.num_images |
|
|
| class SwappingDataset(data.Dataset): |
| """Dataset class for the Artworks dataset and content dataset.""" |
|
|
| def __init__(self, |
| image_dir, |
| img_transform, |
| subffix='jpg', |
| random_seed=1234): |
| """Initialize and preprocess the Swapping dataset.""" |
| self.image_dir = image_dir |
| self.img_transform = img_transform |
| self.subffix = subffix |
| self.dataset = [] |
| self.random_seed = random_seed |
| self.preprocess() |
| self.num_images = len(self.dataset) |
|
|
| def preprocess(self): |
| """Preprocess the Swapping dataset.""" |
| print("processing Swapping dataset images...") |
|
|
| temp_path = os.path.join(self.image_dir,'*/') |
| pathes = glob.glob(temp_path) |
| self.dataset = [] |
| for dir_item in pathes: |
| join_path = glob.glob(os.path.join(dir_item,'*.jpg')) |
| print("processing %s"%dir_item,end='\r') |
| temp_list = [] |
| for item in join_path: |
| temp_list.append(item) |
| self.dataset.append(temp_list) |
| random.seed(self.random_seed) |
| random.shuffle(self.dataset) |
| print('Finished preprocessing the Swapping dataset, total dirs number: %d...'%len(self.dataset)) |
| |
| def __getitem__(self, index): |
| """Return two src domain images and two dst domain images.""" |
| dir_tmp1 = self.dataset[index] |
| dir_tmp1_len = len(dir_tmp1) |
|
|
| filename1 = dir_tmp1[random.randint(0,dir_tmp1_len-1)] |
| filename2 = dir_tmp1[random.randint(0,dir_tmp1_len-1)] |
| image1 = self.img_transform(Image.open(filename1)) |
| image2 = self.img_transform(Image.open(filename2)) |
| return image1, image2 |
| |
| def __len__(self): |
| """Return the number of images.""" |
| return self.num_images |
|
|
| def GetLoader( dataset_roots, |
| batch_size=16, |
| dataloader_workers=8, |
| random_seed = 1234 |
| ): |
| """Build and return a data loader.""" |
| |
| num_workers = dataloader_workers |
| data_root = dataset_roots |
| random_seed = random_seed |
| |
| c_transforms = [] |
| |
| c_transforms.append(T.ToTensor()) |
| c_transforms = T.Compose(c_transforms) |
|
|
| content_dataset = SwappingDataset( |
| data_root, |
| c_transforms, |
| "jpg", |
| random_seed) |
| content_data_loader = data.DataLoader(dataset=content_dataset,batch_size=batch_size, |
| drop_last=True,shuffle=True,num_workers=num_workers,pin_memory=True) |
| prefetcher = data_prefetcher(content_data_loader) |
| return prefetcher |
|
|
| def denorm(x): |
| out = (x + 1) / 2 |
| return out.clamp_(0, 1) |
