199 lines
8.9 KiB
Python
199 lines
8.9 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
from core.base_mesh import MediapipeBaseLandmarks
|
|
|
|
|
|
class GPUDeformer:
|
|
"""
|
|
Class for handling face warping operations on GPU using CUDA acceleration.
|
|
"""
|
|
|
|
@staticmethod
|
|
def warp_face(image, source_landmarks, target_landmarks):
|
|
"""
|
|
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
|
|
|
|
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
|