def initialize(self, opt):
self.opt = opt
do_reenactment = opt.do_reenactment if not opt.isTrain else False
prefix = 'source_' if do_reenactment else ''
source_name = opt.source_name if not opt.isTrain else None
# Get dataset directories.
self.dir_nmfc_video = os.path.join(opt.dataroot, self.opt.phase,
prefix + 'nmfcs')
self.nmfc_video_paths = make_video_dataset(self.dir_nmfc_video,
opt.target_name,
source_name,
opt.max_n_sequences)
self.dir_rgb_video = os.path.join(opt.dataroot, self.opt.phase,
prefix + 'images')
self.rgb_video_paths = make_video_dataset(self.dir_rgb_video,
opt.target_name, source_name,
opt.max_n_sequences)
assert_valid_pairs(self.nmfc_video_paths, self.rgb_video_paths)
if not opt.no_eye_gaze or (not self.opt.no_mouth_D and
self.opt.isTrain) or (self.opt.use_eyes_D
and self.opt.isTrain):
self.dir_landmark_video = os.path.join(opt.dataroot,
self.opt.phase,
prefix + 'landmarks70')
self.landmark_video_paths = make_video_dataset(
self.dir_landmark_video, opt.target_name, source_name,
opt.max_n_sequences)
assert_valid_pairs(self.landmark_video_paths, self.rgb_video_paths)
self.n_of_seqs = len(self.nmfc_video_paths)
self.seq_len_max = max([len(A) for A in self.nmfc_video_paths])
self.init_frame_index(self.nmfc_video_paths)
def __init__(self, opt):
"""Initialize this dataset class.
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseDataset.__init__(self, opt)
self.isTrain = opt.isTrain
self.dir_frame = os.path.join(opt.dataroot, 'frame')
self.dir_flow = os.path.join(opt.dataroot, 'flow')
self.dir_mask = os.path.join(opt.dataroot, 'mask')
self.dir_style = os.path.join(opt.dataroot, 'style')
self.dir_last_fake = os.path.join(opt.dataroot,
'last_fake' + opt.suffix)
self.frames, self.flows, self.masks, self.styles, self.lasts_fake = make_video_dataset(
self.dir_frame, self.dir_flow, self.dir_mask, self.dir_style,
self.dir_last_fake)
assert opt.output_nc == 3 and opt.input_nc == 3, 'Numbers of input and output channels must be 3'
self.transform = get_transform(self.opt)
self.lengths = [len(video) for video in self.frames]