def do_transform(img, mask): warp = (img_type == SPTF.IMG_WARPED or img_type == SPTF.IMG_WARPED_TRANSFORMED) transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED or img_type == SPTF.IMG_TRANSFORMED) flip = img_type != SPTF.IMG_WARPED img = imagelib.warp_by_params(params, img, warp, transform, flip, border_replicate) if mask is not None: mask = imagelib.warp_by_params( params, mask, warp, transform, flip, False) if len(mask.shape) == 2: mask = mask[..., np.newaxis] return img, mask
def batch_func(self, param ): pickled_samples, resolution, face_type, data_format = param samples = pickle.loads(pickled_samples) shuffle_idxs = [] idxs = [*range(len(samples))] random_flip = True rotation_range=[-10,10] scale_range=[-0.05, 0.05] tx_range=[-0.05, 0.05] ty_range=[-0.05, 0.05] random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25,75 motion_blur_chance, motion_blur_mb_max_size = 25, 5 gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5 bs = self.batch_size while True: batches = [ [], [] ] n_batch = 0 while n_batch < bs: try: if len(shuffle_idxs) == 0: shuffle_idxs = idxs.copy() np.random.shuffle(shuffle_idxs) idx = shuffle_idxs.pop() sample = samples[idx] img = sample.load_bgr() h,w,c = img.shape mask = np.zeros ((h,w,1), dtype=np.float32) sample.seg_ie_polys.overlay_mask(mask) warp_params = imagelib.gen_warp_params(resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range ) if face_type == sample.face_type: if w != resolution: img = cv2.resize( img, (resolution, resolution), cv2.INTER_LANCZOS4 ) mask = cv2.resize( mask, (resolution, resolution), cv2.INTER_LANCZOS4 ) else: mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, face_type) img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 ) mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 ) if len(mask.shape) == 2: mask = mask[...,None] img = imagelib.warp_by_params (warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) mask = imagelib.warp_by_params (warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) img = np.clip(img.astype(np.float32), 0, 1) mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 mask = np.clip(mask, 0, 1) if np.random.randint(2) == 0: img = imagelib.apply_random_hsv_shift(img, mask=sd.random_circle_faded ([resolution,resolution])) else: img = imagelib.apply_random_rgb_levels(img, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_motion_blur( img, motion_blur_chance, motion_blur_mb_max_size, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_gaussian_blur( img, gaussian_blur_chance, gaussian_blur_kernel_max_size, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_bilinear_resize( img, random_bilinear_resize_chance, random_bilinear_resize_max_size_per, mask=sd.random_circle_faded ([resolution,resolution])) if data_format == "NCHW": img = np.transpose(img, (2,0,1) ) mask = np.transpose(mask, (2,0,1) ) batches[0].append ( img ) batches[1].append ( mask ) n_batch += 1 except: io.log_err ( traceback.format_exc() ) yield [ np.array(batch) for batch in batches]
def process(samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample in samples: sample_bgr = sample.load_bgr() ct_sample_bgr = None h, w, c = sample_bgr.shape is_face_sample = sample.landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks, (0, 1, 0)) params = imagelib.gen_warp_params( sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range, rnd_seed=sample_rnd_seed) outputs_sample = [] for opts in output_sample_types: resolution = opts.get('resolution', 0) types = opts.get('types', []) motion_blur = opts.get('motion_blur', None) gaussian_blur = opts.get('gaussian_blur', None) ct_mode = opts.get('ct_mode', 'None') normalize_tanh = opts.get('normalize_tanh', False) data_format = opts.get('data_format', 'NHWC') img_type = SPTF.NONE target_face_type = SPTF.NONE mode_type = SPTF.NONE for t in types: if t >= SPTF.IMG_TYPE_BEGIN and t < SPTF.IMG_TYPE_END: img_type = t elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END: target_face_type = t elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END: mode_type = t if mode_type == SPTF.MODE_FACE_MASK_HULL and not is_face_sample: raise ValueError( "MODE_FACE_MASK_HULL applicable only for face samples") if mode_type == SPTF.MODE_FACE_MASK_STRUCT and not is_face_sample: raise ValueError( "MODE_FACE_MASK_STRUCT applicable only for face samples" ) if is_face_sample: if target_face_type == SPTF.NONE: raise ValueError( "target face type must be defined for face samples" ) can_warp = (img_type == SPTF.IMG_WARPED or img_type == SPTF.IMG_WARPED_TRANSFORMED) can_transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED or img_type == SPTF.IMG_TRANSFORMED) if img_type == SPTF.NONE: raise ValueError('expected IMG_ type') if img_type == SPTF.IMG_LANDMARKS_ARRAY: l = sample.landmarks l = np.concatenate([ np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1) ], -1) l = np.clip(l, 0.0, 1.0) out_sample = l elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch_yaw_roll = sample.get_pitch_yaw_roll() if params['flip']: yaw = -yaw if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1) yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1) roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1) out_sample = (pitch, yaw, roll) else: if mode_type == SPTF.NONE: raise ValueError('expected MODE_ type') if mode_type == SPTF.MODE_FACE_MASK_HULL: if sample.eyebrows_expand_mod is not None: img = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) else: img = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample.landmarks) if sample.ie_polys is not None: sample.ie_polys.overlay_mask(img) elif mode_type == SPTF.MODE_FACE_MASK_STRUCT: if sample.eyebrows_expand_mod is not None: img = LandmarksProcessor.get_face_struct_mask( sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) else: img = LandmarksProcessor.get_face_struct_mask( sample_bgr.shape, sample.landmarks) else: img = sample_bgr if motion_blur is not None: chance, mb_max_size = motion_blur chance = np.clip(chance, 0, 100) if np.random.randint(100) < chance: img = imagelib.LinearMotionBlur( img, np.random.randint(mb_max_size) + 1, np.random.randint(360)) if gaussian_blur is not None: chance, kernel_max_size = gaussian_blur chance = np.clip(chance, 0, 100) if np.random.randint(100) < chance: img = cv2.GaussianBlur( img, (np.random.randint(kernel_max_size) * 2 + 1, ) * 2, 0) if is_face_sample: target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[ target_face_type] if target_ft > sample.face_type: raise Exception( 'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_ft)) if sample.face_type == FaceType.MARK_ONLY: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, sample.shape[0], target_ft) if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT: img = cv2.warpAffine( img, mat, (sample.shape[0], sample.shape[0]), flags=cv2.INTER_CUBIC) img = imagelib.warp_by_params( params, img, can_warp, can_transform, can_flip=True, border_replicate=False) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC)[..., None] else: img = cv2.warpAffine( img, mat, (sample.shape[0], sample.shape[0]), flags=cv2.INTER_CUBIC) img = imagelib.warp_by_params( params, img, can_warp, can_transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) else: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, target_ft) if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT: img = imagelib.warp_by_params( params, img, can_warp, can_transform, can_flip=True, border_replicate=False) img = cv2.warpAffine( img, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_CUBIC)[..., None] else: img = imagelib.warp_by_params( params, img, can_warp, can_transform, can_flip=True, border_replicate=True) img = cv2.warpAffine( img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) else: img = imagelib.warp_by_params(params, img, can_warp, can_transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT: out_sample = np.clip(img.astype(np.float32), 0, 1) else: img = np.clip(img.astype(np.float32), 0, 1) if ct_mode is not None and ct_sample is not None: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() img = imagelib.color_transfer( ct_mode, img, cv2.resize(ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR)) if mode_type == SPTF.MODE_BGR: out_sample = img elif mode_type == SPTF.MODE_BGR_SHUFFLE: rnd_state = np.random.RandomState(sample_rnd_seed) out_sample = np.take(img, rnd_state.permutation( img.shape[-1]), axis=-1) elif mode_type == SPTF.MODE_BGR_RANDOM_HSV_SHIFT: rnd_state = np.random.RandomState(sample_rnd_seed) hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) h, s, v = cv2.split(hsv) h = (h + rnd_state.randint(360)) % 360 s = np.clip(s + rnd_state.random() - 0.5, 0, 1) v = np.clip(v + rnd_state.random() - 0.5, 0, 1) hsv = cv2.merge([h, s, v]) out_sample = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1) elif mode_type == SPTF.MODE_G: out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[..., None] elif mode_type == SPTF.MODE_GGG: out_sample = np.repeat( np.expand_dims( cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) if not debug: if normalize_tanh: out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0, 1.0) if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) outputs_sample.append(out_sample) outputs += [outputs_sample] return outputs
def process(samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPST = SampleProcessor.SampleType SPCT = SampleProcessor.ChannelType SPFMT = SampleProcessor.FaceMaskType sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample in samples: sample_face_type = sample.face_type sample_bgr = sample.load_bgr() sample_landmarks = sample.landmarks ct_sample_bgr = None h, w, c = sample_bgr.shape def get_full_face_mask(): if sample.eyebrows_expand_mod is not None: full_face_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) else: full_face_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample_landmarks) return np.clip(full_face_mask, 0, 1) def get_eyes_mask(): eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample_landmarks) return np.clip(eyes_mask, 0, 1) is_face_sample = sample_landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks, (0, 1, 0)) params_per_resolution = {} warp_rnd_state = np.random.RandomState(sample_rnd_seed - 1) for opts in output_sample_types: resolution = opts.get('resolution', None) if resolution is None: continue params_per_resolution[resolution] = imagelib.gen_warp_params( resolution, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range, rnd_state=warp_rnd_state) outputs_sample = [] for opts in output_sample_types: sample_type = opts.get('sample_type', SPST.NONE) channel_type = opts.get('channel_type', SPCT.NONE) resolution = opts.get('resolution', 0) warp = opts.get('warp', False) transform = opts.get('transform', False) motion_blur = opts.get('motion_blur', None) gaussian_blur = opts.get('gaussian_blur', None) random_bilinear_resize = opts.get('random_bilinear_resize', None) random_rgb_levels = opts.get('random_rgb_levels', False) random_hsv_shift = opts.get('random_hsv_shift', False) random_circle_mask = opts.get('random_circle_mask', False) normalize_tanh = opts.get('normalize_tanh', False) ct_mode = opts.get('ct_mode', None) data_format = opts.get('data_format', 'NHWC') if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE: border_replicate = False elif sample_type == SPST.FACE_IMAGE: border_replicate = True border_replicate = opts.get('border_replicate', border_replicate) borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: if not is_face_sample: raise ValueError( "face_samples should be provided for sample_type FACE_*" ) if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: face_type = opts.get('face_type', None) face_mask_type = opts.get('face_mask_type', SPFMT.NONE) if face_type is None: raise ValueError( "face_type must be defined for face samples") if face_type > sample.face_type: raise Exception( 'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, face_type)) if sample_type == SPST.FACE_MASK: if face_mask_type == SPFMT.FULL_FACE: img = get_full_face_mask() elif face_mask_type == SPFMT.EYES: img = get_eyes_mask() elif face_mask_type == SPFMT.FULL_FACE_EYES: img = get_full_face_mask() + get_eyes_mask() else: img = np.zeros(sample_bgr.shape[0:2] + (1, ), dtype=np.float32) if sample.ie_polys is not None: sample.ie_polys.overlay_mask(img) if sample_face_type == FaceType.MARK_ONLY: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, warp_resolution, face_type) img = cv2.warpAffine( img, mat, (warp_resolution, warp_resolution), flags=cv2.INTER_LINEAR) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) img = cv2.resize(img, (resolution, resolution), cv2.INTER_LINEAR) else: if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_LINEAR) else: if w != resolution: img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) if len(img.shape) == 2: img = img[..., None] if channel_type == SPCT.G: out_sample = img.astype(np.float32) else: raise ValueError( "only channel_type.G supported for the mask") elif sample_type == SPST.FACE_IMAGE: img = sample_bgr if random_rgb_levels: random_mask = sd.random_circle_faded( [w, w], rnd_state=np.random.RandomState( sample_rnd_seed) ) if random_circle_mask else None img = imagelib.apply_random_rgb_levels( img, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed)) if random_hsv_shift: random_mask = sd.random_circle_faded( [w, w], rnd_state=np.random.RandomState( sample_rnd_seed + 1)) if random_circle_mask else None img = imagelib.apply_random_hsv_shift( img, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 1)) if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_CUBIC) else: if w != resolution: img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) # Apply random color transfer if ct_mode is not None and ct_sample is not None: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() img = imagelib.color_transfer( ct_mode, img, cv2.resize(ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR)) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate) img = np.clip(img.astype(np.float32), 0, 1) if motion_blur is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 2)) if random_circle_mask else None img = imagelib.apply_random_motion_blur( img, *motion_blur, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 2)) if gaussian_blur is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 3)) if random_circle_mask else None img = imagelib.apply_random_gaussian_blur( img, *gaussian_blur, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 3)) if random_bilinear_resize is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 4)) if random_circle_mask else None img = imagelib.apply_random_bilinear_resize( img, *random_bilinear_resize, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 4)) # Transform from BGR to desired channel_type if channel_type == SPCT.BGR: out_sample = img elif channel_type == SPCT.G: out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[..., None] elif channel_type == SPCT.GGG: out_sample = np.repeat( np.expand_dims( cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) # Final transformations if not debug: if normalize_tanh: out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0, 1.0) if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.IMAGE: img = sample_bgr img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) out_sample = img if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.LANDMARKS_ARRAY: l = sample_landmarks l = np.concatenate([ np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1) ], -1) l = np.clip(l, 0.0, 1.0) out_sample = l elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch, yaw, roll = sample.get_pitch_yaw_roll() if params_per_resolution[resolution]['flip']: yaw = -yaw if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1) yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1) roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1) out_sample = (pitch, yaw) else: raise ValueError('expected sample_type') outputs_sample.append(out_sample) outputs += [outputs_sample] return outputs
def batch_func(self, param ): samples, seg_sample_idxs, resolution, face_type, data_format = param shuffle_idxs = [] bg_shuffle_idxs = [] random_flip = True rotation_range=[-10,10] scale_range=[-0.05, 0.05] tx_range=[-0.05, 0.05] ty_range=[-0.05, 0.05] random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25,75 motion_blur_chance, motion_blur_mb_max_size = 25, 5 gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5 def gen_img_mask(sample): img = sample.load_bgr() h,w,c = img.shape if sample.seg_ie_polys.has_polys(): mask = np.zeros ((h,w,1), dtype=np.float32) sample.seg_ie_polys.overlay_mask(mask) elif sample.has_xseg_mask(): mask = sample.get_xseg_mask() mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 else: raise Exception(f'no mask in sample {sample.filename}') if face_type == sample.face_type: if w != resolution: img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 ) mask = cv2.resize( mask, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 ) else: mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, face_type) img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 ) mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 ) if len(mask.shape) == 2: mask = mask[...,None] return img, mask bs = self.batch_size while True: batches = [ [], [] ] n_batch = 0 while n_batch < bs: try: if len(shuffle_idxs) == 0: shuffle_idxs = seg_sample_idxs.copy() np.random.shuffle(shuffle_idxs) sample = samples[shuffle_idxs.pop()] img, mask = gen_img_mask(sample) if np.random.randint(2) == 0: if len(bg_shuffle_idxs) == 0: bg_shuffle_idxs = seg_sample_idxs.copy() np.random.shuffle(bg_shuffle_idxs) bg_sample = samples[bg_shuffle_idxs.pop()] bg_img, bg_mask = gen_img_mask(bg_sample) bg_wp = imagelib.gen_warp_params(resolution, True, rotation_range=[-180,180], scale_range=[-0.10, 0.10], tx_range=[-0.10, 0.10], ty_range=[-0.10, 0.10] ) bg_img = imagelib.warp_by_params (bg_wp, bg_img, can_warp=False, can_transform=True, can_flip=True, border_replicate=False) bg_mask = imagelib.warp_by_params (bg_wp, bg_mask, can_warp=False, can_transform=True, can_flip=True, border_replicate=False) c_mask = (1-bg_mask) * (1-mask) img = img*(1-c_mask) + bg_img * c_mask warp_params = imagelib.gen_warp_params(resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range ) img = imagelib.warp_by_params (warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) mask = imagelib.warp_by_params (warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) img = np.clip(img.astype(np.float32), 0, 1) mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 mask = np.clip(mask, 0, 1) if np.random.randint(2) == 0: img = imagelib.apply_random_hsv_shift(img, mask=sd.random_circle_faded ([resolution,resolution])) else: img = imagelib.apply_random_rgb_levels(img, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_motion_blur( img, motion_blur_chance, motion_blur_mb_max_size, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_gaussian_blur( img, gaussian_blur_chance, gaussian_blur_kernel_max_size, mask=sd.random_circle_faded ([resolution,resolution])) img = imagelib.apply_random_bilinear_resize( img, random_bilinear_resize_chance, random_bilinear_resize_max_size_per, mask=sd.random_circle_faded ([resolution,resolution])) if data_format == "NCHW": img = np.transpose(img, (2,0,1) ) mask = np.transpose(mask, (2,0,1) ) batches[0].append ( img ) batches[1].append ( mask ) n_batch += 1 except: io.log_err ( traceback.format_exc() ) yield [ np.array(batch) for batch in batches]
def process(samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPST = SampleProcessor.SampleType SPCT = SampleProcessor.ChannelType SPFMT = SampleProcessor.FaceMaskType sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample in samples: sample_bgr = sample.load_bgr() ct_sample_bgr = None h, w, c = sample_bgr.shape is_face_sample = sample.landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks, (0, 1, 0)) params = imagelib.gen_warp_params( sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range) outputs_sample = [] for opts in output_sample_types: sample_type = opts.get('sample_type', SPST.NONE) channel_type = opts.get('channel_type', SPCT.NONE) resolution = opts.get('resolution', 0) warp = opts.get('warp', False) transform = opts.get('transform', False) motion_blur = opts.get('motion_blur', None) gaussian_blur = opts.get('gaussian_blur', None) normalize_tanh = opts.get('normalize_tanh', False) ct_mode = opts.get('ct_mode', 'None') data_format = opts.get('data_format', 'NHWC') if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: if not is_face_sample: raise ValueError( "face_samples should be provided for sample_type FACE_*" ) if is_face_sample: face_type = opts.get('face_type', None) face_mask_type = opts.get('face_mask_type', SPFMT.NONE) if face_type is None: raise ValueError( "face_type must be defined for face samples") if face_type > sample.face_type: raise Exception( 'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_ft)) if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: if sample_type == SPST.FACE_MASK: if face_mask_type == SPFMT.ALL_HULL or \ face_mask_type == SPFMT.EYES_HULL or \ face_mask_type == SPFMT.ALL_EYES_HULL: if face_mask_type == SPFMT.ALL_HULL or \ face_mask_type == SPFMT.ALL_EYES_HULL: if sample.eyebrows_expand_mod is not None: all_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample. eyebrows_expand_mod) else: all_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample.landmarks) all_mask = np.clip(all_mask, 0, 1) if face_mask_type == SPFMT.EYES_HULL or \ face_mask_type == SPFMT.ALL_EYES_HULL: eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample.landmarks) eyes_mask = np.clip(eyes_mask, 0, 1) if face_mask_type == SPFMT.ALL_HULL: img = all_mask elif face_mask_type == SPFMT.EYES_HULL: img = eyes_mask elif face_mask_type == SPFMT.ALL_EYES_HULL: img = all_mask + eyes_mask elif face_mask_type == SPFMT.STRUCT: if sample.eyebrows_expand_mod is not None: img = LandmarksProcessor.get_face_struct_mask( sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample. eyebrows_expand_mod) else: img = LandmarksProcessor.get_face_struct_mask( sample_bgr.shape, sample.landmarks) if sample.ie_polys is not None: sample.ie_polys.overlay_mask(img) if sample.face_type == FaceType.MARK_ONLY: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, sample.shape[0], face_type) img = cv2.warpAffine( img, mat, (sample.shape[0], sample.shape[0]), flags=cv2.INTER_LINEAR) img = imagelib.warp_by_params( params, img, warp, transform, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LINEAR) img = cv2.resize(img, (resolution, resolution), cv2.INTER_LINEAR)[..., None] else: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, face_type) img = imagelib.warp_by_params( params, img, warp, transform, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LINEAR) img = cv2.warpAffine( img, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LINEAR)[..., None] if channel_type == SPCT.G: out_sample = img.astype(np.float32) else: raise ValueError( "only channel_type.G supported for the mask") elif sample_type == SPST.FACE_IMAGE: img = sample_bgr if motion_blur is not None: chance, mb_max_size = motion_blur chance = np.clip(chance, 0, 100) l_rnd_state = np.random.RandomState( sample_rnd_seed) mblur_rnd_chance = l_rnd_state.randint(100) mblur_rnd_kernel = l_rnd_state.randint( mb_max_size) + 1 mblur_rnd_deg = l_rnd_state.randint(360) if mblur_rnd_chance < chance: img = imagelib.LinearMotionBlur( img, mblur_rnd_kernel, mblur_rnd_deg) if gaussian_blur is not None: chance, kernel_max_size = gaussian_blur chance = np.clip(chance, 0, 100) l_rnd_state = np.random.RandomState( sample_rnd_seed + 1) gblur_rnd_chance = l_rnd_state.randint(100) gblur_rnd_kernel = l_rnd_state.randint( kernel_max_size) * 2 + 1 if gblur_rnd_chance < chance: img = cv2.GaussianBlur( img, (gblur_rnd_kernel, ) * 2, 0) if sample.face_type == FaceType.MARK_ONLY: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, sample.shape[0], face_type) img = cv2.warpAffine( img, mat, (sample.shape[0], sample.shape[0]), flags=cv2.INTER_CUBIC) img = imagelib.warp_by_params( params, img, warp, transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) else: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, face_type) img = imagelib.warp_by_params( params, img, warp, transform, can_flip=True, border_replicate=True) img = cv2.warpAffine( img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) img = np.clip(img.astype(np.float32), 0, 1) # Apply random color transfer if ct_mode is not None and ct_sample is not None: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() img = imagelib.color_transfer( ct_mode, img, cv2.resize(ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR)) # Transform from BGR to desired channel_type if channel_type == SPCT.BGR: out_sample = img elif channel_type == SPCT.BGR_SHUFFLE: l_rnd_state = np.random.RandomState( sample_rnd_seed) out_sample = np.take(img, l_rnd_state.permutation( img.shape[-1]), axis=-1) elif channel_type == SPCT.BGR_RANDOM_HSV_SHIFT: l_rnd_state = np.random.RandomState( sample_rnd_seed) hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) h, s, v = cv2.split(hsv) h = (h + l_rnd_state.randint(360)) % 360 s = np.clip(s + l_rnd_state.random() - 0.5, 0, 1) v = np.clip(v + l_rnd_state.random() - 0.5, 0, 1) hsv = cv2.merge([h, s, v]) out_sample = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1) elif channel_type == SPCT.BGR_RANDOM_RGB_LEVELS: l_rnd_state = np.random.RandomState( sample_rnd_seed) np_rnd = l_rnd_state.rand inBlack = np.array([ np_rnd() * 0.25, np_rnd() * 0.25, np_rnd() * 0.25 ], dtype=np.float32) inWhite = np.array([ 1.0 - np_rnd() * 0.25, 1.0 - np_rnd() * 0.25, 1.0 - np_rnd() * 0.25 ], dtype=np.float32) inGamma = np.array([ 0.5 + np_rnd(), 0.5 + np_rnd(), 0.5 + np_rnd() ], dtype=np.float32) outBlack = np.array([0.0, 0.0, 0.0], dtype=np.float32) outWhite = np.array([1.0, 1.0, 1.0], dtype=np.float32) out_sample = np.clip( (img - inBlack) / (inWhite - inBlack), 0, 1) out_sample = (out_sample**(1 / inGamma)) * ( outWhite - outBlack) + outBlack out_sample = np.clip(out_sample, 0, 1) elif channel_type == SPCT.G: out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[..., None] elif channel_type == SPCT.GGG: out_sample = np.repeat( np.expand_dims( cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) # Final transformations if not debug: if normalize_tanh: out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0, 1.0) if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) #else: # img = imagelib.warp_by_params (params, img, warp, transform, can_flip=True, border_replicate=True) # img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) elif sample_type == SPST.LANDMARKS_ARRAY: l = sample.landmarks l = np.concatenate([ np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1) ], -1) l = np.clip(l, 0.0, 1.0) out_sample = l elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch_yaw_roll = sample.get_pitch_yaw_roll() if params['flip']: yaw = -yaw if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1) yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1) roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1) out_sample = (pitch, yaw, roll) else: raise ValueError('expected sample_type') outputs_sample.append(out_sample) outputs += [outputs_sample] return outputs
def process(samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPST = SampleProcessor.SampleType SPCT = SampleProcessor.ChannelType SPFMT = SampleProcessor.FaceMaskType sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample in samples: sample_face_type = sample.face_type sample_bgr = sample.load_bgr() sample_landmarks = sample.landmarks ct_sample_bgr = None h, w, c = sample_bgr.shape def get_full_face_mask(): xseg_mask = sample.get_xseg_mask() if xseg_mask is not None: if xseg_mask.shape[0] != h or xseg_mask.shape[1] != w: xseg_mask = cv2.resize(xseg_mask, (w, h), interpolation=cv2.INTER_CUBIC) xseg_mask = imagelib.normalize_channels(xseg_mask, 1) return np.clip(xseg_mask, 0, 1) else: full_face_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) return np.clip(full_face_mask, 0, 1) def get_eyes_mask(): eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample_landmarks) # set eye masks to 1-2 clip = np.clip(eyes_mask, 0, 1) clip[clip > 0.1] += 1 return clip def get_mouth_mask(): mouth_mask = LandmarksProcessor.get_image_mouth_mask( sample_bgr.shape, sample_landmarks) # set eye masks to 2-3 clip = np.clip(mouth_mask, 0, 1) clip[clip > 0.1] += 2 return clip is_face_sample = sample_landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks, (0, 1, 0)) params_per_resolution = {} warp_rnd_state = np.random.RandomState(sample_rnd_seed - 1) for opts in output_sample_types: resolution = opts.get('resolution', None) if resolution is None: continue params_per_resolution[resolution] = imagelib.gen_warp_params( resolution, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range, rnd_state=warp_rnd_state) outputs_sample = [] for opts in output_sample_types: sample_type = opts.get('sample_type', SPST.NONE) channel_type = opts.get('channel_type', SPCT.NONE) resolution = opts.get('resolution', 0) nearest_resize_to = opts.get('nearest_resize_to', None) warp = opts.get('warp', False) transform = opts.get('transform', False) random_downsample = opts.get('random_downsample', False) random_noise = opts.get('random_noise', False) random_blur = opts.get('random_blur', False) random_jpeg = opts.get('random_jpeg', False) motion_blur = opts.get('motion_blur', None) gaussian_blur = opts.get('gaussian_blur', None) random_bilinear_resize = opts.get('random_bilinear_resize', None) random_rgb_levels = opts.get('random_rgb_levels', False) random_hsv_shift = opts.get('random_hsv_shift', False) random_circle_mask = opts.get('random_circle_mask', False) normalize_tanh = opts.get('normalize_tanh', False) ct_mode = opts.get('ct_mode', None) data_format = opts.get('data_format', 'NHWC') if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE: border_replicate = False elif sample_type == SPST.FACE_IMAGE: border_replicate = True border_replicate = opts.get('border_replicate', border_replicate) borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: if not is_face_sample: raise ValueError( "face_samples should be provided for sample_type FACE_*" ) if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: face_type = opts.get('face_type', None) face_mask_type = opts.get('face_mask_type', SPFMT.NONE) if face_type is None: raise ValueError( "face_type must be defined for face samples") if sample_type == SPST.FACE_MASK: if face_mask_type == SPFMT.FULL_FACE: img = get_full_face_mask() elif face_mask_type == SPFMT.EYES: img = get_eyes_mask() elif face_mask_type == SPFMT.FULL_FACE_EYES: # sets both eyes and mouth mask parts img = get_full_face_mask() mask = img.copy() mask[mask != 0.0] = 1.0 eye_mask = get_eyes_mask() * mask img = np.where(eye_mask > 1, eye_mask, img) mouth_mask = get_mouth_mask() * mask img = np.where(mouth_mask > 2, mouth_mask, img) else: img = np.zeros(sample_bgr.shape[0:2] + (1, ), dtype=np.float32) if sample_face_type == FaceType.MARK_ONLY: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, warp_resolution, face_type) img = cv2.warpAffine( img, mat, (warp_resolution, warp_resolution), flags=cv2.INTER_LINEAR) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_LINEAR) else: if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_LINEAR) else: if w != resolution: img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LINEAR) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) if len(img.shape) == 2: img = img[..., None] if channel_type == SPCT.G: out_sample = img.astype(np.float32) else: raise ValueError( "only channel_type.G supported for the mask") elif sample_type == SPST.FACE_IMAGE: img = sample_bgr if random_rgb_levels: random_mask = sd.random_circle_faded( [w, w], rnd_state=np.random.RandomState( sample_rnd_seed) ) if random_circle_mask else None img = imagelib.apply_random_rgb_levels( img, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed)) if random_hsv_shift: random_mask = sd.random_circle_faded( [w, w], rnd_state=np.random.RandomState( sample_rnd_seed + 1)) if random_circle_mask else None img = imagelib.apply_random_hsv_shift( img, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 1)) if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_CUBIC) else: if w != resolution: img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) # Apply random color transfer if ct_mode is not None and ct_sample is not None or ct_mode == 'fs-aug': if ct_mode == 'fs-aug': img = imagelib.color_augmentation( img, sample_rnd_seed) else: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() img = imagelib.color_transfer( ct_mode, img, cv2.resize(ct_sample_bgr, (resolution, resolution), interpolation=cv2.INTER_LINEAR)) randomization_order = ['blur', 'noise', 'jpeg', 'down'] np.random.shuffle(randomization_order) for random_distortion in randomization_order: # Apply random blur if random_distortion == 'blur' and random_blur: blur_type = np.random.choice( ['motion', 'gaussian']) if blur_type == 'motion': blur_k = np.random.randint(10, 20) blur_angle = 360 * np.random.random() img = LinearMotionBlur( img, blur_k, blur_angle) elif blur_type == 'gaussian': blur_sigma = 5 * np.random.random() + 3 if blur_sigma < 5.0: kernel_size = 2.9 * blur_sigma # 97% of weight else: kernel_size = 2.6 * blur_sigma # 95% of weight kernel_size = int(kernel_size) kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size img = cv2.GaussianBlur( img, (kernel_size, kernel_size), blur_sigma) # Apply random noise if random_distortion == 'noise' and random_noise: noise_type = np.random.choice( ['gaussian', 'laplace', 'poisson']) noise_scale = (20 * np.random.random() + 20) if noise_type == 'gaussian': noise = np.random.normal(scale=noise_scale, size=img.shape) img += noise / 255.0 elif noise_type == 'laplace': noise = np.random.laplace( scale=noise_scale, size=img.shape) img += noise / 255.0 elif noise_type == 'poisson': noise_lam = (15 * np.random.random() + 15) noise = np.random.poisson(lam=noise_lam, size=img.shape) img += noise / 255.0 # Apply random jpeg compression if random_distortion == 'jpeg' and random_jpeg: img = np.clip(img * 255, 0, 255).astype(np.uint8) jpeg_compression_level = np.random.randint( 50, 85) encode_param = [ int(cv2.IMWRITE_JPEG_QUALITY), jpeg_compression_level ] _, enc_img = cv2.imencode( '.jpg', img, encode_param) img = cv2.imdecode( enc_img, cv2.IMREAD_UNCHANGED).astype( np.float32) / 255.0 # Apply random downsampling if random_distortion == 'down' and random_downsample: down_res = np.random.randint( int(0.125 * resolution), int(0.25 * resolution)) img = cv2.resize(img, (down_res, down_res), interpolation=cv2.INTER_CUBIC) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate) img = np.clip(img.astype(np.float32), 0, 1) if motion_blur is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 2)) if random_circle_mask else None img = imagelib.apply_random_motion_blur( img, *motion_blur, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 2)) if gaussian_blur is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 3)) if random_circle_mask else None img = imagelib.apply_random_gaussian_blur( img, *gaussian_blur, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 3)) if random_bilinear_resize is not None: random_mask = sd.random_circle_faded( [resolution, resolution], rnd_state=np.random.RandomState( sample_rnd_seed + 4)) if random_circle_mask else None img = imagelib.apply_random_bilinear_resize( img, *random_bilinear_resize, mask=random_mask, rnd_state=np.random.RandomState( sample_rnd_seed + 4)) # Transform from BGR to desired channel_type if channel_type == SPCT.BGR: out_sample = img elif channel_type == SPCT.LAB_RAND_TRANSFORM: out_sample = random_lab_rotation( img, sample_rnd_seed) elif channel_type == SPCT.G: out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[..., None] elif channel_type == SPCT.GGG: out_sample = np.repeat( np.expand_dims( cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) # Final transformations if nearest_resize_to is not None: out_sample = cv2_resize( out_sample, (nearest_resize_to, nearest_resize_to), interpolation=cv2.INTER_NEAREST) if not debug: if normalize_tanh: out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0, 1.0) if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.IMAGE: img = sample_bgr img = imagelib.warp_by_params( params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) out_sample = img if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.LANDMARKS_ARRAY: l = sample_landmarks l = np.concatenate([ np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1) ], -1) l = np.clip(l, 0.0, 1.0) out_sample = l elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch, yaw, roll = sample.get_pitch_yaw_roll() if params_per_resolution[resolution]['flip']: yaw = -yaw if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1) yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1) roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1) out_sample = (pitch, yaw) else: raise ValueError('expected sample_type') outputs_sample.append(out_sample) outputs += [outputs_sample] return outputs
def process(samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPST = SampleProcessor.SampleType SPCT = SampleProcessor.ChannelType SPFMT = SampleProcessor.FaceMaskType outputs = [] for sample in samples: sample_rnd_seed = np.random.randint(0x80000000) sample_face_type = sample.face_type sample_bgr = sample.load_bgr() sample_landmarks = sample.landmarks ct_sample_bgr = None h, w, c = sample_bgr.shape def get_full_face_mask(): xseg_mask = sample.get_xseg_mask() if xseg_mask is not None: if xseg_mask.shape[0] != h or xseg_mask.shape[1] != w: xseg_mask = cv2.resize(xseg_mask, (w, h), interpolation=cv2.INTER_CUBIC) xseg_mask = imagelib.normalize_channels(xseg_mask, 1) return np.clip(xseg_mask, 0, 1) else: full_face_mask = LandmarksProcessor.get_image_hull_mask( sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) return np.clip(full_face_mask, 0, 1) def get_eyes_mask(): eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample_landmarks) return np.clip(eyes_mask, 0, 1) def get_eyes_mouth_mask(): eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample_landmarks) mouth_mask = LandmarksProcessor.get_image_mouth_mask( sample_bgr.shape, sample_landmarks) mask = eyes_mask + mouth_mask return np.clip(mask, 0, 1) is_face_sample = sample_landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks, (0, 1, 0)) outputs_sample = [] for opts in output_sample_types: resolution = opts.get('resolution', 0) sample_type = opts.get('sample_type', SPST.NONE) channel_type = opts.get('channel_type', SPCT.NONE) nearest_resize_to = opts.get('nearest_resize_to', None) warp = opts.get('warp', False) transform = opts.get('transform', False) random_hsv_shift_amount = opts.get('random_hsv_shift_amount', 0) normalize_tanh = opts.get('normalize_tanh', False) ct_mode = opts.get('ct_mode', None) data_format = opts.get('data_format', 'NHWC') rnd_seed_shift = opts.get('rnd_seed_shift', 0) warp_rnd_seed_shift = opts.get('warp_rnd_seed_shift', rnd_seed_shift) rnd_state = np.random.RandomState(sample_rnd_seed + rnd_seed_shift) warp_rnd_state = np.random.RandomState(sample_rnd_seed + warp_rnd_seed_shift) warp_params = imagelib.gen_warp_params( resolution, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range, rnd_state=rnd_state, warp_rnd_state=warp_rnd_state, ) if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE: border_replicate = False elif sample_type == SPST.FACE_IMAGE: border_replicate = True border_replicate = opts.get('border_replicate', border_replicate) borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: if not is_face_sample: raise ValueError( "face_samples should be provided for sample_type FACE_*" ) if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK: face_type = opts.get('face_type', None) face_mask_type = opts.get('face_mask_type', SPFMT.NONE) if face_type is None: raise ValueError( "face_type must be defined for face samples") if sample_type == SPST.FACE_MASK: if face_mask_type == SPFMT.FULL_FACE: img = get_full_face_mask() elif face_mask_type == SPFMT.EYES: img = get_eyes_mask() elif face_mask_type == SPFMT.EYES_MOUTH: mask = get_full_face_mask().copy() mask[mask != 0.0] = 1.0 img = get_eyes_mouth_mask() * mask else: img = np.zeros(sample_bgr.shape[0:2] + (1, ), dtype=np.float32) if sample_face_type == FaceType.MARK_ONLY: raise NotImplementedError() mat = LandmarksProcessor.get_transform_mat( sample_landmarks, warp_resolution, face_type) img = cv2.warpAffine( img, mat, (warp_resolution, warp_resolution), flags=cv2.INTER_LINEAR) img = imagelib.warp_by_params( warp_params, img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_LINEAR) else: if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_LINEAR) else: if w != resolution: img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LINEAR) img = imagelib.warp_by_params( warp_params, img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR) if face_mask_type == SPFMT.EYES_MOUTH: div = img.max() if div != 0.0: img = img / div # normalize to 1.0 after warp if len(img.shape) == 2: img = img[..., None] if channel_type == SPCT.G: out_sample = img.astype(np.float32) else: raise ValueError( "only channel_type.G supported for the mask") elif sample_type == SPST.FACE_IMAGE: img = sample_bgr if face_type != sample_face_type: mat = LandmarksProcessor.get_transform_mat( sample_landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=borderMode, flags=cv2.INTER_CUBIC) else: if w != resolution: img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) # Apply random color transfer if ct_mode is not None and ct_sample is not None: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() img = imagelib.color_transfer( ct_mode, img, cv2.resize(ct_sample_bgr, (resolution, resolution), interpolation=cv2.INTER_LINEAR)) if random_hsv_shift_amount != 0: a = random_hsv_shift_amount h_amount = max(1, int(360 * a * 0.5)) img_h, img_s, img_v = cv2.split( cv2.cvtColor(img, cv2.COLOR_BGR2HSV)) img_h = (img_h + rnd_state.randint( -h_amount, h_amount + 1)) % 360 img_s = np.clip( img_s + (rnd_state.random() - 0.5) * a, 0, 1) img_v = np.clip( img_v + (rnd_state.random() - 0.5) * a, 0, 1) img = np.clip( cv2.cvtColor(cv2.merge([img_h, img_s, img_v]), cv2.COLOR_HSV2BGR), 0, 1) img = imagelib.warp_by_params( warp_params, img, warp, transform, can_flip=True, border_replicate=border_replicate) img = np.clip(img.astype(np.float32), 0, 1) # Transform from BGR to desired channel_type if channel_type == SPCT.BGR: out_sample = img elif channel_type == SPCT.G: out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[..., None] elif channel_type == SPCT.GGG: out_sample = np.repeat( np.expand_dims( cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) # Final transformations if nearest_resize_to is not None: out_sample = cv2_resize( out_sample, (nearest_resize_to, nearest_resize_to), interpolation=cv2.INTER_NEAREST) if not debug: if normalize_tanh: out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0, 1.0) if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.IMAGE: img = sample_bgr img = imagelib.warp_by_params(warp_params, img, warp, transform, can_flip=True, border_replicate=True) img = cv2.resize(img, (resolution, resolution), interpolation=cv2.INTER_CUBIC) out_sample = img if data_format == "NCHW": out_sample = np.transpose(out_sample, (2, 0, 1)) elif sample_type == SPST.LANDMARKS_ARRAY: l = sample_landmarks l = np.concatenate([ np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1) ], -1) l = np.clip(l, 0.0, 1.0) out_sample = l elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch, yaw, roll = sample.get_pitch_yaw_roll() if warp_params['flip']: yaw = -yaw if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID: pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1) yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1) roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1) out_sample = (pitch, yaw) else: raise ValueError('expected sample_type') outputs_sample.append(out_sample) outputs += [outputs_sample] return outputs
def batch_func(self, param): images_path, masks_path, mask_file_id_hash, data_format = param file_ids = list(mask_file_id_hash.keys()) shuffle_file_ids = [] resolution = 256 random_flip = True rotation_range = [-15, 15] scale_range = [-0.10, 0.95] tx_range = [-0.3, 0.3] ty_range = [-0.3, 0.3] random_bilinear_resize = (25, 75) motion_blur = (25, 5) gaussian_blur = (25, 5) bs = self.batch_size while True: batches = None n_batch = 0 while n_batch < bs: try: if len(shuffle_file_ids) == 0: shuffle_file_ids = file_ids.copy() np.random.shuffle(shuffle_file_ids) file_id = shuffle_file_ids.pop() masks = mask_file_id_hash[file_id] image_path = images_path / f'{file_id}.jpg' skin_path = masks.get(MaskType.skin, None) hair_path = masks.get(MaskType.hair, None) hat_path = masks.get(MaskType.hat, None) #neck_path = masks.get(MaskType.neck, None) img = cv2_imread(image_path).astype(np.float32) / 255.0 mask = cv2_imread(masks_path / skin_path)[..., 0:1].astype( np.float32) / 255.0 if hair_path is not None: hair_path = masks_path / hair_path if hair_path.exists(): hair = cv2_imread(hair_path)[..., 0:1].astype( np.float32) / 255.0 mask *= (1 - hair) if hat_path is not None: hat_path = masks_path / hat_path if hat_path.exists(): hat = cv2_imread(hat_path)[..., 0:1].astype( np.float32) / 255.0 mask *= (1 - hat) #if neck_path is not None: # neck_path = masks_path / neck_path # if neck_path.exists(): # neck = cv2_imread(neck_path)[...,0:1].astype(np.float32) / 255.0 # mask = np.clip(mask+neck, 0, 1) warp_params = imagelib.gen_warp_params( resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range) img = cv2.resize(img, (resolution, resolution), cv2.INTER_LANCZOS4) h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV)) h = (h + np.random.randint(360)) % 360 s = np.clip(s + np.random.random() - 0.5, 0, 1) v = np.clip(v + np.random.random() / 2 - 0.25, 0, 1) img = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR), 0, 1) if motion_blur is not None: chance, mb_max_size = motion_blur chance = np.clip(chance, 0, 100) mblur_rnd_chance = np.random.randint(100) mblur_rnd_kernel = np.random.randint(mb_max_size) + 1 mblur_rnd_deg = np.random.randint(360) if mblur_rnd_chance < chance: img = imagelib.LinearMotionBlur( img, mblur_rnd_kernel, mblur_rnd_deg) img = imagelib.warp_by_params(warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LANCZOS4) if gaussian_blur is not None: chance, kernel_max_size = gaussian_blur chance = np.clip(chance, 0, 100) gblur_rnd_chance = np.random.randint(100) gblur_rnd_kernel = np.random.randint( kernel_max_size) * 2 + 1 if gblur_rnd_chance < chance: img = cv2.GaussianBlur(img, (gblur_rnd_kernel, ) * 2, 0) if random_bilinear_resize is not None: chance, max_size_per = random_bilinear_resize chance = np.clip(chance, 0, 100) pick_chance = np.random.randint(100) resize_to = resolution - int( np.random.rand() * int(resolution * (max_size_per / 100.0))) img = cv2.resize(img, (resize_to, resize_to), cv2.INTER_LINEAR) img = cv2.resize(img, (resolution, resolution), cv2.INTER_LINEAR) mask = cv2.resize(mask, (resolution, resolution), cv2.INTER_LANCZOS4)[..., None] mask = imagelib.warp_by_params( warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LANCZOS4) mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 mask = np.clip(mask, 0, 1) if data_format == "NCHW": img = np.transpose(img, (2, 0, 1)) mask = np.transpose(mask, (2, 0, 1)) if batches is None: batches = [[], []] batches[0].append(img) batches[1].append(mask) n_batch += 1 except: io.log_err(traceback.format_exc()) yield [np.array(batch) for batch in batches]
def batch_func(self, param): pickled_samples, obstructions_images_paths, resolution, face_type, data_format = param samples = pickle.loads(pickled_samples) obstructions_images_paths_len = len(obstructions_images_paths) shuffle_o_idxs = [] o_idxs = [*range(obstructions_images_paths_len)] shuffle_idxs = [] idxs = [*range(len(samples))] random_flip = True rotation_range = [-10, 10] scale_range = [-0.05, 0.05] tx_range = [-0.05, 0.05] ty_range = [-0.05, 0.05] o_random_flip = True o_rotation_range = [-180, 180] o_scale_range = [-0.5, 0.05] o_tx_range = [-0.5, 0.5] o_ty_range = [-0.5, 0.5] random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25, 75 motion_blur_chance, motion_blur_mb_max_size = 25, 5 gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5 bs = self.batch_size while True: batches = [[], []] n_batch = 0 while n_batch < bs: try: if len(shuffle_idxs) == 0: shuffle_idxs = idxs.copy() np.random.shuffle(shuffle_idxs) idx = shuffle_idxs.pop() sample = samples[idx] img = sample.load_bgr() h, w, c = img.shape mask = np.zeros((h, w, 1), dtype=np.float32) sample.ie_polys.overlay_mask(mask) warp_params = imagelib.gen_warp_params( resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range) if face_type == sample.face_type: if w != resolution: img = cv2.resize(img, (resolution, resolution), cv2.INTER_LANCZOS4) mask = cv2.resize(mask, (resolution, resolution), cv2.INTER_LANCZOS4) else: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4) mask = cv2.warpAffine(mask, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4) if len(mask.shape) == 2: mask = mask[..., None] if obstructions_images_paths_len != 0: # apply obstruction if len(shuffle_o_idxs) == 0: shuffle_o_idxs = o_idxs.copy() np.random.shuffle(shuffle_o_idxs) o_idx = shuffle_o_idxs.pop() o_img = cv2_imread( obstructions_images_paths[o_idx]).astype( np.float32) / 255.0 oh, ow, oc = o_img.shape if oc == 4: ohw = max(oh, ow) scale = resolution / ohw #o_img = cv2.resize (o_img, ( int(ow*rate), int(oh*rate), ), cv2.INTER_CUBIC) mat = cv2.getRotationMatrix2D( (ow / 2, oh / 2), np.random.uniform(o_rotation_range[0], o_rotation_range[1]), 1.0) mat += np.float32([[0, 0, -ow / 2], [0, 0, -oh / 2]]) mat *= scale * np.random.uniform( 1 + o_scale_range[0], 1 + o_scale_range[1]) mat += np.float32( [[ 0, 0, resolution / 2 + resolution * np.random.uniform( o_tx_range[0], o_tx_range[1]) ], [ 0, 0, resolution / 2 + resolution * np.random.uniform( o_ty_range[0], o_ty_range[1]) ]]) o_img = cv2.warpAffine( o_img, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4) if o_random_flip and np.random.randint(10) < 4: o_img = o_img[:, ::-1, ...] o_mask = o_img[..., 3:4] o_mask[o_mask > 0] = 1.0 o_mask = cv2.erode(o_mask, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (5, 5)), iterations=1) o_mask = cv2.GaussianBlur(o_mask, (5, 5), 0)[..., None] img = img * (1 - o_mask) + o_img[..., 0:3] * o_mask o_mask[o_mask < 0.5] = 0.0 #import code #code.interact(local=dict(globals(), **locals())) mask *= (1 - o_mask) #cv2.imshow ("", np.clip(o_img*255, 0,255).astype(np.uint8) ) #cv2.waitKey(0) img = imagelib.warp_by_params(warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) mask = imagelib.warp_by_params(warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) img = np.clip(img.astype(np.float32), 0, 1) mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 mask = np.clip(mask, 0, 1) img = imagelib.apply_random_hsv_shift( img, mask=sd.random_circle_faded([resolution, resolution])) img = imagelib.apply_random_motion_blur( img, motion_blur_chance, motion_blur_mb_max_size, mask=sd.random_circle_faded([resolution, resolution])) img = imagelib.apply_random_gaussian_blur( img, gaussian_blur_chance, gaussian_blur_kernel_max_size, mask=sd.random_circle_faded([resolution, resolution])) img = imagelib.apply_random_bilinear_resize( img, random_bilinear_resize_chance, random_bilinear_resize_max_size_per, mask=sd.random_circle_faded([resolution, resolution])) if data_format == "NCHW": img = np.transpose(img, (2, 0, 1)) mask = np.transpose(mask, (2, 0, 1)) batches[0].append(img) batches[1].append(mask) n_batch += 1 except: io.log_err(traceback.format_exc()) yield [np.array(batch) for batch in batches]
def batch_func(self, param): samples, kf_idxs, resolution, face_type, data_format = param kf_idxs_len = len(kf_idxs) shuffle_idxs = [] idxs = [*range(len(samples))] random_flip = True rotation_range = [-10, 10] scale_range = [-0.05, 0.05] tx_range = [-0.05, 0.05] ty_range = [-0.05, 0.05] bs = self.batch_size while True: batches = [[], [], [], [], [], []] n_batch = 0 while n_batch < bs: try: if len(shuffle_idxs) == 0: shuffle_idxs = idxs.copy() np.random.shuffle(shuffle_idxs) idx = shuffle_idxs.pop() key_idx, key_chain_idx, chain_idxs = kf_idxs[ np.random.randint(kf_idxs_len)] key_sample = samples[key_idx] key_chain_sample = samples[key_chain_idx] chain_sample = samples[chain_idxs[np.random.randint( len(chain_idxs))]] #print('==========') #print(key_sample.filename) #print(key_chain_sample.filename) #print(chain_sample.filename) sample = samples[idx] img = sample.load_bgr() key_img = key_sample.load_bgr() key_chain_img = key_chain_sample.load_bgr() chain_img = chain_sample.load_bgr() h, w, c = img.shape mask = LandmarksProcessor.get_image_hull_mask( img.shape, sample.landmarks) mask = np.clip(mask, 0, 1) warp_params = imagelib.gen_warp_params( resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range) if face_type == sample.face_type: if w != resolution: img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) key_img = cv2.resize(key_img, (resolution, resolution), cv2.INTER_CUBIC) key_chain_img = cv2.resize( key_chain_img, (resolution, resolution), cv2.INTER_CUBIC) chain_img = cv2.resize(chain_img, (resolution, resolution), cv2.INTER_CUBIC) mask = cv2.resize(mask, (resolution, resolution), cv2.INTER_CUBIC) else: mat = LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, face_type) img = cv2.warpAffine(img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) key_img = cv2.warpAffine( key_img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) key_chain_img = cv2.warpAffine( key_chain_img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) chain_img = cv2.warpAffine( chain_img, mat, (resolution, resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC) mask = cv2.warpAffine(mask, mat, (resolution, resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_CUBIC) if len(mask.shape) == 2: mask = mask[..., None] img_warped = imagelib.warp_by_params(warp_params, img, can_warp=True, can_transform=True, can_flip=True, border_replicate=True) img_transformed = imagelib.warp_by_params( warp_params, img, can_warp=False, can_transform=True, can_flip=True, border_replicate=True) mask = imagelib.warp_by_params(warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False) key_img = imagelib.warp_by_params(warp_params, key_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) key_chain_img = imagelib.warp_by_params( warp_params, key_chain_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) chain_img = imagelib.warp_by_params(warp_params, chain_img, can_warp=False, can_transform=False, can_flip=False, border_replicate=True) img_warped = np.clip(img_warped.astype(np.float32), 0, 1) img_transformed = np.clip( img_transformed.astype(np.float32), 0, 1) mask[mask < 0.5] = 0.0 mask[mask >= 0.5] = 1.0 mask = np.clip(mask, 0, 1) if data_format == "NCHW": img_warped = np.transpose(img_warped, (2, 0, 1)) img_transformed = np.transpose(img_transformed, (2, 0, 1)) mask = np.transpose(mask, (2, 0, 1)) key_img = np.transpose(key_img, (2, 0, 1)) key_chain_img = np.transpose(key_chain_img, (2, 0, 1)) chain_img = np.transpose(chain_img, (2, 0, 1)) batches[0].append(img_warped) batches[1].append(img_transformed) batches[2].append(mask) batches[3].append(key_img) batches[4].append(key_chain_img) batches[5].append(chain_img) n_batch += 1 except: io.log_err(traceback.format_exc()) yield [np.array(batch) for batch in batches]