125 lines
5.3 KiB
Python
125 lines
5.3 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
from core.base_mesh import MediapipeBaseLandmarks
|
|
|
|
class CPUDeformer:
|
|
"""
|
|
Class for handling face warping operations on CPU using triangle-based warping.
|
|
"""
|
|
|
|
@staticmethod
|
|
def warp_face(image, source_landmarks, target_landmarks):
|
|
"""
|
|
Warp the face image to match the base landmarks using triangle-based warping on CPU.
|
|
Now includes boundary triangles for complete image warping.
|
|
|
|
Args:
|
|
image: PIL Image to warp
|
|
source_landmarks: Source face landmarks
|
|
target_landmarks: Target face landmarks
|
|
|
|
Returns:
|
|
PIL Image: Warped image
|
|
"""
|
|
try:
|
|
# Get image dimensions
|
|
w, h = image.size
|
|
print(f"Image size: {w}x{h}")
|
|
|
|
# Convert landmarks to numpy arrays if they aren't already
|
|
source_landmarks = np.array(source_landmarks)
|
|
target_landmarks = np.array(target_landmarks)
|
|
|
|
# Get triangulation and all landmarks (including boundary points)
|
|
print("Getting triangulation with boundary triangles...")
|
|
triangles, all_target_points = MediapipeBaseLandmarks.get_face_triangles(size=(w, h))
|
|
|
|
# Create arrays for source and target points
|
|
n_base_landmarks = len(source_landmarks)
|
|
n_all_points = len(all_target_points)
|
|
n_boundary_points = n_all_points - n_base_landmarks
|
|
|
|
print(f"Base landmarks: {n_base_landmarks}, Total points: {n_all_points}")
|
|
|
|
# Create complete source and target point arrays
|
|
source_points = np.zeros((n_all_points, 2), dtype=np.float32)
|
|
target_points = np.zeros((n_all_points, 2), dtype=np.float32)
|
|
|
|
# Fill in face landmarks
|
|
source_points[:n_base_landmarks] = source_landmarks
|
|
target_points[:n_base_landmarks] = target_landmarks
|
|
|
|
# Fill in boundary points (use the same points for both source and target)
|
|
if n_boundary_points > 0:
|
|
boundary_points = all_target_points[n_base_landmarks:]
|
|
source_points[n_base_landmarks:] = boundary_points
|
|
target_points[n_base_landmarks:] = boundary_points
|
|
|
|
print(f"Processing {len(triangles)} triangles")
|
|
|
|
# Convert PIL image to numpy array
|
|
img_np = np.array(image)
|
|
output = np.zeros_like(img_np)
|
|
|
|
# Process each triangle
|
|
for i, triangle in enumerate(triangles):
|
|
# Get triangle vertices
|
|
src_tri = source_points[triangle]
|
|
dst_tri = target_points[triangle]
|
|
|
|
# Calculate bounding box for target triangle
|
|
min_x = max(0, int(np.min(dst_tri[:, 0])))
|
|
min_y = max(0, int(np.min(dst_tri[:, 1])))
|
|
max_x = min(w - 1, int(np.ceil(np.max(dst_tri[:, 0]))))
|
|
max_y = min(h - 1, int(np.ceil(np.max(dst_tri[:, 1]))))
|
|
|
|
# Skip if triangle is outside image bounds
|
|
if min_x >= max_x or min_y >= max_y:
|
|
continue
|
|
|
|
# Calculate target triangle matrix for barycentric coordinates
|
|
dst_matrix = np.vstack([
|
|
dst_tri[1] - dst_tri[0],
|
|
dst_tri[2] - dst_tri[0]
|
|
]).T
|
|
|
|
# Calculate inverse matrix for barycentric coordinates if possible
|
|
try:
|
|
dst_matrix_inv = np.linalg.inv(dst_matrix)
|
|
except np.linalg.LinAlgError:
|
|
continue # Skip degenerate triangles
|
|
|
|
# Process each pixel in target triangle's bounding box
|
|
for y in range(min_y, max_y + 1):
|
|
for x in range(min_x, max_x + 1):
|
|
# Calculate barycentric coordinates
|
|
point = np.array([x, y]) - dst_tri[0]
|
|
bary = np.dot(dst_matrix_inv, point)
|
|
alpha = 1.0 - bary[0] - bary[1]
|
|
|
|
# Check if point is inside triangle
|
|
eps = 1e-5
|
|
if (alpha >= -eps and bary[0] >= -eps and bary[1] >= -eps and
|
|
alpha <= 1 + eps and bary[0] <= 1 + eps and bary[1] <= 1 + eps):
|
|
|
|
# Calculate source pixel position using barycentric coordinates
|
|
src_x = int(alpha * src_tri[0, 0] + bary[0] * src_tri[1, 0] +
|
|
bary[1] * src_tri[2, 0] + 0.5)
|
|
src_y = int(alpha * src_tri[0, 1] + bary[0] * src_tri[1, 1] +
|
|
bary[1] * src_tri[2, 1] + 0.5)
|
|
|
|
# Copy pixel if within bounds
|
|
if (0 <= src_x < w and 0 <= src_y < h):
|
|
output[y, x] = img_np[src_y, src_x]
|
|
|
|
# Print progress every 10%
|
|
if (i + 1) % max(1, len(triangles) // 10) == 0:
|
|
print(f"Processed {i + 1}/{len(triangles)} triangles ({(i + 1) * 100 / len(triangles):.1f}%)")
|
|
|
|
return Image.fromarray(output)
|
|
|
|
except Exception as e:
|
|
print(f"Error during CPU warping: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return image |