def resample_mask(id: int):
"""
Resamples the mask of image_id with similar face masks (face + hair is kept from original)
:param id: id of mask image in celebaHQ
:return: resampled mask
"""
# get real parts we want to keep similar
skin = create_mask_from_id(id,
skin=True, nose=False, glasses=False, eyes=False, brows=False, ears=True,
mouth=False, u_lip=False, l_lip=False, hair=False, neck=True, misc=False)
hair = create_mask_from_id(id,
skin=False, nose=False, glasses=False, eyes=False, brows=False, ears=True,
mouth=False, u_lip=False, l_lip=False, hair=True, neck=False, misc=True)
parts_to_resample = create_mask_from_id(id,
skin=False, nose=True, glasses=True, eyes=True, brows=True, ears=False,
mouth=True, u_lip=True, l_lip=True, hair=False, neck=False, misc=False)
# get similar masks
max_score = -1
for i in range(40):
# get random id
rand_id = id
while rand_id is id:
# dont sample from training data?
# 0-6000 is test that I use
# 6001-25200 is training
# 25201-30000 is val that I never touched
rand_id = random.randint(0, 6000)
random_parts = create_mask_from_id(rand_id,
skin=False, nose=True, glasses=True, eyes=True, brows=True, ears=False,
mouth=True, u_lip=True, l_lip=True, hair=False, neck=False, misc=False)
# compute difference of masks
(score, _) = structural_similarity(np.array(parts_to_resample), np.array(random_parts),
full=True, multichannel=True)
# keep highest score
if score > max_score:
resampled_parts = random_parts
max_score = score
new_id = rand_id
print(f"Image {id} resampled to {new_id}")
# combine real and random masks into one mask
add_real_background = True
if add_real_background:
resampled_mask = Image.open('/home/mo/datasets/CelebAMask-HQ/CelebA-HQ-img/' + str(id) +".jpg").resize((256, 256), Image.NEAREST).convert("RGBA")
skin = make_color_transparent(skin, (0, 0, 0))
else:
resampled_mask = Image.new("RGBA", (256, 256), "BLACK")
resampled_mask.paste(skin, (0, 0), skin)
resampled_parts = make_color_transparent(resampled_parts, (0, 0, 0))
resampled_mask.paste(resampled_parts, (0, 0), resampled_parts)
hair = make_color_transparent(hair, (0, 0, 0))
resampled_mask.paste(hair, (0, 0), hair)
return resampled_mask.convert("RGB")
def render_frames(frames: list):
"""
Renders list of PIL or openCV images
:type frames: object
:return merged image of same type as imput image (PIL or openCV)
"""
opencv = False
try:
# assume Opencv image
shape = (frames[0].shape[0], frames[0].shape[1])
opencv = True
except AttributeError:
# we dont have a opencv image
pass
if not opencv:
# so we have a PIL image
shape = frames[0].size
background = Image.new("RGBA", shape, "BLACK")
for frame in frames:
if opencv:
frame = Image.fromarray(frame)
frame = make_color_transparent(frame, (0, 0, 0))
background.paste(frame, (0, 0), frame)
if opencv:
return np.array(background)
else:
return background
def run_segmentation_model(model, opt, root_folder, warp=False):
print("\n---- Running Face2Mask Model ----\n")
dataset = create_dataset(opt)
for i, data in tqdm(enumerate(dataset)):
if i >= opt.num_test:
break
# run model and get processed image
model.set_input(data)
model.test()
# get mask
visuals_mask = model.get_current_visuals()
mask = tensor2im(visuals_mask['fake_B'])
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)
mask = clean_mask(mask)
mask = Image.fromarray(mask)
mask = make_color_transparent(mask, (0, 0, 0), tolerance=50)
# open original
img_path = model.get_image_paths()[0]
img_name = img_path.split("/")[-1][:-4]
image = Image.open(img_path)
# add background to mask
overlayed_mask = overlay_two_images(image, mask)
# save file
overlayed_mask.save(f'{root_folder}/{img_name}.png')
print("SAVED GENERATED MASKS\n")
def face_mask_from_keypoints(keypoints, img):
"""
Converts keypoints to face mask
:param keypoints: tuple of keypoints
:param img: image
:return: mask
"""
# define mask regions
upper_keypoints = keypoints[18:27].copy()
# invert direction
upper_keypoints = upper_keypoints[::-1]
# add padding to eyebrows
padded_upper_keypoints = apply_padding_in_x(upper_keypoints)
# add padded keypoints to make face larger
skin_keypoints = np.concatenate((keypoints[0:17], padded_upper_keypoints))
nose_keypoints = np.concatenate((keypoints[27:28], keypoints[31:36]))
keypoint_dict = {
'skin': {
'keypoints': skin_keypoints, 'color': (204, 0, 0)},
'eyebrow_l': {
'keypoints': keypoints[18:22], 'color': (0, 255, 255)},
'eyebrow_r': {
'keypoints': keypoints[22:27], 'color': (0, 255, 255)},
'nose': {
'keypoints': nose_keypoints, 'color': (76, 153, 0)},
'eye_l': {
'keypoints': keypoints[36:42], 'color': (51, 51, 255)},
'eye_r': {
'keypoints': keypoints[42:48], 'color': (51, 51, 255)},
'lips': {
'keypoints': keypoints[48:60], 'color': (255, 255, 0)},
'mouth': {
'keypoints': keypoints[60:69], 'color': (102, 204, 0)}
}
mask = []
for key, value in keypoint_dict.items():
dilate_radius = 1
if 'eyebrow' in key:
dilate_radius = 2
mask.append(keypoints_to_mask(value['keypoints'], img.shape, dilate_radius, value['color']))
for i in range(len(mask) - 1):
if i is 0:
back = Image.fromarray(mask[i]).convert("RGBA")
front = Image.fromarray(mask[i + 1]).convert("RGBA")
front = make_color_transparent(front, (0, 0, 0))
# back is always assigned as i starts from 0
back.paste(front, (0, 0), front)
return np.array(back.convert("RGB"))
def insert_face(img, face, add_border=True):
"""
inserts face onto image
:param add_border: add black border around face
:param img: Image to paste onto
:param face: face image with removed background
:return: merged images
"""
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
face = make_color_transparent(face, (0, 0, 0))
if add_border:
black = np.array(face)
black[..., :3] = [0, 0, 0]
dilated = dilate(black, 7)
dilated = Image.fromarray(dilated)
dilated.paste(face, (0, 0), face)
face = dilated
img = Image.fromarray(img)
img.paste(face, (0, 0), face)
return img
def preprocess_images(data_points: int, data_path: str, processed_path: str,
train: float, save_a_b: bool, no_misc: bool,
no_hair: bool, no_neck: bool, no_ears: bool,
overlay_mask: bool):
"""
Preprocess images according to pix2pix requirements
"""
print(
f'Processing with following options:\n'
f'Overlay Mask: {overlay_mask}, save_a_b: {save_a_b}, train-percentage: {train * 100}%\n'
f'Misc: {not no_misc}, Hair: {not no_hair}, Neck: {not no_neck}, Ears {not no_ears}'
)
data_path = Path(data_path)
processed_path = Path(processed_path)
if save_a_b:
create_folder_structure(processed_path / 'A')
create_folder_structure(processed_path / 'B')
else:
create_folder_structure(processed_path)
test = 1 - train
for i in tqdm(range(data_points)):
# load and resize 512x512 px mask image
mask_img = create_mask_from_id(
i,
skin=True,
nose=True,
glasses=True,
eyes=True,
brows=True,
ears=not no_ears,
mouth=True,
u_lip=True,
l_lip=True,
hair=not no_hair,
neck=not no_neck,
misc=not no_misc,
)
# load and resize corresponding 1028x1028px image
normal_img = im_resize(Image.open(data_path / f"{i}.jpg"))
# split dataset
if i < test * data_points:
folder = "test"
elif i < (test + train * train) * data_points:
folder = "train"
else:
# using the val images as train as well as we dont really have validation with pix2pix
folder = "train"
if overlay_mask:
mask_img = make_color_transparent(mask_img, (0, 0, 0))
mask_background = normal_img.copy()
mask_background.paste(mask_img, (0, 0), mask_img)
mask_background.convert("RGB")
mask_img = mask_background
# save images
if save_a_b:
normal_img.save(processed_path / "A" / folder / f"{i}.png",
format='PNG',
subsampling=0,
quality=100)
mask_img.save(processed_path / "B" / folder / f"{i}.png",
format='PNG',
subsampling=0,
quality=100)
else:
# merge the two images into one 256x512px image
output = merge_images_side_by_side(im_resize(normal_img),
im_resize(mask_img))
# output.save(processed_path / folder / f"{i}.png", format='PNG', subsampling=0, quality=100)
output.save(processed_path / folder / f"{i}.png",
format='PNG',
subsampling=0,
quality=100)
def create_mask_from_id(
pic_id: int,
face_only_segmentation=False,
warp=False,
showGUI=False,
data_path=Path('/home/mo/datasets/CelebAMask-HQ/CelebAMask-HQ-mask-anno'),
skin=True,
nose=True,
glasses=True,
eyes=True,
brows=True,
ears=True,
mouth=True,
u_lip=True,
l_lip=True,
hair=True,
neck=True,
misc=True,
):
"""Creates one image with all segmentation masks
:param showGUI: displays hui
:type data_path: Path
:param warp: warps the masks randomly
:param face_only_segmentation: only face
:param pic_id: id of image
:param data_path: path object of data root containing masks
:return: the merged Image object
"""
sub_folder = str(pic_id // 2000)
pic_id = str(pic_id).zfill(5)
segments = get_segments()
unprocessed_segments = 0
frames = []
warped_frames = []
# seed for warping
if warp:
deg = random.randint(0, 10)
if random.getrandbits(1):
deg = 360 - deg
seed = [
random.getrandbits(1),
random.getrandbits(1),
random.getrandbits(1), deg
]
# render loop
for key, segment in segments.items():
if face_only_segmentation:
if "skin" in key:
segment["color"] = (255, 255, 255)
else:
# stop to add other segmentation masks
break
try:
for name_id in range(len(segment["name"])):
if not skin and "skin" in key:
break
if not nose and "nose" in key:
break
if not glasses and "glasses" in key:
break
if not eyes and "eye" in key:
break
if not brows and "brows" in key:
break
if not ears and "ear" in key:
break
if not mouth and "mouth" in key:
break
if not u_lip and "u_lip" in key:
break
if not l_lip and "l_lip" in key:
break
if not hair and "hair" in key:
# we want to remove the regions of masks that are below the hair mask
pass
if not neck and "neck" in key:
break
if not misc and "misc" in key:
break
# open segmentation mask
img = Image.open(data_path / sub_folder /
(pic_id + segment["name"][name_id]))
# recolor white to respective color of segment
img = recolor_to_rgba(img, (255, 255, 255), segment["color"])
# make black transparent
img = make_color_transparent(img, (0, 0, 0))
# save individual frames for later modification
frames.append(img)
# warp image
if warp:
img = add_random_warp(img, seed=seed)
warped_frames.append(img)
except FileNotFoundError:
unprocessed_segments = unprocessed_segments + 1
if unprocessed_segments == len(segments):
print("No segmentation maps for file " + pic_id +
" found! Stopping.")
return
# render frames
if warp:
result = render_frames(warped_frames)
else:
result = render_frames(frames)
if showGUI:
result = render_gui(frames)
if not hair:
# remove hair region again, this ensures no mask sticks out "below" the hair
result = make_color_transparent(result, segments['hair']["color"])
return im_resize(result)
def cutout_face(img_path: str, mask_path: str, processed_path: str,
face_mask_path: str):
"""Cuts out face of person
:param face_mask_path:
:param processed_path:
:param img_path: path of folder containing images
:param mask_path: path of folder containing masks
"""
# get all filenames in the folders
img_file_list = get_files(img_path, '.jpg')
mask_file_list = get_files(mask_path, '.png')
if img_file_list is None or mask_file_list is None:
print("No image files found. Exiting")
return
if not len(img_file_list) == len(mask_file_list):
print("Not the same amount og image files and mask files. Exiting")
return
# create folder for processed data
processed_path = Path(processed_path)
img_path = Path(img_path)
face_mask_path = Path(face_mask_path)
mask_path = Path(mask_path)
processed_path.mkdir(exist_ok=True)
for file in tqdm(img_file_list):
# Load images
face_img = Image.open(img_path / file.name)
face_mask_img = Image.open(face_mask_path / f"{file.stem}.png")
mask_img = Image.open(mask_path / f"{file.stem}.png")
# Create dilated face mask as transparent, background is black
face_mask_img_dilated = cv.imread(
cv.samples.findFile(str(face_mask_path / f"{file.stem}.png")))
face_mask_img_dilated = dilate(face_mask_img_dilated, 20)
face_mask_img_dilated = cv.cvtColor(face_mask_img_dilated,
cv.COLOR_BGR2RGB)
face_mask_img_dilated = Image.fromarray(face_mask_img_dilated)
face_mask_img_dilated = make_color_transparent(face_mask_img_dilated,
(255, 255, 255))
# Make face mask transparent, background is black
face_mask_img = make_color_transparent(face_mask_img, (255, 255, 255))
# copy face mask onto the mask image
overlaid_mask = overlay_two_images(mask_img, face_mask_img)
# Make background transparent
overlaid_mask = make_color_transparent(overlaid_mask, (0, 0, 0))
# overlay dilated face onto image
overlaid_img = overlay_two_images(face_img, face_mask_img_dilated)
# overlay mask onto image
overlaid_mask = overlay_two_images(overlaid_img, overlaid_mask)
# save result
merge_images_side_by_side(overlaid_mask, overlaid_img).save(
processed_path / file.name,
format='PNG',
subsampling=0,
quality=100)