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
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]
def process(samples,
sample_process_options,
output_sample_types,
debug,
ct_sample=None):
SPTF = SampleProcessor.Types
sample_rnd_seed = np.random.randint(0x80000000)
outputs = []
for sample in samples:
sample_bgr = sample.load_bgr()
ct_sample_bgr = None
h, w, c = sample_bgr.shape
is_face_sample = sample.landmarks is not None
if debug and is_face_sample:
LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks,
(0, 1, 0))
params = imagelib.gen_warp_params(
sample_bgr,
sample_process_options.random_flip,
rotation_range=sample_process_options.rotation_range,
scale_range=sample_process_options.scale_range,
tx_range=sample_process_options.tx_range,
ty_range=sample_process_options.ty_range,
rnd_seed=sample_rnd_seed)
outputs_sample = []
for opts in output_sample_types:
resolution = opts.get('resolution', 0)
types = opts.get('types', [])
motion_blur = opts.get('motion_blur', None)
gaussian_blur = opts.get('gaussian_blur', None)
ct_mode = opts.get('ct_mode', 'None')
normalize_tanh = opts.get('normalize_tanh', False)
data_format = opts.get('data_format', 'NHWC')
img_type = SPTF.NONE
target_face_type = SPTF.NONE
mode_type = SPTF.NONE
for t in types:
if t >= SPTF.IMG_TYPE_BEGIN and t < SPTF.IMG_TYPE_END:
img_type = t
elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END:
target_face_type = t
elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END:
mode_type = t
if mode_type == SPTF.MODE_FACE_MASK_HULL and not is_face_sample:
raise ValueError(
"MODE_FACE_MASK_HULL applicable only for face samples")
if mode_type == SPTF.MODE_FACE_MASK_STRUCT and not is_face_sample:
raise ValueError(
"MODE_FACE_MASK_STRUCT applicable only for face samples"
)
if is_face_sample:
if target_face_type == SPTF.NONE:
raise ValueError(
"target face type must be defined for face samples"
)
can_warp = (img_type == SPTF.IMG_WARPED
or img_type == SPTF.IMG_WARPED_TRANSFORMED)
can_transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED
or img_type == SPTF.IMG_TRANSFORMED)
if img_type == SPTF.NONE:
raise ValueError('expected IMG_ type')
if img_type == SPTF.IMG_LANDMARKS_ARRAY:
l = sample.landmarks
l = np.concatenate([
np.expand_dims(l[:, 0] / w, -1),
np.expand_dims(l[:, 1] / h, -1)
], -1)
l = np.clip(l, 0.0, 1.0)
out_sample = l
elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
pitch_yaw_roll = sample.get_pitch_yaw_roll()
if params['flip']:
yaw = -yaw
if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)
out_sample = (pitch, yaw, roll)
else:
if mode_type == SPTF.NONE:
raise ValueError('expected MODE_ type')
if mode_type == SPTF.MODE_FACE_MASK_HULL:
if sample.eyebrows_expand_mod is not None:
img = LandmarksProcessor.get_image_hull_mask(
sample_bgr.shape,
sample.landmarks,
eyebrows_expand_mod=sample.eyebrows_expand_mod)
else:
img = LandmarksProcessor.get_image_hull_mask(
sample_bgr.shape, sample.landmarks)
if sample.ie_polys is not None:
sample.ie_polys.overlay_mask(img)
elif mode_type == SPTF.MODE_FACE_MASK_STRUCT:
if sample.eyebrows_expand_mod is not None:
img = LandmarksProcessor.get_face_struct_mask(
sample_bgr.shape,
sample.landmarks,
eyebrows_expand_mod=sample.eyebrows_expand_mod)
else:
img = LandmarksProcessor.get_face_struct_mask(
sample_bgr.shape, sample.landmarks)
else:
img = sample_bgr
if motion_blur is not None:
chance, mb_max_size = motion_blur
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
img = imagelib.LinearMotionBlur(
img,
np.random.randint(mb_max_size) + 1,
np.random.randint(360))
if gaussian_blur is not None:
chance, kernel_max_size = gaussian_blur
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
img = cv2.GaussianBlur(
img,
(np.random.randint(kernel_max_size) * 2 +
1, ) * 2, 0)
if is_face_sample:
target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[
target_face_type]
if target_ft > sample.face_type:
raise Exception(
'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.'
%
(sample.filename, sample.face_type, target_ft))
if sample.face_type == FaceType.MARK_ONLY:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, sample.shape[0], target_ft)
if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
img = cv2.warpAffine(
img,
mat, (sample.shape[0], sample.shape[0]),
flags=cv2.INTER_CUBIC)
img = imagelib.warp_by_params(
params,
img,
can_warp,
can_transform,
can_flip=True,
border_replicate=False)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)[..., None]
else:
img = cv2.warpAffine(
img,
mat, (sample.shape[0], sample.shape[0]),
flags=cv2.INTER_CUBIC)
img = imagelib.warp_by_params(
params,
img,
can_warp,
can_transform,
can_flip=True,
border_replicate=True)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, target_ft)
if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
img = imagelib.warp_by_params(
params,
img,
can_warp,
can_transform,
can_flip=True,
border_replicate=False)
img = cv2.warpAffine(
img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_CUBIC)[..., None]
else:
img = imagelib.warp_by_params(
params,
img,
can_warp,
can_transform,
can_flip=True,
border_replicate=True)
img = cv2.warpAffine(
img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
else:
img = imagelib.warp_by_params(params,
img,
can_warp,
can_transform,
can_flip=True,
border_replicate=True)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)
if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
out_sample = np.clip(img.astype(np.float32), 0, 1)
else:
img = np.clip(img.astype(np.float32), 0, 1)
if ct_mode is not None and ct_sample is not None:
if ct_sample_bgr is None:
ct_sample_bgr = ct_sample.load_bgr()
img = imagelib.color_transfer(
ct_mode, img,
cv2.resize(ct_sample_bgr,
(resolution, resolution),
cv2.INTER_LINEAR))
if mode_type == SPTF.MODE_BGR:
out_sample = img
elif mode_type == SPTF.MODE_BGR_SHUFFLE:
rnd_state = np.random.RandomState(sample_rnd_seed)
out_sample = np.take(img,
rnd_state.permutation(
img.shape[-1]),
axis=-1)
elif mode_type == SPTF.MODE_BGR_RANDOM_HSV_SHIFT:
rnd_state = np.random.RandomState(sample_rnd_seed)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
h, s, v = cv2.split(hsv)
h = (h + rnd_state.randint(360)) % 360
s = np.clip(s + rnd_state.random() - 0.5, 0, 1)
v = np.clip(v + rnd_state.random() - 0.5, 0, 1)
hsv = cv2.merge([h, s, v])
out_sample = np.clip(
cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1)
elif mode_type == SPTF.MODE_G:
out_sample = cv2.cvtColor(img,
cv2.COLOR_BGR2GRAY)[...,
None]
elif mode_type == SPTF.MODE_GGG:
out_sample = np.repeat(
np.expand_dims(
cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
(3, ), -1)
if not debug:
if normalize_tanh:
out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
1.0)
if data_format == "NCHW":
out_sample = np.transpose(out_sample, (2, 0, 1))
outputs_sample.append(out_sample)
outputs += [outputs_sample]
return outputs
def process(samples,
sample_process_options,
output_sample_types,
debug,
ct_sample=None):
SPST = SampleProcessor.SampleType
SPCT = SampleProcessor.ChannelType
SPFMT = SampleProcessor.FaceMaskType
sample_rnd_seed = np.random.randint(0x80000000)
outputs = []
for sample in samples:
sample_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)
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 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 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),
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),
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),
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
motion_blur_chance, motion_blur_mb_max_size = 25, 5
gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5
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=False)
bg_mask = imagelib.warp_by_params (bg_wp, bg_mask, can_warp=False, can_transform=True, can_flip=True, border_replicate=False)
c_mask = (1-bg_mask) * (1-mask)
img = img*(1-c_mask) + bg_img * c_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 )
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]
def process(samples,
sample_process_options,
output_sample_types,
debug,
ct_sample=None):
SPST = SampleProcessor.SampleType
SPCT = SampleProcessor.ChannelType
SPFMT = SampleProcessor.FaceMaskType
sample_rnd_seed = np.random.randint(0x80000000)
outputs = []
for sample in samples:
sample_bgr = sample.load_bgr()
ct_sample_bgr = None
h, w, c = sample_bgr.shape
is_face_sample = sample.landmarks is not None
if debug and is_face_sample:
LandmarksProcessor.draw_landmarks(sample_bgr, sample.landmarks,
(0, 1, 0))
params = imagelib.gen_warp_params(
sample_bgr,
sample_process_options.random_flip,
rotation_range=sample_process_options.rotation_range,
scale_range=sample_process_options.scale_range,
tx_range=sample_process_options.tx_range,
ty_range=sample_process_options.ty_range)
outputs_sample = []
for opts in output_sample_types:
sample_type = opts.get('sample_type', SPST.NONE)
channel_type = opts.get('channel_type', SPCT.NONE)
resolution = opts.get('resolution', 0)
warp = opts.get('warp', False)
transform = opts.get('transform', False)
motion_blur = opts.get('motion_blur', None)
gaussian_blur = opts.get('gaussian_blur', None)
normalize_tanh = opts.get('normalize_tanh', False)
ct_mode = opts.get('ct_mode', 'None')
data_format = opts.get('data_format', 'NHWC')
if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
if not is_face_sample:
raise ValueError(
"face_samples should be provided for sample_type FACE_*"
)
if is_face_sample:
face_type = opts.get('face_type', None)
face_mask_type = opts.get('face_mask_type', SPFMT.NONE)
if face_type is None:
raise ValueError(
"face_type must be defined for face samples")
if face_type > sample.face_type:
raise Exception(
'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.'
% (sample.filename, sample.face_type, target_ft))
if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
if sample_type == SPST.FACE_MASK:
if face_mask_type == SPFMT.ALL_HULL or \
face_mask_type == SPFMT.EYES_HULL or \
face_mask_type == SPFMT.ALL_EYES_HULL:
if face_mask_type == SPFMT.ALL_HULL or \
face_mask_type == SPFMT.ALL_EYES_HULL:
if sample.eyebrows_expand_mod is not None:
all_mask = LandmarksProcessor.get_image_hull_mask(
sample_bgr.shape,
sample.landmarks,
eyebrows_expand_mod=sample.
eyebrows_expand_mod)
else:
all_mask = LandmarksProcessor.get_image_hull_mask(
sample_bgr.shape, sample.landmarks)
all_mask = np.clip(all_mask, 0, 1)
if face_mask_type == SPFMT.EYES_HULL or \
face_mask_type == SPFMT.ALL_EYES_HULL:
eyes_mask = LandmarksProcessor.get_image_eye_mask(
sample_bgr.shape, sample.landmarks)
eyes_mask = np.clip(eyes_mask, 0, 1)
if face_mask_type == SPFMT.ALL_HULL:
img = all_mask
elif face_mask_type == SPFMT.EYES_HULL:
img = eyes_mask
elif face_mask_type == SPFMT.ALL_EYES_HULL:
img = all_mask + eyes_mask
elif face_mask_type == SPFMT.STRUCT:
if sample.eyebrows_expand_mod is not None:
img = LandmarksProcessor.get_face_struct_mask(
sample_bgr.shape,
sample.landmarks,
eyebrows_expand_mod=sample.
eyebrows_expand_mod)
else:
img = LandmarksProcessor.get_face_struct_mask(
sample_bgr.shape, sample.landmarks)
if sample.ie_polys is not None:
sample.ie_polys.overlay_mask(img)
if sample.face_type == FaceType.MARK_ONLY:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, sample.shape[0], face_type)
img = cv2.warpAffine(
img,
mat, (sample.shape[0], sample.shape[0]),
flags=cv2.INTER_LINEAR)
img = imagelib.warp_by_params(
params,
img,
warp,
transform,
can_flip=True,
border_replicate=False,
cv2_inter=cv2.INTER_LINEAR)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_LINEAR)[..., None]
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, face_type)
img = imagelib.warp_by_params(
params,
img,
warp,
transform,
can_flip=True,
border_replicate=False,
cv2_inter=cv2.INTER_LINEAR)
img = cv2.warpAffine(
img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_LINEAR)[..., None]
if channel_type == SPCT.G:
out_sample = img.astype(np.float32)
else:
raise ValueError(
"only channel_type.G supported for the mask")
elif sample_type == SPST.FACE_IMAGE:
img = sample_bgr
if motion_blur is not None:
chance, mb_max_size = motion_blur
chance = np.clip(chance, 0, 100)
l_rnd_state = np.random.RandomState(
sample_rnd_seed)
mblur_rnd_chance = l_rnd_state.randint(100)
mblur_rnd_kernel = l_rnd_state.randint(
mb_max_size) + 1
mblur_rnd_deg = l_rnd_state.randint(360)
if mblur_rnd_chance < chance:
img = imagelib.LinearMotionBlur(
img, mblur_rnd_kernel, mblur_rnd_deg)
if gaussian_blur is not None:
chance, kernel_max_size = gaussian_blur
chance = np.clip(chance, 0, 100)
l_rnd_state = np.random.RandomState(
sample_rnd_seed + 1)
gblur_rnd_chance = l_rnd_state.randint(100)
gblur_rnd_kernel = l_rnd_state.randint(
kernel_max_size) * 2 + 1
if gblur_rnd_chance < chance:
img = cv2.GaussianBlur(
img, (gblur_rnd_kernel, ) * 2, 0)
if sample.face_type == FaceType.MARK_ONLY:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, sample.shape[0], face_type)
img = cv2.warpAffine(
img,
mat, (sample.shape[0], sample.shape[0]),
flags=cv2.INTER_CUBIC)
img = imagelib.warp_by_params(
params,
img,
warp,
transform,
can_flip=True,
border_replicate=True)
img = cv2.resize(img, (resolution, resolution),
cv2.INTER_CUBIC)
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, face_type)
img = imagelib.warp_by_params(
params,
img,
warp,
transform,
can_flip=True,
border_replicate=True)
img = cv2.warpAffine(
img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
img = np.clip(img.astype(np.float32), 0, 1)
# Apply random color transfer
if ct_mode is not None and ct_sample is not None:
if ct_sample_bgr is None:
ct_sample_bgr = ct_sample.load_bgr()
img = imagelib.color_transfer(
ct_mode, img,
cv2.resize(ct_sample_bgr,
(resolution, resolution),
cv2.INTER_LINEAR))
# Transform from BGR to desired channel_type
if channel_type == SPCT.BGR:
out_sample = img
elif channel_type == SPCT.BGR_SHUFFLE:
l_rnd_state = np.random.RandomState(
sample_rnd_seed)
out_sample = np.take(img,
l_rnd_state.permutation(
img.shape[-1]),
axis=-1)
elif channel_type == SPCT.BGR_RANDOM_HSV_SHIFT:
l_rnd_state = np.random.RandomState(
sample_rnd_seed)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
h, s, v = cv2.split(hsv)
h = (h + l_rnd_state.randint(360)) % 360
s = np.clip(s + l_rnd_state.random() - 0.5, 0, 1)
v = np.clip(v + l_rnd_state.random() - 0.5, 0, 1)
hsv = cv2.merge([h, s, v])
out_sample = np.clip(
cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR), 0, 1)
elif channel_type == SPCT.BGR_RANDOM_RGB_LEVELS:
l_rnd_state = np.random.RandomState(
sample_rnd_seed)
np_rnd = l_rnd_state.rand
inBlack = np.array([
np_rnd() * 0.25,
np_rnd() * 0.25,
np_rnd() * 0.25
],
dtype=np.float32)
inWhite = np.array([
1.0 - np_rnd() * 0.25, 1.0 - np_rnd() * 0.25,
1.0 - np_rnd() * 0.25
],
dtype=np.float32)
inGamma = np.array([
0.5 + np_rnd(), 0.5 + np_rnd(), 0.5 + np_rnd()
],
dtype=np.float32)
outBlack = np.array([0.0, 0.0, 0.0],
dtype=np.float32)
outWhite = np.array([1.0, 1.0, 1.0],
dtype=np.float32)
out_sample = np.clip(
(img - inBlack) / (inWhite - inBlack), 0, 1)
out_sample = (out_sample**(1 / inGamma)) * (
outWhite - outBlack) + outBlack
out_sample = np.clip(out_sample, 0, 1)
elif channel_type == SPCT.G:
out_sample = cv2.cvtColor(img,
cv2.COLOR_BGR2GRAY)[...,
None]
elif channel_type == SPCT.GGG:
out_sample = np.repeat(
np.expand_dims(
cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), -1),
(3, ), -1)
# Final transformations
if not debug:
if normalize_tanh:
out_sample = np.clip(out_sample * 2.0 - 1.0, -1.0,
1.0)
if data_format == "NCHW":
out_sample = np.transpose(out_sample, (2, 0, 1))
#else:
# img = imagelib.warp_by_params (params, img, warp, transform, can_flip=True, border_replicate=True)
# img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
elif sample_type == SPST.LANDMARKS_ARRAY:
l = sample.landmarks
l = np.concatenate([
np.expand_dims(l[:, 0] / w, -1),
np.expand_dims(l[:, 1] / h, -1)
], -1)
l = np.clip(l, 0.0, 1.0)
out_sample = l
elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
pitch_yaw_roll = sample.get_pitch_yaw_roll()
if params['flip']:
yaw = -yaw
if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)
out_sample = (pitch, yaw, roll)
else:
raise ValueError('expected sample_type')
outputs_sample.append(out_sample)
outputs += [outputs_sample]
return outputs
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)
random_downsample = opts.get('random_downsample', False)
random_noise = opts.get('random_noise', False)
random_blur = opts.get('random_blur', False)
random_jpeg = opts.get('random_jpeg', 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))
randomization_order = ['blur', 'noise', 'jpeg', 'down']
np.random.shuffle(randomization_order)
for random_distortion in randomization_order:
# Apply random blur
if random_distortion == 'blur' and random_blur:
blur_type = np.random.choice(
['motion', 'gaussian'])
if blur_type == 'motion':
blur_k = np.random.randint(10, 20)
blur_angle = 360 * np.random.random()
img = LinearMotionBlur(
img, blur_k, blur_angle)
elif blur_type == 'gaussian':
blur_sigma = 5 * np.random.random() + 3
if blur_sigma < 5.0:
kernel_size = 2.9 * blur_sigma # 97% of weight
else:
kernel_size = 2.6 * blur_sigma # 95% of weight
kernel_size = int(kernel_size)
kernel_size = kernel_size + 1 if kernel_size % 2 == 0 else kernel_size
img = cv2.GaussianBlur(
img, (kernel_size, kernel_size),
blur_sigma)
# Apply random noise
if random_distortion == 'noise' and random_noise:
noise_type = np.random.choice(
['gaussian', 'laplace', 'poisson'])
noise_scale = (20 * np.random.random() + 20)
if noise_type == 'gaussian':
noise = np.random.normal(scale=noise_scale,
size=img.shape)
img += noise / 255.0
elif noise_type == 'laplace':
noise = np.random.laplace(
scale=noise_scale, size=img.shape)
img += noise / 255.0
elif noise_type == 'poisson':
noise_lam = (15 * np.random.random() + 15)
noise = np.random.poisson(lam=noise_lam,
size=img.shape)
img += noise / 255.0
# Apply random jpeg compression
if random_distortion == 'jpeg' and random_jpeg:
img = np.clip(img * 255, 0,
255).astype(np.uint8)
jpeg_compression_level = np.random.randint(
50, 85)
encode_param = [
int(cv2.IMWRITE_JPEG_QUALITY),
jpeg_compression_level
]
_, enc_img = cv2.imencode(
'.jpg', img, encode_param)
img = cv2.imdecode(
enc_img, cv2.IMREAD_UNCHANGED).astype(
np.float32) / 255.0
# Apply random downsampling
if random_distortion == 'down' and random_downsample:
down_res = np.random.randint(
int(0.125 * resolution),
int(0.25 * resolution))
img = cv2.resize(img, (down_res, down_res),
interpolation=cv2.INTER_CUBIC)
img = cv2.resize(img, (resolution, resolution),
interpolation=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)
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 process(samples,
sample_process_options,
output_sample_types,
debug,
ct_sample=None):
SPST = SampleProcessor.SampleType
SPCT = SampleProcessor.ChannelType
SPFMT = SampleProcessor.FaceMaskType
outputs = []
for sample in samples:
sample_rnd_seed = np.random.randint(0x80000000)
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)
def get_eyes_mouth_mask():
eyes_mask = LandmarksProcessor.get_image_eye_mask(
sample_bgr.shape, sample_landmarks)
mouth_mask = LandmarksProcessor.get_image_mouth_mask(
sample_bgr.shape, sample_landmarks)
mask = eyes_mask + mouth_mask
return np.clip(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))
outputs_sample = []
for opts in output_sample_types:
resolution = opts.get('resolution', 0)
sample_type = opts.get('sample_type', SPST.NONE)
channel_type = opts.get('channel_type', SPCT.NONE)
nearest_resize_to = opts.get('nearest_resize_to', None)
warp = opts.get('warp', False)
transform = opts.get('transform', False)
random_hsv_shift_amount = opts.get('random_hsv_shift_amount',
0)
normalize_tanh = opts.get('normalize_tanh', False)
ct_mode = opts.get('ct_mode', None)
data_format = opts.get('data_format', 'NHWC')
rnd_seed_shift = opts.get('rnd_seed_shift', 0)
warp_rnd_seed_shift = opts.get('warp_rnd_seed_shift',
rnd_seed_shift)
rnd_state = np.random.RandomState(sample_rnd_seed +
rnd_seed_shift)
warp_rnd_state = np.random.RandomState(sample_rnd_seed +
warp_rnd_seed_shift)
warp_params = 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=rnd_state,
warp_rnd_state=warp_rnd_state,
)
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.EYES_MOUTH:
mask = get_full_face_mask().copy()
mask[mask != 0.0] = 1.0
img = get_eyes_mouth_mask() * mask
else:
img = np.zeros(sample_bgr.shape[0:2] + (1, ),
dtype=np.float32)
if sample_face_type == FaceType.MARK_ONLY:
raise NotImplementedError()
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(
warp_params,
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(
warp_params,
img,
warp,
transform,
can_flip=True,
border_replicate=border_replicate,
cv2_inter=cv2.INTER_LINEAR)
if face_mask_type == SPFMT.EYES_MOUTH:
div = img.max()
if div != 0.0:
img = img / div # normalize to 1.0 after warp
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),
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))
if random_hsv_shift_amount != 0:
a = random_hsv_shift_amount
h_amount = max(1, int(360 * a * 0.5))
img_h, img_s, img_v = cv2.split(
cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
img_h = (img_h + rnd_state.randint(
-h_amount, h_amount + 1)) % 360
img_s = np.clip(
img_s + (rnd_state.random() - 0.5) * a, 0, 1)
img_v = np.clip(
img_v + (rnd_state.random() - 0.5) * a, 0, 1)
img = np.clip(
cv2.cvtColor(cv2.merge([img_h, img_s, img_v]),
cv2.COLOR_HSV2BGR), 0, 1)
img = imagelib.warp_by_params(
warp_params,
img,
warp,
transform,
can_flip=True,
border_replicate=border_replicate)
img = np.clip(img.astype(np.float32), 0, 1)
# 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 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(warp_params,
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 warp_params['flip']:
yaw = -yaw
if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
pitch = np.clip((pitch / math.pi) / 2.0 + 0.5, 0, 1)
yaw = np.clip((yaw / math.pi) / 2.0 + 0.5, 0, 1)
roll = np.clip((roll / math.pi) / 2.0 + 0.5, 0, 1)
out_sample = (pitch, yaw)
else:
raise ValueError('expected sample_type')
outputs_sample.append(out_sample)
outputs += [outputs_sample]
return outputs
def batch_func(self, param):
images_path, masks_path, mask_file_id_hash, data_format = param
file_ids = list(mask_file_id_hash.keys())
shuffle_file_ids = []
resolution = 256
random_flip = True
rotation_range = [-15, 15]
scale_range = [-0.10, 0.95]
tx_range = [-0.3, 0.3]
ty_range = [-0.3, 0.3]
random_bilinear_resize = (25, 75)
motion_blur = (25, 5)
gaussian_blur = (25, 5)
bs = self.batch_size
while True:
batches = None
n_batch = 0
while n_batch < bs:
try:
if len(shuffle_file_ids) == 0:
shuffle_file_ids = file_ids.copy()
np.random.shuffle(shuffle_file_ids)
file_id = shuffle_file_ids.pop()
masks = mask_file_id_hash[file_id]
image_path = images_path / f'{file_id}.jpg'
skin_path = masks.get(MaskType.skin, None)
hair_path = masks.get(MaskType.hair, None)
hat_path = masks.get(MaskType.hat, None)
#neck_path = masks.get(MaskType.neck, None)
img = cv2_imread(image_path).astype(np.float32) / 255.0
mask = cv2_imread(masks_path / skin_path)[..., 0:1].astype(
np.float32) / 255.0
if hair_path is not None:
hair_path = masks_path / hair_path
if hair_path.exists():
hair = cv2_imread(hair_path)[..., 0:1].astype(
np.float32) / 255.0
mask *= (1 - hair)
if hat_path is not None:
hat_path = masks_path / hat_path
if hat_path.exists():
hat = cv2_imread(hat_path)[..., 0:1].astype(
np.float32) / 255.0
mask *= (1 - hat)
#if neck_path is not None:
# neck_path = masks_path / neck_path
# if neck_path.exists():
# neck = cv2_imread(neck_path)[...,0:1].astype(np.float32) / 255.0
# mask = np.clip(mask+neck, 0, 1)
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 = cv2.resize(img, (resolution, resolution),
cv2.INTER_LANCZOS4)
h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
h = (h + np.random.randint(360)) % 360
s = np.clip(s + np.random.random() - 0.5, 0, 1)
v = np.clip(v + np.random.random() / 2 - 0.25, 0, 1)
img = np.clip(
cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR),
0, 1)
if motion_blur is not None:
chance, mb_max_size = motion_blur
chance = np.clip(chance, 0, 100)
mblur_rnd_chance = np.random.randint(100)
mblur_rnd_kernel = np.random.randint(mb_max_size) + 1
mblur_rnd_deg = np.random.randint(360)
if mblur_rnd_chance < chance:
img = imagelib.LinearMotionBlur(
img, mblur_rnd_kernel, mblur_rnd_deg)
img = imagelib.warp_by_params(warp_params,
img,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=False,
cv2_inter=cv2.INTER_LANCZOS4)
if gaussian_blur is not None:
chance, kernel_max_size = gaussian_blur
chance = np.clip(chance, 0, 100)
gblur_rnd_chance = np.random.randint(100)
gblur_rnd_kernel = np.random.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:
chance, max_size_per = random_bilinear_resize
chance = np.clip(chance, 0, 100)
pick_chance = np.random.randint(100)
resize_to = resolution - int(
np.random.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)
mask = cv2.resize(mask, (resolution, resolution),
cv2.INTER_LANCZOS4)[..., None]
mask = imagelib.warp_by_params(
warp_params,
mask,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=False,
cv2_inter=cv2.INTER_LANCZOS4)
mask[mask < 0.5] = 0.0
mask[mask >= 0.5] = 1.0
mask = np.clip(mask, 0, 1)
if data_format == "NCHW":
img = np.transpose(img, (2, 0, 1))
mask = np.transpose(mask, (2, 0, 1))
if batches is None:
batches = [[], []]
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, 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 batch_func(self, param):
samples, kf_idxs, resolution, face_type, data_format = param
kf_idxs_len = len(kf_idxs)
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]
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()
key_idx, key_chain_idx, chain_idxs = kf_idxs[
np.random.randint(kf_idxs_len)]
key_sample = samples[key_idx]
key_chain_sample = samples[key_chain_idx]
chain_sample = samples[chain_idxs[np.random.randint(
len(chain_idxs))]]
#print('==========')
#print(key_sample.filename)
#print(key_chain_sample.filename)
#print(chain_sample.filename)
sample = samples[idx]
img = sample.load_bgr()
key_img = key_sample.load_bgr()
key_chain_img = key_chain_sample.load_bgr()
chain_img = chain_sample.load_bgr()
h, w, c = img.shape
mask = LandmarksProcessor.get_image_hull_mask(
img.shape, sample.landmarks)
mask = np.clip(mask, 0, 1)
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_CUBIC)
key_img = cv2.resize(key_img,
(resolution, resolution),
cv2.INTER_CUBIC)
key_chain_img = cv2.resize(
key_chain_img, (resolution, resolution),
cv2.INTER_CUBIC)
chain_img = cv2.resize(chain_img,
(resolution, resolution),
cv2.INTER_CUBIC)
mask = cv2.resize(mask, (resolution, resolution),
cv2.INTER_CUBIC)
else:
mat = LandmarksProcessor.get_transform_mat(
sample.landmarks, resolution, face_type)
img = cv2.warpAffine(img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
key_img = cv2.warpAffine(
key_img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
key_chain_img = cv2.warpAffine(
key_chain_img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
chain_img = cv2.warpAffine(
chain_img,
mat, (resolution, resolution),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
mask = cv2.warpAffine(mask,
mat, (resolution, resolution),
borderMode=cv2.BORDER_CONSTANT,
flags=cv2.INTER_CUBIC)
if len(mask.shape) == 2:
mask = mask[..., None]
img_warped = imagelib.warp_by_params(warp_params,
img,
can_warp=True,
can_transform=True,
can_flip=True,
border_replicate=True)
img_transformed = imagelib.warp_by_params(
warp_params,
img,
can_warp=False,
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)
key_img = imagelib.warp_by_params(warp_params,
key_img,
can_warp=False,
can_transform=False,
can_flip=False,
border_replicate=True)
key_chain_img = imagelib.warp_by_params(
warp_params,
key_chain_img,
can_warp=False,
can_transform=False,
can_flip=False,
border_replicate=True)
chain_img = imagelib.warp_by_params(warp_params,
chain_img,
can_warp=False,
can_transform=False,
can_flip=False,
border_replicate=True)
img_warped = np.clip(img_warped.astype(np.float32), 0, 1)
img_transformed = np.clip(
img_transformed.astype(np.float32), 0, 1)
mask[mask < 0.5] = 0.0
mask[mask >= 0.5] = 1.0
mask = np.clip(mask, 0, 1)
if data_format == "NCHW":
img_warped = np.transpose(img_warped, (2, 0, 1))
img_transformed = np.transpose(img_transformed,
(2, 0, 1))
mask = np.transpose(mask, (2, 0, 1))
key_img = np.transpose(key_img, (2, 0, 1))
key_chain_img = np.transpose(key_chain_img, (2, 0, 1))
chain_img = np.transpose(chain_img, (2, 0, 1))
batches[0].append(img_warped)
batches[1].append(img_transformed)
batches[2].append(mask)
batches[3].append(key_img)
batches[4].append(key_chain_img)
batches[5].append(chain_img)
n_batch += 1
except:
io.log_err(traceback.format_exc())
yield [np.array(batch) for batch in batches]
