def apply_xseg(input_path, model_path): if not input_path.exists(): raise ValueError(f'{input_path} not found. Please ensure it exists.') if not model_path.exists(): raise ValueError(f'{model_path} not found. Please ensure it exists.') face_type = io.input_str( "XSeg model face type", 'same', ['h', 'mf', 'f', 'wf', 'head', 'same'], help_message= "Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'" ).lower() if face_type == 'same': face_type = None else: face_type = { 'h': FaceType.HALF, 'mf': FaceType.MID_FULL, 'f': FaceType.FULL, 'wf': FaceType.WHOLE_FACE, 'head': FaceType.HEAD }[face_type] io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.') device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True) nn.initialize(device_config) xseg = XSegNet(name='XSeg', load_weights=True, weights_file_root=model_path, data_format=nn.data_format, raise_on_no_model_files=True) xseg_res = xseg.get_resolution() images_paths = pathex.get_image_paths(input_path, return_Path_class=True) for filepath in io.progress_bar_generator(images_paths, "Processing"): dflimg = DFLIMG.load(filepath) if dflimg is None or not dflimg.has_data(): io.log_info(f'{filepath} is not a DFLIMG') continue img = cv2_imread(filepath).astype(np.float32) / 255.0 h, w, c = img.shape img_face_type = FaceType.fromString(dflimg.get_face_type()) if face_type is not None and img_face_type != face_type: lmrks = dflimg.get_source_landmarks() fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type) imat = LandmarksProcessor.get_transform_mat( lmrks, w, img_face_type) g_p = LandmarksProcessor.transform_points( np.float32([(0, 0), (w, 0), (0, w)]), fmat, True) g_p2 = LandmarksProcessor.transform_points(g_p, imat) mat = cv2.getAffineTransform(g_p2, np.float32([(0, 0), (w, 0), (0, w)])) img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4) img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) else: if w != xseg_res: img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) if len(img.shape) == 2: img = img[..., None] mask = xseg.extract(img) if face_type is not None and img_face_type != face_type: mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4) mask = cv2.warpAffine(mask, mat, (w, w), np.zeros((h, w, c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4) mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) mask[mask < 0.5] = 0 mask[mask >= 0.5] = 1 dflimg.set_xseg_mask(mask) dflimg.save()
def get_pitch_yaw_roll(self): if self.pitch_yaw_roll is None: self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll( self.landmarks, size=self.shape[1]) return self.pitch_yaw_roll
def process_data(self, data): filename_path = Path(data.filename) filename_path_str = str(filename_path) if self.cached_image[0] == filename_path_str: image = self.cached_image[ 1] #cached image for manual extractor else: image = cv2_imread(filename_path_str) if image is None: self.log_err( 'Failed to extract %s, reason: cv2_imread() fail.' % (str(filename_path))) return data image_shape = image.shape if len(image_shape) == 2: h, w = image.shape image = image[:, :, np.newaxis] ch = 1 else: h, w, ch = image.shape if ch == 1: image = np.repeat(image, 3, -1) elif ch == 4: image = image[:, :, 0:3] wm, hm = w % 2, h % 2 if wm + hm != 0: #fix odd image image = image[0:h - hm, 0:w - wm, :] self.cached_image = (filename_path_str, image) src_dflimg = None h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? if filename_path.suffix == '.png': src_dflimg = DFLPNG.load(str(filename_path)) if filename_path.suffix == '.jpg': src_dflimg = DFLJPG.load(str(filename_path)) if 'rects' in self.type: if min(w, h) < 128: self.log_err('Image is too small %s : [%d, %d]' % (str(filename_path), w, h)) data.rects = [] else: for rot in ([0, 90, 270, 180]): data.rects_rotation = rot if rot == 0: rotated_image = image elif rot == 90: rotated_image = image.swapaxes(0, 1)[:, ::-1, :] elif rot == 180: rotated_image = image[::-1, ::-1, :] elif rot == 270: rotated_image = image.swapaxes(0, 1)[::-1, :, :] rects = data.rects = self.e.extract(rotated_image, is_bgr=True) if len(rects) != 0: break return data elif self.type == 'landmarks': if data.rects_rotation == 0: rotated_image = image elif data.rects_rotation == 90: rotated_image = image.swapaxes(0, 1)[:, ::-1, :] elif data.rects_rotation == 180: rotated_image = image[::-1, ::-1, :] elif data.rects_rotation == 270: rotated_image = image.swapaxes(0, 1)[::-1, :, :] data.landmarks = self.e.extract( rotated_image, data.rects, self.second_pass_e if (src_dflimg is None and data.landmarks_accurate) else None, is_bgr=True) if data.rects_rotation != 0: for i, (rect, lmrks) in enumerate(zip(data.rects, data.landmarks)): new_rect, new_lmrks = rect, lmrks (l, t, r, b) = rect if data.rects_rotation == 90: new_rect = (t, h - l, b, h - r) if lmrks is not None: new_lmrks = lmrks[:, ::-1].copy() new_lmrks[:, 1] = h - new_lmrks[:, 1] elif data.rects_rotation == 180: if lmrks is not None: new_rect = (w - l, h - t, w - r, h - b) new_lmrks = lmrks.copy() new_lmrks[:, 0] = w - new_lmrks[:, 0] new_lmrks[:, 1] = h - new_lmrks[:, 1] elif data.rects_rotation == 270: new_rect = (w - b, l, w - t, r) if lmrks is not None: new_lmrks = lmrks[:, ::-1].copy() new_lmrks[:, 0] = w - new_lmrks[:, 0] data.rects[i], data.landmarks[i] = new_rect, new_lmrks return data elif self.type == 'final': data.final_output_files = [] rects = data.rects landmarks = data.landmarks if self.debug_dir is not None: debug_output_file = str( Path(self.debug_dir) / (filename_path.stem + '.jpg')) debug_image = image.copy() if src_dflimg is not None and len(rects) != 1: #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it print("src_dflimg is not None and len(rects) != 1", str(filename_path)) output_file = str(self.final_output_path / filename_path.name) if str(filename_path) != str(output_file): shutil.copy(str(filename_path), str(output_file)) data.final_output_files.append(output_file) else: face_idx = 0 for rect, image_landmarks in zip(rects, landmarks): if src_dflimg is not None and face_idx > 1: #cannot extract more than 1 face from dflimg break if image_landmarks is None: continue rect = np.array(rect) if self.face_type == FaceType.MARK_ONLY: face_image = image face_image_landmarks = image_landmarks else: image_to_face_mat = LandmarksProcessor.get_transform_mat( image_landmarks, self.image_size, self.face_type) face_image = cv2.warpAffine( image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4) face_image_landmarks = LandmarksProcessor.transform_points( image_landmarks, image_to_face_mat) landmarks_bbox = LandmarksProcessor.transform_points( [(0, 0), (0, self.image_size - 1), (self.image_size - 1, self.image_size - 1), (self.image_size - 1, 0)], image_to_face_mat, True) rect_area = mathlib.polygon_area( np.array(rect[[0, 2, 2, 0]]), np.array(rect[[1, 1, 3, 3]])) landmarks_area = mathlib.polygon_area( landmarks_bbox[:, 0], landmarks_bbox[:, 1]) if landmarks_area > 4 * rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area continue if self.debug_dir is not None: LandmarksProcessor.draw_rect_landmarks( debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True) if src_dflimg is not None and filename_path.suffix == '.jpg': #if extracting from dflimg and jpg copy it in order not to lose quality output_file = str(self.final_output_path / filename_path.name) if str(filename_path) != str(output_file): shutil.copy(str(filename_path), str(output_file)) else: output_file = '{}_{}{}'.format( str(self.final_output_path / filename_path.stem), str(face_idx), '.jpg') cv2_imwrite(output_file, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), 85]) DFLJPG.embed_data( output_file, face_type=FaceType.toString(self.face_type), landmarks=face_image_landmarks.tolist(), source_filename=filename_path.name, source_rect=rect, source_landmarks=image_landmarks.tolist(), image_to_face_mat=image_to_face_mat, pitch_yaw_roll=data.pitch_yaw_roll) data.final_output_files.append(output_file) face_idx += 1 data.faces_detected = face_idx if self.debug_dir is not None: cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50]) return data elif self.type == 'fanseg': if src_dflimg is not None: fanseg_mask = self.e.extract(image / 255.0) src_dflimg.embed_and_set( filename_path_str, fanseg_mask=fanseg_mask, #fanseg_mask_ver=FANSegmentator.VERSION, )
def dfl_estimate_pitch_yaw_roll(dfl_img): from facelib import LandmarksProcessor return LandmarksProcessor.estimate_pitch_yaw_roll(dfl_img.get_landmarks())
def get_data(self, host_dict): if self.type == 'landmarks-manual': need_remark_face = False while len(self.input_data) > 0: data = self.input_data[0] filepath, data_rects, data_landmarks = data.filepath, data.rects, data.landmarks is_frame_done = False if self.image_filepath != filepath: self.image_filepath = filepath if self.cache_original_image[0] == filepath: self.original_image = self.cache_original_image[1] else: self.original_image = imagelib.normalize_channels( cv2_imread(filepath), 3) self.cache_original_image = (filepath, self.original_image) (h, w, c) = self.original_image.shape self.view_scale = 1.0 if self.manual_window_size == 0 else self.manual_window_size / ( h * (16.0 / 9.0)) if self.cache_image[0] == (h, w, c) + (self.view_scale, filepath): self.image = self.cache_image[1] else: self.image = cv2.resize(self.original_image, (int( w * self.view_scale), int(h * self.view_scale)), interpolation=cv2.INTER_LINEAR) self.cache_image = ((h, w, c) + (self.view_scale, filepath), self.image) (h, w, c) = self.image.shape sh = (0, 0, w, min(100, h)) if self.cache_text_lines_img[0] == sh: self.text_lines_img = self.cache_text_lines_img[1] else: self.text_lines_img = (imagelib.get_draw_text_lines( self.image, sh, [ '[L Mouse click] - lock/unlock selection. [Mouse wheel] - change rect', '[R Mouse Click] - manual face rectangle', '[Enter] / [Space] - confirm / skip frame', '[,] [.]- prev frame, next frame. [Q] - skip remaining frames', '[a] - accuracy on/off (more fps)', '[h] - hide this help' ], (1, 1, 1)) * 255).astype(np.uint8) self.cache_text_lines_img = (sh, self.text_lines_img) if need_remark_face: # need remark image from input data that already has a marked face? need_remark_face = False if len( data_rects ) != 0: # If there was already a face then lock the rectangle to it until the mouse is clicked self.rect = data_rects.pop() self.landmarks = data_landmarks.pop() data_rects.clear() data_landmarks.clear() self.rect_locked = True self.rect_size = (self.rect[2] - self.rect[0]) / 2 self.x = (self.rect[0] + self.rect[2]) / 2 self.y = (self.rect[1] + self.rect[3]) / 2 self.redraw() if len(data_rects) == 0: (h, w, c) = self.image.shape while True: io.process_messages(0.0001) if not self.force_landmarks: new_x = self.x new_y = self.y new_rect_size = self.rect_size mouse_events = io.get_mouse_events(self.wnd_name) for ev in mouse_events: (x, y, ev, flags) = ev if ev == io.EVENT_MOUSEWHEEL and not self.rect_locked: mod = 1 if flags > 0 else -1 diff = 1 if new_rect_size <= 40 else np.clip( new_rect_size / 10, 1, 10) new_rect_size = max(5, new_rect_size + diff * mod) elif ev == io.EVENT_LBUTTONDOWN: if self.force_landmarks: self.x = new_x self.y = new_y self.force_landmarks = False self.rect_locked = True self.redraw() else: self.rect_locked = not self.rect_locked self.extract_needed = True elif ev == io.EVENT_RBUTTONDOWN: self.force_landmarks = not self.force_landmarks if self.force_landmarks: self.rect_locked = False elif not self.rect_locked: new_x = np.clip(x, 0, w - 1) / self.view_scale new_y = np.clip(y, 0, h - 1) / self.view_scale key_events = io.get_key_events(self.wnd_name) key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[ -1] if len(key_events) > 0 else (0, 0, False, False, False) if key == ord('\r') or key == ord('\n'): #confirm frame is_frame_done = True data_rects.append(self.rect) data_landmarks.append(self.landmarks) break elif key == ord(' '): #confirm skip frame is_frame_done = True break elif key == ord(',') and len(self.result) > 0: #go prev frame if self.rect_locked: self.rect_locked = False # Only save the face if the rect is still locked data_rects.append(self.rect) data_landmarks.append(self.landmarks) self.input_data.insert(0, self.result.pop()) io.progress_bar_inc(-1) need_remark_face = True break elif key == ord('.'): #go next frame if self.rect_locked: self.rect_locked = False # Only save the face if the rect is still locked data_rects.append(self.rect) data_landmarks.append(self.landmarks) need_remark_face = True is_frame_done = True break elif key == ord('q'): #skip remaining if self.rect_locked: self.rect_locked = False data_rects.append(self.rect) data_landmarks.append(self.landmarks) while len(self.input_data) > 0: self.result.append(self.input_data.pop(0)) io.progress_bar_inc(1) break elif key == ord('h'): self.hide_help = not self.hide_help break elif key == ord('a'): self.landmarks_accurate = not self.landmarks_accurate break if self.force_landmarks: pt2 = np.float32([new_x, new_y]) pt1 = np.float32([self.x, self.y]) pt_vec_len = npla.norm(pt2 - pt1) pt_vec = pt2 - pt1 if pt_vec_len != 0: pt_vec /= pt_vec_len self.rect_size = pt_vec_len self.rect = (int(self.x - self.rect_size), int(self.y - self.rect_size), int(self.x + self.rect_size), int(self.y + self.rect_size)) if pt_vec_len > 0: lmrks = np.concatenate( (np.zeros((17, 2), np.float32), LandmarksProcessor.landmarks_2D), axis=0) lmrks -= lmrks[30:31, :] mat = cv2.getRotationMatrix2D( (0, 0), -np.arctan2(pt_vec[1], pt_vec[0]) * 180 / math.pi, pt_vec_len) mat[:, 2] += (self.x, self.y) self.landmarks = LandmarksProcessor.transform_points( lmrks, mat) self.redraw() elif self.x != new_x or \ self.y != new_y or \ self.rect_size != new_rect_size or \ self.extract_needed: self.x = new_x self.y = new_y self.rect_size = new_rect_size self.rect = (int(self.x - self.rect_size), int(self.y - self.rect_size), int(self.x + self.rect_size), int(self.y + self.rect_size)) return ExtractSubprocessor.Data( filepath, rects=[self.rect], landmarks_accurate=self.landmarks_accurate) else: is_frame_done = True if is_frame_done: self.result.append(data) self.input_data.pop(0) io.progress_bar_inc(1) self.extract_needed = True self.rect_locked = False else: if len(self.input_data) > 0: return self.input_data.pop(0) return None
def process(sample, sample_process_options, output_sample_types, debug): source = sample.load_bgr() h, w, c = source.shape is_face_sample = sample.landmarks is not None if debug and is_face_sample: LandmarksProcessor.draw_landmarks(source, sample.landmarks, (0, 1, 0)) params = image_utils.gen_warp_params( source, 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) images = [[None] * 3 for _ in range(4)] sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample_type in output_sample_types: f = sample_type[0] size = sample_type[1] random_sub_size = 0 if len(sample_type) < 3 else min( sample_type[2], size) if f & SampleProcessor.TypeFlags.SOURCE != 0: img_type = 0 elif f & SampleProcessor.TypeFlags.WARPED != 0: img_type = 1 elif f & SampleProcessor.TypeFlags.WARPED_TRANSFORMED != 0: img_type = 2 elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0: img_type = 3 else: raise ValueError('expected SampleTypeFlags type') face_mask_type = 0 if f & SampleProcessor.TypeFlags.FACE_MASK_FULL != 0: face_mask_type = 1 elif f & SampleProcessor.TypeFlags.FACE_MASK_EYES != 0: face_mask_type = 2 target_face_type = -1 if f & SampleProcessor.TypeFlags.FACE_ALIGN_HALF != 0: target_face_type = FaceType.HALF elif f & SampleProcessor.TypeFlags.FACE_ALIGN_FULL != 0: target_face_type = FaceType.FULL elif f & SampleProcessor.TypeFlags.FACE_ALIGN_HEAD != 0: target_face_type = FaceType.HEAD elif f & SampleProcessor.TypeFlags.FACE_ALIGN_AVATAR != 0: target_face_type = FaceType.AVATAR if images[img_type][face_mask_type] is None: img = source if is_face_sample: if face_mask_type == 1: img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask( source, sample.landmarks)), -1) elif face_mask_type == 2: mask = LandmarksProcessor.get_image_eye_mask( source, sample.landmarks) mask = np.expand_dims( cv2.blur(mask, (w // 32, w // 32)), -1) mask[mask > 0.0] = 1.0 img = np.concatenate((img, mask), -1) images[img_type][face_mask_type] = image_utils.warp_by_params( params, img, (img_type == 1 or img_type == 2), (img_type == 2 or img_type == 3), img_type != 0, face_mask_type == 0) img = images[img_type][face_mask_type] if is_face_sample and target_face_type != -1: if target_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_face_type)) img = cv2.warpAffine(img, LandmarksProcessor.get_transform_mat( sample.landmarks, size, target_face_type), (size, size), flags=cv2.INTER_LANCZOS4) else: img = cv2.resize(img, (size, size), cv2.INTER_LANCZOS4) if random_sub_size != 0: sub_size = size - random_sub_size rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_size) start_x = rnd_state.randint(sub_size + 1) start_y = rnd_state.randint(sub_size + 1) img = img[start_y:start_y + sub_size, start_x:start_x + sub_size, :] img_bgr = img[..., 0:3] img_mask = img[..., 3:4] if f & SampleProcessor.TypeFlags.MODE_BGR != 0: img = img elif f & SampleProcessor.TypeFlags.MODE_BGR_SHUFFLE != 0: img_bgr = np.take(img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1) img = np.concatenate((img_bgr, img_mask), -1) elif f & SampleProcessor.TypeFlags.MODE_G != 0: img = np.concatenate((np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), img_mask), -1) elif f & SampleProcessor.TypeFlags.MODE_GGG != 0: img = np.concatenate((np.repeat( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), (3, ), -1), img_mask), -1) elif is_face_sample and f & SampleProcessor.TypeFlags.MODE_M != 0: if face_mask_type == 0: raise ValueError('no face_mask_type defined') img = img_mask else: raise ValueError('expected SampleTypeFlags mode') if not debug and sample_process_options.normalize_tanh: img = img * 2.0 - 1.0 outputs.append(img) if debug: result = [] for output in outputs: if output.shape[2] < 4: result += [ output, ] elif output.shape[2] == 4: result += [ output[..., 0:3] * output[..., 3:4], ] return result else: return outputs
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 is_face_sample: if target_face_type == SPTF.NONE: raise ValueError("target face type must be defined for face samples") else: if mode_type == SPTF.MODE_FACE_MASK_ALL_HULL: raise ValueError("MODE_FACE_MASK_ALL_HULL applicable only for face samples") if mode_type == SPTF.MODE_FACE_MASK_EYES_HULL: raise ValueError("MODE_FACE_MASK_EYES_HULL applicable only for face samples") if mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL: raise ValueError("MODE_FACE_MASK_ALL_EYES_HULL applicable only for face samples") if mode_type == SPTF.MODE_FACE_MASK_STRUCT: raise ValueError("MODE_FACE_MASK_STRUCT applicable only 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_ALL_HULL or \ mode_type == SPTF.MODE_FACE_MASK_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL: if mode_type == SPTF.MODE_FACE_MASK_ALL_HULL or \ mode_type == SPTF.MODE_FACE_MASK_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 mode_type == SPTF.MODE_FACE_MASK_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL: eyes_mask = LandmarksProcessor.get_image_eye_mask (sample_bgr.shape, sample.landmarks) eyes_mask = np.clip(eyes_mask, 0, 1) if mode_type == SPTF.MODE_FACE_MASK_ALL_HULL: img = all_mask elif mode_type == SPTF.MODE_FACE_MASK_EYES_HULL: img = eyes_mask elif mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL: img = all_mask + eyes_mask 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) rnd_state = np.random.RandomState (sample_rnd_seed) mblur_rnd_chance = rnd_state.randint(100) mblur_rnd_kernel = rnd_state.randint(mb_max_size)+1 mblur_rnd_deg = 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) rnd_state = np.random.RandomState (sample_rnd_seed+1) gblur_rnd_chance = rnd_state.randint(100) gblur_rnd_kernel = rnd_state.randint(kernel_max_size)*2+1 if gblur_rnd_chance < chance: img = cv2.GaussianBlur(img, (gblur_rnd_kernel,) *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_ALL_HULL or \ mode_type == SPTF.MODE_FACE_MASK_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_STRUCT: img = cv2.warpAffine( img, mat, (sample.shape[0],sample.shape[0]), flags=cv2.INTER_LINEAR ) img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LINEAR) img = cv2.resize( img, (resolution,resolution), cv2.INTER_LINEAR )[...,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_ALL_HULL or \ mode_type == SPTF.MODE_FACE_MASK_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_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, cv2_inter=cv2.INTER_LINEAR) img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LINEAR )[...,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_ALL_HULL or \ mode_type == SPTF.MODE_FACE_MASK_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_ALL_EYES_HULL or \ mode_type == SPTF.MODE_FACE_MASK_STRUCT: out_sample = img.astype(np.float32) 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_BGR_RANDOM_RGB_LEVELS: rnd_state = np.random.RandomState (sample_rnd_seed) np_rnd = 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 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 MergeMaskedFace_test(path, cfg): data = np.load(path, allow_pickle=True) [ img_bgr, predictor_input_shape, frame_info, img_face_landmarks, prd_face_bgr, prd_face_mask_a_0, prd_face_dst_mask_a_0, img_bgr_uint8 ] = data img_path = glob.glob(datadir() + '/data_dst/aligned/' + path.split('/')[-1].split('_')[1] + '*')[0] dflimg = DFLIMG.load(Path(img_path)) face_type_ = dflimg.get_face_type() cfg.face_type = FaceType.WHOLE_FACE if face_type_ == 'head': cfg.face_type = FaceType.HEAD elif face_type_ == 'f': cfg.face_type = FaceType.FULL else: cfg.face_type = FaceType.WHOLE_FACE out_img = None if cfg.show_mode == 1 or cfg.show_mode == 3: img_size = img_bgr.shape[1], img_bgr.shape[0] img_face_mask_a = LandmarksProcessor.get_image_hull_mask( img_bgr.shape, img_face_landmarks) input_size = predictor_input_shape[0] mask_subres_size = input_size * 4 output_size = input_size if cfg.super_resolution_power != 0: output_size *= 4 face_mat = LandmarksProcessor.get_transform_mat( img_face_landmarks, output_size, face_type=cfg.face_type) face_output_mat = LandmarksProcessor.get_transform_mat( img_face_landmarks, output_size, face_type=cfg.face_type, scale=1.0 + 0.01 * cfg.output_face_scale) if mask_subres_size == output_size: face_mask_output_mat = face_output_mat else: face_mask_output_mat = LandmarksProcessor.get_transform_mat( img_face_landmarks, mask_subres_size, face_type=cfg.face_type, scale=1.0 + 0.01 * cfg.output_face_scale) dst_face_bgr = cv2.warpAffine(img_bgr, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC) dst_face_bgr = np.clip(dst_face_bgr, 0, 1) dst_face_mask_a_0 = cv2.warpAffine(img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC) dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1) predictor_input_bgr = cv2.resize(dst_face_bgr, (input_size, input_size)) if cfg.mask_mode == 1: #dst wrk_face_mask_a_0 = cv2.resize(dst_face_mask_a_0, (output_size, output_size), interpolation=cv2.INTER_CUBIC) elif cfg.mask_mode == 2: #learned-prd wrk_face_mask_a_0 = prd_face_mask_a_0 elif cfg.mask_mode == 3: #learned-dst wrk_face_mask_a_0 = prd_face_dst_mask_a_0 elif cfg.mask_mode == 4: #learned-prd*learned-dst wrk_face_mask_a_0 = prd_face_mask_a_0 * prd_face_dst_mask_a_0 elif cfg.mask_mode == 5: #learned-prd+learned-dst wrk_face_mask_a_0 = np.clip( prd_face_mask_a_0 + prd_face_dst_mask_a_0, 0, 1) elif cfg.mask_mode >= 6 and cfg.mask_mode <= 9: #XSeg modes if cfg.mask_mode == 6: #'XSeg-prd' wrk_face_mask_a_0 = X_prd_face_mask_a_0 elif cfg.mask_mode == 7: #'XSeg-dst' wrk_face_mask_a_0 = X_dst_face_mask_a_0 elif cfg.mask_mode == 8: #'XSeg-prd*XSeg-dst' wrk_face_mask_a_0 = X_prd_face_mask_a_0 * X_dst_face_mask_a_0 elif cfg.mask_mode == 9: #learned-prd*learned-dst*XSeg-prd*XSeg-dst wrk_face_mask_a_0 = prd_face_mask_a_0 * prd_face_dst_mask_a_0 * X_prd_face_mask_a_0 * X_dst_face_mask_a_0 wrk_face_mask_a_0[wrk_face_mask_a_0 < ( 1.0 / 255.0)] = 0.0 # get rid of noise # resize to mask_subres_size if wrk_face_mask_a_0.shape[0] != mask_subres_size: wrk_face_mask_a_0 = cv2.resize( wrk_face_mask_a_0, (mask_subres_size, mask_subres_size), interpolation=cv2.INTER_CUBIC) # process mask in local predicted space if 'raw' not in cfg.mode: # add zero pad wrk_face_mask_a_0 = np.pad(wrk_face_mask_a_0, input_size) ero = cfg.erode_mask_modifier blur = cfg.blur_mask_modifier if ero > 0: wrk_face_mask_a_0 = cv2.erode(wrk_face_mask_a_0, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (ero, ero)), iterations=1) elif ero < 0: wrk_face_mask_a_0 = cv2.dilate(wrk_face_mask_a_0, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (-ero, -ero)), iterations=1) # clip eroded/dilated mask in actual predict area # pad with half blur size in order to accuratelly fade to zero at the boundary clip_size = input_size + blur // 2 wrk_face_mask_a_0[:clip_size, :] = 0 wrk_face_mask_a_0[-clip_size:, :] = 0 wrk_face_mask_a_0[:, :clip_size] = 0 wrk_face_mask_a_0[:, -clip_size:] = 0 if blur > 0: blur = blur + (1 - blur % 2) wrk_face_mask_a_0 = cv2.GaussianBlur(wrk_face_mask_a_0, (blur, blur), 0) wrk_face_mask_a_0 = wrk_face_mask_a_0[input_size:-input_size, input_size:-input_size] wrk_face_mask_a_0 = np.clip(wrk_face_mask_a_0, 0, 1) img_face_mask_a = cv2.warpAffine(wrk_face_mask_a_0, face_mask_output_mat, img_size, np.zeros(img_bgr.shape[0:2], dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC)[..., None] img_face_mask_a = np.clip(img_face_mask_a, 0.0, 1.0) img_face_mask_a[img_face_mask_a < (1.0 / 255.0)] = 0.0 # get rid of noise if wrk_face_mask_a_0.shape[0] != output_size: wrk_face_mask_a_0 = cv2.resize(wrk_face_mask_a_0, (output_size, output_size), interpolation=cv2.INTER_CUBIC) wrk_face_mask_a = wrk_face_mask_a_0[..., None] out_merging_mask_a = None if cfg.mode == 'original': return img_bgr, img_face_mask_a elif 'raw' in cfg.mode: if cfg.mode == 'raw-rgb': out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_merging_mask_a = img_face_mask_a elif cfg.mode == 'raw-predict': out_img = prd_face_bgr out_merging_mask_a = wrk_face_mask_a else: raise ValueError(f"undefined raw type {cfg.mode}") out_img = np.clip(out_img, 0.0, 1.0) else: # Process if the mask meets minimum size maxregion = np.argwhere(img_face_mask_a >= 0.1) if maxregion.size != 0: miny, minx = maxregion.min(axis=0)[:2] maxy, maxx = maxregion.max(axis=0)[:2] lenx = maxx - minx leny = maxy - miny if min(lenx, leny) >= 4: wrk_face_mask_area_a = wrk_face_mask_a.copy() wrk_face_mask_area_a[wrk_face_mask_area_a > 0] = 1.0 if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0: if cfg.color_transfer_mode == 1: #rct prd_face_bgr = imagelib.reinhard_color_transfer( np.clip( prd_face_bgr * wrk_face_mask_area_a * 255, 0, 255).astype(np.uint8), np.clip( dst_face_bgr * wrk_face_mask_area_a * 255, 0, 255).astype(np.uint8), ) prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) elif cfg.color_transfer_mode == 2: #lct prd_face_bgr = imagelib.linear_color_transfer( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 3: #mkl prd_face_bgr = imagelib.color_transfer_mkl( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 4: #mkl-m prd_face_bgr = imagelib.color_transfer_mkl( prd_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) elif cfg.color_transfer_mode == 5: #idt prd_face_bgr = imagelib.color_transfer_idt( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m prd_face_bgr = imagelib.color_transfer_idt( prd_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) elif cfg.color_transfer_mode == 7: #sot-m prd_face_bgr = imagelib.color_transfer_sot( prd_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a, steps=10, batch_size=30) prd_face_bgr = np.clip(prd_face_bgr, 0.0, 1.0) elif cfg.color_transfer_mode == 8: #mix-m prd_face_bgr = imagelib.color_transfer_mix( prd_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) if cfg.mode == 'hist-match': hist_mask_a = np.ones(prd_face_bgr.shape[:2] + (1, ), dtype=np.float32) if cfg.masked_hist_match: hist_mask_a *= wrk_face_mask_area_a white = (1.0 - hist_mask_a) * np.ones( prd_face_bgr.shape[:2] + (1, ), dtype=np.float32) hist_match_1 = prd_face_bgr * hist_mask_a + white hist_match_1[hist_match_1 > 1.0] = 1.0 hist_match_2 = dst_face_bgr * hist_mask_a + white hist_match_2[hist_match_1 > 1.0] = 1.0 prd_face_bgr = imagelib.color_hist_match( hist_match_1, hist_match_2, cfg.hist_match_threshold).astype(dtype=np.float32) if 'seamless' in cfg.mode: #mask used for cv2.seamlessClone img_face_seamless_mask_a = None for i in range(1, 10): a = img_face_mask_a > i / 10.0 if len(np.argwhere(a)) == 0: continue img_face_seamless_mask_a = img_face_mask_a.copy() img_face_seamless_mask_a[a] = 1.0 img_face_seamless_mask_a[ img_face_seamless_mask_a <= i / 10.0] = 0.0 break out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_img = np.clip(out_img, 0.0, 1.0) if 'seamless' in cfg.mode: try: #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering) l, t, w, h = cv2.boundingRect( (img_face_seamless_mask_a * 255).astype( np.uint8)) s_maskx, s_masky = int(l + w / 2), int(t + h / 2) out_img = cv2.seamlessClone( (out_img * 255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a * 255).astype(np.uint8), (s_maskx, s_masky), cv2.NORMAL_CLONE) out_img = out_img.astype(dtype=np.float32) / 255.0 except Exception as e: #seamlessClone may fail in some cases e_str = traceback.format_exc() if 'MemoryError' in e_str: raise Exception( "Seamless fail: " + e_str ) #reraise MemoryError in order to reprocess this data by other processes else: print("Seamless fail: " + e_str) cfg_mp = cfg.motion_blur_power / 100.0 out_img = geometric_transformation_of_mask( img_bgr, out_img, img_face_mask_a, H=cfg.horizontal_shear, V=cfg.vertical_shear, PH=cfg.horizontal_shift, PV=cfg.vertical_shift) if ('seamless' in cfg.mode and cfg.color_transfer_mode != 0) or \ cfg.mode == 'seamless-hist-match' or \ cfg_mp != 0 or \ cfg.blursharpen_amount != 0 or \ cfg.image_denoise_power != 0 or \ cfg.bicubic_degrade_power != 0: out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC) if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0: if cfg.color_transfer_mode == 1: out_face_bgr = imagelib.reinhard_color_transfer( np.clip( out_face_bgr * wrk_face_mask_area_a * 255, 0, 255).astype(np.uint8), np.clip( dst_face_bgr * wrk_face_mask_area_a * 255, 0, 255).astype(np.uint8)) out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) elif cfg.color_transfer_mode == 2: #lct out_face_bgr = imagelib.linear_color_transfer( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 3: #mkl out_face_bgr = imagelib.color_transfer_mkl( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 4: #mkl-m out_face_bgr = imagelib.color_transfer_mkl( out_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) elif cfg.color_transfer_mode == 5: #idt out_face_bgr = imagelib.color_transfer_idt( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m out_face_bgr = imagelib.color_transfer_idt( out_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) elif cfg.color_transfer_mode == 7: #sot-m out_face_bgr = imagelib.color_transfer_sot( out_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a, steps=10, batch_size=30) out_face_bgr = np.clip(out_face_bgr, 0.0, 1.0) elif cfg.color_transfer_mode == 8: #mix-m out_face_bgr = imagelib.color_transfer_mix( out_face_bgr * wrk_face_mask_area_a, dst_face_bgr * wrk_face_mask_area_a) if cfg.mode == 'seamless-hist-match': out_face_bgr = imagelib.color_hist_match( out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) if cfg_mp != 0: k_size = int(frame_info.motion_power * cfg_mp) if k_size >= 1: k_size = np.clip(k_size + 1, 2, 50) if cfg.super_resolution_power != 0: k_size *= 2 out_face_bgr = imagelib.LinearMotionBlur( out_face_bgr, k_size, frame_info.motion_deg) if cfg.blursharpen_amount != 0: out_face_bgr = imagelib.blursharpen( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount) if cfg.image_denoise_power != 0: n = cfg.image_denoise_power while n > 0: img_bgr_denoised = cv2.medianBlur(img_bgr, 5) if int(n / 100) != 0: img_bgr = img_bgr_denoised else: pass_power = (n % 100) / 100.0 img_bgr = img_bgr * ( 1.0 - pass_power ) + img_bgr_denoised * pass_power n = max(n - 10, 0) if cfg.bicubic_degrade_power != 0: p = 1.0 - cfg.bicubic_degrade_power / 101.0 img_bgr_downscaled = cv2.resize( img_bgr, (int(img_size[0] * p), int(img_size[1] * p)), interpolation=cv2.INTER_CUBIC) img_bgr = cv2.resize(img_bgr_downscaled, img_size, interpolation=cv2.INTER_CUBIC) new_out = cv2.warpAffine( out_face_bgr, face_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_img = np.clip( img_bgr * (1 - img_face_mask_a) + (new_out * img_face_mask_a), 0, 1.0) if cfg.color_degrade_power != 0: out_img_reduced = imagelib.reduce_colors(out_img, 256) if cfg.color_degrade_power == 100: out_img = out_img_reduced else: alpha = cfg.color_degrade_power / 100.0 out_img = (out_img * (1.0 - alpha) + out_img_reduced * alpha) out_merging_mask_a = img_face_mask_a if out_img is None: out_img = img_bgr.copy() if cfg.show_mode == 3: g = dflimg.get_source_rect() a, b = 30, int(30 * (g[3] - g[1]) / (g[2] - g[0])) out = out_img[g[1] - int(1.5 * a):g[3] + a, g[0] - b:g[2] + b] if cfg.show_mode == 4: g = dflimg.get_source_rect() a, b = 30, int(30 * (g[3] - g[1]) / (g[2] - g[0])) out = img_bgr[g[1] - int(1.5 * a):g[3] + a, g[0] - b:g[2] + b] if cfg.show_mode == 2: out = img_bgr if cfg.show_mode == 1: out = out_img return out
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 convert_face (self, img_bgr, img_face_landmarks, debug): if self.over_res != 1: img_bgr = cv2.resize ( img_bgr, ( img_bgr.shape[1]*self.over_res, img_bgr.shape[0]*self.over_res ) ) img_face_landmarks = img_face_landmarks*self.over_res if debug: debugs = [img_bgr.copy()] img_size = img_bgr.shape[1], img_bgr.shape[0] img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type) face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type, scale=self.output_face_scale) dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size)) predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1) predicted_bgra = self.predictor_func ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) ) prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 ) prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0) if self.mask_mode == 2: #dst prd_face_mask_a_0 = predictor_input_mask_a_0 elif self.mask_mode == 3: #FAN-prd prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( np.expand_dims(prd_face_bgr_256,0) ) [0] prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (self.predictor_input_size, self.predictor_input_size)) elif self.mask_mode == 4: #FAN-dst face_256_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL) dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 ) dst_face_256_mask = self.fan_seg.extract_from_bgr( np.expand_dims(dst_face_256_bgr,0) ) [0] prd_face_mask_a_0 = cv2.resize (dst_face_256_mask, (self.predictor_input_size, self.predictor_input_size)) prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 prd_face_mask_a = np.expand_dims (prd_face_mask_a_0, axis=-1) prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1) img_face_mask_aaa = cv2.warpAffine( prd_face_mask_aaa, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) img_face_mask_aaa = np.clip (img_face_mask_aaa, 0.0, 1.0) img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0 #get rid of noise if debug: debugs += [img_face_mask_aaa.copy()] if 'seamless' in self.mode: #mask used for cv2.seamlessClone img_face_seamless_mask_aaa = None for i in range(9, 0, -1): a = img_face_mask_aaa > i / 10.0 if len(np.argwhere(a)) == 0: continue img_face_seamless_mask_aaa = img_face_mask_aaa.copy() img_face_seamless_mask_aaa[a] = 1.0 img_face_seamless_mask_aaa[img_face_seamless_mask_aaa <= i / 10.0] = 0.0 out_img = img_bgr.copy() if self.mode == 'raw': if self.raw_mode == 'rgb' or self.raw_mode == 'rgb-mask': out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) if self.raw_mode == 'rgb-mask': out_img = np.concatenate ( [out_img, np.expand_dims (img_face_mask_aaa[:,:,0],-1)], -1 ) if self.raw_mode == 'mask-only': out_img = img_face_mask_aaa if self.raw_mode == 'predicted-only': out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(out_img.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) elif ('seamless' not in self.mode) or (img_face_seamless_mask_aaa is not None): #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask ar = [] for i in range(1, 10): maxregion = np.argwhere( img_face_mask_aaa > i / 10.0 ) if maxregion.size != 0: miny,minx = maxregion.min(axis=0)[:2] maxy,maxx = maxregion.max(axis=0)[:2] lenx = maxx - minx leny = maxy - miny maskx = ( minx+(lenx/2) ) masky = ( miny+(leny/2) ) if lenx >= 4 and leny >= 4: ar += [ [ lenx, leny, maskx, masky] ] if len(ar) > 0: lenx, leny, maskx, masky = np.mean ( ar, axis=0 ) if debug: io.log_info ("lenx/leny:(%d/%d) maskx/masky:(%f/%f)" % (lenx, leny, maskx, masky ) ) maskx = int( maskx ) masky = int( masky ) lowest_len = min (lenx, leny) if debug: io.log_info ("lowest_len = %f" % (lowest_len) ) img_mask_blurry_aaa = img_face_mask_aaa if self.erode_mask_modifier != 0: ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier ) if debug: io.log_info ("erode_size = %d" % (ero) ) if ero > 0: img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) elif ero < 0: img_mask_blurry_aaa = cv2.dilate(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) if self.seamless_erode_mask_modifier != 0: ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.seamless_erode_mask_modifier ) if debug: io.log_info ("seamless_erode_size = %d" % (ero) ) if ero > 0: img_face_seamless_mask_aaa = cv2.erode(img_face_seamless_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) elif ero < 0: img_face_seamless_mask_aaa = cv2.dilate(img_face_seamless_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) img_face_seamless_mask_aaa = np.clip (img_face_seamless_mask_aaa, 0, 1) if self.clip_hborder_mask_per > 0: #clip hborder before blur prd_hborder_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=np.float32) prd_border_size = int ( prd_hborder_rect_mask_a.shape[1] * self.clip_hborder_mask_per ) prd_hborder_rect_mask_a[:,0:prd_border_size,:] = 0 prd_hborder_rect_mask_a[:,-prd_border_size:,:] = 0 prd_hborder_rect_mask_a = np.expand_dims(cv2.blur(prd_hborder_rect_mask_a, (prd_border_size, prd_border_size) ),-1) img_prd_hborder_rect_mask_a = cv2.warpAffine( prd_hborder_rect_mask_a, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) img_prd_hborder_rect_mask_a = np.expand_dims (img_prd_hborder_rect_mask_a, -1) img_mask_blurry_aaa *= img_prd_hborder_rect_mask_a img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) if debug: debugs += [img_mask_blurry_aaa.copy()] if self.blur_mask_modifier > 0: blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier ) if debug: io.log_info ("blur_size = %d" % (blur) ) if blur > 0: img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) ) img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) if debug: debugs += [img_mask_blurry_aaa.copy()] if self.color_transfer_mode is not None: if self.color_transfer_mode == 'rct': if debug: debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] prd_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (prd_face_bgr*255).astype(np.uint8), 0, 255), np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), source_mask=prd_face_mask_a, target_mask=prd_face_mask_a) prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) if debug: debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] elif self.color_transfer_mode == 'lct': if debug: debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr) prd_face_bgr = np.clip( prd_face_bgr, 0.0, 1.0) if debug: debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] if self.mode == 'hist-match-bw': prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 ) if self.mode == 'hist-match' or self.mode == 'hist-match-bw': if debug: debugs += [ cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) ] hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) if self.masked_hist_match: hist_mask_a *= prd_face_mask_a hist_match_1 = prd_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) hist_match_1[ hist_match_1 > 1.0 ] = 1.0 hist_match_2 = dst_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) hist_match_2[ hist_match_1 > 1.0 ] = 1.0 prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold ) if self.mode == 'hist-match-bw': prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) out_img = np.clip(out_img, 0.0, 1.0) if debug: debugs += [out_img.copy()] if self.mode == 'overlay': pass if 'seamless' in self.mode: try: out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_seamless_mask_aaa*255).astype(np.uint8), (maskx,masky) , cv2.NORMAL_CLONE ) out_img = out_img.astype(dtype=np.float32) / 255.0 except Exception as e: #seamlessClone may fail in some cases e_str = traceback.format_exc() if 'MemoryError' in e_str: raise Exception("Seamless fail: " + e_str) #reraise MemoryError in order to reprocess this data by other processes else: print ("Seamless fail: " + e_str) if debug: debugs += [out_img.copy()] out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 ) if self.mode == 'seamless-hist-match': out_face_bgr = cv2.warpAffine( out_img, face_mat, (self.output_size, self.output_size) ) new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold) new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) if self.final_image_color_degrade_power != 0: if debug: debugs += [out_img.copy()] out_img_reduced = imagelib.reduce_colors(out_img, 256) if self.final_image_color_degrade_power == 100: out_img = out_img_reduced else: alpha = self.final_image_color_degrade_power / 100.0 out_img = (out_img*(1.0-alpha) + out_img_reduced*alpha) if self.alpha: out_img = np.concatenate ( [out_img, np.expand_dims (img_mask_blurry_aaa[:,:,0],-1)], -1 ) if self.over_res != 1: out_img = cv2.resize ( out_img, ( img_bgr.shape[1] // self.over_res, img_bgr.shape[0] // self.over_res ) ) out_img = np.clip (out_img, 0.0, 1.0 ) if debug: debugs += [out_img.copy()] return debugs if debug else out_img
def batch_func(self, generator_id): gen_sq = self.generators_sq[generator_id] if self.generators_random_seed is not None: np.random.seed(self.generators_random_seed[generator_id]) samples = self.samples samples_len = len(samples) samples_idxs = [*range(samples_len) ][generator_id::self.generators_count] repeat_samples_idxs = [] if len(samples_idxs) == 0: raise ValueError('No training data provided.') if self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: if all([samples[idx] == None for idx in samples_idxs]): raise ValueError( 'Not enough training data. Gather more faces!') if self.sample_type == SampleType.FACE: shuffle_idxs = [] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: shuffle_idxs = [] shuffle_idxs_2D = [[]] * samples_len while True: while not gen_sq.empty(): idxs = gen_sq.get() for idx in idxs: if idx in samples_idxs: repeat_samples_idxs.append(idx) batches = None for n_batch in range(self.batch_size): while True: sample = None if len(repeat_samples_idxs) > 0: idx = repeat_samples_idxs.pop() if self.sample_type == SampleType.FACE: sample = samples[idx] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF] else: if self.sample_type == SampleType.FACE: if len(shuffle_idxs) == 0: shuffle_idxs = samples_idxs.copy() np.random.shuffle(shuffle_idxs) idx = shuffle_idxs.pop() sample = samples[idx] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: if len(shuffle_idxs) == 0: shuffle_idxs = samples_idxs.copy() np.random.shuffle(shuffle_idxs) idx = shuffle_idxs.pop() if samples[idx] != None: if len(shuffle_idxs_2D[idx]) == 0: a = shuffle_idxs_2D[idx] = [ *range(len(samples[idx])) ] np.random.shuffle(a) idx2 = shuffle_idxs_2D[idx].pop() sample = samples[idx][idx2] idx = (idx << 16) | (idx2 & 0xFFFF) if sample is not None: try: x = SampleProcessor.process( sample, self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception( "Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc())) if type(x) != tuple and type(x) != list: raise Exception( 'SampleProcessor.process returns NOT tuple/list' ) if batches is None: batches = [[] for _ in range(len(x))] if self.add_sample_idx: batches += [[]] i_sample_idx = len(batches) - 1 if self.add_pitch: batches += [[]] i_pitch = len(batches) - 1 if self.add_yaw: batches += [[]] i_yaw = len(batches) - 1 for i in range(len(x)): batches[i].append(x[i]) if self.add_sample_idx: batches[i_sample_idx].append(idx) if self.add_pitch or self.add_yaw: pitch, yaw = LandmarksProcessor.estimate_pitch_yaw( sample.landmarks) if self.add_pitch: batches[i_pitch].append([pitch]) if self.add_yaw: batches[i_yaw].append([yaw]) break yield [np.array(batch) for batch in batches]
def main(model_class_name=None, saved_models_path=None, training_data_src_path=None, force_model_name=None, input_path=None, output_path=None, output_mask_path=None, aligned_path=None, force_gpu_idxs=None, cpu_only=None): io.log_info("Running merger.\r\n") try: if not input_path.exists(): io.log_err('Input directory not found. Please ensure it exists.') return if not output_path.exists(): output_path.mkdir(parents=True, exist_ok=True) if not output_mask_path.exists(): output_mask_path.mkdir(parents=True, exist_ok=True) if not saved_models_path.exists(): io.log_err('Model directory not found. Please ensure it exists.') return # Initialize model import models model = models.import_model(model_class_name)( is_training=False, saved_models_path=saved_models_path, force_gpu_idxs=force_gpu_idxs, cpu_only=cpu_only) predictor_func, predictor_input_shape, cfg = model.get_MergerConfig() # Preparing MP functions predictor_func = MPFunc(predictor_func) run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0 xseg_256_extract_func = MPClassFuncOnDemand( XSegNet, 'extract', name='XSeg', resolution=256, weights_file_root=saved_models_path, place_model_on_cpu=True, run_on_cpu=run_on_cpu) face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer, 'enhance', place_model_on_cpu=True, run_on_cpu=run_on_cpu) is_interactive = io.input_bool("Use interactive merger?", True) if not io.is_colab() else False if not is_interactive: cfg.ask_settings() input_path_image_paths = pathex.get_image_paths(input_path) if cfg.type == MergerConfig.TYPE_MASKED: if not aligned_path.exists(): io.log_err( 'Aligned directory not found. Please ensure it exists.') return packed_samples = None try: packed_samples = samplelib.PackedFaceset.load(aligned_path) except: io.log_err( f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}" ) if packed_samples is not None: io.log_info("Using packed faceset.") def generator(): for sample in io.progress_bar_generator( packed_samples, "Collecting alignments"): filepath = Path(sample.filename) yield filepath, DFLIMG.load( filepath, loader_func=lambda x: sample.read_raw_file()) else: def generator(): for filepath in io.progress_bar_generator( pathex.get_image_paths(aligned_path), "Collecting alignments"): filepath = Path(filepath) yield filepath, DFLIMG.load(filepath) alignments = {} multiple_faces_detected = False for filepath, dflimg in generator(): if dflimg is None or not dflimg.has_data(): io.log_err(f"{filepath.name} is not a dfl image file") continue source_filename = dflimg.get_source_filename() if source_filename is None: continue source_filepath = Path(source_filename) source_filename_stem = source_filepath.stem if source_filename_stem not in alignments.keys(): alignments[source_filename_stem] = [] alignments_ar = alignments[source_filename_stem] alignments_ar.append( (dflimg.get_source_landmarks(), filepath, source_filepath)) if len(alignments_ar) > 1: multiple_faces_detected = True if multiple_faces_detected: io.log_info("") io.log_info( "Warning: multiple faces detected. Only one alignment file should refer one source file." ) io.log_info("") for a_key in list(alignments.keys()): a_ar = alignments[a_key] if len(a_ar) > 1: for _, filepath, source_filepath in a_ar: io.log_info( f"alignment {filepath.name} refers to {source_filepath.name} " ) io.log_info("") alignments[a_key] = [a[0] for a in a_ar] if multiple_faces_detected: io.log_info( "It is strongly recommended to process the faces separatelly." ) io.log_info( "Use 'recover original filename' to determine the exact duplicates." ) io.log_info("") frames = [ InteractiveMergerSubprocessor.Frame(frame_info=FrameInfo( filepath=Path(p), landmarks_list=alignments.get(Path(p).stem, None))) for p in input_path_image_paths ] if multiple_faces_detected: io.log_info( "Warning: multiple faces detected. Motion blur will not be used." ) io.log_info("") else: s = 256 local_pts = [(s // 2 - 1, s // 2 - 1), (s // 2 - 1, 0)] #center+up frames_len = len(frames) for i in io.progress_bar_generator(range(len(frames)), "Computing motion vectors"): fi_prev = frames[max(0, i - 1)].frame_info fi = frames[i].frame_info fi_next = frames[min(i + 1, frames_len - 1)].frame_info if len(fi_prev.landmarks_list) == 0 or \ len(fi.landmarks_list) == 0 or \ len(fi_next.landmarks_list) == 0: continue mat_prev = LandmarksProcessor.get_transform_mat( fi_prev.landmarks_list[0], s, face_type=FaceType.FULL) mat = LandmarksProcessor.get_transform_mat( fi.landmarks_list[0], s, face_type=FaceType.FULL) mat_next = LandmarksProcessor.get_transform_mat( fi_next.landmarks_list[0], s, face_type=FaceType.FULL) pts_prev = LandmarksProcessor.transform_points( local_pts, mat_prev, True) pts = LandmarksProcessor.transform_points( local_pts, mat, True) pts_next = LandmarksProcessor.transform_points( local_pts, mat_next, True) prev_vector = pts[0] - pts_prev[0] next_vector = pts_next[0] - pts[0] motion_vector = pts_next[0] - pts_prev[0] fi.motion_power = npla.norm(motion_vector) motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array( [0, 0], dtype=np.float32) fi.motion_deg = -math.atan2( motion_vector[1], motion_vector[0]) * 180 / math.pi if len(frames) == 0: io.log_info("No frames to merge in input_dir.") else: if False: pass else: InteractiveMergerSubprocessor( is_interactive=is_interactive, merger_session_filepath=model.get_strpath_storage_for_file( 'merger_session.dat'), predictor_func=predictor_func, predictor_input_shape=predictor_input_shape, face_enhancer_func=face_enhancer_func, xseg_256_extract_func=xseg_256_extract_func, merger_config=cfg, frames=frames, frames_root_path=input_path, output_path=output_path, output_mask_path=output_mask_path, model_iter=model.get_iter()).run() model.finalize() except Exception as e: print(traceback.format_exc())
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) 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 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) 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) # 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)) 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 random_bilinear_resize is not None: l_rnd_state = np.random.RandomState( sample_rnd_seed + 2) chance, max_size_per = random_bilinear_resize chance = np.clip(chance, 0, 100) pick_chance = l_rnd_state.randint(100) resize_to = resolution - int( l_rnd_state.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) # 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() / 2 - 0.25, 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)) 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 convert_face (self, img_bgr, img_face_landmarks, debug): if debug: debugs = [img_bgr.copy()] img_size = img_bgr.shape[1], img_bgr.shape[0] img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr, img_face_landmarks) face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type) dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size)) predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1) predicted_bgra = self.predictor ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) ) prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 ) prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0) prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 prd_face_mask_a = np.expand_dims (prd_face_mask_a_0, axis=-1) prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1) img_prd_face_mask_aaa = cv2.warpAffine( prd_face_mask_aaa, face_mat, img_size, np.zeros(img_bgr.shape, dtype=float), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) img_prd_face_mask_aaa = np.clip (img_prd_face_mask_aaa, 0.0, 1.0) img_face_mask_aaa = img_prd_face_mask_aaa if debug: debugs += [img_face_mask_aaa.copy()] img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0 img_face_mask_flatten_aaa = img_face_mask_aaa.copy() img_face_mask_flatten_aaa[img_face_mask_flatten_aaa > 0.9] = 1.0 maxregion = np.argwhere(img_face_mask_flatten_aaa==1.0) out_img = img_bgr.copy() if maxregion.size != 0: miny,minx = maxregion.min(axis=0)[:2] maxy,maxx = maxregion.max(axis=0)[:2] lenx = maxx - minx leny = maxy - miny masky = int(minx+(lenx//2)) maskx = int(miny+(leny//2)) lowest_len = min (lenx, leny) if debug: print ("lowest_len = %f" % (lowest_len) ) ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier ) blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier ) if debug: print ("ero = %d, blur = %d" % (ero, blur) ) img_mask_blurry_aaa = img_face_mask_aaa if self.erode_mask: if ero > 0: img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) elif ero < 0: img_mask_blurry_aaa = cv2.dilate(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) if self.blur_mask and blur > 0: img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) ) img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) if self.clip_border_mask_per > 0: prd_border_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=prd_face_mask_a.dtype) prd_border_size = int ( prd_border_rect_mask_a.shape[1] * self.clip_border_mask_per ) prd_border_rect_mask_a[0:prd_border_size,:,:] = 0 prd_border_rect_mask_a[-prd_border_size:,:,:] = 0 prd_border_rect_mask_a[:,0:prd_border_size,:] = 0 prd_border_rect_mask_a[:,-prd_border_size:,:] = 0 prd_border_rect_mask_a = np.expand_dims(cv2.blur(prd_border_rect_mask_a, (prd_border_size, prd_border_size) ),-1) if self.mode == 'hist-match-bw': prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 ) if self.mode == 'hist-match' or self.mode == 'hist-match-bw': if debug: debugs += [ cv2.warpAffine( prd_face_bgr, face_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) ] hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=prd_face_bgr.dtype) if self.masked_hist_match: hist_mask_a *= prd_face_mask_a new_prd_face_bgr = image_utils.color_hist_match(prd_face_bgr*hist_mask_a, dst_face_bgr*hist_mask_a ) prd_face_bgr = new_prd_face_bgr if self.mode == 'hist-match-bw': prd_face_bgr = prd_face_bgr.astype(np.float32) out_img = cv2.warpAffine( prd_face_bgr, face_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) if debug: debugs += [out_img.copy()] debugs += [img_mask_blurry_aaa.copy()] if self.mode == 'seamless' or self.mode == 'seamless-hist-match': out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 ) if debug: debugs += [out_img.copy()] out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_mask_flatten_aaa*255).astype(np.uint8), (masky,maskx) , cv2.NORMAL_CLONE ) out_img = out_img.astype(np.float32) / 255.0 if debug: debugs += [out_img.copy()] if self.clip_border_mask_per > 0: img_prd_border_rect_mask_a = cv2.warpAffine( prd_border_rect_mask_a, face_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) img_prd_border_rect_mask_a = np.expand_dims (img_prd_border_rect_mask_a, -1) out_img = out_img * img_prd_border_rect_mask_a + img_bgr * (1.0 - img_prd_border_rect_mask_a) img_mask_blurry_aaa *= img_prd_border_rect_mask_a out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 ) if self.mode == 'seamless-hist-match': out_face_bgr = cv2.warpAffine( out_img, face_mat, (self.output_size, self.output_size) ) new_out_face_bgr = image_utils.color_hist_match(out_face_bgr, dst_face_bgr ) new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) if debug: debugs += [out_img.copy()] return debugs if debug else out_img
def process_data(self, data): filename_path = Path( data[0] ) filename_path_str = str(filename_path) if self.cached_image[0] == filename_path_str: image = self.cached_image[1] #cached image for manual extractor else: image = cv2_imread( filename_path_str ) if image is None: self.log_err ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) ) return None image_shape = image.shape if len(image_shape) == 2: h, w = image.shape ch = 1 else: h, w, ch = image.shape if ch == 1: image = np.repeat ( image [:,:,np.newaxis], 3, -1 ) elif ch == 4: image = image[:,:,0:3] wm = w % 2 hm = h % 2 if wm + hm != 0: #fix odd image image = image[0:h-hm,0:w-wm,:] self.cached_image = ( filename_path_str, image ) src_dflimg = None h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? if filename_path.suffix == '.jpg': src_dflimg = DFLJPG.load ( str(filename_path) ) if self.type == 'rects': if min(w,h) < 128: self.log_err ( 'Image is too small %s : [%d, %d]' % ( str(filename_path), w, h ) ) rects = [] else: rects = self.e.extract_from_bgr (image) return [str(filename_path), rects] elif self.type == 'landmarks': rects = data[1] if rects is None: landmarks = None else: landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None) return [str(filename_path), landmarks] elif self.type == 'final': result = [] faces = data[1] if self.debug_dir is not None: debug_output_file = str( Path(self.debug_dir) / (filename_path.stem+'.jpg') ) debug_image = image.copy() if src_dflimg is not None and len(faces) != 1: #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it print("src_dflimg is not None and len(faces) != 1", str(filename_path) ) output_file = str(self.output_path / filename_path.name) if str(filename_path) != str(output_file): shutil.copy ( str(filename_path), str(output_file) ) result.append (output_file) else: face_idx = 0 for face in faces: rect = np.array(face[0]) image_landmarks = face[1] if image_landmarks is None: continue image_landmarks = np.array(image_landmarks) if self.face_type == FaceType.MARK_ONLY: face_image = image face_image_landmarks = image_landmarks else: image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type) face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4) face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat) landmarks_bbox = LandmarksProcessor.transform_points ( [ (0,0), (0,self.image_size-1), (self.image_size-1, self.image_size-1), (self.image_size-1,0) ], image_to_face_mat, True) rect_area = mathlib.polygon_area(np.array(rect[[0,2,2,0]]), np.array(rect[[1,1,3,3]])) landmarks_area = mathlib.polygon_area(landmarks_bbox[:,0], landmarks_bbox[:,1] ) if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area continue if self.debug_dir is not None: LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True) if src_dflimg is not None: #if extracting from dflimg copy it in order not to lose quality output_file = str(self.output_path / filename_path.name) if str(filename_path) != str(output_file): shutil.copy ( str(filename_path), str(output_file) ) else: output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.jpg') cv2_imwrite(output_file, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), 85] ) DFLJPG.embed_data(output_file, face_type=FaceType.toString(self.face_type), landmarks=face_image_landmarks.tolist(), source_filename=filename_path.name, source_rect=rect, source_landmarks=image_landmarks.tolist(), image_to_face_mat=image_to_face_mat ) result.append (output_file) face_idx += 1 if self.debug_dir is not None: cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) return result
def apply_celebamaskhq(input_dir): input_path = Path(input_dir) img_path = input_path / 'aligned' mask_path = input_path / 'mask' if not img_path.exists(): raise ValueError( f'{str(img_path)} directory not found. Please ensure it exists.') CelebAMASKHQSubprocessor(pathex.get_image_paths(img_path), pathex.get_image_paths(mask_path, subdirs=True)).run() return paths_to_extract = [] for filename in io.progress_bar_generator(pathex.get_image_paths(img_path), desc="Processing"): filepath = Path(filename) dflimg = DFLIMG.load(filepath) if dflimg is not None: paths_to_extract.append(filepath) image_to_face_mat = dflimg.get_image_to_face_mat() src_filename = dflimg.get_source_filename() #img = cv2_imread(filename) h, w, c = dflimg.get_shape() fanseg_mask = LandmarksProcessor.get_image_hull_mask( (h, w, c), dflimg.get_landmarks()) idx_name = '%.5d' % int(src_filename.split('.')[0]) idx_files = [x for x in masks_files if idx_name in x] skin_files = [x for x in idx_files if 'skin' in x] eye_glass_files = [x for x in idx_files if 'eye_g' in x] for files, is_invert in [(skin_files, False), (eye_glass_files, True)]: if len(files) > 0: mask = cv2_imread(files[0]) mask = mask[..., 0] mask[mask == 255] = 1 mask = mask.astype(np.float32) mask = cv2.resize(mask, (1024, 1024)) mask = cv2.warpAffine(mask, image_to_face_mat, (w, h), cv2.INTER_LANCZOS4) if not is_invert: fanseg_mask *= mask[..., None] else: fanseg_mask *= (1 - mask[..., None]) #cv2.imshow("", (fanseg_mask*255).astype(np.uint8) ) #cv2.waitKey(0) dflimg.embed_and_set(filename, fanseg_mask=fanseg_mask)
def process(sample, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types sample_bgr = sample.load_bgr() ct_sample_bgr = None ct_sample_mask = 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) cached_images = collections.defaultdict(dict) sample_rnd_seed = np.random.randint(0x80000000) SPTF_FACETYPE_TO_FACETYPE = { SPTF.FACE_TYPE_HALF: FaceType.HALF, SPTF.FACE_TYPE_FULL: FaceType.FULL, SPTF.FACE_TYPE_HEAD: FaceType.HEAD, SPTF.FACE_TYPE_AVATAR: FaceType.AVATAR } outputs = [] for opts in output_sample_types: resolution = opts.get('resolution', 0) types = opts.get('types', []) random_sub_res = opts.get('random_sub_res', 0) normalize_std_dev = opts.get('normalize_std_dev', False) normalize_vgg = opts.get('normalize_vgg', False) motion_blur = opts.get('motion_blur', None) apply_ct = opts.get('apply_ct', False) normalize_tanh = opts.get('normalize_tanh', False) img_type = SPTF.NONE target_face_type = SPTF.NONE face_mask_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 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) img = l elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch_yaw_roll = sample.pitch_yaw_roll if pitch_yaw_roll is not None: pitch, yaw, roll = pitch_yaw_roll else: pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll( sample.landmarks) if params['flip']: yaw = -yaw if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch = (pitch + 1.0) / 2.0 yaw = (yaw + 1.0) / 2.0 roll = (roll + 1.0) / 2.0 img = (pitch, yaw, roll) else: if mode_type == SPTF.NONE: raise ValueError('expected MODE_ type') img = cached_images.get(img_type, None) if img is None: img = sample_bgr mask = None cur_sample = sample if is_face_sample: if motion_blur is not None: chance, mb_range = motion_blur chance = np.clip(chance, 0, 100) if np.random.randint(100) < chance: mb_range = [3, 5, 7, 9][:np.clip(mb_range, 0, 3) + 1] dim = mb_range[np.random.randint( len(mb_range))] img = imagelib.LinearMotionBlur( img, dim, np.random.randint(180)) mask = cur_sample.load_fanseg_mask( ) #using fanseg_mask if exist if mask is None: mask = LandmarksProcessor.get_image_hull_mask( img.shape, cur_sample.landmarks) if cur_sample.ie_polys is not None: cur_sample.ie_polys.overlay_mask(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, True) if mask is not None: mask = imagelib.warp_by_params(params, mask, warp, transform, flip, False)[..., np.newaxis] img = np.concatenate((img, mask), -1) cached_images[img_type] = img if is_face_sample and target_face_type != SPTF.NONE: ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type] if 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, ft)) img = cv2.warpAffine(img, LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, ft), (resolution, resolution), flags=cv2.INTER_CUBIC) else: img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) if random_sub_res != 0: sub_size = resolution - random_sub_res rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_res) start_x = rnd_state.randint(sub_size + 1) start_y = rnd_state.randint(sub_size + 1) img = img[start_y:start_y + sub_size, start_x:start_x + sub_size, :] img = np.clip(img, 0, 1) img_bgr = img[..., 0:3] img_mask = img[..., 3:4] if apply_ct: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR) img_bgr = imagelib.linear_color_transfer( img_bgr, ct_sample_bgr_resized) img_bgr = np.clip(img_bgr, 0.0, 1.0) if normalize_std_dev: img_bgr = (img_bgr - img_bgr.mean((0, 1))) / img_bgr.std( (0, 1)) elif normalize_vgg: img_bgr = np.clip(img_bgr * 255, 0, 255) img_bgr[:, :, 0] -= 103.939 img_bgr[:, :, 1] -= 116.779 img_bgr[:, :, 2] -= 123.68 if mode_type == SPTF.MODE_BGR: img = img_bgr elif mode_type == SPTF.MODE_BGR_SHUFFLE: rnd_state = np.random.RandomState(sample_rnd_seed) img = np.take(img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1) elif mode_type == SPTF.MODE_G: img = np.concatenate((np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), img_mask), -1) elif mode_type == SPTF.MODE_GGG: img = np.concatenate((np.repeat( np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), (3, ), -1), img_mask), -1) elif mode_type == SPTF.MODE_M and is_face_sample: img = img_mask if not debug: if normalize_tanh: img = np.clip(img * 2.0 - 1.0, -1.0, 1.0) else: img = np.clip(img, 0.0, 1.0) outputs.append(img) if debug: result = [] for output in outputs: if output.shape[2] < 4: result += [ output, ] elif output.shape[2] == 4: result += [ output[..., 0:3] * output[..., 3:4], ] return result else: return outputs
def dev_test_68(input_dir): # process 68 landmarks dataset with .pts files input_path = Path(input_dir) if not input_path.exists(): raise ValueError('input_dir not found. Please ensure it exists.') output_path = input_path.parent / (input_path.name + '_aligned') io.log_info(f'Output dir is % {output_path}') if output_path.exists(): output_images_paths = pathex.get_image_paths(output_path) if len(output_images_paths) > 0: io.input_bool( "WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False) for filename in output_images_paths: Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) images_paths = pathex.get_image_paths(input_path) for filepath in io.progress_bar_generator(images_paths, "Processing"): filepath = Path(filepath) pts_filepath = filepath.parent / (filepath.stem + '.pts') if pts_filepath.exists(): pts = pts_filepath.read_text() pts_lines = pts.split('\n') lmrk_lines = None for pts_line in pts_lines: if pts_line == '{': lmrk_lines = [] elif pts_line == '}': break else: if lmrk_lines is not None: lmrk_lines.append(pts_line) if lmrk_lines is not None and len(lmrk_lines) == 68: try: lmrks = [ np.array(lmrk_line.strip().split(' ')).astype( np.float32).tolist() for lmrk_line in lmrk_lines ] except Exception as e: print(e) print(filepath) continue rect = LandmarksProcessor.get_rect_from_landmarks(lmrks) output_filepath = output_path / (filepath.stem + '.jpg') img = cv2_imread(filepath) img = imagelib.normalize_channels(img, 3) cv2_imwrite(output_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) DFLJPG.embed_data(output_filepath, face_type=FaceType.toString( FaceType.MARK_ONLY), landmarks=lmrks, source_filename=filepath.name, source_rect=rect, source_landmarks=lmrks) io.log_info("Done.")
def process(sample, sample_process_options, output_sample_types, debug): sample_bgr = sample.load_bgr() 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)) close_sample = sample.close_target_list[np.random.randint( 0, len(sample.close_target_list) )] if sample.close_target_list is not None else None close_sample_bgr = close_sample.load_bgr( ) if close_sample is not None else None if debug and close_sample_bgr is not None: LandmarksProcessor.draw_landmarks(close_sample_bgr, close_sample.landmarks, (0, 1, 0)) params = image_utils.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) images = [[None] * 3 for _ in range(30)] sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for sample_type in output_sample_types: f = sample_type[0] size = sample_type[1] random_sub_size = 0 if len(sample_type) < 3 else min( sample_type[2], size) if f & SampleProcessor.TypeFlags.SOURCE != 0: img_type = 0 elif f & SampleProcessor.TypeFlags.WARPED != 0: img_type = 1 elif f & SampleProcessor.TypeFlags.WARPED_TRANSFORMED != 0: img_type = 2 elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0: img_type = 3 elif f & SampleProcessor.TypeFlags.LANDMARKS_ARRAY != 0: img_type = 4 else: raise ValueError('expected SampleTypeFlags type') if f & SampleProcessor.TypeFlags.RANDOM_CLOSE != 0: img_type += 10 elif f & SampleProcessor.TypeFlags.MORPH_TO_RANDOM_CLOSE != 0: img_type += 20 face_mask_type = 0 if f & SampleProcessor.TypeFlags.FACE_MASK_FULL != 0: face_mask_type = 1 elif f & SampleProcessor.TypeFlags.FACE_MASK_EYES != 0: face_mask_type = 2 target_face_type = -1 if f & SampleProcessor.TypeFlags.FACE_ALIGN_HALF != 0: target_face_type = FaceType.HALF elif f & SampleProcessor.TypeFlags.FACE_ALIGN_FULL != 0: target_face_type = FaceType.FULL elif f & SampleProcessor.TypeFlags.FACE_ALIGN_HEAD != 0: target_face_type = FaceType.HEAD elif f & SampleProcessor.TypeFlags.FACE_ALIGN_AVATAR != 0: target_face_type = FaceType.AVATAR if img_type == 4: 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) img = l else: if images[img_type][face_mask_type] is None: if img_type >= 10 and img_type <= 19: #RANDOM_CLOSE img_type -= 10 img = close_sample_bgr cur_sample = close_sample elif img_type >= 20 and img_type <= 29: #MORPH_TO_RANDOM_CLOSE img_type -= 20 res = sample.shape[0] s_landmarks = sample.landmarks.copy() d_landmarks = close_sample.landmarks.copy() idxs = list(range(len(s_landmarks))) #remove landmarks near boundaries for i in idxs[:]: s_l = s_landmarks[i] d_l = d_landmarks[i] if s_l[0] < 5 or s_l[1] < 5 or s_l[0] >= res-5 or s_l[1] >= res-5 or \ d_l[0] < 5 or d_l[1] < 5 or d_l[0] >= res-5 or d_l[1] >= res-5: idxs.remove(i) #remove landmarks that close to each other in 5 dist for landmarks in [s_landmarks, d_landmarks]: for i in idxs[:]: s_l = landmarks[i] for j in idxs[:]: if i == j: continue s_l_2 = landmarks[j] diff_l = np.abs(s_l - s_l_2) if np.sqrt(diff_l.dot(diff_l)) < 5: idxs.remove(i) break s_landmarks = s_landmarks[idxs] d_landmarks = d_landmarks[idxs] s_landmarks = np.concatenate([ s_landmarks, [[0, 0], [res // 2, 0], [res - 1, 0], [0, res // 2], [res - 1, res // 2], [0, res - 1], [res // 2, res - 1], [res - 1, res - 1]] ]) d_landmarks = np.concatenate([ d_landmarks, [[0, 0], [res // 2, 0], [res - 1, 0], [0, res // 2], [res - 1, res // 2], [0, res - 1], [res // 2, res - 1], [res - 1, res - 1]] ]) img = image_utils.morph_by_points( sample_bgr, s_landmarks, d_landmarks) cur_sample = close_sample else: img = sample_bgr cur_sample = sample if is_face_sample: if face_mask_type == 1: img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask( img.shape, cur_sample.landmarks)), -1) elif face_mask_type == 2: mask = LandmarksProcessor.get_image_eye_mask( img.shape, cur_sample.landmarks) mask = np.expand_dims( cv2.blur(mask, (w // 32, w // 32)), -1) mask[mask > 0.0] = 1.0 img = np.concatenate((img, mask), -1) images[img_type][ face_mask_type] = image_utils.warp_by_params( params, img, (img_type == 1 or img_type == 2), (img_type == 2 or img_type == 3), img_type != 0, face_mask_type == 0) img = images[img_type][face_mask_type] if is_face_sample and target_face_type != -1: if target_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_face_type)) img = cv2.warpAffine(img, LandmarksProcessor.get_transform_mat( sample.landmarks, size, target_face_type), (size, size), flags=cv2.INTER_CUBIC) else: img = cv2.resize(img, (size, size), cv2.INTER_CUBIC) if random_sub_size != 0: sub_size = size - random_sub_size rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_size) start_x = rnd_state.randint(sub_size + 1) start_y = rnd_state.randint(sub_size + 1) img = img[start_y:start_y + sub_size, start_x:start_x + sub_size, :] img_bgr = img[..., 0:3] img_mask = img[..., 3:4] if f & SampleProcessor.TypeFlags.MODE_BGR != 0: img = img elif f & SampleProcessor.TypeFlags.MODE_BGR_SHUFFLE != 0: img_bgr = np.take(img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1) img = np.concatenate((img_bgr, img_mask), -1) elif f & SampleProcessor.TypeFlags.MODE_G != 0: img = np.concatenate((np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), img_mask), -1) elif f & SampleProcessor.TypeFlags.MODE_GGG != 0: img = np.concatenate((np.repeat( np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), (3, ), -1), img_mask), -1) elif is_face_sample and f & SampleProcessor.TypeFlags.MODE_M != 0: if face_mask_type == 0: raise ValueError('no face_mask_type defined') img = img_mask else: raise ValueError('expected SampleTypeFlags mode') if not debug and sample_process_options.normalize_tanh: img = img * 2.0 - 1.0 outputs.append(img) if debug: result = [] for output in outputs: if output.shape[2] < 4: result += [ output, ] elif output.shape[2] == 4: result += [ output[..., 0:3] * output[..., 3:4], ] return result else: return outputs
def dev_segmented_extract(input_dir, output_dir): # extract and merge .json labelme files within the faces device_config = nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(suggest_all_gpu=True)) input_path = Path(input_dir) if not input_path.exists(): raise ValueError('input_dir not found. Please ensure it exists.') output_path = Path(output_dir) io.log_info("Performing extract segmented faces.") io.log_info(f'Output dir is {output_path}') if output_path.exists(): output_images_paths = pathex.get_image_paths(output_path) if len(output_images_paths) > 0: io.input_bool( "WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False) for filename in output_images_paths: Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) images_paths = pathex.get_image_paths(input_path) extract_data = [] images_jsons = {} images_processed = 0 for filepath in io.progress_bar_generator(images_paths, "Processing"): filepath = Path(filepath) json_filepath = filepath.parent / (filepath.stem + '.json') if json_filepath.exists(): try: json_dict = json.loads(json_filepath.read_text()) images_jsons[filepath] = json_dict total_points = [[x, y] for shape in json_dict['shapes'] for x, y in shape['points']] total_points = np.array(total_points) if len(total_points) == 0: io.log_info( f"No points found in {json_filepath}, skipping.") continue l, r = int(total_points[:, 0].min()), int(total_points[:, 0].max()) t, b = int(total_points[:, 1].min()), int(total_points[:, 1].max()) extract_data.append( ExtractSubprocessor.Data(filepath, rects=[[l, t, r, b]])) images_processed += 1 except: io.log_err(f"err {filepath}") return image_size = 1024 face_type = FaceType.HEAD extract_data = ExtractSubprocessor(extract_data, 'landmarks', image_size, face_type, device_config=device_config).run() extract_data = ExtractSubprocessor(extract_data, 'final', image_size, face_type, final_output_path=output_path, device_config=device_config).run() for data in extract_data: filepath = output_path / (data.filepath.stem + '_0.jpg') dflimg = DFLIMG.load(filepath) image_to_face_mat = dflimg.get_image_to_face_mat() json_dict = images_jsons[data.filepath] ie_polys = IEPolys() for shape in json_dict['shapes']: ie_poly = ie_polys.add(1) points = np.array([[x, y] for x, y in shape['points']]) points = LandmarksProcessor.transform_points( points, image_to_face_mat) for x, y in points: ie_poly.add(int(x), int(y)) dflimg.embed_and_set(filepath, ie_polys=ie_polys) io.log_info(f"Images found: {len(images_paths)}") io.log_info(f"Images processed: {images_processed}")
def process_data(self, filepath): try: dflimg = DFLIMG.load (filepath) if dflimg is None or not dflimg.has_data(): self.log_err (f"{filepath.name} is not a dfl image file") else: img = cv2_imread(filepath) h,w = img.shape[:2] if h != w: raise Exception(f'w != h in {filepath}') image_size = self.image_size face_type = self.face_type output_filepath = self.output_dirpath / filepath.name if face_type is not None: lmrks = dflimg.get_landmarks() mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type) img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 ) img = np.clip(img, 0, 255).astype(np.uint8) cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) dfl_dict = dflimg.get_dict() dflimg = DFLIMG.load (output_filepath) dflimg.set_dict(dfl_dict) xseg_mask = dflimg.get_xseg_mask() if xseg_mask is not None: xseg_res = 256 xseg_lmrks = lmrks.copy() xseg_lmrks *= (xseg_res / w) xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type) xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 ) xseg_mask[xseg_mask < 0.5] = 0 xseg_mask[xseg_mask >= 0.5] = 1 dflimg.set_xseg_mask(xseg_mask) seg_ie_polys = dflimg.get_seg_ie_polys() for poly in seg_ie_polys.get_polys(): poly_pts = poly.get_pts() poly_pts = LandmarksProcessor.transform_points(poly_pts, mat) poly.set_points(poly_pts) dflimg.set_seg_ie_polys(seg_ie_polys) lmrks = LandmarksProcessor.transform_points(lmrks, mat) dflimg.set_landmarks(lmrks) image_to_face_mat = dflimg.get_image_to_face_mat() if image_to_face_mat is not None: image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type ) dflimg.set_image_to_face_mat(image_to_face_mat) dflimg.set_face_type( FaceType.toString(face_type) ) dflimg.save() else: dfl_dict = dflimg.get_dict() scale = w / image_size img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4) cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) dflimg = DFLIMG.load (output_filepath) dflimg.set_dict(dfl_dict) lmrks = dflimg.get_landmarks() lmrks /= scale dflimg.set_landmarks(lmrks) seg_ie_polys = dflimg.get_seg_ie_polys() seg_ie_polys.mult_points( 1.0 / scale) dflimg.set_seg_ie_polys(seg_ie_polys) image_to_face_mat = dflimg.get_image_to_face_mat() if image_to_face_mat is not None: face_type = FaceType.fromString ( dflimg.get_face_type() ) image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type ) dflimg.set_image_to_face_mat(image_to_face_mat) dflimg.save() return (1, filepath, output_filepath) except: self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}") return (0, filepath, None)
def get_eyes_mask(): eyes_mask = LandmarksProcessor.get_image_eye_mask( sample_bgr.shape, sample_landmarks) return np.clip(eyes_mask, 0, 1)
def final_stage( data, image, face_type, image_size, jpeg_quality, extract_from_dflimg=False, output_debug_path=None, final_output_path=None, ): data.final_output_files = [] filepath = data.filepath rects = data.rects landmarks = data.landmarks if output_debug_path is not None: debug_image = image.copy() if extract_from_dflimg and len(rects) != 1: #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it print("extract_from_dflimg and len(rects) != 1", filepath) output_filepath = final_output_path / filepath.name if filepath != str(output_file): shutil.copy(str(filepath), str(output_filepath)) data.final_output_files.append(output_filepath) else: face_idx = 0 for rect, image_landmarks in zip(rects, landmarks): if extract_from_dflimg and face_idx > 1: #cannot extract more than 1 face from dflimg break if image_landmarks is None: continue rect = np.array(rect) if face_type == FaceType.MARK_ONLY: image_to_face_mat = None face_image = image face_image_landmarks = image_landmarks else: image_to_face_mat = LandmarksProcessor.get_transform_mat( image_landmarks, image_size, face_type) face_image = cv2.warpAffine(image, image_to_face_mat, (image_size, image_size), cv2.INTER_LANCZOS4) face_image_landmarks = LandmarksProcessor.transform_points( image_landmarks, image_to_face_mat) landmarks_bbox = LandmarksProcessor.transform_points( [(0, 0), (0, image_size - 1), (image_size - 1, image_size - 1), (image_size - 1, 0)], image_to_face_mat, True) rect_area = mathlib.polygon_area( np.array(rect[[0, 2, 2, 0]]).astype(np.float32), np.array(rect[[1, 1, 3, 3]]).astype(np.float32)) landmarks_area = mathlib.polygon_area( landmarks_bbox[:, 0].astype(np.float32), landmarks_bbox[:, 1].astype(np.float32)) if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4 * rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area continue if output_debug_path is not None: LandmarksProcessor.draw_rect_landmarks( debug_image, rect, image_landmarks, face_type, image_size, transparent_mask=True) output_path = final_output_path if data.force_output_path is not None: output_path = data.force_output_path if extract_from_dflimg and filepath.suffix == '.jpg': #if extracting from dflimg and jpg copy it in order not to lose quality output_filepath = output_path / filepath.name if filepath != output_filepath: shutil.copy(str(filepath), str(output_filepath)) else: output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg" cv2_imwrite( output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]) dflimg = DFLJPG.load(output_filepath) dflimg.set_face_type(FaceType.toString(face_type)) dflimg.set_landmarks(face_image_landmarks.tolist()) dflimg.set_source_filename(filepath.name) dflimg.set_source_rect(rect) dflimg.set_source_landmarks(image_landmarks.tolist()) dflimg.set_image_to_face_mat(image_to_face_mat) dflimg.save() data.final_output_files.append(output_filepath) face_idx += 1 data.faces_detected = face_idx if output_debug_path is not None: cv2_imwrite(output_debug_path / (filepath.stem + '.jpg'), debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50]) return data
def main(input_dir, output_dir): input_path = Path(input_dir) output_path = Path(output_dir) if not input_path.exists(): raise ValueError('Input directory not found. Please ensure it exists.') if not output_path.exists(): output_path.mkdir(parents=True) wnd_name = "Labeling tool" io.named_window(wnd_name) io.capture_mouse(wnd_name) io.capture_keys(wnd_name) #for filename in io.progress_bar_generator (Path_utils.get_image_paths(input_path), desc="Labeling"): for filename in Path_utils.get_image_paths(input_path): filepath = Path(filename) if filepath.suffix == '.png': dflimg = DFLPNG.load(str(filepath)) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load(str(filepath)) else: dflimg = None if dflimg is None: io.log_err("%s is not a dfl image file" % (filepath.name)) continue lmrks = dflimg.get_landmarks() lmrks_list = lmrks.tolist() orig_img = cv2_imread(str(filepath)) h, w, c = orig_img.shape mask_orig = LandmarksProcessor.get_image_hull_mask( orig_img.shape, lmrks).astype(np.uint8)[:, :, 0] ero_dil_rate = w // 8 mask_ero = cv2.erode( mask_orig, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ero_dil_rate, ero_dil_rate)), iterations=1) mask_dil = cv2.dilate(mask_orig, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (ero_dil_rate, ero_dil_rate)), iterations=1) #mask_bg = np.zeros(orig_img.shape[:2],np.uint8) mask_bg = 1 - mask_dil mask_bgp = np.ones(orig_img.shape[:2], np.uint8) #default - all background possible mask_fg = np.zeros(orig_img.shape[:2], np.uint8) mask_fgp = np.zeros(orig_img.shape[:2], np.uint8) img = orig_img.copy() l_thick = 2 def draw_4_lines(masks_out, pts, thickness=1): fgp, fg, bg, bgp = masks_out h, w = fg.shape fgp_pts = [] fg_pts = np.array([pts[i:i + 2] for i in range(len(pts) - 1)]) bg_pts = [] bgp_pts = [] for i in range(len(fg_pts)): a, b = line = fg_pts[i] ba = b - a v = ba / npl.norm(ba) ccpv = np.array([v[1], -v[0]]) cpv = np.array([-v[1], v[0]]) step = 1 / max(np.abs(cpv)) fgp_pts.append( np.clip(line + ccpv * step * thickness, 0, w - 1).astype(np.int)) bg_pts.append( np.clip(line + cpv * step * thickness, 0, w - 1).astype(np.int)) bgp_pts.append( np.clip(line + cpv * step * thickness * 2, 0, w - 1).astype(np.int)) fgp_pts = np.array(fgp_pts) bg_pts = np.array(bg_pts) bgp_pts = np.array(bgp_pts) cv2.polylines(fgp, fgp_pts, False, (1, ), thickness=thickness) cv2.polylines(fg, fg_pts, False, (1, ), thickness=thickness) cv2.polylines(bg, bg_pts, False, (1, ), thickness=thickness) cv2.polylines(bgp, bgp_pts, False, (1, ), thickness=thickness) def draw_lines(masks_steps, pts, thickness=1): lines = np.array([pts[i:i + 2] for i in range(len(pts) - 1)]) for mask, step in masks_steps: h, w = mask.shape mask_lines = [] for i in range(len(lines)): a, b = line = lines[i] ba = b - a ba_len = npl.norm(ba) if ba_len != 0: v = ba / ba_len pv = np.array([-v[1], v[0]]) pv_inv_max = 1 / max(np.abs(pv)) mask_lines.append( np.clip(line + pv * pv_inv_max * thickness * step, 0, w - 1).astype(np.int)) else: mask_lines.append(np.array(line, dtype=np.int)) cv2.polylines(mask, mask_lines, False, (1, ), thickness=thickness) def draw_fill_convex(mask_out, pts, scale=1.0): hull = cv2.convexHull(np.array(pts)) if scale != 1.0: pts_count = hull.shape[0] sum_x = np.sum(hull[:, 0, 0]) sum_y = np.sum(hull[:, 0, 1]) hull_center = np.array([sum_x / pts_count, sum_y / pts_count]) hull = hull_center + (hull - hull_center) * scale hull = hull.astype(pts.dtype) cv2.fillConvexPoly(mask_out, hull, (1, )) def get_gc_mask_bgr(gc_mask): h, w = gc_mask.shape bgr = np.zeros((h, w, 3), dtype=np.uint8) bgr[gc_mask == 0] = (0, 0, 0) bgr[gc_mask == 1] = (255, 255, 255) bgr[gc_mask == 2] = (0, 0, 255) #RED bgr[gc_mask == 3] = (0, 255, 0) #GREEN return bgr def get_gc_mask_result(gc_mask): return np.where((gc_mask == 1) + (gc_mask == 3), 1, 0).astype(np.int) #convex inner of right chin to end of right eyebrow #draw_fill_convex ( mask_fgp, lmrks_list[8:17]+lmrks_list[26:27] ) #convex inner of start right chin to right eyebrow #draw_fill_convex ( mask_fgp, lmrks_list[8:9]+lmrks_list[22:27] ) #convex inner of nose draw_fill_convex(mask_fgp, lmrks[27:36]) #convex inner of nose half draw_fill_convex(mask_fg, lmrks[27:36], scale=0.5) #left corner of mouth to left corner of nose #draw_lines ( [ (mask_fg,0), ], lmrks_list[49:50]+lmrks_list[32:33], l_thick) #convex inner: right corner of nose to centers of eyebrows #draw_fill_convex ( mask_fgp, lmrks_list[35:36]+lmrks_list[19:20]+lmrks_list[24:25]) #right corner of mouth to right corner of nose #draw_lines ( [ (mask_fg,0), ], lmrks_list[54:55]+lmrks_list[35:36], l_thick) #left eye #draw_fill_convex ( mask_fg, lmrks_list[36:40] ) #right eye #draw_fill_convex ( mask_fg, lmrks_list[42:48] ) #right chin draw_lines([ (mask_bg, 0), (mask_fg, -1), ], lmrks[8:17], l_thick) #left eyebrow center to right eyeprow center draw_lines([ (mask_bg, -1), (mask_fg, 0), ], lmrks_list[19:20] + lmrks_list[24:25], l_thick) # #draw_lines ( [ (mask_bg,-1), (mask_fg,0), ], lmrks_list[24:25] + lmrks_list[19:17:-1], l_thick) #half right eyebrow to end of right chin draw_lines([ (mask_bg, -1), (mask_fg, 0), ], lmrks_list[24:27] + lmrks_list[16:17], l_thick) #import code #code.interact(local=dict(globals(), **locals())) #compose mask layers gc_mask = np.zeros(orig_img.shape[:2], np.uint8) gc_mask[mask_bgp == 1] = 2 gc_mask[mask_fgp == 1] = 3 gc_mask[mask_bg == 1] = 0 gc_mask[mask_fg == 1] = 1 gc_bgr_before = get_gc_mask_bgr(gc_mask) #io.show_image (wnd_name, gc_mask ) ##points, hierarcy = cv2.findContours(original_mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) ##gc_mask = ( (1-erode_mask)*2 + erode_mask )# * dilate_mask #gc_mask = (1-erode_mask)*2 + erode_mask #cv2.addWeighted( #gc_mask = mask_0_27 + (1-mask_0_27)*2 # ##import code ##code.interact(local=dict(globals(), **locals())) # #rect = (1,1,img.shape[1]-2,img.shape[0]-2) # # cv2.grabCut(img, gc_mask, None, np.zeros((1, 65), np.float64), np.zeros((1, 65), np.float64), 5, cv2.GC_INIT_WITH_MASK) gc_bgr = get_gc_mask_bgr(gc_mask) gc_mask_result = get_gc_mask_result(gc_mask) gc_mask_result_1 = gc_mask_result[:, :, np.newaxis] #import code #code.interact(local=dict(globals(), **locals())) orig_img_gc_layers_masked = (0.5 * orig_img + 0.5 * gc_bgr).astype( np.uint8) orig_img_gc_before_layers_masked = (0.5 * orig_img + 0.5 * gc_bgr_before).astype( np.uint8) pink_bg = np.full(orig_img.shape, (255, 0, 255), dtype=np.uint8) orig_img_result = orig_img * gc_mask_result_1 orig_img_result_pinked = orig_img_result + pink_bg * (1 - gc_mask_result_1) #io.show_image (wnd_name, blended_img) ##gc_mask, bgdModel, fgdModel = # #mask2 = np.where((gc_mask==1) + (gc_mask==3),255,0).astype('uint8')[:,:,np.newaxis] #mask2 = np.repeat(mask2, (3,), -1) # ##mask2 = np.where(gc_mask!=0,255,0).astype('uint8') #blended_img = orig_img #-\ # #0.3 * np.full(original_img.shape, (50,50,50)) * (1-mask_0_27)[:,:,np.newaxis] # #0.3 * np.full(original_img.shape, (50,50,50)) * (1-dilate_mask)[:,:,np.newaxis] +\ # #0.3 * np.full(original_img.shape, (50,50,50)) * (erode_mask)[:,:,np.newaxis] #blended_img = np.clip(blended_img, 0, 255).astype(np.uint8) ##import code ##code.interact(local=dict(globals(), **locals())) orig_img_lmrked = orig_img.copy() LandmarksProcessor.draw_landmarks(orig_img_lmrked, lmrks, transparent_mask=True) screen = np.concatenate([ orig_img_gc_before_layers_masked, orig_img_gc_layers_masked, orig_img, orig_img_lmrked, orig_img_result_pinked, orig_img_result, ], axis=1) io.show_image(wnd_name, screen.astype(np.uint8)) while True: io.process_messages() for (x, y, ev, flags) in io.get_mouse_events(wnd_name): pass #print (x,y,ev,flags) key_events = [ev for ev, in io.get_key_events(wnd_name)] for key in key_events: if key == ord('1'): pass if key == ord('2'): pass if key == ord('3'): pass if ord(' ') in key_events: break import code code.interact(local=dict(globals(), **locals()))
def process(sample, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types sample_bgr = sample.load_bgr() ct_sample_bgr = None ct_sample_mask = 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) cached_images = collections.defaultdict(dict) sample_rnd_seed = np.random.randint(0x80000000) outputs = [] for opts in output_sample_types: resolution = opts.get('resolution', 0) types = opts.get('types', []) border_replicate = opts.get('border_replicate', True) random_sub_res = opts.get('random_sub_res', 0) normalize_std_dev = opts.get('normalize_std_dev', False) normalize_vgg = opts.get('normalize_vgg', False) motion_blur = opts.get('motion_blur', None) apply_ct = opts.get('apply_ct', False) normalize_tanh = opts.get('normalize_tanh', False) img_type = SPTF.NONE target_face_type = SPTF.NONE face_mask_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 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) img = l elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch_yaw_roll = sample.pitch_yaw_roll if pitch_yaw_roll is not None: pitch, yaw, roll = pitch_yaw_roll else: pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll( sample.landmarks) if params['flip']: yaw = -yaw if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: pitch = (pitch + 1.0) / 2.0 yaw = (yaw + 1.0) / 2.0 roll = (roll + 1.0) / 2.0 img = (pitch, yaw, roll) else: if mode_type == SPTF.NONE: raise ValueError('expected MODE_ type') 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] img = np.concatenate((img, mask), -1) return img img = sample_bgr ### Prepare a mask mask = None if is_face_sample: mask = sample.load_fanseg_mask( ) #using fanseg_mask if exist if mask is None: if sample.eyebrows_expand_mod is not None: mask = LandmarksProcessor.get_image_hull_mask( img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod) else: mask = LandmarksProcessor.get_image_hull_mask( img.shape, sample.landmarks) if sample.ie_polys is not None: sample.ie_polys.overlay_mask(mask) ################## 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 is_face_sample and target_face_type != SPTF.NONE: 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: #first warp to target facetype img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat( sample.landmarks, sample.shape[0], target_ft), (sample.shape[0], sample.shape[0]), flags=cv2.INTER_CUBIC) mask = cv2.warpAffine( mask, LandmarksProcessor.get_transform_mat( sample.landmarks, sample.shape[0], target_ft), (sample.shape[0], sample.shape[0]), flags=cv2.INTER_CUBIC) #then apply transforms img = do_transform(img, mask) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) else: img = do_transform(img, mask) img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat( sample.landmarks, resolution, target_ft), (resolution, resolution), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC) else: img = do_transform(img, mask) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC) if random_sub_res != 0: sub_size = resolution - random_sub_res rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_res) start_x = rnd_state.randint(sub_size + 1) start_y = rnd_state.randint(sub_size + 1) img = img[start_y:start_y + sub_size, start_x:start_x + sub_size, :] img = np.clip(img, 0, 1) img_bgr = img[..., 0:3] img_mask = img[..., 3:4] if apply_ct and ct_sample is not None: if ct_sample_bgr is None: ct_sample_bgr = ct_sample.load_bgr() ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR) img_bgr = imagelib.linear_color_transfer( img_bgr, ct_sample_bgr_resized) img_bgr = np.clip(img_bgr, 0.0, 1.0) if normalize_std_dev: img_bgr = (img_bgr - img_bgr.mean((0, 1))) / img_bgr.std( (0, 1)) elif normalize_vgg: img_bgr = np.clip(img_bgr * 255, 0, 255) img_bgr[:, :, 0] -= 103.939 img_bgr[:, :, 1] -= 116.779 img_bgr[:, :, 2] -= 123.68 if mode_type == SPTF.MODE_BGR: img = img_bgr elif mode_type == SPTF.MODE_BGR_SHUFFLE: rnd_state = np.random.RandomState(sample_rnd_seed) img = np.take(img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1) elif mode_type == SPTF.MODE_BGR_RANDOM_HUE_SHIFT: rnd_state = np.random.RandomState(sample_rnd_seed) hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) h, s, v = cv2.split(hsv) h = (h + rnd_state.randint(360)) % 360 hsv = cv2.merge([h, s, v]) img = np.clip(cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1) elif mode_type == SPTF.MODE_G: img = np.concatenate((np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), img_mask), -1) elif mode_type == SPTF.MODE_GGG: img = np.repeat( np.expand_dims( cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), (3, ), -1) elif mode_type == SPTF.MODE_M and is_face_sample: img = img_mask if not debug: if normalize_tanh: img = np.clip(img * 2.0 - 1.0, -1.0, 1.0) else: img = np.clip(img, 0.0, 1.0) outputs.append(img) if debug: result = [] for output in outputs: if output.shape[2] < 4: result += [ output, ] elif output.shape[2] == 4: result += [ output[..., 0:3] * output[..., 3:4], ] return result else: return outputs
def MergeMaskedFace(predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks, img_landmarks_prev, img_landmarks_next): img_size = img_bgr.shape[1], img_bgr.shape[0] img_face_mask_a = LandmarksProcessor.get_image_hull_mask( img_bgr.shape, img_face_landmarks) if cfg.mode == 'original': return img_bgr, img_face_mask_a out_img = img_bgr.copy() out_merging_mask_a = None input_size = predictor_input_shape[0] mask_subres_size = input_size * 4 output_size = input_size if cfg.super_resolution_power != 0 or cfg.smooth_rect: output_size *= 4 if cfg.smooth_rect: average_frame_count = 5 average_center_frame_count = 1 else: average_frame_count = 0 average_center_frame_count = 0 def get_transform_mat(*args, **kwargs): return LandmarksProcessor.get_averaged_transform_mat( img_face_landmarks, img_landmarks_prev, img_landmarks_next, average_frame_count, average_center_frame_count, *args, **kwargs) face_mat = get_transform_mat(output_size, face_type=cfg.face_type) face_output_mat = get_transform_mat(output_size, face_type=cfg.face_type, scale=1.0 + 0.01 * cfg.output_face_scale) if mask_subres_size == output_size: face_mask_output_mat = face_output_mat else: face_mask_output_mat = get_transform_mat(mask_subres_size, face_type=cfg.face_type, scale=1.0 + 0.01 * cfg.output_face_scale) dst_face_bgr = cv2.warpAffine(img_bgr, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC) dst_face_bgr = np.clip(dst_face_bgr, 0, 1) dst_face_mask_a_0 = cv2.warpAffine(img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC) dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1) predictor_input_bgr = cv2.resize(dst_face_bgr, (input_size, input_size)) predicted = predictor_func(predictor_input_bgr) if isinstance(predicted, tuple): #merger return bgr,mask prd_face_bgr = np.clip(predicted[0], 0, 1.0) prd_face_mask_a_0 = np.clip(predicted[1], 0, 1.0) predictor_masked = True else: #merger return bgr only, using dst mask prd_face_bgr = np.clip(predicted, 0, 1.0) prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0, (input_size, input_size)) predictor_masked = False if cfg.super_resolution_power != 0: prd_face_bgr_enhanced = cfg.superres_func(prd_face_bgr) mod = cfg.super_resolution_power / 100.0 prd_face_bgr = cv2.resize(prd_face_bgr, (output_size, output_size)) * ( 1.0 - mod) + prd_face_bgr_enhanced * mod prd_face_bgr = np.clip(prd_face_bgr, 0, 1) elif cfg.smooth_rect: prd_face_bgr = cv2.resize(prd_face_bgr, (output_size, output_size), cv2.INTER_CUBIC) prd_face_bgr = np.clip(prd_face_bgr, 0, 1) if cfg.super_resolution_power != 0 or cfg.smooth_rect: if predictor_masked: prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) else: prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) if cfg.mask_mode == 2: #dst prd_face_mask_a_0 = cv2.resize(dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) elif cfg.mask_mode >= 3 and cfg.mask_mode <= 8: if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6: prd_face_fanseg_bgr = cv2.resize(prd_face_bgr, (cfg.fanseg_input_size, ) * 2) prd_face_fanseg_mask = cfg.fanseg_extract_func( FaceType.FULL, prd_face_fanseg_bgr) FAN_prd_face_mask_a_0 = cv2.resize(prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC) if cfg.mask_mode >= 4 and cfg.mask_mode <= 7: full_face_fanseg_mat = get_transform_mat(cfg.fanseg_input_size, face_type=FaceType.FULL) dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (cfg.fanseg_input_size, ) * 2, flags=cv2.INTER_CUBIC) dst_face_fanseg_mask = cfg.fanseg_extract_func( FaceType.FULL, dst_face_fanseg_bgr) if cfg.face_type == FaceType.FULL: FAN_dst_face_mask_a_0 = cv2.resize(dst_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC) else: face_fanseg_mat = get_transform_mat(cfg.fanseg_input_size, face_type=cfg.face_type) fanseg_rect_corner_pts = np.array( [[0, 0], [cfg.fanseg_input_size - 1, 0], [0, cfg.fanseg_input_size - 1]], dtype=np.float32) a = LandmarksProcessor.transform_points(fanseg_rect_corner_pts, face_fanseg_mat, invert=True) b = LandmarksProcessor.transform_points( a, full_face_fanseg_mat) m = cv2.getAffineTransform(b, fanseg_rect_corner_pts) FAN_dst_face_mask_a_0 = cv2.warpAffine( dst_face_fanseg_mask, m, (cfg.fanseg_input_size, ) * 2, flags=cv2.INTER_CUBIC) FAN_dst_face_mask_a_0 = cv2.resize(FAN_dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) if cfg.mask_mode == 3: #FAN-prd prd_face_mask_a_0 = FAN_prd_face_mask_a_0 elif cfg.mask_mode == 4: #FAN-dst prd_face_mask_a_0 = FAN_dst_face_mask_a_0 elif cfg.mask_mode == 5: prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 elif cfg.mask_mode == 6: prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 elif cfg.mask_mode == 7: prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0 prd_face_mask_a_0[prd_face_mask_a_0 < (1.0 / 255.0)] = 0.0 # get rid of noise # resize to mask_subres_size if prd_face_mask_a_0.shape[0] != mask_subres_size: prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0, (mask_subres_size, mask_subres_size), cv2.INTER_CUBIC) # process mask in local predicted space if 'raw' not in cfg.mode: # add zero pad prd_face_mask_a_0 = np.pad(prd_face_mask_a_0, input_size) ero = cfg.erode_mask_modifier blur = cfg.blur_mask_modifier if ero > 0: prd_face_mask_a_0 = cv2.erode(prd_face_mask_a_0, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (ero, ero)), iterations=1) elif ero < 0: prd_face_mask_a_0 = cv2.dilate(prd_face_mask_a_0, cv2.getStructuringElement( cv2.MORPH_ELLIPSE, (-ero, -ero)), iterations=1) # clip eroded/dilated mask in actual predict area # pad with half blur size in order to accuratelly fade to zero at the boundary clip_size = input_size + blur // 2 prd_face_mask_a_0[:clip_size, :] = 0 prd_face_mask_a_0[-clip_size:, :] = 0 prd_face_mask_a_0[:, :clip_size] = 0 prd_face_mask_a_0[:, -clip_size:] = 0 if blur > 0: blur = blur + (1 - blur % 2) prd_face_mask_a_0 = cv2.GaussianBlur(prd_face_mask_a_0, (blur, blur), 0) prd_face_mask_a_0 = prd_face_mask_a_0[input_size:-input_size, input_size:-input_size] prd_face_mask_a_0 = np.clip(prd_face_mask_a_0, 0, 1) img_face_mask_a = cv2.warpAffine(prd_face_mask_a_0, face_mask_output_mat, img_size, np.zeros(img_bgr.shape[0:2], dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC)[..., None] img_face_mask_a = np.clip(img_face_mask_a, 0.0, 1.0) img_face_mask_a[img_face_mask_a < (1.0 / 255.0)] = 0.0 # get rid of noise if prd_face_mask_a_0.shape[0] != output_size: prd_face_mask_a_0 = cv2.resize(prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) prd_face_mask_a = prd_face_mask_a_0[..., None] prd_face_mask_area_a = prd_face_mask_a.copy() prd_face_mask_area_a[prd_face_mask_area_a > 0] = 1.0 if 'raw' in cfg.mode: if cfg.mode == 'raw-rgb': out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_merging_mask_a = img_face_mask_a out_img = np.clip(out_img, 0.0, 1.0) else: #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask ar = [] for i in range(1, 10): maxregion = np.argwhere(img_face_mask_a > i / 10.0) if maxregion.size != 0: miny, minx = maxregion.min(axis=0)[:2] maxy, maxx = maxregion.max(axis=0)[:2] lenx = maxx - minx leny = maxy - miny if min(lenx, leny) >= 4: ar += [[lenx, leny]] if len(ar) > 0: if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0: if cfg.color_transfer_mode == 1: #rct prd_face_bgr = imagelib.reinhard_color_transfer( np.clip(prd_face_bgr * 255, 0, 255).astype(np.uint8), np.clip(dst_face_bgr * 255, 0, 255).astype(np.uint8), source_mask=prd_face_mask_area_a, target_mask=prd_face_mask_area_a) prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) elif cfg.color_transfer_mode == 2: #lct prd_face_bgr = imagelib.linear_color_transfer( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 3: #mkl prd_face_bgr = imagelib.color_transfer_mkl( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 4: #mkl-m prd_face_bgr = imagelib.color_transfer_mkl( prd_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) elif cfg.color_transfer_mode == 5: #idt prd_face_bgr = imagelib.color_transfer_idt( prd_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m prd_face_bgr = imagelib.color_transfer_idt( prd_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) elif cfg.color_transfer_mode == 7: #sot-m prd_face_bgr = imagelib.color_transfer_sot( prd_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) prd_face_bgr = np.clip(prd_face_bgr, 0.0, 1.0) elif cfg.color_transfer_mode == 8: #mix-m prd_face_bgr = imagelib.color_transfer_mix( prd_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) if cfg.mode == 'hist-match': hist_mask_a = np.ones(prd_face_bgr.shape[:2] + (1, ), dtype=np.float32) if cfg.masked_hist_match: hist_mask_a *= prd_face_mask_area_a white = (1.0 - hist_mask_a) * np.ones( prd_face_bgr.shape[:2] + (1, ), dtype=np.float32) hist_match_1 = prd_face_bgr * hist_mask_a + white hist_match_1[hist_match_1 > 1.0] = 1.0 hist_match_2 = dst_face_bgr * hist_mask_a + white hist_match_2[hist_match_1 > 1.0] = 1.0 prd_face_bgr = imagelib.color_hist_match( hist_match_1, hist_match_2, cfg.hist_match_threshold).astype(dtype=np.float32) if 'seamless' in cfg.mode: #mask used for cv2.seamlessClone img_face_seamless_mask_a = None for i in range(1, 10): a = img_face_mask_a > i / 10.0 if len(np.argwhere(a)) == 0: continue img_face_seamless_mask_a = img_face_mask_a.copy() img_face_seamless_mask_a[a] = 1.0 img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0 break out_img = cv2.warpAffine(prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_img = np.clip(out_img, 0.0, 1.0) if 'seamless' in cfg.mode: try: #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering) l, t, w, h = cv2.boundingRect( (img_face_seamless_mask_a * 255).astype(np.uint8)) s_maskx, s_masky = int(l + w / 2), int(t + h / 2) out_img = cv2.seamlessClone( (out_img * 255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a * 255).astype(np.uint8), (s_maskx, s_masky), cv2.NORMAL_CLONE) out_img = out_img.astype(dtype=np.float32) / 255.0 except Exception as e: #seamlessClone may fail in some cases e_str = traceback.format_exc() if 'MemoryError' in e_str: raise Exception( "Seamless fail: " + e_str ) #reraise MemoryError in order to reprocess this data by other processes else: print("Seamless fail: " + e_str) out_img = img_bgr * (1 - img_face_mask_a) + (out_img * img_face_mask_a) out_face_bgr = cv2.warpAffine(out_img, face_mat, (output_size, output_size)) if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0: if cfg.color_transfer_mode == 1: face_mask_a = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size))[..., None] out_face_bgr = imagelib.reinhard_color_transfer( (out_face_bgr * 255).astype(np.uint8), (dst_face_bgr * 255).astype(np.uint8), source_mask=face_mask_a, target_mask=face_mask_a) out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) elif cfg.color_transfer_mode == 2: #lct out_face_bgr = imagelib.linear_color_transfer( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 3: #mkl out_face_bgr = imagelib.color_transfer_mkl( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 4: #mkl-m out_face_bgr = imagelib.color_transfer_mkl( out_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) elif cfg.color_transfer_mode == 5: #idt out_face_bgr = imagelib.color_transfer_idt( out_face_bgr, dst_face_bgr) elif cfg.color_transfer_mode == 6: #idt-m out_face_bgr = imagelib.color_transfer_idt( out_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) elif cfg.color_transfer_mode == 7: #sot-m out_face_bgr = imagelib.color_transfer_sot( out_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) out_face_bgr = np.clip(out_face_bgr, 0.0, 1.0) elif cfg.color_transfer_mode == 8: #mix-m out_face_bgr = imagelib.color_transfer_mix( out_face_bgr * prd_face_mask_area_a, dst_face_bgr * prd_face_mask_area_a) if cfg.mode == 'seamless-hist-match': out_face_bgr = imagelib.color_hist_match( out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) cfg_mp = cfg.motion_blur_power / 100.0 if cfg_mp != 0: k_size = int(frame_info.motion_power * cfg_mp) if k_size >= 1: k_size = np.clip(k_size + 1, 2, 50) if cfg.super_resolution_power != 0 or cfg.smooth_rect: k_size *= 2 out_face_bgr = imagelib.LinearMotionBlur( out_face_bgr, k_size, frame_info.motion_deg) if cfg.blursharpen_amount != 0: out_face_bgr = cfg.blursharpen_func(out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount) if cfg.image_denoise_power != 0: n = cfg.image_denoise_power while n > 0: img_bgr_denoised = cv2.medianBlur(img_bgr, 5) if int(n / 100) != 0: img_bgr = img_bgr_denoised else: pass_power = (n % 100) / 100.0 img_bgr = img_bgr * ( 1.0 - pass_power) + img_bgr_denoised * pass_power n = max(n - 10, 0) if cfg.bicubic_degrade_power != 0: p = 1.0 - cfg.bicubic_degrade_power / 101.0 img_bgr_downscaled = cv2.resize( img_bgr, (int(img_size[0] * p), int(img_size[1] * p)), cv2.INTER_CUBIC) img_bgr = cv2.resize(img_bgr_downscaled, img_size, cv2.INTER_CUBIC) new_out = cv2.warpAffine(out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT) out_img = np.clip( img_bgr * (1 - img_face_mask_a) + (new_out * img_face_mask_a), 0, 1.0) if cfg.color_degrade_power != 0: out_img_reduced = imagelib.reduce_colors(out_img, 256) if cfg.color_degrade_power == 100: out_img = out_img_reduced else: alpha = cfg.color_degrade_power / 100.0 out_img = (out_img * (1.0 - alpha) + out_img_reduced * alpha) out_merging_mask_a = img_face_mask_a return out_img, out_merging_mask_a
def main (args, device_args): io.log_info ("Running converter.\r\n") aligned_dir = args.get('aligned_dir', None) avaperator_aligned_dir = args.get('avaperator_aligned_dir', None) try: input_path = Path(args['input_dir']) output_path = Path(args['output_dir']) model_path = Path(args['model_dir']) if not input_path.exists(): io.log_err('Input directory not found. Please ensure it exists.') return if output_path.exists(): for filename in Path_utils.get_image_paths(output_path): Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) if not model_path.exists(): io.log_err('Model directory not found. Please ensure it exists.') return is_interactive = io.input_bool ("Use interactive converter? (y/n skip:y) : ", True) if not io.is_colab() else False import models model = models.import_model( args['model_name'] )(model_path, device_args=device_args) cfg = model.get_ConverterConfig() if not is_interactive: cfg.ask_settings() input_path_image_paths = Path_utils.get_image_paths(input_path) if cfg.type == ConverterConfig.TYPE_MASKED: if aligned_dir is None: io.log_err('Aligned directory not found. Please ensure it exists.') return aligned_path = Path(aligned_dir) if not aligned_path.exists(): io.log_err('Aligned directory not found. Please ensure it exists.') return alignments = {} multiple_faces_detected = False aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"): filepath = Path(filepath) if filepath.suffix == '.png': dflimg = DFLPNG.load( str(filepath) ) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load ( str(filepath) ) else: dflimg = None if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue source_filename_stem = Path( dflimg.get_source_filename() ).stem if source_filename_stem not in alignments.keys(): alignments[ source_filename_stem ] = [] alignments_ar = alignments[ source_filename_stem ] alignments_ar.append (dflimg.get_source_landmarks()) if len(alignments_ar) > 1: multiple_faces_detected = True if multiple_faces_detected: io.log_info ("Warning: multiple faces detected. Strongly recommended to process them separately.") frames = [ ConvertSubprocessor.Frame( frame_info=FrameInfo(filename=p, landmarks_list=alignments.get(Path(p).stem, None))) for p in input_path_image_paths ] if multiple_faces_detected: io.log_info ("Warning: multiple faces detected. Motion blur will not be used.") else: s = 256 local_pts = [ (s//2-1, s//2-1), (s//2-1,0) ] #center+up frames_len = len(frames) for i in io.progress_bar_generator( range(len(frames)) , "Computing motion vectors"): fi_prev = frames[max(0, i-1)].frame_info fi = frames[i].frame_info fi_next = frames[min(i+1, frames_len-1)].frame_info if len(fi_prev.landmarks_list) == 0 or \ len(fi.landmarks_list) == 0 or \ len(fi_next.landmarks_list) == 0: continue mat_prev = LandmarksProcessor.get_transform_mat ( fi_prev.landmarks_list[0], s, face_type=FaceType.FULL) mat = LandmarksProcessor.get_transform_mat ( fi.landmarks_list[0] , s, face_type=FaceType.FULL) mat_next = LandmarksProcessor.get_transform_mat ( fi_next.landmarks_list[0], s, face_type=FaceType.FULL) pts_prev = LandmarksProcessor.transform_points (local_pts, mat_prev, True) pts = LandmarksProcessor.transform_points (local_pts, mat, True) pts_next = LandmarksProcessor.transform_points (local_pts, mat_next, True) prev_vector = pts[0]-pts_prev[0] next_vector = pts_next[0]-pts[0] motion_vector = pts_next[0] - pts_prev[0] fi.motion_power = npla.norm(motion_vector) motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array([0,0],dtype=np.float32) fi.motion_deg = -math.atan2(motion_vector[1],motion_vector[0])*180 / math.pi elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR: filesdata = [] for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"): filepath = Path(filepath) if filepath.suffix == '.png': dflimg = DFLPNG.load( str(filepath) ) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load ( str(filepath) ) else: dflimg = None if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue filesdata += [ ( FrameInfo(filename=str(filepath), landmarks_list=[dflimg.get_landmarks()] ), dflimg.get_source_filename() ) ] filesdata = sorted(filesdata, key=operator.itemgetter(1)) #sort by filename frames = [] filesdata_len = len(filesdata) for i in range(len(filesdata)): frame_info = filesdata[i][0] prev_temporal_frame_infos = [] next_temporal_frame_infos = [] for t in range (cfg.temporal_face_count): prev_frame_info = filesdata[ max(i -t, 0) ][0] next_frame_info = filesdata[ min(i +t, filesdata_len-1 )][0] prev_temporal_frame_infos.insert (0, prev_frame_info ) next_temporal_frame_infos.append ( next_frame_info ) frames.append ( ConvertSubprocessor.Frame(prev_temporal_frame_infos=prev_temporal_frame_infos, frame_info=frame_info, next_temporal_frame_infos=next_temporal_frame_infos) ) if len(frames) == 0: io.log_info ("No frames to convert in input_dir.") else: ConvertSubprocessor ( is_interactive = is_interactive, converter_config = cfg, frames = frames, output_path = output_path, ).run() model.finalize() except Exception as e: print ( 'Error: %s' % (str(e))) traceback.print_exc()
def get_transform_mat(*args, **kwargs): return LandmarksProcessor.get_averaged_transform_mat( img_face_landmarks, img_landmarks_prev, img_landmarks_next, average_frame_count, average_center_frame_count, *args, **kwargs)
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) 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)) 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 extract(self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False): if len(rects) == 0: return [] if is_bgr: input_image = input_image[:, :, ::-1] is_bgr = False (h, w, ch) = input_image.shape landmarks = [] for (left, top, right, bottom) in rects: scale = (right - left + bottom - top) / 195.0 center = np.array([(left + right) / 2.0, (top + bottom) / 2.0]) centers = [center] if multi_sample: centers += [ center + [-1, -1], center + [1, -1], center + [1, 1], center + [-1, 1], ] images = [] ptss = [] try: for c in centers: images += [self.crop(input_image, c, scale)] images = np.stack(images) images = images.astype(np.float32) / 255.0 predicted = [] for i in range(len(images)): predicted += [ self.model.predict(images[i][None, ...]).transpose( 0, 3, 1, 2)[0] ] predicted = np.stack(predicted) for i, pred in enumerate(predicted): ptss += [ self.get_pts_from_predict(pred, centers[i], scale) ] pts_img = np.mean(np.array(ptss), 0) landmarks.append(pts_img) except: landmarks.append(None) if second_pass_extractor is not None: for i, lmrks in enumerate(landmarks): try: if lmrks is not None: image_to_face_mat = LandmarksProcessor.get_transform_mat( lmrks, 256, FaceType.FULL) face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC) rects2 = second_pass_extractor.extract(face_image, is_bgr=is_bgr) if len( rects2 ) == 1: #dont do second pass if faces != 1 detected in cropped image lmrks2 = self.extract(face_image, [rects2[0]], is_bgr=is_bgr, multi_sample=True)[0] landmarks[i] = LandmarksProcessor.transform_points( lmrks2, image_to_face_mat, True) except: pass return landmarks