def initialize(self, opt):
self.opt = opt
root = opt.dataroot
if opt.isTrain:
self.L_paths = sorted(make_grouped_dataset(path.join(root, 'train_keypoints')))
self.I_paths = sorted(make_grouped_dataset(path.join(root, 'train_images')))
check_path_valid(self.L_paths, self.I_paths)
else:
self.L_paths = sorted(make_dataset(opt.seq_path.replace('images', 'keypoints')))
self.I_paths = sorted(make_dataset(opt.seq_path))
self.ref_L_paths = sorted(make_dataset(opt.ref_img_path.replace('images', 'keypoints')))
self.ref_I_paths = sorted(make_dataset(opt.ref_img_path))
self.n_of_seqs = len(self.I_paths) # number of sequences to train
if opt.isTrain: print('%d sequences' % self.n_of_seqs)
# mapping from keypoints to face part
self.add_upper_face = not opt.no_upper_face
self.part_list = [[list(range(0, 17)) + ((list(range(68, 83)) + [0]) if self.add_upper_face else [])], # face
[range(17, 22)], # right eyebrow
[range(22, 27)], # left eyebrow
[[28, 31], range(31, 36), [35, 28]], # nose
[[36,37,38,39], [39,40,41,36]], # right eye
[[42,43,44,45], [45,46,47,42]], # left eye
[range(48, 55), [54,55,56,57,58,59,48], range(60, 65), [64,65,66,67,60]], # mouth and tongue
]
self.ref_dist_x, self.ref_dist_y = [None] * 83, [None] * 83
self.dist_scale_x, self.dist_scale_y = [None] * 83, [None] * 83
self.fix_crop_pos = True
def get_paths(self, opt):
root = opt.dataroot
phase = opt.eval_set +'_256'
dir_A = os.path.join(opt.dataroot, phase, 'A')
dir_B_noise = os.path.join(opt.dataroot, phase, self.opt.pose_extractor)
A_paths = sorted(make_grouped_dataset(dir_A))
B_paths_noise = sorted(make_grouped_dataset(dir_B_noise))
check_path_valid(A_paths, B_paths_noise)
return A_paths, B_paths_noise
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + '_A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + '_B')
self.use_real = opt.use_real_img
self.A_is_label = self.opt.label_nc != 0
self.A_paths = sorted(make_grouped_dataset(self.dir_A))
if self.use_real:
self.B_paths = sorted(make_grouped_dataset(self.dir_B))
check_path_valid(self.A_paths, self.B_paths)
if self.opt.use_instance:
self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
self.I_paths = sorted(make_grouped_dataset(self.dir_inst))
check_path_valid(self.A_paths, self.I_paths)
if opt.custom_data_root is not None:
self.A_paths2 = sorted(make_grouped_dataset(opt.custom_data_root))
def kitti_path_match(A_paths2, dataroot):
B_paths2 = []
for batch_dir in A_paths2:
B_paths2_seq = []
for sample_dir in batch_dir:
recording, color_mask_name = sample_dir.strip('/').split('/')[-2:]
frame_name = color_mask_name.split('_')[-1]
B_path2 = os.path.join(dataroot, 'test_B', recording, frame_name)
if os.path.exists(B_path2):
B_paths2_seq.append(B_path2)
else:
B_paths2_seq.append(B_paths2_seq[-1])
B_paths2.append(B_paths2_seq)
return B_paths2
def cityscapes_path_match(A_paths2, dataroot):
B_paths2 = []
for batch_dir in A_paths2:
B_paths2_seq = []
for sample_dir in batch_dir:
recording, color_mask_name = sample_dir.strip('/').split('/')[-2:]
frame_name = color_mask_name.replace('color_mask_', '')
B_path2 = os.path.join(dataroot, 'test_B', recording, frame_name)
if os.path.exists(B_path2):
B_paths2_seq.append(B_path2)
else:
B_paths2_seq.append(B_paths2_seq[-1])
B_paths2.append(B_paths2_seq)
return B_paths2
if 'kitti' in opt.dataroot.lower():
path_match = kitti_path_match
elif 'cityscapes' in opt.dataroot.lower():
path_match = cityscapes_path_match
self.B_paths2 = path_match(self.A_paths2, opt.dataroot)
self.A_paths = self.A_paths2
self.B_paths = self.B_paths2
self.init_frame_idx(self.A_paths)
def get_paths(self, opt):
root = opt.dataroot
phase = 'test' if opt.phase == 'val' else opt.phase
dir_A = os.path.join(opt.dataroot, phase + '_keypoints')
dir_B = os.path.join(opt.dataroot, phase + '_data')
A_paths = sorted(make_grouped_dataset(dir_A))
B_paths = sorted(make_grouped_dataset(dir_B))
check_path_valid(A_paths, B_paths)
return A_paths, B_paths, None
def __init__(self, data_root, phase, label_nc):
self.data_dir = os.path.join(data_root, 'images')
self.annotations_dir = os.path.join(data_root, 'annotations')
self.ir_is_label = label_nc != 0
self.data_paths = sorted(make_grouped_dataset(self.data_dir))
self.ir_paths = sorted([path for path in self.data_paths if path[0].find("lwir") != -1])
self.rgb_paths = sorted([path for path in self.data_paths if path[0].find("visible") != -1])
self.annotations_path = sorted(make_grouped_dataset(self.annotations_dir))
check_path_valid(self.ir_paths, self.rgb_paths, self.annotations_path)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + '_keypoints')
self.dir_B = os.path.join(opt.dataroot, opt.phase + '_img')
self.A_paths = sorted(make_grouped_dataset(self.dir_A))
self.B_paths = sorted(make_grouped_dataset(self.dir_B))
check_path_valid(self.A_paths, self.B_paths)
self.init_frame_idx(self.A_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + '_keypoints')
self.dir_B = os.path.join(opt.dataroot, opt.phase + '_img')
self.A_paths = sorted(make_grouped_dataset(self.dir_A))
self.B_paths = sorted(make_grouped_dataset(self.dir_B))
check_path_valid(self.A_paths, self.B_paths)
self.init_frame_idx(self.A_paths)
self.scale_ratio = np.array([[0.9, 1], [1, 1], [0.9, 1], [1, 1.1], [0.9, 0.9], [0.9, 0.9]])#np.random.uniform(0.9, 1.1, size=[6, 2])
self.scale_ratio_sym = np.array([[1, 1], [0.9, 1], [1, 1], [0.9, 1], [1, 1], [1, 1]]) #np.random.uniform(0.9, 1.1, size=[6, 2])
self.scale_shift = np.zeros((6, 2)) #np.random.uniform(-5, 5, size=[6, 2])
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_dp = os.path.join(opt.dataroot, opt.phase + '_densepose')
self.dir_op = os.path.join(opt.dataroot, opt.phase + '_openpose')
self.dir_img = os.path.join(opt.dataroot, opt.phase + '_img')
self.img_paths = sorted(make_grouped_dataset(self.dir_img))
if not opt.openpose_only:
self.dp_paths = sorted(make_grouped_dataset(self.dir_dp))
check_path_valid(self.dp_paths, self.img_paths)
if not opt.densepose_only:
self.op_paths = sorted(make_grouped_dataset(self.dir_op))
check_path_valid(self.op_paths, self.img_paths)
self.init_frame_idx(self.img_paths)
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + '_A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + '_B')
self.A_is_label = self.opt.label_nc != 0
self.A_paths = sorted(make_grouped_dataset(self.dir_A))
self.B_paths = sorted(make_grouped_dataset(self.dir_B))
check_path_valid(self.A_paths, self.B_paths)
if opt.use_instance:
self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
self.I_paths = sorted(make_grouped_dataset(self.dir_inst))
check_path_valid(self.A_paths, self.I_paths)
self.n_of_seqs = len(self.A_paths) # number of sequences to train
self.seq_len_max = max([len(A) for A in self.A_paths])
self.n_frames_total = self.opt.n_frames_total # current number of frames to train in a single iteration
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + '_A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + '_B')
self.use_real = opt.use_real_img
self.A_is_label = self.opt.label_nc != 0
self.A_paths = sorted(make_grouped_dataset(self.dir_A))
if self.use_real:
self.B_paths = sorted(make_grouped_dataset(self.dir_B))
check_path_valid(self.A_paths, self.B_paths)
if self.opt.use_instance:
self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
self.I_paths = sorted(make_grouped_dataset(self.dir_inst))
check_path_valid(self.A_paths, self.I_paths)
self.init_frame_idx(self.A_paths)
def initialize(self, opt):
self.opt = opt
root = opt.dataroot
self.L_is_label = self.opt.label_nc != 0
if opt.isTrain:
self.L_paths = sorted(
make_grouped_dataset(path.join(root, 'train_labels')))
self.I_paths = sorted(
make_grouped_dataset(path.join(root, 'train_images')))
check_path_valid(self.L_paths, self.I_paths)
self.n_of_seqs = len(self.L_paths)
print('%d sequences' % self.n_of_seqs)
else:
self.I_paths = sorted(make_dataset(opt.seq_path))
self.L_paths = sorted(
make_dataset(opt.seq_path.replace('images', 'labels')))
self.ref_I_paths = sorted(make_dataset(opt.ref_img_path))
self.ref_L_paths = sorted(
make_dataset(opt.ref_img_path.replace('images', 'labels')))
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_dp = os.path.join(opt.dataroot, opt.phase + '_densepose')
if opt.phase == 'test' and opt.custom_data_root is not None:
self.dir_dp = os.path.join(opt.custom_data_root)
self.dir_img = os.path.join(opt.dataroot, opt.phase + '_img')
self.img_paths = sorted(make_grouped_dataset(self.dir_img))
self.dp_paths = sorted(make_grouped_dataset(self.dir_dp))
#check_path_valid(self.dp_paths, self.img_paths)
if opt.phase == 'test' and opt.custom_data_root is not None and opt.no_first_img:
self.img_paths = self.dp_paths
elif opt.phase == 'test' and opt.custom_data_root is not None and opt.use_real_img:
def pose_paths_match(dp_paths, dataroot):
img_paths = []
for batch_dir in dp_paths:
img_paths_seq = []
for top_dir in batch_dir:
recording, color_mask_name = top_dir.strip('/').split(
'/')[-2:]
frame_name = color_mask_name.replace('_IUV', '')
frame_root = list(
Path(dataroot).glob('*_img/%s' % recording))[0]
img_path = frame_root / frame_name
if img_path.exists():
img_paths_seq.append(str(img_path))
else:
img_paths_seq.append(img_paths_seq[-1])
img_paths.append(img_paths_seq)
return img_paths
self.img_paths = pose_paths_match(self.dp_paths, opt.alldataroot)
else:
check_path_valid(self.dp_paths, self.img_paths)
self.init_frame_idx(self.img_paths)
def __init__(self, opt):
super(IR2RGBTestVideoDataset, self).__init__(opt)
self.use_real = opt.use_real_img
self.A_is_label = self.opt.label_nc != 0
self.A_paths = sorted(make_grouped_dataset(self.ir_dir))
if self.use_real:
self.B_paths = sorted(make_grouped_dataset(self.rgb_dir))
check_path_valid(self.A_paths, self.B_paths)
# Init frame idx:
self.n_of_seqs = min(len(self.A_paths), self.opt.max_dataset_size) # number of sequences to train
self.seq_len_max = max([len(A) for A in self.A_paths]) # max number of frames in the training sequences
self.seq_idx = 0 # index for current sequence
self.frame_idx = self.opt.start_frame if not self.opt.is_train else 0 # index for current frame in the sequence
self.frames_count = [] # number of frames in each sequence
for path in self.A_paths:
self.frames_count.append(len(path) - self.opt.n_input_gen_frames + 1)
self.folder_prob = [count / sum(self.frames_count) for count in self.frames_count]
self.n_frames_total = self.opt.n_frames_total if self.opt.is_train else 1
self.A, self.B = None, None
def get_paths(self, opt):
root = opt.dataroot
phase = opt.phase
phase_dir = phase + '_256'
self.phase_dir = phase_dir
dir_A = os.path.join(opt.dataroot, phase_dir, 'train_A')
dir_B_clean = os.path.join(opt.dataroot, phase_dir,
'train_' + 'video2d')
dir_B_noise = os.path.join(opt.dataroot, phase_dir,
'train_' + 'alphapose')
A_paths = sorted(make_grouped_dataset(dir_A))
B_paths_clean = sorted(make_grouped_dataset(dir_B_clean))
B_paths_noise = sorted(make_grouped_dataset(dir_B_noise))
check_path_valid(A_paths, B_paths_clean)
check_path_valid(A_paths, B_paths_noise)
if phase == 'train':
if self.opt.use_mask:
dir_C = os.path.join(opt.dataroot, phase_dir, 'train_C')
C_paths = sorted(make_grouped_dataset(dir_C))
check_path_valid(A_paths, C_paths)
C_paths = self.split_ref_gen(C_paths, self.opt.sub_dataset)
else:
C_paths = None
A_paths = self.split_ref_gen(A_paths, self.opt.sub_dataset)
B_paths_clean = self.split_ref_gen(B_paths_clean,
self.opt.sub_dataset)
B_paths_noise = self.split_ref_gen(B_paths_noise,
self.opt.sub_dataset)
else:
assert self.opt.use_mask == False
C_paths = None
if opt.test_list is not None:
scv_path = os.path.join(opt.dataroot, opt.test_list)
file = pd.read_csv(scv_path)
A_paths = file['A_paths'].map(ast.literal_eval)
B_paths_noise = file['B_paths_noise'].map(ast.literal_eval)
B_paths_clean = file['B_paths_clean'].map(ast.literal_eval)
else:
A_paths = self.split_ref_gen(A_paths, self.opt.sub_dataset)
B_paths_clean = self.split_ref_gen(B_paths_clean,
self.opt.sub_dataset)
B_paths_noise = self.split_ref_gen(B_paths_noise,
self.opt.sub_dataset)
A_paths, B_paths_clean, B_paths_noise = self.pad_for_latest_frames(
A_paths, B_paths_clean, B_paths_noise)
return A_paths, B_paths_clean, B_paths_noise, C_paths
