From aaca513c6fabc33cd03a4f64c13f34af4ae8cebc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Sun, 26 Jan 2025 13:04:14 +0100 Subject: [PATCH] WIP cpu and gpu deformer --- core/cpu_deformer.py | 125 +++++++++++++++++++++++++++ core/gpu_deformer.py | 199 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 324 insertions(+) create mode 100644 core/cpu_deformer.py create mode 100644 core/gpu_deformer.py diff --git a/core/cpu_deformer.py b/core/cpu_deformer.py new file mode 100644 index 0000000..cbf140f --- /dev/null +++ b/core/cpu_deformer.py @@ -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 \ No newline at end of file diff --git a/core/gpu_deformer.py b/core/gpu_deformer.py new file mode 100644 index 0000000..6517f21 --- /dev/null +++ b/core/gpu_deformer.py @@ -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