def upgradeToFaceSamples ( samples ):
sample_list = []
for s in io.progress_bar_generator(samples, "Loading"):
s_filename_path = Path(s.filename)
try:
if s_filename_path.suffix == '.png':
dflimg = DFLPNG.load ( str(s_filename_path) )
elif s_filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(s_filename_path) )
else:
dflimg = None
if dflimg is None:
print ("%s is not a dfl image file required for training" % (s_filename_path.name) )
continue
landmarks = dflimg.get_landmarks()
pitch_yaw_roll = dflimg.get_pitch_yaw_roll()
if pitch_yaw_roll is None:
pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks)
sample_list.append( s.copy_and_set(sample_type=SampleType.FACE,
face_type=FaceType.fromString (dflimg.get_face_type()),
shape=dflimg.get_shape(),
landmarks=landmarks,
ie_polys=dflimg.get_ie_polys(),
pitch_yaw_roll=pitch_yaw_roll,
source_filename=dflimg.get_source_filename(),
fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) )
except:
print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
return sample_list
def sort_by_face_pitch(input_path):
io.log_info("Sorting by face pitch...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
filepath = Path(filepath)
dflimg = DFLIMG.load(filepath)
if dflimg is None:
io.log_err("%s is not a dfl image file" % (filepath.name))
trash_img_list.append([str(filepath)])
continue
pitch_yaw_roll = dflimg.get_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(
dflimg.get_landmarks())
img_list.append([str(filepath), pitch])
io.log_info("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def process_data(self, data):
filepath = Path(data[0])
try:
dflimg = DFLIMG.load(filepath)
if dflimg is None:
self.log_err("%s is not a dfl image file" %
(filepath.name))
return [1, [str(filepath)]]
bgr = cv2_imread(str(filepath))
if bgr is None:
raise Exception("Unable to load %s" % (filepath.name))
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
sharpness = estimate_sharpness(
gray) if self.include_by_blur else 0
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(
dflimg.get_landmarks())
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
self.log_err(e)
return [1, [str(filepath)]]
return [0, [str(filepath), sharpness, hist, yaw, pitch]]
def process_data(self, data):
filepath = Path(data[0])
try:
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:
self.log_err("%s 不是DeepFaceLab的图片格式，请使用DeepFaceLab提取脸图" %
(filepath.name))
return [1, [str(filepath)]]
bgr = cv2_imread(str(filepath))
if bgr is None:
raise Exception("无法加载 %s" % (filepath.name))
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
sharpness = estimate_sharpness(
gray) if self.include_by_blur else 0
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(
dflimg.get_landmarks())
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
self.log_err(e)
return [1, [str(filepath)]]
return [0, [str(filepath), sharpness, hist, yaw]]
def process_data(self, data):
filepath = Path(data[0])
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")
return [1, [str(filepath)]]
bgr = cv2_imread(str(filepath))
if bgr is None:
raise Exception("Unable to load %s" % (filepath.name))
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
if self.faster:
source_rect = dflimg.get_source_rect()
sharpness = mathlib.polygon_area(
np.array(source_rect[[0, 2, 2, 0]]).astype(np.float32),
np.array(source_rect[[1, 1, 3, 3]]).astype(np.float32))
else:
face_mask = LandmarksProcessor.get_image_hull_mask(
gray.shape, dflimg.get_landmarks())
sharpness = estimate_sharpness(
(gray[..., None] * face_mask).astype(np.uint8))
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(
dflimg.get_landmarks(), size=dflimg.get_shape()[1])
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
self.log_err(e)
return [1, [str(filepath)]]
return [0, [str(filepath), sharpness, hist, yaw, pitch]]
def sort_by_face_pitch(input_path):
io.log_info("根据[pitch]排序... ")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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 不是DeepFaceLab的图片格式，请使用DeepFaceLab提取脸图" %
(filepath.name))
trash_img_list.append([str(filepath)])
continue
pitch_yaw_roll = dflimg.get_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(
dflimg.get_landmarks())
img_list.append([str(filepath), pitch])
io.log_info("排序...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def get_pitch_yaw_roll(input_path, r=0.05):
import os
import numpy as np
import cv2
from shutil import copyfile
from pathlib import Path
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from facelib import LandmarksProcessor
from joblib import Subprocessor
import multiprocessing
from interact import interact as io
from imagelib import estimate_sharpness
io.log_info("Sorting by face yaw...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator(
Path_utils.get_image_paths(input_path), "Loading"):
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))
trash_img_list.append([str(filepath)])
continue
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(
dflimg.get_landmarks())
img_list.append([str(filepath), pitch, yaw, roll])
img_list.sort(key=lambda item: item[1])
with open(os.path.join(input_path, "_pitch_yaw_roll.csv"), "w") as f:
for i in img_list:
f.write("%s,%f,%f,%f\n" %
(os.path.basename(i[0]), i[1], i[2], i[3]))
import cv
width = 800
img = cv.cv_new((width, width))
xs = [i[1] for i in img_list]
ys = [i[2] for i in img_list]
cs = [(128, 128, 128)] * len(xs)
rs = [int(r * width / 2)] * len(xs)
cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs)
cs = [(0xcc, 0x66, 0x33)] * len(xs)
rs = [2] * len(xs)
cv.cv_scatter(img, xs, ys, [-1, 1], [-1, 1], cs, rs)
cv.cv_save(img, os.path.join(input_path, "_pitch_yaw_roll.bmp"))
return img_list
def sort_by_face_pitch(input_path):
io.log_info ("Sorting by face pitch...")
img_list = []
trash_img_list = []
for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
filepath = Path(filepath)
dflimg = DFLIMG.load (filepath)
if dflimg is None or not dflimg.has_data():
io.log_err (f"{filepath.name} is not a dfl image file")
trash_img_list.append ( [str(filepath)] )
continue
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )
img_list.append( [str(filepath), pitch ] )
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def get_pitch_yaw_roll(self):
if self.pitch_yaw_roll is None:
self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(
landmarks)
return self.pitch_yaw_roll
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_FULL_NO_ALIGN: FaceType.FULL_NO_ALIGN
}
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')
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, True)
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 = 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)
if sample.face_type == FaceType.MARK_ONLY:
if mask is not None:
img = np.concatenate((img, mask), -1)
else:
img = do_transform(img, mask)
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))
if sample.face_type == FaceType.MARK_ONLY:
img = cv2.warpAffine(
img,
LandmarksProcessor.get_transform_mat(
sample.landmarks, sample.shape[0],
ft), (sample.shape[0], sample.shape[0]),
flags=cv2.INTER_CUBIC)
mask = img[..., 3:4] if img.shape[2] > 3 else None
img = img[..., 0:3]
img = do_transform(img, mask)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)
else:
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 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_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 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(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
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,
rnd_seed=sample_rnd_seed)
outputs_sample = []
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)
gaussian_blur = opts.get('gaussian_blur', None)
ct_mode = opts.get('ct_mode', 'None')
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]
return img, mask
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 gaussian_blur is not None:
chance, kernel_max_size = gaussian_blur
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
img = cv2.GaussianBlur(
img,
(np.random.randint(kernel_max_size) * 2 + 1, )
* 2, 0)
if is_face_sample 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, mask = do_transform(img, mask)
img = np.concatenate((img, mask), -1)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)
else:
img, mask = do_transform(img, mask)
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, target_ft)
img = cv2.warpAffine(
img,
mat, (resolution, resolution),
borderMode=(cv2.BORDER_REPLICATE
if border_replicate else
cv2.BORDER_CONSTANT),
flags=cv2.INTER_CUBIC)
mask = cv2.warpAffine(
mask,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_CUBIC)
img = np.concatenate((img, mask[..., None]), -1)
else:
img, mask = do_transform(img, mask)
img = np.concatenate((img, mask), -1)
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).astype(np.float32)
img_bgr = img[..., 0:3]
img_mask = img[..., 3:4]
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()
ct_sample_bgr_resized = cv2.resize(
ct_sample_bgr, (resolution, resolution),
cv2.INTER_LINEAR)
if ct_mode == 'lct':
img_bgr = imagelib.linear_color_transfer(
img_bgr, ct_sample_bgr_resized)
img_bgr = np.clip(img_bgr, 0.0, 1.0)
elif ct_mode == 'rct':
img_bgr = imagelib.reinhard_color_transfer(
np.clip((img_bgr * 255).astype(np.uint8), 0,
255),
np.clip((ct_sample_bgr_resized * 255).astype(
np.uint8), 0, 255))
img_bgr = np.clip(
img_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
elif ct_mode == 'mkl':
img_bgr = imagelib.color_transfer_mkl(
img_bgr, ct_sample_bgr_resized)
elif ct_mode == 'idt':
img_bgr = imagelib.color_transfer_idt(
img_bgr, ct_sample_bgr_resized)
elif ct_mode == 'sot':
img_bgr = imagelib.color_transfer_sot(
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_HSV_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
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])
img = np.clip(cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0,
1)
elif mode_type == SPTF.MODE_G:
img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)[...,
None]
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_sample.append(img)
outputs += [outputs_sample]
return outputs
def dfl_estimate_pitch_yaw_roll(dfl_img):
from facelib import LandmarksProcessor
return LandmarksProcessor.estimate_pitch_yaw_roll(dfl_img.get_landmarks())
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', ColorTransferMode.NONE)
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()
if apply_ct == ColorTransferMode.LCT:
img_bgr = imagelib.linear_color_transfer(
img_bgr, ct_sample_bgr)
elif ColorTransferMode.RCT <= apply_ct <= ColorTransferMode.MASKED_RCT_PAPER_CLIP:
ct_options = {
ColorTransferMode.RCT: (False, False, False),
ColorTransferMode.RCT_CLIP: (False, False, True),
ColorTransferMode.RCT_PAPER: (False, True, False),
ColorTransferMode.RCT_PAPER_CLIP:
(False, True, True),
ColorTransferMode.MASKED_RCT: (True, False, False),
ColorTransferMode.MASKED_RCT_CLIP:
(True, False, True),
ColorTransferMode.MASKED_RCT_PAPER:
(True, True, False),
ColorTransferMode.MASKED_RCT_PAPER_CLIP:
(True, True, True),
}
use_masks, use_paper, use_clip = ct_options[apply_ct]
if not use_masks:
img_bgr = imagelib.reinhard_color_transfer(
img_bgr,
ct_sample_bgr,
clip=use_clip,
preserve_paper=use_paper)
else:
if ct_sample_mask is None:
ct_sample_mask = ct_sample.load_mask()
img_bgr = imagelib.reinhard_color_transfer(
img_bgr,
ct_sample_bgr,
clip=use_clip,
preserve_paper=use_paper,
source_mask=img_mask,
target_mask=ct_sample_mask)
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_LAB_RAND_TRANSFORM:
rnd_state = np.random.RandomState(sample_rnd_seed)
img = random_color_transform(img_bgr, rnd_state)
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
