222 lines
7.5 KiB
Python
222 lines
7.5 KiB
Python
import os
|
|
import dlib
|
|
import cv2
|
|
import copy
|
|
import numpy as np
|
|
|
|
from PIL import Image, ImageDraw
|
|
from .constants import script_models_dir
|
|
from .utils import expand_polygon, calculate_distance
|
|
from .modules import depthmap
|
|
|
|
try:
|
|
from .modules import mmdetdd
|
|
except Exception as e:
|
|
pass
|
|
|
|
depthmap_generator = depthmap.SimpleDepthMapGenerator()
|
|
|
|
try:
|
|
landmark_detector = dlib.shape_predictor(
|
|
os.path.join(script_models_dir, 'shape_predictor_68_face_landmarks.dat')
|
|
)
|
|
except Exception as e:
|
|
# when reloading the module, landmark_detector is already loaded
|
|
# and the file cant be opened again
|
|
pass
|
|
|
|
def _get_image_mask(image):
|
|
width, height = image.size
|
|
return np.full((height, width, 3), 0, dtype=np.uint8)
|
|
|
|
def _get_detected_faces_dlib(image):
|
|
|
|
# Load the pre-trained model for detecting facial landmarks
|
|
face_detector = dlib.get_frontal_face_detector()
|
|
|
|
# Convert the image to grayscale
|
|
frame = np.array(image)
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Detect faces in the image
|
|
return face_detector(gray, 1), gray
|
|
|
|
def _calculate_mask_padding(p1, p2, percents):
|
|
distance = calculate_distance((p1.x, p1.y), (p2.x, p2.y))
|
|
value = int(distance * (percents / 100))
|
|
return value
|
|
|
|
def get_eyes_mask_dlib(init_image, padding=20, in_pixels=False):
|
|
|
|
mask = _get_image_mask(init_image)
|
|
faces, gray = _get_detected_faces_dlib(init_image)
|
|
padding_original = padding
|
|
|
|
if not len(faces):
|
|
return Image.fromarray(mask), False
|
|
|
|
for face in faces:
|
|
# Use the landmark detector to find facial landmarks
|
|
landmarks = landmark_detector(gray, face)
|
|
# Extract the coordinates of the left and right eyes
|
|
left_eye_x = [landmarks.part(i).x for i in range(36, 42)]
|
|
left_eye_y = [landmarks.part(i).y for i in range(36, 42)]
|
|
right_eye_x = [landmarks.part(i).x for i in range(42, 48)]
|
|
right_eye_y = [landmarks.part(i).y for i in range(42, 48)]
|
|
# Calculate mask padding
|
|
if not in_pixels:
|
|
padding = _calculate_mask_padding(landmarks.part(36), landmarks.part(39), padding_original)
|
|
# Draw a filled white polygon around the left eye
|
|
left_eye_points = []
|
|
for i in range(len(left_eye_x)):
|
|
left_eye_points.append([left_eye_x[i], left_eye_y[i]])
|
|
left_eye_points = expand_polygon(left_eye_points, padding)
|
|
left_eye_points = np.array(left_eye_points, np.int32)
|
|
left_eye_points = left_eye_points.reshape((-1, 1, 2))
|
|
cv2.fillPoly(mask, [left_eye_points], (255, 255, 255))
|
|
# Calculate mask padding
|
|
if not in_pixels:
|
|
padding = _calculate_mask_padding(landmarks.part(42), landmarks.part(45), padding_original)
|
|
# Draw a filled white polygon around the right eye
|
|
right_eye_points = []
|
|
for i in range(len(right_eye_x)):
|
|
right_eye_points.append([right_eye_x[i], right_eye_y[i]])
|
|
right_eye_points = expand_polygon(right_eye_points, padding)
|
|
right_eye_points = np.array(right_eye_points, np.int32)
|
|
right_eye_points = right_eye_points.reshape((-1, 1, 2))
|
|
cv2.fillPoly(mask, [right_eye_points], (255, 255, 255))
|
|
|
|
return Image.fromarray(mask), True
|
|
|
|
def get_face_mask_dlib(init_image, padding=20, in_pixels=False):
|
|
|
|
mask = _get_image_mask(init_image)
|
|
faces, gray = _get_detected_faces_dlib(init_image)
|
|
padding_original = padding
|
|
|
|
if not len(faces):
|
|
return Image.fromarray(mask), False
|
|
|
|
for face in faces:
|
|
# Use the landmark detector to find facial landmarks
|
|
landmarks = landmark_detector(gray, face)
|
|
face_x = [landmarks.part(i).x for i in range(0, 17)] + [landmarks.part(i).x for i in reversed(range(17, 27))]
|
|
face_y = [landmarks.part(i).y for i in range(0, 17)] + [landmarks.part(i).y for i in reversed(range(17, 27))]
|
|
# Calculate mask padding
|
|
if not in_pixels:
|
|
padding = _calculate_mask_padding(landmarks.part(0), landmarks.part(16), padding_original)
|
|
# Draw a filled white polygon around the face
|
|
face_points = []
|
|
for i in range(len(face_x)):
|
|
face_points.append([face_x[i], face_y[i]])
|
|
face_points = expand_polygon(face_points, padding)
|
|
face_points = np.array(face_points, np.int32)
|
|
face_points = face_points.reshape((-1, 1, 2))
|
|
cv2.fillPoly(mask, [face_points], (255, 255, 255))
|
|
|
|
return Image.fromarray(mask), True
|
|
|
|
def get_face_mask_depth(init_image):
|
|
|
|
mask = _get_image_mask(init_image)
|
|
faces, gray = _get_detected_faces_dlib(init_image)
|
|
|
|
if len(faces) != 1:
|
|
return Image.fromarray(mask), False
|
|
|
|
body_mask, body_mask_success = get_body_mask_depth(init_image)
|
|
|
|
if not body_mask_success:
|
|
return Image.fromarray(mask), False
|
|
|
|
face = faces[0]
|
|
landmarks = landmark_detector(gray, face)
|
|
lowest_point = None
|
|
|
|
for i in range(0, 17):
|
|
point = landmarks.part(i)
|
|
if lowest_point is None or point.y > lowest_point.y:
|
|
lowest_point = point
|
|
|
|
width, height = body_mask.size
|
|
|
|
for x in range(0, width):
|
|
for y in range(lowest_point.y, height):
|
|
body_mask.putpixel((x, y), (0, 0, 0))
|
|
|
|
return body_mask, True
|
|
|
|
def get_body_mask_depth(init_image):
|
|
|
|
def remap_range(value, minIn, MaxIn, minOut, maxOut):
|
|
if value > MaxIn: value = MaxIn;
|
|
if value < minIn: value = minIn;
|
|
finalValue = ((value - minIn) / (MaxIn - minIn)) * (maxOut - minOut) + minOut
|
|
return finalValue
|
|
|
|
def create_depth_mask_from_depth_map(img, treshold, clean_cut):
|
|
img = copy.deepcopy(img.convert("RGBA"))
|
|
mask_img = copy.deepcopy(img.convert("L"))
|
|
mask_datas = mask_img.getdata()
|
|
datas = img.getdata()
|
|
newData = []
|
|
maxD = max(mask_datas)
|
|
if clean_cut and treshold == 0:
|
|
treshold = 128
|
|
for i in range(len(mask_datas)):
|
|
if clean_cut and mask_datas[i] > treshold:
|
|
newrgb = 255
|
|
elif mask_datas[i] > treshold and not clean_cut:
|
|
newrgb = int(remap_range(mask_datas[i],treshold,255,0,255))
|
|
else:
|
|
newrgb = 0
|
|
newData.append((newrgb,newrgb,newrgb,255))
|
|
img.putdata(newData)
|
|
return img
|
|
|
|
try:
|
|
d_m = depthmap_generator.calculate_depth_maps(init_image, init_image.width, init_image.height, 1, False)
|
|
d_m = create_depth_mask_from_depth_map(d_m, 128, True)
|
|
return d_m, True
|
|
except Exception as e:
|
|
print(e)
|
|
return _get_image_mask(init_image), False
|
|
|
|
def get_face_mask_mmdet(init_image):
|
|
|
|
mask = _get_image_mask(init_image)
|
|
faces, gray = _get_detected_faces_dlib(init_image)
|
|
|
|
if len(faces) != 1:
|
|
return Image.fromarray(mask), False
|
|
|
|
body_mask, body_mask_success = get_body_mask_mmdet(init_image)
|
|
|
|
if not body_mask_success:
|
|
return Image.fromarray(mask), False
|
|
|
|
face = faces[0]
|
|
landmarks = landmark_detector(gray, face)
|
|
lowest_point = None
|
|
|
|
for i in range(0, 17):
|
|
point = landmarks.part(i)
|
|
if lowest_point is None or point.y > lowest_point.y:
|
|
lowest_point = point
|
|
|
|
width, height = body_mask.size
|
|
|
|
for x in range(0, width):
|
|
for y in range(lowest_point.y, height):
|
|
body_mask.putpixel((x, y), 0)
|
|
|
|
return body_mask, True
|
|
|
|
def get_body_mask_mmdet(init_image):
|
|
try:
|
|
mask = mmdetdd.get_person_mask(init_image)
|
|
return mask, True
|
|
except Exception as e:
|
|
print(e)
|
|
return _get_image_mask(init_image), False
|