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_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