File size: 4,763 Bytes
caa56d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
from enum import Enum
from functools import reduce
import cv2
import numpy as np
from scipy.ndimage import binary_dilation
from .DeepFakeMask import Mask
def dist(a, b):
x1, y1 = a
x2, y2 = b
return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
# return np.linalg.norm(a-b)
def get_five_key(landmarks_68):
# get the five key points by using the landmarks
leye_center = (landmarks_68[36] + landmarks_68[39]) * 0.5
reye_center = (landmarks_68[42] + landmarks_68[45]) * 0.5
nose = landmarks_68[33]
lmouth = landmarks_68[48]
rmouth = landmarks_68[54]
leye_left = landmarks_68[36]
leye_right = landmarks_68[39]
reye_left = landmarks_68[42]
reye_right = landmarks_68[45]
out = [
tuple(x.astype("int32"))
for x in [
leye_center,
reye_center,
nose,
lmouth,
rmouth,
leye_left,
leye_right,
reye_left,
reye_right,
]
]
return out
def remove_eyes(image, landmarks, opt):
##l: left eye; r: right eye, b: both eye
if opt == "l":
(x1, y1), (x2, y2) = landmarks[5:7]
elif opt == "r":
(x1, y1), (x2, y2) = landmarks[7:9]
elif opt == "b":
(x1, y1), (x2, y2) = landmarks[:2]
else:
print("wrong region")
mask = np.zeros_like(image[..., 0])
line = cv2.line(mask, (x1, y1), (x2, y2), color=(1), thickness=2)
w = dist((x1, y1), (x2, y2))
dilation = int(w // 4)
if opt != "b":
dilation *= 4
line = binary_dilation(line, iterations=dilation)
return line
def remove_nose(image, landmarks):
(x1, y1), (x2, y2) = landmarks[:2]
x3, y3 = landmarks[2]
mask = np.zeros_like(image[..., 0])
x4 = int((x1 + x2) / 2)
y4 = int((y1 + y2) / 2)
line = cv2.line(mask, (x3, y3), (x4, y4), color=(1), thickness=2)
w = dist((x1, y1), (x2, y2))
dilation = int(w // 4)
line = binary_dilation(line, iterations=dilation)
return line
def remove_mouth(image, landmarks):
(x1, y1), (x2, y2) = landmarks[3:5]
mask = np.zeros_like(image[..., 0])
line = cv2.line(mask, (x1, y1), (x2, y2), color=(1), thickness=2)
w = dist((x1, y1), (x2, y2))
dilation = int(w // 3)
line = binary_dilation(line, iterations=dilation)
return line
class SladdRegion(Enum):
left_eye = 0
right_eye = 1
nose = 2
mouth = 3
# composition
both_eyes = left_eye + right_eye # 4
class SladdMasking(Mask):
# [0, 1, 2, 3, (0, 1), (0, 2), (1, 2), (2, 3), (0, 1, 2), (0, 1, 2, 3)]
# left-eye, right-eye, nose, mouth, ...
ALL_REGIONS = [
SladdRegion.left_eye,
SladdRegion.right_eye,
SladdRegion.nose,
SladdRegion.mouth,
]
REGIONS = [
[SladdRegion.left_eye],
[SladdRegion.right_eye],
[SladdRegion.nose],
[SladdRegion.mouth],
[SladdRegion.left_eye, SladdRegion.right_eye],
[SladdRegion.left_eye, SladdRegion.nose],
[SladdRegion.right_eye, SladdRegion.nose],
[SladdRegion.nose, SladdRegion.mouth],
[SladdRegion.left_eye, SladdRegion.right_eye, SladdRegion.nose],
ALL_REGIONS,
]
def init(self, compose: bool = False, single: bool = True, **kwargs):
# super().__init__(**kwargs)
self.compose = compose
if compose:
self.regions = SladdMasking.REGIONS
else:
self.regions = [reg for reg in SladdMasking.REGIONS if len(reg) == 1]
if single:
self.regions = [self.ALL_REGIONS]
@property
def total(self) -> int:
return len(self.regions)
@staticmethod
def parse(img, reg, landmarks) -> np.ndarray:
five_key = get_five_key(landmarks)
if reg is SladdRegion.left_eye:
mask = remove_eyes(img, five_key, "l")
elif reg is SladdRegion.right_eye:
mask = remove_eyes(img, five_key, "r")
elif reg is SladdRegion.nose:
mask = remove_nose(img, five_key)
elif reg is SladdRegion.mouth:
mask = remove_mouth(img, five_key)
else:
raise ValueError("Invalid region")
# elif reg == SladdRegion4:
# mask = remove_eyes(img, five_key, "b")
return mask
def build_mask(self) -> np.ndarray:
self.init()
h, w = self.face.shape[:2]
# print(len(self.regions))
regs = [self.regions[0][self.idx]]
# if isinstance(reg, int):
# mask = parse(img, reg, landmarks)
masks = [SladdMasking.parse(self.face, reg, self.landmarks) for reg in regs]
mask = reduce(np.maximum, masks)
mask = mask.reshape([mask.shape[0],mask.shape[1], 1])
return mask
|