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
