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)
return np.clip(eyes_mask, 0, 1)
is_face_sample = sample_landmarks is not None
if debug and is_face_sample:
LandmarksProcessor.draw_landmarks(sample_bgr, sample_landmarks,
(0, 1, 0))
params_per_resolution = {}
warp_rnd_state = np.random.RandomState(sample_rnd_seed - 1)
for opts in output_sample_types:
resolution = opts.get('resolution', None)
if resolution is None:
continue
params_per_resolution[resolution] = imagelib.gen_warp_params(
resolution,
sample_process_options.random_flip,
rotation_range=sample_process_options.rotation_range,
scale_range=sample_process_options.scale_range,
tx_range=sample_process_options.tx_range,
ty_range=sample_process_options.ty_range,
rnd_state=warp_rnd_state)
outputs_sample = []
for opts in output_sample_types:
sample_type = opts.get('sample_type', SPST.NONE)
channel_type = opts.get('channel_type', SPCT.NONE)
resolution = opts.get('resolution', 0)
warp = opts.get('warp', False)
transform = opts.get('transform', False)
motion_blur = opts.get('motion_blur', None)
gaussian_blur = opts.get('gaussian_blur', None)
random_bilinear_resize = opts.get('random_bilinear_resize',
None)
random_rgb_levels = opts.get('random_rgb_levels', False)
random_hsv_shift = opts.get('random_hsv_shift', False)
random_circle_mask = opts.get('random_circle_mask', False)
normalize_tanh = opts.get('normalize_tanh', False)
ct_mode = opts.get('ct_mode', None)
data_format = opts.get('data_format', 'NHWC')
if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE:
border_replicate = False
elif sample_type == SPST.FACE_IMAGE:
border_replicate = True
border_replicate = opts.get('border_replicate',
border_replicate)
borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT
if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
if not is_face_sample:
raise ValueError(
"face_samples should be provided for sample_type FACE_*"
)
if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
face_type = opts.get('face_type', None)
face_mask_type = opts.get('face_mask_type', SPFMT.NONE)
if face_type is None:
raise ValueError(
"face_type must be defined for face samples")
if 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()
img += get_eyes_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:
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.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),
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 batch_func(self, param):
samples, seg_sample_idxs, resolution, face_type, data_format = param
shuffle_idxs = []
bg_shuffle_idxs = []
random_flip = True
rotation_range = [-10, 10]
scale_range = [-0.05, 0.05]
tx_range = [-0.05, 0.05]
ty_range = [-0.05, 0.05]
random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25, 75
sharpen_chance, sharpen_kernel_max_size = 25, 5
motion_blur_chance, motion_blur_mb_max_size = 25, 5
gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5
random_jpeg_compress_chance = 25
def gen_img_mask(sample):
img = sample.load_bgr()
h, w, c = img.shape
if sample.seg_ie_polys.has_polys():
mask = np.zeros((h, w, 1), dtype=np.float32)
sample.seg_ie_polys.overlay_mask(mask)
elif sample.has_xseg_mask():
mask = sample.get_xseg_mask()
mask[mask < 0.5] = 0.0
mask[mask >= 0.5] = 1.0
else:
raise Exception(f'no mask in sample {sample.filename}')
if face_type == sample.face_type:
if w != resolution:
img = cv2.resize(img, (resolution, resolution),
interpolation=cv2.INTER_LANCZOS4)
mask = cv2.resize(mask, (resolution, resolution),
interpolation=cv2.INTER_LANCZOS4)
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, face_type)
img = cv2.warpAffine(img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_LANCZOS4)
mask = cv2.warpAffine(mask,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_LANCZOS4)
if len(mask.shape) == 2:
mask = mask[..., None]
return img, mask
bs = self.batch_size
while True:
batches = [[], []]
n_batch = 0
while n_batch < bs:
try:
if len(shuffle_idxs) == 0:
shuffle_idxs = seg_sample_idxs.copy()
np.random.shuffle(shuffle_idxs)
sample = samples[shuffle_idxs.pop()]
img, mask = gen_img_mask(sample)
if np.random.randint(2) == 0:
if len(bg_shuffle_idxs) == 0:
bg_shuffle_idxs = seg_sample_idxs.copy()
np.random.shuffle(bg_shuffle_idxs)
bg_sample = samples[bg_shuffle_idxs.pop()]
bg_img, bg_mask = gen_img_mask(bg_sample)
bg_wp = imagelib.gen_warp_params(
resolution,
True,
rotation_range=[-180, 180],
scale_range=[-0.10, 0.10],
tx_range=[-0.10, 0.10],
ty_range=[-0.10, 0.10])
bg_img = imagelib.warp_by_params(bg_wp,
bg_img,
can_warp=False,
can_transform=True,
can_flip=True,
border_replicate=True)
bg_mask = imagelib.warp_by_params(
bg_wp,
bg_mask,
can_warp=False,
can_transform=True,
can_flip=True,
border_replicate=False)
bg_img = bg_img * (1 - bg_mask)
if np.random.randint(2) == 0:
bg_img = imagelib.apply_random_hsv_shift(bg_img)
else:
bg_img = imagelib.apply_random_rgb_levels(bg_img)
c_mask = 1.0 - (1 - bg_mask) * (1 - mask)
rnd = 0.15 + np.random.uniform() * 0.85
img = img * (c_mask) + img * (
1 - c_mask) * rnd + bg_img * (1 - c_mask) * (1 -
rnd)
warp_params = imagelib.gen_warp_params(
resolution,
random_flip,
rotation_range=rotation_range,
scale_range=scale_range,
tx_range=tx_range,
ty_range=ty_range)
img = imagelib.warp_by_params(warp_params,
img,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=True)
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_overlay_triangle(
img,
max_alpha=0.25,
mask=sd.random_circle_faded([resolution, resolution]))
if np.random.randint(2) == 0:
img = imagelib.apply_random_hsv_shift(
img,
mask=sd.random_circle_faded(
[resolution, resolution]))
else:
img = imagelib.apply_random_rgb_levels(
img,
mask=sd.random_circle_faded(
[resolution, resolution]))
if np.random.randint(2) == 0:
# random face flare
krn = np.random.randint(resolution // 4, resolution)
krn = krn - krn % 2 + 1
img = img + cv2.GaussianBlur(img * mask, (krn, krn), 0)
if np.random.randint(2) == 0:
img = imagelib.apply_random_sharpen(
img,
sharpen_chance,
sharpen_kernel_max_size,
mask=sd.random_circle_faded(
[resolution, resolution]))
else:
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]))
if np.random.randint(2) == 0:
img = imagelib.apply_random_nearest_resize(
img,
random_bilinear_resize_chance,
random_bilinear_resize_max_size_per,
mask=sd.random_circle_faded(
[resolution, resolution]))
else:
img = imagelib.apply_random_bilinear_resize(
img,
random_bilinear_resize_chance,
random_bilinear_resize_max_size_per,
mask=sd.random_circle_faded(
[resolution, resolution]))
img = np.clip(img, 0, 1)
img = imagelib.apply_random_jpeg_compress(
img,
random_jpeg_compress_chance,
mask=sd.random_circle_faded([resolution, resolution]))
if data_format == "NCHW":
img = np.transpose(img, (2, 0, 1))
mask = np.transpose(mask, (2, 0, 1))
batches[0].append(img)
batches[1].append(mask)
n_batch += 1
except:
io.log_err(traceback.format_exc())
yield [np.array(batch) for batch in batches]
def batch_func(self, param):
pickled_samples, resolution, face_type, data_format = param
samples = pickle.loads(pickled_samples)
shuffle_idxs = []
idxs = [*range(len(samples))]
random_flip = True
rotation_range = [-10, 10]
scale_range = [-0.05, 0.05]
tx_range = [-0.05, 0.05]
ty_range = [-0.05, 0.05]
random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25, 75
motion_blur_chance, motion_blur_mb_max_size = 25, 5
gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5
bs = self.batch_size
while True:
batches = [[], []]
n_batch = 0
while n_batch < bs:
try:
if len(shuffle_idxs) == 0:
shuffle_idxs = idxs.copy()
np.random.shuffle(shuffle_idxs)
idx = shuffle_idxs.pop()
sample = samples[idx]
img = sample.load_bgr()
h, w, c = img.shape
mask = np.zeros((h, w, 1), dtype=np.float32)
sample.seg_ie_polys.overlay_mask(mask)
warp_params = imagelib.gen_warp_params(
resolution,
random_flip,
rotation_range=rotation_range,
scale_range=scale_range,
tx_range=tx_range,
ty_range=ty_range)
if face_type == sample.face_type:
if w != resolution:
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_LANCZOS4)
mask = cv2.resize(mask, (resolution, resolution),
cv2.INTER_LANCZOS4)
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, face_type)
img = cv2.warpAffine(img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_LANCZOS4)
mask = cv2.warpAffine(mask,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_LANCZOS4)
if len(mask.shape) == 2:
mask = mask[..., None]
img = imagelib.warp_by_params(warp_params,
img,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=False)
mask = imagelib.warp_by_params(warp_params,
mask,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=False)
img = np.clip(img.astype(np.float32), 0, 1)
mask[mask < 0.5] = 0.0
mask[mask >= 0.5] = 1.0
mask = np.clip(mask, 0, 1)
if np.random.randint(2) == 0:
img = imagelib.apply_random_hsv_shift(
img,
mask=sd.random_circle_faded(
[resolution, resolution]))
else:
img = imagelib.apply_random_rgb_levels(
img,
mask=sd.random_circle_faded(
[resolution, resolution]))
img = imagelib.apply_random_motion_blur(
img,
motion_blur_chance,
motion_blur_mb_max_size,
mask=sd.random_circle_faded([resolution, resolution]))
img = imagelib.apply_random_gaussian_blur(
img,
gaussian_blur_chance,
gaussian_blur_kernel_max_size,
mask=sd.random_circle_faded([resolution, resolution]))
img = imagelib.apply_random_bilinear_resize(
img,
random_bilinear_resize_chance,
random_bilinear_resize_max_size_per,
mask=sd.random_circle_faded([resolution, resolution]))
if data_format == "NCHW":
img = np.transpose(img, (2, 0, 1))
mask = np.transpose(mask, (2, 0, 1))
batches[0].append(img)
batches[1].append(mask)
n_batch += 1
except:
io.log_err(traceback.format_exc())
yield [np.array(batch) for batch in batches]
