WIP
cpu and gpu deformer
This commit is contained in:
Rafał Tłuchowski
@@ -0,0 +1,125 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from core.base_landmarks 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
|
||||
@@ -0,0 +1,199 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from core.base_landmarks import MediapipeBaseLandmarks
|
||||
|
||||
|
||||
class GPUDeformer:
|
||||
"""
|
||||
Class for handling face warping operations on GPU using CUDA acceleration.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def warp_face(image, source_landmarks, target_landmarks, device='cuda'):
|
||||
"""
|
||||
Warp the face image using parallel GPU processing with reverse mapping.
|
||||
Now includes boundary triangles for complete image warping.
|
||||
|
||||
Args:
|
||||
image: Input PIL Image or numpy array
|
||||
source_landmarks: Detected face landmarks
|
||||
target_landmarks: Target base landmarks
|
||||
device: 'cuda' or 'cpu' (default: 'cuda')
|
||||
|
||||
Returns:
|
||||
PIL Image: Warped image
|
||||
"""
|
||||
try:
|
||||
import cupy as cp
|
||||
print("Starting GPU warping process with boundary triangles...")
|
||||
|
||||
# Convert and validate input data
|
||||
source_landmarks = np.asarray(source_landmarks, dtype=np.float32)
|
||||
target_landmarks = np.asarray(target_landmarks, dtype=np.float32)
|
||||
|
||||
# Debug landmark ranges
|
||||
print("\nLandmark ranges:")
|
||||
print(f"Source landmarks - X: [{source_landmarks[:, 0].min():.1f}, {source_landmarks[:, 0].max():.1f}], "
|
||||
f"Y: [{source_landmarks[:, 1].min():.1f}, {source_landmarks[:, 1].max():.1f}]")
|
||||
print(f"Target landmarks - X: [{target_landmarks[:, 0].min():.1f}, {target_landmarks[:, 0].max():.1f}], "
|
||||
f"Y: [{target_landmarks[:, 1].min():.1f}, {target_landmarks[:, 1].max():.1f}]")
|
||||
|
||||
# Convert image to numpy array if needed
|
||||
if isinstance(image, Image.Image):
|
||||
img_np = np.array(image, dtype=np.uint8)
|
||||
else:
|
||||
img_np = np.asarray(image, dtype=np.uint8)
|
||||
|
||||
h, w = img_np.shape[:2]
|
||||
channels = img_np.shape[2] if len(img_np.shape) > 2 else 1
|
||||
print(f"\nImage shape: {img_np.shape}")
|
||||
|
||||
# Get triangulation including boundary triangles
|
||||
triangles, all_points = MediapipeBaseLandmarks.get_face_triangles(size=(w, h))
|
||||
print(f"Loaded {len(triangles)} triangles (including boundary triangles)")
|
||||
|
||||
# Create complete source and target point arrays
|
||||
n_base_landmarks = len(source_landmarks)
|
||||
n_all_points = len(all_points)
|
||||
print(f"Base landmarks: {n_base_landmarks}, Total points: {n_all_points}")
|
||||
|
||||
# Initialize complete arrays for source and target points
|
||||
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
|
||||
n_base = 468 # Hardcode original landmark count
|
||||
source_points[:n_base] = source_landmarks[:n_base]
|
||||
target_points[:n_base] = target_landmarks[:n_base]
|
||||
|
||||
# Boundary points remain identical
|
||||
source_points[n_base:] = all_points[n_base:]
|
||||
target_points[n_base:] = all_points[n_base:]
|
||||
|
||||
# CUDA kernel for parallel triangle processing with reverse mapping
|
||||
cuda_kernel = cp.RawKernel(r'''
|
||||
extern "C" __global__
|
||||
void process_triangles(
|
||||
const float* source_points,
|
||||
const float* target_points,
|
||||
const int* triangles,
|
||||
const unsigned char* input_image,
|
||||
unsigned char* output_image,
|
||||
const int num_triangles,
|
||||
const int width,
|
||||
const int height,
|
||||
const int channels
|
||||
) {
|
||||
// Get triangle index
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (tid >= num_triangles) return;
|
||||
|
||||
// Get triangle indices
|
||||
const int idx0 = triangles[tid * 3];
|
||||
const int idx1 = triangles[tid * 3 + 1];
|
||||
const int idx2 = triangles[tid * 3 + 2];
|
||||
|
||||
// Get source points
|
||||
const float src_x0 = source_points[idx0 * 2];
|
||||
const float src_y0 = source_points[idx0 * 2 + 1];
|
||||
const float src_x1 = source_points[idx1 * 2];
|
||||
const float src_y1 = source_points[idx1 * 2 + 1];
|
||||
const float src_x2 = source_points[idx2 * 2];
|
||||
const float src_y2 = source_points[idx2 * 2 + 1];
|
||||
|
||||
// Get target points
|
||||
const float dst_x0 = target_points[idx0 * 2];
|
||||
const float dst_y0 = target_points[idx0 * 2 + 1];
|
||||
const float dst_x1 = target_points[idx1 * 2];
|
||||
const float dst_y1 = target_points[idx1 * 2 + 1];
|
||||
const float dst_x2 = target_points[idx2 * 2];
|
||||
const float dst_y2 = target_points[idx2 * 2 + 1];
|
||||
|
||||
// Calculate target triangle area
|
||||
const float dst_area = (dst_x1 - dst_x0) * (dst_y2 - dst_y0) -
|
||||
(dst_x2 - dst_x0) * (dst_y1 - dst_y0);
|
||||
|
||||
const float abs_area = abs(dst_area);
|
||||
if (abs_area < 1e-6f) return;
|
||||
|
||||
// Calculate target triangle bounding box with border check
|
||||
const int min_x = max(0, (int)min(min(dst_x0, dst_x1), dst_x2));
|
||||
const int min_y = max(0, (int)min(min(dst_y0, dst_y1), dst_y2));
|
||||
const int max_x = min(width - 1, (int)(max(max(dst_x0, dst_x1), dst_x2) + 0.5f));
|
||||
const int max_y = min(height - 1, (int)(max(max(dst_y0, dst_y1), dst_y2) + 0.5f));
|
||||
|
||||
// Process each pixel in target triangle's bounding box
|
||||
for (int y = min_y; y <= max_y; y++) {
|
||||
for (int x = min_x; x <= max_x; x++) {
|
||||
// Calculate barycentric coordinates
|
||||
float lambda1 = ((y - dst_y2) * (dst_x0 - dst_x2) +
|
||||
(dst_x2 - x) * (dst_y0 - dst_y2)) / dst_area;
|
||||
float lambda2 = ((y - dst_y0) * (dst_x1 - dst_x0) +
|
||||
(dst_x0 - x) * (dst_y1 - dst_y0)) / dst_area;
|
||||
float lambda0 = 1.0f - lambda1 - lambda2;
|
||||
|
||||
// Check if point is inside triangle with tolerance
|
||||
const float eps = 1e-5f;
|
||||
if (lambda0 >= -eps && lambda1 >= -eps && lambda2 >= -eps &&
|
||||
lambda0 <= 1+eps && lambda1 <= 1+eps && lambda2 <= 1+eps) {
|
||||
|
||||
// Calculate source position using barycentric coordinates
|
||||
float source_x = lambda0 * src_x0 + lambda1 * src_x1 + lambda2 * src_x2;
|
||||
float source_y = lambda0 * src_y0 + lambda1 * src_y1 + lambda2 * src_y2;
|
||||
|
||||
// Round to nearest pixel
|
||||
int sx = (int)(source_x + 0.5f);
|
||||
int sy = (int)(source_y + 0.5f);
|
||||
|
||||
// Copy pixel if within bounds
|
||||
if (sx >= 0 && sx < width && sy >= 0 && sy < height) {
|
||||
for (int c = 0; c < channels; c++) {
|
||||
output_image[(y * width + x) * channels + c] =
|
||||
input_image[(sy * width + sx) * channels + c];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
''', 'process_triangles')
|
||||
|
||||
# Move data to GPU
|
||||
source_points_gpu = cp.asarray(source_points.ravel())
|
||||
target_points_gpu = cp.asarray(target_points.ravel())
|
||||
triangles_gpu = cp.asarray(triangles.ravel().astype(np.int32))
|
||||
img_gpu = cp.asarray(img_np)
|
||||
output_gpu = cp.zeros_like(img_gpu)
|
||||
|
||||
# Configure CUDA grid
|
||||
threadsPerBlock = 256
|
||||
blocksPerGrid = (len(triangles) + threadsPerBlock - 1) // threadsPerBlock
|
||||
print(f"CUDA config: {blocksPerGrid} blocks, {threadsPerBlock} threads per block")
|
||||
|
||||
# Execute kernel
|
||||
print(f"Processing {len(triangles)} triangles on GPU...")
|
||||
cuda_kernel(
|
||||
(blocksPerGrid,), (threadsPerBlock,),
|
||||
(source_points_gpu, target_points_gpu, triangles_gpu,
|
||||
img_gpu, output_gpu, len(triangles),
|
||||
w, h, channels)
|
||||
)
|
||||
cp.cuda.stream.get_current_stream().synchronize()
|
||||
print("CUDA kernel execution completed")
|
||||
print("Flushing CUDA cache...")
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Transfer result back to CPU and convert to PIL
|
||||
output_np = cp.asnumpy(output_gpu)
|
||||
return Image.fromarray(output_np)
|
||||
|
||||
except ImportError as e:
|
||||
print(f"CuPy not installed or CUDA not available: {str(e)}")
|
||||
return image
|
||||
except Exception as e:
|
||||
print(f"Error during GPU warping: {str(e)}")
|
||||
print(f"Error details: {str(e.__class__.__name__)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return image
|
||||
Reference in New Issue
Block a user