import os import cv2 import random import torch import numpy as np from PIL import Image import torch.nn.functional as F import torch.utils.data as data import torchvision.transforms as transforms class MakeupDataset(data.Dataset): def __init__(self, opts): self.opts = opts self.phase = opts.phase self.data_root = opts.data_root self.resize_size = opts.resize_size self.crop_size = opts.crop_size self.flip = opts.flip self.semantic_dim = opts.semantic_dim non_makeup_name = os.listdir(os.path.join(self.data_root, 'non_makeup')) self.non_makeup_path = [os.path.join(self.data_root, 'non_makeup', x) for x in non_makeup_name] makeup_name = os.listdir(os.path.join(self.data_root, 'makeup')) self.makeup_path = [os.path.join(self.data_root, 'makeup', x) for x in makeup_name] self.non_makeup_size = len(self.non_makeup_path) self.makeup_size = len(self.makeup_path) if self.phase == 'train' or self.phase=='1000_test': self.data_size = self.non_makeup_size # self.data_size = 10 else: self.data_size = self.non_makeup_size * self.makeup_size # Color perturbation has already been done here, so it's okay not to add color contrast loss self.color_transform = transforms.Compose( [transforms.transforms.ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.25), transforms.ToTensor()]) self.type_transform = transforms.Compose( [transforms.ToTensor()]) self.image_resnet_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) print('The size of dataset: ',self.data_size) def load_img(self, path, angle=0): img = cv2.imread(path) img_color = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img_color = self.rotate(img_color, angle) img_gray = self.rotate(img_gray, angle) return img_color, img_gray def load_parse(self, path, angle=0): parse = cv2.imread(path, cv2.IMREAD_GRAYSCALE) parse = self.rotate(parse, angle) return parse def rotate(self, img, angle): img = Image.fromarray(img) img = img.rotate(angle) img = np.array(img) return img def preprocessing(self, color_img, gray_img, parse): color_img = cv2.resize(color_img, (self.resize_size, self.resize_size)) gray_img = cv2.resize(gray_img, (self.resize_size, self.resize_size)) parse = cv2.resize(parse, (self.resize_size, self.resize_size), interpolation=cv2.INTER_NEAREST) if np.random.random() > 0.5: h1 = int(np.ceil(np.random.uniform(1e-2, self.resize_size - self.crop_size))) w1 = int(np.ceil(np.random.uniform(1e-2, self.resize_size - self.crop_size))) color_img = color_img[h1:h1 + self.crop_size, w1:w1 + self.crop_size] gray_img = gray_img[h1:h1 + self.crop_size, w1:w1 + self.crop_size] parse = parse[h1:h1 + self.crop_size, w1:w1 + self.crop_size] if self.flip: if np.random.random() > 0.5: color_img = np.fliplr(color_img) gray_img = np.fliplr(gray_img) parse = np.fliplr(parse) color_img = cv2.resize(color_img, (self.crop_size, self.crop_size)) gray_img = cv2.resize(gray_img, (self.crop_size, self.crop_size)) parse = cv2.resize(parse, (self.crop_size, self.crop_size), interpolation=cv2.INTER_NEAREST) return color_img, gray_img, parse def split_parse(self, parse): # 0背景；1面部；2左眉毛；3右眉毛；4左眼；5右眼；6眼镜 # 7左耳朵；8右耳朵；9耳环；10鼻子；11口内部；12上嘴唇 # 13下嘴唇；14脖子；15项链；16衣服；17头发；18帽子 h, w = parse.shape result = np.zeros([h, w, self.semantic_dim]) # 与面部无关的都是背景 result[:, :, 0][np.where(parse == 0)] = 1 result[:, :, 0][np.where(parse == 16)] = 1 result[:, :, 0][np.where(parse == 17)] = 1 result[:, :, 0][np.where(parse == 18)] = 1 result[:, :, 0][np.where(parse == 9)] = 1 result[:, :, 1][np.where(parse == 1)] = 1 result[:, :, 2][np.where(parse == 2)] = 1 result[:, :, 2][np.where(parse == 3)] = 1 result[:, :, 3][np.where(parse == 4)] = 1 result[:, :, 3][np.where(parse == 5)] = 1 result[:, :, 1][np.where(parse == 6)] = 1 result[:, :, 4][np.where(parse == 7)] = 1 result[:, :, 4][np.where(parse == 8)] = 1 # result[:, :, 4][np.where(parse == 9)] = 1 result[:, :, 5][np.where(parse == 10)] = 1 result[:, :, 6][np.where(parse == 11)] = 1 result[:, :, 7][np.where(parse == 12)] = 1 result[:, :, 8][np.where(parse == 13)] = 1 result[:, :, 9][np.where(parse == 14)] = 1 result[:, :, 9][np.where(parse == 15)] = 1 result = np.array(result) return result def local_masks(self, split_parse): h, w, c = split_parse.shape all_mask = np.zeros([h, w]) all_mask[np.where(split_parse[:, :, 0] == 0)] = 1 all_mask[np.where(split_parse[:, :, 3] == 1)] = 0 all_mask[np.where(split_parse[:, :, 6] == 1)] = 0 brow_mask = np.zeros([h, w]) brow_mask[np.where(split_parse[:, :, 2] == 1)] = 1 eye_mask = np.zeros([h, w]) eye_mask[np.where(split_parse[:, :, 3] == 1)] = 1 kernel = np.ones((30, 30), np.uint8) if self.crop_size==512: kernel = np.ones((60, 60), np.uint8) if self.crop_size==1024: kernel = np.ones((80, 80), np.uint8) eye_maskall = cv2.dilate(eye_mask, kernel, iterations=1) eye_mask = eye_maskall - eye_mask eye_mask[np.where(all_mask == 0)] = 0 eye_mask[np.where(brow_mask == 1)] = 0 lip_mask = np.zeros([h, w]) lip_mask[np.where(split_parse[:, :, 7] == 1)] = 1 lip_mask[np.where(split_parse[:, :, 8] == 1)] = 1 face_mask = np.zeros([h, w]) face_mask[np.where(all_mask == 1)] = 1 face_mask[np.where(brow_mask == 1)] = 0 face_mask[np.where(eye_mask == 1)] = 0 face_mask[np.where(lip_mask == 1)] = 0 # 去掉脖子和耳朵 face_mask[np.where(split_parse[:, :, 4] == 1)] = 0 face_mask[np.where(split_parse[:, :, 9] == 1)] = 0 all_mask = np.expand_dims(all_mask, axis=2) # Expansion of the dimension face_mask = np.expand_dims(face_mask, axis=2) brow_mask = np.expand_dims(brow_mask, axis=2) eye_mask = np.expand_dims(eye_mask, axis=2) lip_mask = np.expand_dims(lip_mask, axis=2) all_mask = np.concatenate((all_mask, all_mask, all_mask), axis=2) face_mask = np.concatenate((face_mask, face_mask, face_mask), axis=2) brow_mask = np.concatenate((brow_mask, brow_mask, brow_mask), axis=2) eye_mask = np.concatenate((eye_mask, eye_mask, eye_mask), axis=2) lip_mask = np.concatenate((lip_mask, lip_mask, lip_mask), axis=2) return all_mask, face_mask, brow_mask, eye_mask, lip_mask def affine_transform(self, x, theta): theta = theta.view(-1, 2, 3) grid = F.affine_grid(theta, x.size(), align_corners=True) x = F.grid_sample(x, grid, mode='nearest', align_corners=True) return x def test_preprocessing(self, color_img, gray_img, parse): color_img = cv2.resize(color_img, (self.resize_size, self.resize_size)) gray_img = cv2.resize(gray_img, (self.resize_size, self.resize_size)) parse = cv2.resize(parse, (self.resize_size, self.resize_size), interpolation=cv2.INTER_NEAREST) h1 = (self.resize_size - self.crop_size) // 2 w1 = (self.resize_size - self.crop_size) // 2 color_img = color_img[h1:h1 + self.crop_size, w1:w1 + self.crop_size] gray_img = gray_img[h1:h1 + self.crop_size, w1:w1 + self.crop_size] parse = parse[h1:h1 + self.crop_size, w1:w1 + self.crop_size] color_img = cv2.resize(color_img, (self.crop_size, self.crop_size)) gray_img = cv2.resize(gray_img, (self.crop_size, self.crop_size)) parse = cv2.resize(parse, (self.crop_size, self.crop_size), interpolation=cv2.INTER_NEAREST) return color_img, gray_img, parse def change_color_processing(self,image): color_img = Image.fromarray(image.astype('uint8')) color_img_change = self.color_transform(color_img).unsqueeze(0) color_img_change = np.array(color_img_change.squeeze(0).data) * 255. color_img_change = np.transpose(color_img_change, (1, 2, 0)) return color_img_change def self_supervised_processing(self, color_img, parse): color_img = Image.fromarray(color_img.astype('uint8')) color_img_change = self.color_transform(color_img).unsqueeze(0) parse = Image.fromarray((parse).astype('uint8')) parse = self.type_transform(parse).unsqueeze(0) # generate reference image theta1 = np.zeros(9) theta1[0:6] = np.random.randn(6) * 0.15 theta1 = theta1 + np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]) affine1 = np.reshape(theta1, (3, 3)) affine1 = np.reshape(affine1, -1)[0:6] affine1 = torch.from_numpy(affine1).type(torch.FloatTensor) # print(reference_img_1.size()) color_img_change_warp = self.affine_transform(color_img_change, affine1) color_img_change = np.array(color_img_change.squeeze(0).data) * 255. color_img_change = np.transpose(color_img_change, (1, 2, 0)) color_img_change_warp = np.array(color_img_change_warp.squeeze(0).data) * 255. color_img_change_warp = np.transpose(color_img_change_warp, (1, 2, 0)) parse_warp = self.affine_transform(parse, affine1) parse_warp = np.array(parse_warp.squeeze(0).data) * 255. parse_warp = np.transpose(parse_warp, (1, 2, 0)) parse_warp = parse_warp.astype('uint8') # 变换到[0,1] return color_img_change, color_img_change_warp, np.round(parse_warp) def __getitem__(self, index): if self.phase == 'train': # 需要加载：1未化妆图像；2未化妆图像的灰度图；3未化妆图像语义解析；4未化妆图像的局部mask # 5化妆图像；6化妆图像的语义解析；7化妆图像的局部mask # 8化妆图像color扰动，9化妆图像color扰动后的warp结果，10化妆图像color扰动后的warp结果的mask # 11未化妆图像归一化图像；12化妆图像归一化图像；13化妆图像color扰动后的warp结果warp结果 non_makeup_index = index makeup_index = random.randint(0, self.makeup_size - 1) # 随机旋转 if random.random() > 0.2: non_makeup_angle = np.random.randint(0, 60) - 30 makeup_angle = np.random.randint(0, 60) - 30 else: non_makeup_angle = 0 makeup_angle = 0 # 加载color和gray图像 non_makeup_color_img, non_makeup_gray_img = self.load_img(self.non_makeup_path[non_makeup_index], non_makeup_angle) makeup_color_img, makeup_gray_img = self.load_img(self.makeup_path[makeup_index], makeup_angle) # 加载parse non_makeup_parse = self.load_parse( self.non_makeup_path[non_makeup_index].replace('images', 'seg1')[:-4] + '.png',non_makeup_angle) makeup_parse = self.load_parse(self.makeup_path[makeup_index].replace('images', 'seg1')[:-4] + '.png',makeup_angle) # 数据增强 non_makeup_color_img, non_makeup_gray_img, non_makeup_parse = self.preprocessing(non_makeup_color_img, non_makeup_gray_img, non_makeup_parse) makeup_color_img, makeup_gray_img, makeup_parse = self.preprocessing(makeup_color_img, makeup_gray_img, makeup_parse) # 数据增强，自监督 makeup_color_change_img, makeup_color_change_warp_img, makeup_warp_parse = self.self_supervised_processing( makeup_color_img, makeup_parse) makeup_warp_parse = np.squeeze(makeup_warp_parse, axis=-1) makeup_color_change_img2 = self.change_color_processing(makeup_color_img) makeup_color_change_img3 = self.change_color_processing(makeup_color_img) makeup_color_change_img4 = self.change_color_processing(makeup_color_img) # if random.random() > 0.5: # makeup_color_change_warp_img = np.fliplr(makeup_color_change_warp_img) # makeup_warp_parse = np.fliplr(makeup_warp_parse) # 处理parse变为通道数为N的数据格式 non_makeup_split_parse = self.split_parse(non_makeup_parse) makeup_split_parse = self.split_parse(makeup_parse) # print(np.unique(makeup_warp_parse)) makeup_split_parse_warp = self.split_parse(makeup_warp_parse) # 加载局部mask non_makeup_all_mask, non_makeup_face_mask, non_makeup_brow_mask, non_makeup_eye_mask, non_makeup_lip_mask = self.local_masks( non_makeup_split_parse) makeup_all_mask, makeup_face_mask, makeup_brow_mask, makeup_eye_mask, makeup_lip_mask = self.local_masks( makeup_split_parse) makeup_all_warp_mask, makeup_face_warp_mask, makeup_brow_warp_mask, makeup_eye_warp_mask, makeup_lip_warp_mask = self.local_masks( makeup_split_parse_warp) # non_makeup_split_parse[:,:,-1]=non_makeup_eye_mask[:,:,0] # makeup_split_parse[:, :, -1] = makeup_eye_mask[:, :, 0] # makeup_split_parse_warp[:, :, -1] = makeup_eye_warp_mask[:, :, 0] # 各种归一化 non_makeup_color_img_norm = Image.fromarray(non_makeup_color_img.astype('uint8')) # print(np.unique(non_makeup_color_img_norm)) non_makeup_color_img_norm = self.image_resnet_transform(non_makeup_color_img_norm) makeup_color_img_norm = Image.fromarray(makeup_color_img.astype('uint8')) # print(np.unique(makeup_color_img_norm)) makeup_color_img_norm = self.image_resnet_transform(makeup_color_img_norm) makeup_color_change_warp_img_norm = Image.fromarray(makeup_color_change_warp_img.astype('uint8')) # print(np.unique(makeup_color_change_warp_img_norm)) makeup_color_change_warp_img_norm = self.image_resnet_transform(makeup_color_change_warp_img_norm) # makeup_color_change_img[np.where(makeup_all_mask == 0)] = 0 makeup_color_change_img = makeup_color_change_img / 127.5 - 1 makeup_color_change_img = np.transpose(makeup_color_change_img, (2, 0, 1)) makeup_color_change_img2 = makeup_color_change_img2 / 127.5 - 1 makeup_color_change_img2 = np.transpose(makeup_color_change_img2, (2, 0, 1)) makeup_color_change_img3 = makeup_color_change_img3 / 127.5 - 1 makeup_color_change_img3 = np.transpose(makeup_color_change_img3, (2, 0, 1)) makeup_color_change_img4 = makeup_color_change_img4 / 127.5 - 1 makeup_color_change_img4 = np.transpose(makeup_color_change_img4, (2, 0, 1)) makeup_color_change_warp_img = makeup_color_change_warp_img / 127.5 - 1 makeup_color_change_warp_img = np.transpose(makeup_color_change_warp_img, (2, 0, 1)) makeup_split_parse_warp = np.transpose(makeup_split_parse_warp, (2, 0, 1)) makeup_all_warp_mask = np.transpose(makeup_all_warp_mask, (2, 0, 1)) makeup_face_warp_mask = np.transpose(makeup_face_warp_mask, (2, 0, 1)) makeup_brow_warp_mask = np.transpose(makeup_brow_warp_mask, (2, 0, 1)) makeup_eye_warp_mask = np.transpose(makeup_eye_warp_mask, (2, 0, 1)) makeup_lip_warp_mask = np.transpose(makeup_lip_warp_mask, (2, 0, 1)) non_makeup_color_img = non_makeup_color_img / 127.5 - 1 non_makeup_color_img = np.transpose(non_makeup_color_img, (2, 0, 1)) # non_makeup_gray_img[np.where(non_makeup_all_mask[:,:,0]==0)]=0 non_makeup_gray_img = np.expand_dims(non_makeup_gray_img, -1) non_makeup_gray_img = non_makeup_gray_img / 127.5 - 1 non_makeup_gray_img = np.transpose(non_makeup_gray_img, (2, 0, 1)) non_makeup_split_parse = np.transpose(non_makeup_split_parse, (2, 0, 1)) non_makeup_all_mask = np.transpose(non_makeup_all_mask, (2, 0, 1)) non_makeup_face_mask = np.transpose(non_makeup_face_mask, (2, 0, 1)) non_makeup_brow_mask = np.transpose(non_makeup_brow_mask, (2, 0, 1)) non_makeup_eye_mask = np.transpose(non_makeup_eye_mask, (2, 0, 1)) non_makeup_lip_mask = np.transpose(non_makeup_lip_mask, (2, 0, 1)) makeup_color_img = makeup_color_img / 127.5 - 1 makeup_color_img = np.transpose(makeup_color_img, (2, 0, 1)) # makeup_gray_img[np.where(makeup_all_mask[:,:,0] == 0)] = 0 makeup_gray_img = np.expand_dims(makeup_gray_img, -1) makeup_gray_img = makeup_gray_img / 127.5 - 1 makeup_gray_img = np.transpose(makeup_gray_img, (2, 0, 1)) makeup_split_parse = np.transpose(makeup_split_parse, (2, 0, 1)) makeup_all_mask = np.transpose(makeup_all_mask, (2, 0, 1)) makeup_face_mask = np.transpose(makeup_face_mask, (2, 0, 1)) makeup_brow_mask = np.transpose(makeup_brow_mask, (2, 0, 1)) makeup_eye_mask = np.transpose(makeup_eye_mask, (2, 0, 1)) makeup_lip_mask = np.transpose(makeup_lip_mask, (2, 0, 1)) data = {'non_makeup_color_img_norm': non_makeup_color_img_norm, 'makeup_color_img_norm': makeup_color_img_norm, 'makeup_color_change_warp_img_norm': makeup_color_change_warp_img_norm, 'non_makeup_color_img': torch.from_numpy(non_makeup_color_img).type(torch.FloatTensor), 'non_makeup_gray_img': torch.from_numpy(non_makeup_gray_img).type(torch.FloatTensor), 'non_makeup_split_parse': torch.from_numpy(non_makeup_split_parse).type(torch.FloatTensor), 'non_makeup_all_mask': torch.from_numpy(non_makeup_all_mask).type(torch.FloatTensor), 'non_makeup_face_mask': torch.from_numpy(non_makeup_face_mask).type(torch.FloatTensor), 'non_makeup_brow_mask': torch.from_numpy(non_makeup_brow_mask).type(torch.FloatTensor), 'non_makeup_eye_mask': torch.from_numpy(non_makeup_eye_mask).type(torch.FloatTensor), 'non_makeup_lip_mask': torch.from_numpy(non_makeup_lip_mask).type(torch.FloatTensor), 'makeup_color_img': torch.from_numpy(makeup_color_img).type(torch.FloatTensor), 'makeup_gray_img': torch.from_numpy(makeup_gray_img).type(torch.FloatTensor), 'makeup_split_parse': torch.from_numpy(makeup_split_parse).type(torch.FloatTensor), 'makeup_all_mask': torch.from_numpy(makeup_all_mask).type(torch.FloatTensor), 'makeup_face_mask': torch.from_numpy(makeup_face_mask).type(torch.FloatTensor), 'makeup_brow_mask': torch.from_numpy(makeup_brow_mask).type(torch.FloatTensor), 'makeup_eye_mask': torch.from_numpy(makeup_eye_mask).type(torch.FloatTensor), 'makeup_lip_mask': torch.from_numpy(makeup_lip_mask).type(torch.FloatTensor), 'makeup_color_change_img': torch.from_numpy(makeup_color_change_img).type(torch.FloatTensor), 'makeup_color_change_warp_img': torch.from_numpy(makeup_color_change_warp_img).type( torch.FloatTensor), 'makeup_all_warp_mask': torch.from_numpy(makeup_all_warp_mask).type(torch.FloatTensor), 'makeup_face_warp_mask': torch.from_numpy(makeup_face_warp_mask).type(torch.FloatTensor), 'makeup_brow_warp_mask': torch.from_numpy(makeup_brow_warp_mask).type(torch.FloatTensor), 'makeup_eye_warp_mask': torch.from_numpy(makeup_eye_warp_mask).type(torch.FloatTensor), 'makeup_lip_warp_mask': torch.from_numpy(makeup_lip_warp_mask).type(torch.FloatTensor), 'makeup_split_parse_warp': torch.from_numpy(makeup_split_parse_warp).type(torch.FloatTensor), 'makeup_color_change_img2': torch.from_numpy(makeup_color_change_img2).type(torch.FloatTensor), 'makeup_color_change_img3': torch.from_numpy(makeup_color_change_img3).type(torch.FloatTensor), 'makeup_color_change_img4': torch.from_numpy(makeup_color_change_img4).type(torch.FloatTensor), } return data elif self.phase == 'test_pair': non_makeup_index = index // self.makeup_size makeup_index = index % self.makeup_size # non_makeup_angle = np.random.randint(0, 60) - 30 # makeup_angle = np.random.randint(0, 60) - 30 non_makeup_color_img, non_makeup_gray_img = self.load_img(self.non_makeup_path[non_makeup_index]) makeup_color_img, makeup_gray_img = self.load_img(self.makeup_path[makeup_index]) # 加载parse non_makeup_parse = self.load_parse( self.non_makeup_path[non_makeup_index].replace('images', 'seg1')[:-4] + '.png') makeup_parse = self.load_parse(self.makeup_path[makeup_index].replace('images', 'seg1')[:-4] + '.png') # # 加载parse # non_makeup_parse = self.load_parse( # self.non_makeup_path[non_makeup_index].replace('images', 'seg1')) # makeup_parse = self.load_parse(self.makeup_path[makeup_index].replace('images', 'seg1')) # 数据增强 non_makeup_color_img, non_makeup_gray_img, non_makeup_parse = self.test_preprocessing(non_makeup_color_img, non_makeup_gray_img, non_makeup_parse) makeup_color_img, makeup_gray_img, makeup_parse = self.test_preprocessing(makeup_color_img, makeup_gray_img, makeup_parse) # 处理parse变为通道数为N的数据格式 non_makeup_split_parse = self.split_parse(non_makeup_parse) makeup_split_parse = self.split_parse(makeup_parse) # 加载局部mask non_makeup_all_mask, non_makeup_face_mask, non_makeup_brow_mask, non_makeup_eye_mask, non_makeup_lip_mask = self.local_masks( non_makeup_split_parse) makeup_all_mask, makeup_face_mask, makeup_brow_mask, makeup_eye_mask, makeup_lip_mask = self.local_masks( makeup_split_parse) # 各种归一化 non_makeup_color_img_norm = Image.fromarray(non_makeup_color_img.astype('uint8')) # print(np.unique(non_makeup_color_img_norm)) non_makeup_color_img_norm = self.image_resnet_transform(non_makeup_color_img_norm) makeup_color_img_norm = Image.fromarray(makeup_color_img.astype('uint8')) # print(np.unique(makeup_color_img_norm)) makeup_color_img_norm = self.image_resnet_transform(makeup_color_img_norm) non_makeup_color_img = non_makeup_color_img / 127.5 - 1 non_makeup_color_img = np.transpose(non_makeup_color_img, (2, 0, 1)) # non_makeup_gray_img[np.where(non_makeup_all_mask[:,:,0]==0)]=0 non_makeup_gray_img = np.expand_dims(non_makeup_gray_img, -1) non_makeup_gray_img = non_makeup_gray_img / 127.5 - 1 non_makeup_gray_img = np.transpose(non_makeup_gray_img, (2, 0, 1)) non_makeup_split_parse = np.transpose(non_makeup_split_parse, (2, 0, 1)) non_makeup_all_mask = np.transpose(non_makeup_all_mask, (2, 0, 1)) non_makeup_face_mask = np.transpose(non_makeup_face_mask, (2, 0, 1)) non_makeup_brow_mask = np.transpose(non_makeup_brow_mask, (2, 0, 1)) non_makeup_eye_mask = np.transpose(non_makeup_eye_mask, (2, 0, 1)) non_makeup_lip_mask = np.transpose(non_makeup_lip_mask, (2, 0, 1)) makeup_color_img = makeup_color_img / 127.5 - 1 makeup_color_img = np.transpose(makeup_color_img, (2, 0, 1)) # makeup_gray_img[np.where(makeup_all_mask[:,:,0] == 0)] = 0 makeup_gray_img = np.expand_dims(makeup_gray_img, -1) makeup_gray_img = makeup_gray_img / 127.5 - 1 makeup_gray_img = np.transpose(makeup_gray_img, (2, 0, 1)) makeup_split_parse = np.transpose(makeup_split_parse, (2, 0, 1)) makeup_all_mask = np.transpose(makeup_all_mask, (2, 0, 1)) makeup_face_mask = np.transpose(makeup_face_mask, (2, 0, 1)) makeup_brow_mask = np.transpose(makeup_brow_mask, (2, 0, 1)) makeup_eye_mask = np.transpose(makeup_eye_mask, (2, 0, 1)) makeup_lip_mask = np.transpose(makeup_lip_mask, (2, 0, 1)) data = {'non_makeup_color_img_norm': non_makeup_color_img_norm, 'makeup_color_img_norm': makeup_color_img_norm, 'non_makeup_color_img': torch.from_numpy(non_makeup_color_img).type(torch.FloatTensor), 'non_makeup_gray_img': torch.from_numpy(non_makeup_gray_img).type(torch.FloatTensor), 'non_makeup_split_parse': torch.from_numpy(non_makeup_split_parse).type(torch.FloatTensor), 'non_makeup_all_mask': torch.from_numpy(non_makeup_all_mask).type(torch.FloatTensor), 'non_makeup_face_mask': torch.from_numpy(non_makeup_face_mask).type(torch.FloatTensor), 'non_makeup_brow_mask': torch.from_numpy(non_makeup_brow_mask).type(torch.FloatTensor), 'non_makeup_eye_mask': torch.from_numpy(non_makeup_eye_mask).type(torch.FloatTensor), 'non_makeup_lip_mask': torch.from_numpy(non_makeup_lip_mask).type(torch.FloatTensor), 'makeup_color_img': torch.from_numpy(makeup_color_img).type(torch.FloatTensor), 'makeup_gray_img': torch.from_numpy(makeup_gray_img).type(torch.FloatTensor), 'makeup_split_parse': torch.from_numpy(makeup_split_parse).type(torch.FloatTensor), 'makeup_all_mask': torch.from_numpy(makeup_all_mask).type(torch.FloatTensor), 'makeup_face_mask': torch.from_numpy(makeup_face_mask).type(torch.FloatTensor), 'makeup_brow_mask': torch.from_numpy(makeup_brow_mask).type(torch.FloatTensor), 'makeup_eye_mask': torch.from_numpy(makeup_eye_mask).type(torch.FloatTensor), 'makeup_lip_mask': torch.from_numpy(makeup_lip_mask).type(torch.FloatTensor) } return data else: raise ValueError(f'This mode {self.opt.phase} is currently not supported') def __len__(self): return self.data_size if __name__ == '__main__': def save_imgs(imgs, names, path): if not os.path.exists(path): os.mkdir(path) for img, name in zip(imgs, names): img = tensor2img(img) img = np.array(img) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(os.path.join(path, name + '.jpg'), img) def tensor2img(img): img = img[0].cpu().float().numpy() if img.shape[0] == 1: img = np.tile(img, (3, 1, 1)) img = (np.transpose(img, (1, 2, 0)) + 1) / 2.0 * 255.0 return img.astype(np.uint8) from options import Options parser = Options() opts = parser.parse() opts.data_root='../MT-Dataset/images' # data loader print('\n--- load dataset ---') dataset = MakeupDataset(opts) train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True, num_workers=opts.nThreads) data_names = ['non_makeup_color_img', 'non_makeup_all_mask', 'non_makeup_face_mask', 'non_makeup_brow_mask', 'non_makeup_eye_mask', 'non_makeup_lip_mask', 'makeup_color_img', 'makeup_all_mask', 'makeup_face_mask', 'makeup_brow_mask', 'makeup_eye_mask', 'makeup_lip_mask', 'makeup_color_change_img', 'makeup_color_change_warp_img', 'makeup_all_warp_mask'] print('dataset size:', len(train_loader)) for i, data in enumerate(train_loader): if i > 5: break print(i) imgs = [data[x] for x in data_names] names = [str(i) + x for x in data_names] save_imgs(imgs, names, path='./debug_load_dataset')