import os.path from data.base_dataset import BaseDataset, get_params, get_transform, normalize from data.image_folder import make_dataset, make_dataset_test from PIL import Image import torch import json import numpy as np import os.path as osp from PIL import ImageDraw class AlignedDataset(BaseDataset): def initialize(self, opt): self.opt = opt self.root = opt.dataroot self.diction = {} self.fine_height = 256 self.fine_width = 192 self.radius = 5 # load data list from pairs file human_names = [] cloth_names = [] with open(os.path.join(opt.dataroot, opt.datapairs), 'r') as f: for line in f.readlines(): h_name, c_name = line.strip().split() human_names.append(h_name) cloth_names.append(c_name) self.human_names = human_names self.cloth_names = cloth_names self.dataset_size = len(human_names) # input A (label maps) dir_A = '_A' if self.opt.label_nc == 0 else '_label' self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A) self.A_paths = sorted(make_dataset(self.dir_A)) self.fine_height = 256 self.fine_width = 192 self.radius = 5 # input A test (label maps) dir_A = '_A' if self.opt.label_nc == 0 else '_label' self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A) self.A_paths = sorted(make_dataset_test(self.dir_A)) # input B (real images) dir_B = '_B' if self.opt.label_nc == 0 else '_img' self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B) self.B_paths = sorted(make_dataset(self.dir_B)) self.dataset_size = len(self.A_paths) self.build_index(self.B_paths) dir_E = '_edge' self.dir_E = os.path.join(opt.dataroot, opt.phase + dir_E) self.E_paths = sorted(make_dataset(self.dir_E)) self.ER_paths = make_dataset(self.dir_E) dir_M = '_mask' self.dir_M = os.path.join(opt.dataroot, opt.phase + dir_M) self.M_paths = sorted(make_dataset(self.dir_M)) self.MR_paths = make_dataset(self.dir_M) dir_MC = '_colormask' self.dir_MC = os.path.join(opt.dataroot, opt.phase + dir_MC) self.MC_paths = sorted(make_dataset(self.dir_MC)) self.MCR_paths = make_dataset(self.dir_MC) dir_C = '_color' self.dir_C = os.path.join(opt.dataroot, opt.phase + dir_C) self.C_paths = sorted(make_dataset(self.dir_C)) self.CR_paths = make_dataset(self.dir_C) # self.build_index(self.C_paths) dir_A = '_A' if self.opt.label_nc == 0 else '_label' self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A) self.A_paths = sorted(make_dataset_test(self.dir_A)) def random_sample(self, item): name = item.split('/')[-1] name = name.split('-')[0] lst = self.diction[name] new_lst = [] for dir in lst: if dir != item: new_lst.append(dir) return new_lst[np.random.randint(len(new_lst))] def build_index(self, dirs): for k, dir in enumerate(dirs): name = dir.split('/')[-1] name = name.split('-')[0] # print(name) for k, d in enumerate(dirs[max(k-20, 0):k+20]): if name in d: if name not in self.diction.keys(): self.diction[name] = [] self.diction[name].append(d) else: self.diction[name].append(d) def __getitem__(self, index): train_mask = 9600 # input A (label maps) box = [] # for k,x in enumerate(self.A_paths): # if '000386' in x : # index=k # break test = np.random.randint(2032) # for k, s in enumerate(self.B_paths): # if '006581' in s: # test = k # break # get names from the pairs file c_name = self.cloth_names[index] h_name = self.human_names[index] # A_path = self.A_paths[index] A_path = osp.join(self.dir_A, h_name.replace(".jpg", ".png")) A = Image.open(A_path).convert('L') params = get_params(self.opt, A.size) if self.opt.label_nc == 0: transform_A = get_transform(self.opt, params) A_tensor = transform_A(A.convert('RGB')) else: transform_A = get_transform( self.opt, params, method=Image.NEAREST, normalize=False) A_tensor = transform_A(A) * 255.0 B_tensor = inst_tensor = feat_tensor = 0 # input B (real images) # B_path = self.B_paths[index] B_path = osp.join(self.dir_B, h_name) name = B_path.split('/')[-1] B = Image.open(B_path).convert('RGB') transform_B = get_transform(self.opt, params) B_tensor = transform_B(B) # input M (masks) M_path = B_path # self.M_paths[np.random.randint(1)] MR_path = B_path # self.MR_paths[np.random.randint(1)] M = Image.open(M_path).convert('L') MR = Image.open(MR_path).convert('L') M_tensor = transform_A(MR) ### input_MC (colorMasks) MC_path = B_path # self.MC_paths[1] MCR_path = B_path # self.MCR_paths[1] MCR = Image.open(MCR_path).convert('L') MC_tensor = transform_A(MCR) ### input_C (color) # print(self.C_paths) # C_path = self.C_paths[test] C_path = osp.join(self.dir_C, c_name) C = Image.open(C_path).convert('RGB') C_tensor = transform_B(C) # Edge # E_path = self.E_paths[test] E_path = osp.join(self.dir_E, c_name) # print(E_path) E = Image.open(E_path).convert('L') E_tensor = transform_A(E) # Pose pose_name = B_path.replace('.jpg', '_keypoints.json').replace('.png', '_keypoints.json').replace( 'test_img', 'test_pose') with open(osp.join(pose_name), 'r') as f: pose_label = json.load(f) pose_data = pose_label['people'][0]['pose_keypoints'] pose_data = np.array(pose_data) pose_data = pose_data.reshape((-1, 3)) point_num = pose_data.shape[0] pose_map = torch.zeros(point_num, self.fine_height, self.fine_width) r = self.radius im_pose = Image.new('L', (self.fine_width, self.fine_height)) pose_draw = ImageDraw.Draw(im_pose) for i in range(point_num): one_map = Image.new('L', (self.fine_width, self.fine_height)) draw = ImageDraw.Draw(one_map) pointx = pose_data[i, 0] pointy = pose_data[i, 1] if pointx > 1 and pointy > 1: draw.rectangle((pointx-r, pointy-r, pointx + r, pointy+r), 'white', 'white') pose_draw.rectangle( (pointx-r, pointy-r, pointx+r, pointy+r), 'white', 'white') one_map = transform_B(one_map.convert('RGB')) pose_map[i] = one_map[0] P_tensor = pose_map input_dict = {'label': A_tensor, 'image': B_tensor, 'path': A_path, 'name': A_path.split("/")[-1], 'edge': E_tensor, 'color': C_tensor, 'mask': M_tensor, 'colormask': MC_tensor, 'pose': P_tensor } return input_dict def __len__(self): return len(self.A_paths) // self.opt.batchSize * self.opt.batchSize def name(self): return 'AlignedDataset'