diff --git a/__init__.py b/__init__.py index 51f749a..df128c6 100644 --- a/__init__.py +++ b/__init__.py @@ -19,7 +19,7 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "FaceFitAndRestore": "Face Fit or Restore", - "FaceWrapper": "Face Landmarks Detector" + "FaceWrapper": "Face Wrapper" } __version__ = "1.0.1" \ No newline at end of file diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 8d3efed..674abb1 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -48,6 +48,10 @@ class FaceWrapper: image_np = self._convert_to_numpy(image) height, width = image_np.shape[:2] + if mode == "Wrap": + return self._wrap_mode(image_np, None, width, height, + device, x_scale, y_transform, processor_settings) + # Detect facial landmarks landmarks_df = self.face_detector.detect_landmarks(image_np) if landmarks_df is None: @@ -59,15 +63,10 @@ class FaceWrapper: return self._debug_mode(image_np, landmarks_df, width, height, show_detection, show_target, landmark_size, show_labels, x_scale, y_transform, processor_settings) - elif mode == "Un-Wrap": return self._unwrap_mode(image_np, landmarks_df, width, height, device, x_scale, y_transform, processor_settings) - elif mode == "Wrap": - print("Wrap mode not implemented yet") - return (image, processor_settings or {}) - def _debug_mode(self, image_np, landmarks_df, width, height, show_detection, show_target, landmark_size, show_labels, x_scale, y_transform, processor_settings): @@ -133,8 +132,7 @@ class FaceWrapper: warped_image = GPUDeformer.warp_face( pil_image, source_landmarks, - base_landmarks, - device='cuda' + base_landmarks ) else: warped_image = CPUDeformer.warp_face( @@ -151,6 +149,43 @@ class FaceWrapper: landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) return (output_image, self._update_settings(processor_settings, landmarks_data)) + def _wrap_mode(self, image_np, landmarks_df, width, height, device, + x_scale, y_transform, processor_settings): + if not processor_settings or 'target_lm' not in processor_settings: + print("No landmarks found in processor settings") + output_image = torch.from_numpy(image_np.astype(np.float32) / 255.0).unsqueeze(0) + return (output_image, {}) + + # Source landmarks from stored base landmarks (target_lm in settings) + source_x = processor_settings['target_lm']['x'] + source_y = processor_settings['target_lm']['y'] + source_landmarks = np.column_stack((source_x, source_y))[:468] + + # Target landmarks from detected face (detected_lm in settings) + detected_x = processor_settings['detected_lm']['x'] + detected_y = processor_settings['detected_lm']['y'] + target_landmarks = np.column_stack((detected_x, detected_y))[:468] + + # Select warping method + pil_image = Image.fromarray(image_np) + if device == "CUDA" and torch.cuda.is_available(): + warped_image = GPUDeformer.warp_face( + pil_image, + source_landmarks, + target_landmarks + ) + else: + warped_image = CPUDeformer.warp_face( + pil_image, + source_landmarks, + target_landmarks + ) + + output_image = np.array(warped_image).astype(np.float32) / 255.0 + output_image = torch.from_numpy(output_image).unsqueeze(0) + + return (output_image, processor_settings) + def _convert_to_numpy(self, image): """Improved image conversion with channel handling""" if torch.is_tensor(image): diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..160486c --- /dev/null +++ b/readme.md @@ -0,0 +1,158 @@ +# Face Processor for ComfyUI + +A custom node collection for ComfyUI that provides advanced face detection, alignment, and transformation capabilities using MediaPipe Face Mesh. + +## Features + +- **Face Detection & Landmark Extraction**: Uses MediaPipe Face Mesh to detect and extract 468 facial landmarks +- **Face Alignment**: Automatic face alignment based on eye positions +- **Face Transformation**: + - Warping between source and target face landmarks + - Scale and translation controls for face shape adjustment + - CPU and CUDA-accelerated processing options +- **Debug Visualization**: + - Visual landmark overlay with customizable parameters + - Support for both detected and target landmark visualization + - Optional landmark labels + +## Installation + +1. Clone this repository into your ComfyUI's `custom_nodes` directory: +```bash +cd ComfyUI/custom_nodes +git clone [repository-url] face_processor +``` + +2. Install required dependencies: +```bash +pip install mediapipe opencv-python numpy pandas pillow torch +``` + +For CUDA acceleration: +1. Install [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) +2. Install [CuPy](https://docs.cupy.dev/en/stable/install.html): +```bash +pip install cupy-cuda11x # Replace with your CUDA version +``` + +## Nodes + +### FaceWrapper +Main node for face detection and transformation operations. + +#### Inputs: +- `image`: Input image (ComfyUI IMAGE type) +- `mode`: Operating mode + - `Debug`: Visualization of detected landmarks + - `Un-Wrap`: Transform face to normalized position + - `Wrap`: Transform normalized face back to original position +- `device`: Processing device (`CPU` or `CUDA`) +- `show_detection`: Toggle detected landmarks visualization +- `show_target`: Toggle target landmarks visualization +- `landmark_size`: Size of landmark points in visualization +- `show_labels`: Toggle landmark index labels +- `x_scale`: Horizontal scaling factor (0.5 to 1.0) +- `y_transform`: Vertical translation (-0.5 to 0.5) +- `processor_settings`: Optional settings dictionary + +#### Outputs: +- `image`: Processed image +- `processor_settings`: Updated settings dictionary + +### FaceFitAndRestore +Node for face cropping and restoration operations. + +#### Inputs: +- `mode`: Operating mode + - `Fit`: Crop and align face + - `Restore`: Place processed face back in original image +- `image`: Input image +- `padding_percent`: Additional padding around face (0.0 to 1.0) +- `bbox_size`: Output size for cropped face (512, 1024, or 2048) +- `processor_settings`: Required for Restore mode + +#### Outputs: +- `image`: Processed image +- `processor_settings`: Updated settings dictionary +- `mask`: Mask indicating face region + +## Technical Details + +### Core Components + +#### Face Detection +- Uses MediaPipe Face Mesh for robust face detection and landmark extraction +- Provides 468 facial landmarks with 3D coordinates +- Supports various input formats (PIL Image, numpy array, torch tensor) + +#### Image Processing +- Automatic face rotation based on eye positions +- Aspect ratio-preserving resizing +- Support for square cropping with configurable padding +- Boundary triangulation for complete face warping + +#### Face Warping +- Triangle-based warping using predefined mesh topology from MediaPipe Face Mesh +- CPU implementation using pure Python/NumPy +- CUDA-accelerated GPU implementation using CuPy +- Handles both forward and inverse warping + +### Performance Considerations + +- GPU acceleration requires CUDA toolkit and CuPy +- CPU fallback available for all operations +- Progressive feedback during long operations +- Memory-efficient processing for large images + +## Example Usage + +Basic face detection and visualization: +```python +face_wrapper = FaceWrapper() +result_image, settings = face_wrapper.detect_face( + image=input_image, + mode="Debug", + device="CPU", + show_detection=True, + show_target=False, + landmark_size=4, + show_labels=True, + x_scale=1.0, + y_transform=0.0 +) +``` + +Face normalization workflow: +1. Detect and normalize face: +```python +# Unwrap face to normalized position +normalized_face, settings = face_wrapper.detect_face( + image=input_image, + mode="Un-Wrap", + device="CUDA", + x_scale=1.0, + y_transform=0.0 +) +``` + +2. Process normalized face with your preferred method + +3. Restore face to original position: +```python +# Wrap processed face back +final_image, _ = face_wrapper.detect_face( + image=processed_face, + mode="Wrap", + device="CUDA", + processor_settings=settings +) +``` + +## License + +[Your License Here] + +## Acknowledgments + +- MediaPipe Face Mesh for facial landmark detection +- ComfyUI project for the node system framework \ No newline at end of file