batch-face-swap/scripts/bfs_utils.py

107 lines
3.7 KiB
Python

from modules import images
from modules.shared import opts
from PIL import Image, ImageOps, ImageChops
import mediapipe as mp
import numpy as np
import math
import cv2
import os
def getFacialLandmarks(image):
height, width, _ = image.shape
mp_face_mesh = mp.solutions.face_mesh
with mp_face_mesh.FaceMesh(static_image_mode=True,max_num_faces=4,min_detection_confidence=0.5) as face_mesh:
height, width, _ = image.shape
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
result = face_mesh.process(image_rgb)
facelandmarks = []
if result.multi_face_landmarks is not None:
for facial_landmarks in result.multi_face_landmarks:
landmarks = []
for i in range(0, 468):
pt1 = facial_landmarks.landmark[i]
x = int(pt1.x * width)
y = int(pt1.y * height)
landmarks.append([x, y])
#cv2.circle(image, (x, y), 2, (100,100,0), -1)
#cv2.imshow("Cropped", image)
facelandmarks.append(np.array(landmarks, np.int32))
return facelandmarks
def apply_overlay(image, paste_loc, imageOriginal, mask):
x, y, w, h = paste_loc
base_image = Image.new('RGBA', (imageOriginal.width, imageOriginal.height))
image = images.resize_image(1, image, w, h)
base_image.paste(image, (x, y))
face = base_image
new_mask = ImageChops.multiply(face.getchannel("A"), mask)
face.putalpha(new_mask)
imageOriginal = imageOriginal.convert('RGBA')
image = Image.alpha_composite(imageOriginal, face)
return image
def composite(image1, image2, mask, visualizationOpacity):
mask_np = np.array(mask)
mask_np = np.where(mask_np == 255, visualizationOpacity, 0)
mask = Image.fromarray(mask_np).convert('L')
image = image2.copy()
image.paste(image1, None, mask)
return image
def findBiggestFace(inputImage):
# Store a copy of the input image:
biggestFace = inputImage.copy()
# Set initial values for the
# largest contour:
largestArea = 0
largestContourIndex = 0
# Find the contours on the binary image:
contours, hierarchy = cv2.findContours(inputImage, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
# Get the largest contour in the contours list:
for i, cc in enumerate(contours):
# Find the area of the contour:
area = cv2.contourArea(cc)
# Store the index of the largest contour:
if area > largestArea:
largestArea = area
largestContourIndex = i
# Once we get the biggest face, paint it black:
tempMat = inputImage.copy()
cv2.drawContours(tempMat, contours, largestContourIndex, (0, 0, 0), -1, 8, hierarchy)
# Erase smaller faces:
biggestFace = biggestFace - tempMat
return biggestFace
def maskResize(mask, maskSize, height):
size = maskSize * (0.011*height)/5
kernel = np.ones((int(math.ceil((0.011*height)+abs(size))),int(math.ceil((0.011*height)+abs(size)))),'uint8')
if size > 0:
mask = cv2.dilate(mask,kernel,iterations=1)
else:
mask = cv2.erode(mask,kernel,iterations=1)
return mask
def listFiles(path, searchSubdir, allFiles):
if searchSubdir:
for root, _, files in os.walk(os.path.abspath(path)):
for file in files:
if file.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.PNG', '.JPG', '.JPEG', '.BMP')):
allFiles.append(os.path.join(root, file))
else:
allFiles = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f)) and f.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.PNG', '.JPG', '.JPEG', '.BMP'))]
return allFiles