├── __init__.py ├── core ├── __init__.py ├── base_mesh.py ├── cpu_deformer.py ├── face_detector.py ├── gpu_deformer.py ├── image_processor.py ├── lm_mapping.py └── resources │ └── model_loader.py ├── nodes ├── __init__.py ├── face_fit_and_restore.py ├── face_tracker.py ├── face_wrapper.py ├── image_feeder.py └── image_filters.py ├── readme.md ├── requirements.txt └── workflow └── FaceProcessor_basic.json /__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from .nodes.image_feeder import ImageFeeder 4 | from .nodes.image_filters import HighPassFilter 5 | from .nodes.face_wrapper import FaceWrapper 6 | from .nodes.face_fit_and_restore import FaceFitAndRestore 7 | from .nodes.face_tracker import FaceTracker 8 | 9 | # Get the path to the current directory 10 | NODE_PATH = os.path.dirname(os.path.realpath(__file__)) 11 | 12 | NODE_CLASS_MAPPINGS = { 13 | "FaceFitAndRestore": FaceFitAndRestore, 14 | "FaceWrapper": FaceWrapper, 15 | "HighPassFilter": HighPassFilter, 16 | "ImageFeeder": ImageFeeder, 17 | "FaceTracker": FaceTracker 18 | } 19 | 20 | NODE_DISPLAY_NAME_MAPPINGS = { 21 | "FaceFitAndRestore": "Face Fit or Restore", 22 | "FaceWrapper": "Face Wrapper", 23 | "HighPassFilter": "High Pass Filter (HPF)", 24 | "ImageFeeder": "Image Feeder", 25 | "FaceTracker": "Face Tracker" 26 | } 27 | 28 | __version__ = "1.2.0" 29 | -------------------------------------------------------------------------------- /core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SykkoAtHome/ComfyUI_FaceProcessor/7687b344c15f8c5c0d8983d60e98447a20d01757/core/__init__.py -------------------------------------------------------------------------------- /core/base_mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from typing import Optional, Tuple, List, Union 3 | 4 | from ..core.resources.model_loader import ModelOBJ 5 | 6 | 7 | class MediapipeBaseLandmarks: 8 | 9 | """ 10 | Class for handling base landmark positions and face topology. 11 | Coordinates are stored in normalized 0-1 range. 12 | """ 13 | _instance: Optional[ModelOBJ] = None 14 | _boundary_faces: Optional[np.ndarray] = None 15 | _boundary_landmarks: Optional[np.ndarray] = None 16 | _current_size: Optional[Union[int, Tuple[int, int]]] = None 17 | 18 | 19 | @classmethod 20 | def _get_model_instance(cls) -> ModelOBJ: 21 | """ 22 | Get or create ModelOBJ instance using singleton pattern 23 | 24 | Returns: 25 | ModelOBJ: Instance of model loader 26 | """ 27 | if cls._instance is None: 28 | cls._instance = ModelOBJ() 29 | return cls._instance 30 | 31 | @classmethod 32 | def create_boundary_triangles(cls, size) -> np.ndarray: 33 | """ 34 | Create additional triangles between face oval and image boundaries 35 | 36 | Args: 37 | size: Target size (int or tuple of width, height) 38 | 39 | Returns: 40 | numpy.ndarray: Array of boundary face indices 41 | """ 42 | if isinstance(size, int): 43 | width = height = size 44 | else: 45 | width, height = size 46 | 47 | print(f"Creating boundary triangles for image size: {width}x{height}") 48 | 49 | boundary_points = [] 50 | boundary_triangles = [] 51 | 52 | edges = { 53 | 'top': { 54 | 'landmarks': [54, 103, 67, 109, 10, 338, 297, 332, 284], 55 | 'start': (0, 0), 56 | 'middle': (0.5, 0), 57 | 'end': (1, 0), 58 | 'segments': 8 59 | }, 60 | 'right': { 61 | 'landmarks': [284, 251, 389, 356, 454, 323, 361, 288, 397, 365], 62 | 'start': (1, 0), 63 | 'middle': (1, 0.5), 64 | 'end': (1, 1), 65 | 'segments': 9 66 | }, 67 | 'bottom': { 68 | 'landmarks': [365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136], 69 | 'start': (1, 1), 70 | 'middle': (0.5, 1), 71 | 'end': (0, 1), 72 | 'segments': 10 73 | }, 74 | 'left': { 75 | 'landmarks': [136, 172, 58, 132, 93, 234, 127, 162, 21, 54], 76 | 'start': (0, 1), 77 | 'middle': (0, 0.5), 78 | 'end': (0, 0), 79 | 'segments': 9 80 | } 81 | } 82 | 83 | current_point_idx = cls._get_model_instance().num_landmarks 84 | 85 | for edge_name, edge_info in edges.items(): 86 | print(f"Processing {edge_name} edge") 87 | 88 | landmarks = edge_info['landmarks'] 89 | n_segments = edge_info['segments'] 90 | print(f"Creating {n_segments} segments for {len(landmarks)} landmarks") 91 | 92 | # Generate normalized edge points 93 | edge_points = cls._interpolate_edge_points( 94 | edge_info['start'], 95 | edge_info['end'], 96 | n_segments 97 | ) 98 | 99 | # Create quads and convert to triangles 100 | for i in range(n_segments): 101 | # Define quad vertices 102 | quad = [ 103 | landmarks[i], # Current landmark 104 | landmarks[i + 1], # Next landmark 105 | current_point_idx + i + 1, # Next edge point 106 | current_point_idx + i # Current edge point 107 | ] 108 | 109 | # Convert quad to two triangles 110 | boundary_triangles.extend([ 111 | [quad[0], quad[1], quad[2]], # First triangle 112 | [quad[0], quad[2], quad[3]] # Second triangle 113 | ]) 114 | 115 | # Add edge points to boundary points list 116 | boundary_points.extend(edge_points) 117 | current_point_idx += len(edge_points) 118 | 119 | print(f"Created {len(boundary_triangles)} triangles from {len(boundary_points)} boundary points") 120 | 121 | cls._boundary_landmarks = np.array(boundary_points) 122 | cls._boundary_faces = np.array(boundary_triangles) 123 | 124 | return cls._boundary_faces 125 | 126 | @staticmethod 127 | def _interpolate_edge_points(start_pos: Tuple[float, float], 128 | end_pos: Tuple[float, float], 129 | n_segments: int) -> List[Tuple[float, float]]: 130 | """ 131 | Interpolate points along an edge 132 | 133 | Args: 134 | start_pos: Start position (x, y) 135 | end_pos: End position (x, y) 136 | n_segments: Number of segments to create 137 | 138 | Returns: 139 | List[Tuple[float, float]]: List of interpolated points 140 | """ 141 | points = [] 142 | for i in range(n_segments + 1): 143 | t = i / n_segments 144 | x = start_pos[0] + t * (end_pos[0] - start_pos[0]) 145 | y = start_pos[1] + t * (end_pos[1] - start_pos[1]) 146 | points.append((x, y)) 147 | return points 148 | 149 | @classmethod 150 | def get_face_triangles(cls, size=None, x_scale: float = 1.0, y_translation: float = 0.0) -> Tuple[ 151 | np.ndarray, np.ndarray]: 152 | """ 153 | Get face triangulation including boundary triangles with optional transformations 154 | 155 | Args: 156 | size: Target size (int or tuple) 157 | x_scale: Horizontal scaling factor (0.5 to 1.0) 158 | y_translation: Vertical translation (-0.5 to 0.5) 159 | 160 | Returns: 161 | Tuple[np.ndarray, np.ndarray]: (triangles array, landmarks array) 162 | """ 163 | model = cls._get_model_instance() 164 | landmarks = model.get_transformed_landmarks(x_scale, y_translation) 165 | faces = model.get_faces() 166 | 167 | # Create boundary triangles if needed (using normalized coordinates) 168 | if cls._boundary_faces is None: 169 | cls.create_boundary_triangles(size) 170 | 171 | if size is not None: 172 | if isinstance(size, tuple): 173 | scaled_landmarks = landmarks * np.array(size) 174 | if cls._boundary_landmarks is not None: 175 | boundary_landmarks = cls._boundary_landmarks * np.array(size) 176 | if y_translation != 0: 177 | boundary_landmarks = boundary_landmarks.copy() 178 | boundary_landmarks[:, 1] += y_translation * size[1] 179 | else: 180 | scaled_landmarks = landmarks * size 181 | if cls._boundary_landmarks is not None: 182 | boundary_landmarks = cls._boundary_landmarks * size 183 | if y_translation != 0: 184 | boundary_landmarks = boundary_landmarks.copy() 185 | boundary_landmarks[:, 1] += y_translation * size 186 | 187 | if cls._boundary_faces is not None: 188 | all_points = np.vstack([scaled_landmarks, boundary_landmarks]) 189 | all_triangles = np.vstack([faces, cls._boundary_faces]) 190 | return all_triangles, all_points 191 | 192 | return faces, landmarks 193 | 194 | @classmethod 195 | def get_base_landmarks(cls, size=None, x_scale: float = 1.0, y_translation: float = 0.0) -> np.ndarray: 196 | """ 197 | Get base landmark positions with optional transformations. 198 | 199 | Args: 200 | size: Target size (int or tuple) 201 | x_scale: Horizontal scaling factor (0.5 to 1.0) 202 | y_translation: Vertical translation (-0.5 to 0.5) 203 | 204 | Returns: 205 | numpy.ndarray: Base landmarks with applied transformations 206 | """ 207 | model = cls._get_model_instance() 208 | landmarks = model.get_transformed_landmarks(x_scale, y_translation) 209 | 210 | if size is not None: 211 | if isinstance(size, tuple): 212 | landmarks = landmarks * np.array(size) 213 | else: 214 | landmarks = landmarks * size 215 | 216 | return landmarks -------------------------------------------------------------------------------- /core/cpu_deformer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | from ..core.base_mesh import MediapipeBaseLandmarks 4 | 5 | class CPUDeformer: 6 | """ 7 | Class for handling face warping operations on CPU using triangle-based warping. 8 | """ 9 | 10 | @staticmethod 11 | def warp_face(image, source_landmarks, target_landmarks): 12 | """ 13 | Warp the face image to match the base landmarks using triangle-based warping on CPU. 14 | Now includes boundary triangles for complete image warping. 15 | 16 | Args: 17 | image: PIL Image to warp 18 | source_landmarks: Source face landmarks 19 | target_landmarks: Target face landmarks 20 | 21 | Returns: 22 | PIL Image: Warped image 23 | """ 24 | try: 25 | # Get image dimensions 26 | w, h = image.size 27 | print(f"Image size: {w}x{h}") 28 | 29 | # Convert landmarks to numpy arrays if they aren't already 30 | source_landmarks = np.array(source_landmarks) 31 | target_landmarks = np.array(target_landmarks) 32 | 33 | # Get triangulation and all landmarks (including boundary points) 34 | print("Getting triangulation with boundary triangles...") 35 | triangles, all_target_points = MediapipeBaseLandmarks.get_face_triangles(size=(w, h)) 36 | 37 | # Create arrays for source and target points 38 | n_base_landmarks = len(source_landmarks) 39 | n_all_points = len(all_target_points) 40 | n_boundary_points = n_all_points - n_base_landmarks 41 | 42 | print(f"Base landmarks: {n_base_landmarks}, Total points: {n_all_points}") 43 | 44 | # Create complete source and target point arrays 45 | source_points = np.zeros((n_all_points, 2), dtype=np.float32) 46 | target_points = np.zeros((n_all_points, 2), dtype=np.float32) 47 | 48 | # Fill in face landmarks 49 | source_points[:n_base_landmarks] = source_landmarks 50 | target_points[:n_base_landmarks] = target_landmarks 51 | 52 | # Fill in boundary points (use the same points for both source and target) 53 | if n_boundary_points > 0: 54 | boundary_points = all_target_points[n_base_landmarks:] 55 | source_points[n_base_landmarks:] = boundary_points 56 | target_points[n_base_landmarks:] = boundary_points 57 | 58 | print(f"Processing {len(triangles)} triangles") 59 | 60 | # Convert PIL image to numpy array 61 | img_np = np.array(image) 62 | output = np.zeros_like(img_np) 63 | 64 | # Process each triangle 65 | for i, triangle in enumerate(triangles): 66 | # Get triangle vertices 67 | src_tri = source_points[triangle] 68 | dst_tri = target_points[triangle] 69 | 70 | # Calculate bounding box for target triangle 71 | min_x = max(0, int(np.min(dst_tri[:, 0]))) 72 | min_y = max(0, int(np.min(dst_tri[:, 1]))) 73 | max_x = min(w - 1, int(np.ceil(np.max(dst_tri[:, 0])))) 74 | max_y = min(h - 1, int(np.ceil(np.max(dst_tri[:, 1])))) 75 | 76 | # Skip if triangle is outside image bounds 77 | if min_x >= max_x or min_y >= max_y: 78 | continue 79 | 80 | # Calculate target triangle matrix for barycentric coordinates 81 | dst_matrix = np.vstack([ 82 | dst_tri[1] - dst_tri[0], 83 | dst_tri[2] - dst_tri[0] 84 | ]).T 85 | 86 | # Calculate inverse matrix for barycentric coordinates if possible 87 | try: 88 | dst_matrix_inv = np.linalg.inv(dst_matrix) 89 | except np.linalg.LinAlgError: 90 | continue # Skip degenerate triangles 91 | 92 | # Process each pixel in target triangle's bounding box 93 | for y in range(min_y, max_y + 1): 94 | for x in range(min_x, max_x + 1): 95 | # Calculate barycentric coordinates 96 | point = np.array([x, y]) - dst_tri[0] 97 | bary = np.dot(dst_matrix_inv, point) 98 | alpha = 1.0 - bary[0] - bary[1] 99 | 100 | # Check if point is inside triangle 101 | eps = 1e-5 102 | if (alpha >= -eps and bary[0] >= -eps and bary[1] >= -eps and 103 | alpha <= 1 + eps and bary[0] <= 1 + eps and bary[1] <= 1 + eps): 104 | 105 | # Calculate source pixel position using barycentric coordinates 106 | src_x = int(alpha * src_tri[0, 0] + bary[0] * src_tri[1, 0] + 107 | bary[1] * src_tri[2, 0] + 0.5) 108 | src_y = int(alpha * src_tri[0, 1] + bary[0] * src_tri[1, 1] + 109 | bary[1] * src_tri[2, 1] + 0.5) 110 | 111 | # Copy pixel if within bounds 112 | if (0 <= src_x < w and 0 <= src_y < h): 113 | output[y, x] = img_np[src_y, src_x] 114 | 115 | # Print progress every 10% 116 | if (i + 1) % max(1, len(triangles) // 10) == 0: 117 | print(f"Processed {i + 1}/{len(triangles)} triangles ({(i + 1) * 100 / len(triangles):.1f}%)") 118 | 119 | return Image.fromarray(output) 120 | 121 | except Exception as e: 122 | print(f"Error during CPU warping: {str(e)}") 123 | import traceback 124 | traceback.print_exc() 125 | return image -------------------------------------------------------------------------------- /core/face_detector.py: -------------------------------------------------------------------------------- 1 | from typing import Union, Optional 2 | 3 | import cv2 4 | import numpy as np 5 | import pandas as pd 6 | import torch 7 | from PIL import Image 8 | from pandas import DataFrame 9 | 10 | from ..core.lm_mapping import LandmarkMappings 11 | from ..core.resources.model_loader import ModelDlib, ModelMediaPipe 12 | from ..core.image_processor import ImageProcessor 13 | 14 | 15 | class FaceDetector: 16 | def __init__(self): 17 | """Initialize the FaceDetector with MediaPipe Face Mesh.""" 18 | # Models will be initialized on demand 19 | self.mediapipe_model = None 20 | self.dlib_model = None 21 | self.image_processor = ImageProcessor() 22 | 23 | def detect_landmarks(self, 24 | image: Union[torch.Tensor, np.ndarray, Image.Image], 25 | refiner: Optional[str] = None 26 | ) -> Union[DataFrame, None]: 27 | """ 28 | Detect facial landmarks using MediaPipe with optional refinement. 29 | 30 | Args: 31 | image: Input image in various formats 32 | refiner: Optional refinement method ('Dlib', 'InsightFace', or None) 33 | 34 | Returns: 35 | DataFrame with landmark coordinates (x, y) and indices or None if no face detected 36 | """ 37 | # Always detect with MediaPipe first 38 | mp_landmarks = self.detect_landmarks_mp(image) 39 | if mp_landmarks is None: 40 | return None 41 | 42 | # Apply refinement if requested 43 | if refiner is not None: 44 | if refiner.lower() == 'dlib': 45 | dlib_landmarks = self.detect_landmarks_dlib(image) 46 | if dlib_landmarks is not None: 47 | return self.landmarks_interpolation(mp_landmarks, dlib_landmarks) 48 | 49 | return mp_landmarks 50 | 51 | def detect_landmarks_mp(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: 52 | """ 53 | Detect facial landmarks in the given image using MediaPipe. 54 | 55 | Args: 56 | image: Input image in various formats 57 | 58 | Returns: 59 | Optional[pd.DataFrame]: DataFrame containing landmark coordinates (x, y) and indices or None if no face detected 60 | """ 61 | try: 62 | # Initialize MediaPipe model if not already done 63 | if self.mediapipe_model is None: 64 | self.mediapipe_model = ModelMediaPipe() 65 | 66 | image_np = self.image_processor.convert_to_numpy(image) 67 | if image_np is None: 68 | return None 69 | 70 | try: 71 | results = self.mediapipe_model.face_mesh.process(image_np) 72 | except RuntimeError as runtime_error: 73 | if "Task runner is currently not running" in str(runtime_error): 74 | # Reset and reinitialize the MediaPipe model if the task has been closed 75 | self.mediapipe_model._face_mesh = None 76 | self.mediapipe_model._model_loaded = False 77 | self.mediapipe_model = ModelMediaPipe() 78 | results = self.mediapipe_model.face_mesh.process(image_np) 79 | else: 80 | raise 81 | 82 | # Results may come from MediaPipe Solutions or Tasks. Try both field names. 83 | face_landmarks_list = getattr(results, 'multi_face_landmarks', None) 84 | if not face_landmarks_list: 85 | face_landmarks_list = getattr(results, 'face_landmarks', None) 86 | if not face_landmarks_list: 87 | print("No face detected in the image") 88 | return None 89 | 90 | # Take first detected face 91 | face_landmarks = face_landmarks_list[0] 92 | 93 | # Prepare data structure for landmarks 94 | landmarks_data = { 95 | 'x': [], 96 | 'y': [], 97 | 'index': [] 98 | } 99 | 100 | # Get image dimensions for coordinate conversion 101 | image_height, image_width = image_np.shape[:2] 102 | 103 | # Extract landmark coordinates. MediaPipe may return a 104 | # NormalizedLandmarkList object (with a `landmark` attribute) 105 | # or simply a list of `Landmark` objects depending on the API 106 | # (Solutions vs Tasks). The previous implementation assumed the 107 | # former and failed with `'list' object has no attribute 108 | # "landmark"' when using the Tasks API. To support both 109 | # structures we access the iterable safely. 110 | landmarks_iter = getattr(face_landmarks, 'landmark', face_landmarks) 111 | 112 | for idx, landmark in enumerate(landmarks_iter): 113 | # Convert normalized coordinates to pixel coordinates 114 | x = landmark.x * image_width 115 | y = landmark.y * image_height 116 | 117 | landmarks_data['x'].append(x) 118 | landmarks_data['y'].append(y) 119 | landmarks_data['index'].append(idx) 120 | 121 | return pd.DataFrame(landmarks_data) 122 | 123 | except Exception as e: 124 | print(f"Error in MediaPipe landmark detection: {str(e)}") 125 | return None 126 | 127 | def detect_landmarks_dlib(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: 128 | """ 129 | Detect facial landmarks using dlib's 68 point predictor. 130 | 131 | Args: 132 | image: Input image in various formats 133 | 134 | Returns: 135 | Optional[pd.DataFrame]: DataFrame containing landmark coordinates (x, y) and indices or None if no face detected 136 | """ 137 | try: 138 | # Initialize Dlib model if not already done 139 | if self.dlib_model is None: 140 | self.dlib_model = ModelDlib() 141 | 142 | # Convert image to numpy format suitable for dlib 143 | image_np = self.image_processor.convert_to_numpy(image) 144 | if image_np is None: 145 | print("Failed to convert image for dlib processing") 146 | return None 147 | 148 | # Convert to grayscale for dlib 149 | gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) 150 | 151 | # Detect faces using Dlib face detector 152 | faces = self.dlib_model.face_detector(gray) 153 | if not faces: 154 | print("No face detected by dlib") 155 | return None 156 | 157 | # Get first face 158 | face = faces[0] 159 | 160 | # Detect landmarks using shape predictor 161 | shape = self.dlib_model.shape_predictor(gray, face) 162 | 163 | # Create landmark data structure 164 | landmarks_data = { 165 | 'x': [], 166 | 'y': [], 167 | 'index': [] 168 | } 169 | 170 | # Extract landmarks 171 | for i in range(68): 172 | point = shape.part(i) 173 | landmarks_data['x'].append(float(point.x)) 174 | landmarks_data['y'].append(float(point.y)) 175 | landmarks_data['index'].append(i + 1) # Adding 1 to match dlib's indexing 176 | 177 | return pd.DataFrame(landmarks_data) 178 | 179 | except Exception as e: 180 | print(f"Error in dlib landmark detection: {str(e)}") 181 | return None 182 | 183 | 184 | def landmarks_interpolation(self, mp_landmarks: pd.DataFrame, dlib_landmarks: pd.DataFrame) -> pd.DataFrame: 185 | """ 186 | Interpolates all MediaPipe landmarks based on dlib reference points using RBF interpolation. 187 | 188 | Args: 189 | mp_landmarks: DataFrame with columns [x, y, index] 190 | dlib_landmarks: DataFrame with columns [x, y, index] 191 | 192 | Returns: 193 | pd.DataFrame: DataFrame with interpolated landmark positions 194 | """ 195 | from scipy.interpolate import RBFInterpolator 196 | 197 | # Validate input DataFrames 198 | required_columns = ['x', 'y', 'index'] 199 | if not all(col in mp_landmarks.columns for col in required_columns): 200 | print("Error: MediaPipe landmarks missing required columns") 201 | return mp_landmarks 202 | 203 | if not all(col in dlib_landmarks.columns for col in required_columns): 204 | print("Error: Dlib landmarks missing required columns") 205 | return mp_landmarks 206 | 207 | try: 208 | # Get control points from Dlib landmarks 209 | control_points_df = LandmarkMappings.get_control_points(dlib_landmarks) 210 | if control_points_df is None: 211 | print("Failed to generate control points, returning original landmarks") 212 | return mp_landmarks 213 | 214 | # Prepare source (MediaPipe) and destination (Dlib) points for RBF 215 | control_src = [] # MediaPipe points 216 | control_dst = [] # Dlib-based points 217 | 218 | for _, control_point in control_points_df.iterrows(): 219 | mp_idx = control_point['index'] 220 | mp_point = mp_landmarks[mp_landmarks['index'] == mp_idx] 221 | 222 | if not mp_point.empty: 223 | control_src.append([mp_point['x'].iloc[0], mp_point['y'].iloc[0]]) 224 | control_dst.append([control_point['x'], control_point['y']]) 225 | 226 | control_src = np.array(control_src) 227 | control_dst = np.array(control_dst) 228 | 229 | print(f"Using {len(control_src)} control points for RBF interpolation") 230 | 231 | # Create and fit RBF interpolation for both x and y coordinates 232 | kernel = 'thin_plate_spline' 233 | rbf_x = RBFInterpolator(control_src, control_dst[:, 0], kernel=kernel) 234 | rbf_y = RBFInterpolator(control_src, control_dst[:, 1], kernel=kernel) 235 | 236 | # Transform all MediaPipe points 237 | all_points = mp_landmarks[['x', 'y']].values 238 | 239 | # Apply transformation 240 | transformed_x = rbf_x(all_points) 241 | transformed_y = rbf_y(all_points) 242 | 243 | # Create result DataFrame 244 | result = mp_landmarks.copy() 245 | result['x'] = transformed_x 246 | result['y'] = transformed_y 247 | 248 | print("Successfully interpolated all landmarks using RBF") 249 | return result 250 | 251 | except Exception as e: 252 | print(f"Error during landmark interpolation: {str(e)}") 253 | return mp_landmarks 254 | -------------------------------------------------------------------------------- /core/gpu_deformer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | import torch 4 | from ..core.base_mesh import MediapipeBaseLandmarks 5 | 6 | 7 | class GPUDeformer: 8 | """ 9 | Class for handling face warping operations on GPU using CUDA acceleration. 10 | """ 11 | 12 | @staticmethod 13 | def warp_face(image, source_landmarks, target_landmarks): 14 | """ 15 | Warp the face image using parallel GPU processing with reverse mapping. 16 | Now includes boundary triangles for complete image warping. 17 | 18 | Args: 19 | image: Input PIL Image or numpy array 20 | source_landmarks: Detected face landmarks 21 | target_landmarks: Target base landmarks 22 | 23 | Returns: 24 | PIL Image: Warped image 25 | """ 26 | try: 27 | import cupy as cp 28 | print("Starting GPU warping process with boundary triangles...") 29 | 30 | # Convert and validate input data 31 | source_landmarks = np.asarray(source_landmarks, dtype=np.float32) 32 | target_landmarks = np.asarray(target_landmarks, dtype=np.float32) 33 | 34 | # Debug landmark ranges 35 | print("\nLandmark ranges:") 36 | print(f"Source landmarks - X: [{source_landmarks[:, 0].min():.1f}, {source_landmarks[:, 0].max():.1f}], " 37 | f"Y: [{source_landmarks[:, 1].min():.1f}, {source_landmarks[:, 1].max():.1f}]") 38 | print(f"Target landmarks - X: [{target_landmarks[:, 0].min():.1f}, {target_landmarks[:, 0].max():.1f}], " 39 | f"Y: [{target_landmarks[:, 1].min():.1f}, {target_landmarks[:, 1].max():.1f}]") 40 | 41 | # Convert image to numpy array if needed 42 | if isinstance(image, Image.Image): 43 | img_np = np.array(image, dtype=np.uint8) 44 | else: 45 | img_np = np.asarray(image, dtype=np.uint8) 46 | 47 | h, w = img_np.shape[:2] 48 | channels = img_np.shape[2] if len(img_np.shape) > 2 else 1 49 | print(f"\nImage shape: {img_np.shape}") 50 | 51 | # Get triangulation including boundary triangles 52 | triangles, all_points = MediapipeBaseLandmarks.get_face_triangles(size=(w, h)) 53 | print(f"Loaded {len(triangles)} triangles (including boundary triangles)") 54 | 55 | # Create complete source and target point arrays 56 | n_base_landmarks = len(source_landmarks) 57 | n_all_points = len(all_points) 58 | print(f"Base landmarks: {n_base_landmarks}, Total points: {n_all_points}") 59 | 60 | # Initialize complete arrays for source and target points 61 | source_points = np.zeros((n_all_points, 2), dtype=np.float32) 62 | target_points = np.zeros((n_all_points, 2), dtype=np.float32) 63 | 64 | # Fill in face landmarks 65 | n_base = 468 # Hardcode original landmark count 66 | source_points[:n_base] = source_landmarks[:n_base] 67 | target_points[:n_base] = target_landmarks[:n_base] 68 | 69 | # Boundary points remain identical 70 | source_points[n_base:] = all_points[n_base:] 71 | target_points[n_base:] = all_points[n_base:] 72 | 73 | # CUDA kernel for parallel triangle processing with reverse mapping 74 | cuda_kernel = cp.RawKernel(r''' 75 | extern "C" __global__ 76 | void process_triangles( 77 | const float* source_points, 78 | const float* target_points, 79 | const int* triangles, 80 | const unsigned char* input_image, 81 | unsigned char* output_image, 82 | const int num_triangles, 83 | const int width, 84 | const int height, 85 | const int channels 86 | ) { 87 | // Get triangle index 88 | const int tid = blockIdx.x * blockDim.x + threadIdx.x; 89 | if (tid >= num_triangles) return; 90 | 91 | // Get triangle indices 92 | const int idx0 = triangles[tid * 3]; 93 | const int idx1 = triangles[tid * 3 + 1]; 94 | const int idx2 = triangles[tid * 3 + 2]; 95 | 96 | // Get source points 97 | const float src_x0 = source_points[idx0 * 2]; 98 | const float src_y0 = source_points[idx0 * 2 + 1]; 99 | const float src_x1 = source_points[idx1 * 2]; 100 | const float src_y1 = source_points[idx1 * 2 + 1]; 101 | const float src_x2 = source_points[idx2 * 2]; 102 | const float src_y2 = source_points[idx2 * 2 + 1]; 103 | 104 | // Get target points 105 | const float dst_x0 = target_points[idx0 * 2]; 106 | const float dst_y0 = target_points[idx0 * 2 + 1]; 107 | const float dst_x1 = target_points[idx1 * 2]; 108 | const float dst_y1 = target_points[idx1 * 2 + 1]; 109 | const float dst_x2 = target_points[idx2 * 2]; 110 | const float dst_y2 = target_points[idx2 * 2 + 1]; 111 | 112 | // Calculate target triangle area 113 | const float dst_area = (dst_x1 - dst_x0) * (dst_y2 - dst_y0) - 114 | (dst_x2 - dst_x0) * (dst_y1 - dst_y0); 115 | 116 | const float abs_area = abs(dst_area); 117 | if (abs_area < 1e-6f) return; 118 | 119 | // Calculate target triangle bounding box with border check 120 | const int min_x = max(0, (int)min(min(dst_x0, dst_x1), dst_x2)); 121 | const int min_y = max(0, (int)min(min(dst_y0, dst_y1), dst_y2)); 122 | const int max_x = min(width - 1, (int)(max(max(dst_x0, dst_x1), dst_x2) + 0.5f)); 123 | const int max_y = min(height - 1, (int)(max(max(dst_y0, dst_y1), dst_y2) + 0.5f)); 124 | 125 | // Process each pixel in target triangle's bounding box 126 | for (int y = min_y; y <= max_y; y++) { 127 | for (int x = min_x; x <= max_x; x++) { 128 | // Calculate barycentric coordinates 129 | float lambda1 = ((y - dst_y2) * (dst_x0 - dst_x2) + 130 | (dst_x2 - x) * (dst_y0 - dst_y2)) / dst_area; 131 | float lambda2 = ((y - dst_y0) * (dst_x1 - dst_x0) + 132 | (dst_x0 - x) * (dst_y1 - dst_y0)) / dst_area; 133 | float lambda0 = 1.0f - lambda1 - lambda2; 134 | 135 | // Check if point is inside triangle with tolerance 136 | const float eps = 1e-5f; 137 | if (lambda0 >= -eps && lambda1 >= -eps && lambda2 >= -eps && 138 | lambda0 <= 1+eps && lambda1 <= 1+eps && lambda2 <= 1+eps) { 139 | 140 | // Calculate source position using barycentric coordinates 141 | float source_x = lambda0 * src_x0 + lambda1 * src_x1 + lambda2 * src_x2; 142 | float source_y = lambda0 * src_y0 + lambda1 * src_y1 + lambda2 * src_y2; 143 | 144 | // Round to nearest pixel 145 | int sx = (int)(source_x + 0.5f); 146 | int sy = (int)(source_y + 0.5f); 147 | 148 | // Copy pixel if within bounds 149 | if (sx >= 0 && sx < width && sy >= 0 && sy < height) { 150 | for (int c = 0; c < channels; c++) { 151 | output_image[(y * width + x) * channels + c] = 152 | input_image[(sy * width + sx) * channels + c]; 153 | } 154 | } 155 | } 156 | } 157 | } 158 | } 159 | ''', 'process_triangles') 160 | 161 | # Move data to GPU 162 | source_points_gpu = cp.asarray(source_points.ravel()) 163 | target_points_gpu = cp.asarray(target_points.ravel()) 164 | triangles_gpu = cp.asarray(triangles.ravel().astype(np.int32)) 165 | img_gpu = cp.asarray(img_np) 166 | output_gpu = cp.zeros_like(img_gpu) 167 | 168 | # Configure CUDA grid 169 | threadsPerBlock = 256 170 | blocksPerGrid = (len(triangles) + threadsPerBlock - 1) // threadsPerBlock 171 | print(f"CUDA config: {blocksPerGrid} blocks, {threadsPerBlock} threads per block") 172 | 173 | # Execute kernel 174 | print(f"Processing {len(triangles)} triangles on GPU...") 175 | cuda_kernel( 176 | (blocksPerGrid,), (threadsPerBlock,), 177 | (source_points_gpu, target_points_gpu, triangles_gpu, 178 | img_gpu, output_gpu, len(triangles), 179 | w, h, channels) 180 | ) 181 | cp.cuda.stream.get_current_stream().synchronize() 182 | print("CUDA kernel execution completed") 183 | print("Flushing CUDA cache...") 184 | torch.cuda.empty_cache() 185 | 186 | # Transfer result back to CPU and convert to PIL 187 | output_np = cp.asnumpy(output_gpu) 188 | return Image.fromarray(output_np) 189 | 190 | except ImportError as e: 191 | print(f"CuPy not installed or CUDA not available: {str(e)}") 192 | return image 193 | except Exception as e: 194 | print(f"Error during GPU warping: {str(e)}") 195 | print(f"Error details: {str(e.__class__.__name__)}") 196 | import traceback 197 | traceback.print_exc() 198 | return image 199 | -------------------------------------------------------------------------------- /core/image_processor.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import pandas as pd 4 | import torch 5 | from PIL import Image 6 | from typing import Union, Optional, Tuple 7 | 8 | class ImageProcessor: 9 | 10 | @staticmethod 11 | def calculate_face_bbox(landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[Tuple[int, int, int, int]]: 12 | """Calculate bounding box for face based on landmarks.""" 13 | if landmarks_df is None or landmarks_df.empty: 14 | return None 15 | 16 | min_x = landmarks_df['x'].min() 17 | max_x = landmarks_df['x'].max() 18 | min_y = landmarks_df['y'].min() 19 | max_y = landmarks_df['y'].max() 20 | 21 | width = max_x - min_x 22 | height = max_y - min_y 23 | 24 | pad_x = width * padding_percent 25 | pad_y = height * padding_percent 26 | 27 | x1 = max(0, min_x - pad_x) 28 | y1 = max(0, min_y - pad_y) 29 | x2 = max_x + pad_x 30 | y2 = max_y + pad_y 31 | 32 | return int(x1), int(y1), int(x2 - x1), int(y2 - y1) 33 | 34 | @staticmethod 35 | def resize_image(image: Union[torch.Tensor, np.ndarray, Image.Image], target_size: int) -> Optional[np.ndarray]: 36 | """Resize image to target size while maintaining aspect ratio and cropping to square.""" 37 | if image is None: 38 | return None 39 | 40 | image_np = ImageProcessor.convert_to_numpy(image) 41 | h, w = image_np.shape[:2] 42 | 43 | # Determine the smaller dimension (height or width) 44 | min_dim = min(h, w) 45 | 46 | # Calculate the crop box to make the image square 47 | start_x = (w - min_dim) // 2 48 | start_y = (h - min_dim) // 2 49 | end_x = start_x + min_dim 50 | end_y = start_y + min_dim 51 | 52 | # Crop the image to a square 53 | cropped_image = image_np[start_y:end_y, start_x:end_x] 54 | 55 | # Resize the cropped image to the target size 56 | resized = cv2.resize(cropped_image, (target_size, target_size), interpolation=cv2.INTER_LANCZOS4) 57 | 58 | return resized 59 | 60 | @staticmethod 61 | def calculate_rotation_angle(landmarks_df: pd.DataFrame) -> float: 62 | """Calculate rotation angle based on eyes position.""" 63 | if landmarks_df is None or landmarks_df.empty: 64 | return 0.0 65 | 66 | LEFT_EYE = 33 # Center of the left eye 67 | RIGHT_EYE = 263 # Center of the right eye 68 | 69 | left_eye = landmarks_df[landmarks_df['index'] == LEFT_EYE].iloc[0] 70 | right_eye = landmarks_df[landmarks_df['index'] == RIGHT_EYE].iloc[0] 71 | 72 | dx = right_eye['x'] - left_eye['x'] 73 | dy = right_eye['y'] - left_eye['y'] 74 | return np.degrees(np.arctan2(dy, dx)) 75 | 76 | @staticmethod 77 | def rotate_image(image: Union[torch.Tensor, np.ndarray, Image.Image], landmarks_df: pd.DataFrame) -> Tuple[Optional[np.ndarray], Optional[pd.DataFrame]]: 78 | """Rotate image based on facial landmarks.""" 79 | image_np = ImageProcessor.convert_to_numpy(image) 80 | 81 | if image_np is None or landmarks_df is None: 82 | return None, None 83 | 84 | angle = ImageProcessor.calculate_rotation_angle(landmarks_df) 85 | height, width = image_np.shape[:2] 86 | center = (width // 2, height // 2) 87 | 88 | rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0) 89 | rotated_image = cv2.warpAffine(image_np, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4) 90 | 91 | # Transform landmarks 92 | ones = np.ones(shape=(len(landmarks_df), 1)) 93 | points = np.hstack([landmarks_df[['x', 'y']].values, ones]) 94 | transformed_points = rotation_matrix.dot(points.T).T 95 | 96 | updated_landmarks = landmarks_df.copy() 97 | updated_landmarks['x'] = transformed_points[:, 0] 98 | updated_landmarks['y'] = transformed_points[:, 1] 99 | 100 | return rotated_image, updated_landmarks 101 | 102 | @staticmethod 103 | def crop_face_to_square(image: np.ndarray, landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Tuple[ 104 | Optional[np.ndarray], Optional[Tuple[int, int, int, int]]]: 105 | """Crop face region to a square (1:1) based on landmarks.""" 106 | bbox = ImageProcessor.calculate_face_bbox(landmarks_df, padding_percent) 107 | if bbox is None: 108 | return None, None 109 | 110 | x, y, w, h = bbox 111 | 112 | # Calculate the center of the bounding box 113 | center_x = x + w // 2 114 | center_y = y + h // 2 115 | 116 | # Determine the size of the square crop 117 | crop_size = max(w, h) 118 | half_size = crop_size // 2 119 | 120 | # Calculate the crop coordinates 121 | x1 = max(0, center_x - half_size) 122 | y1 = max(0, center_y - half_size) 123 | x2 = min(image.shape[1], center_x + half_size) 124 | y2 = min(image.shape[0], center_y + half_size) 125 | 126 | # Adjust if the crop goes out of bounds 127 | if x2 - x1 < crop_size: 128 | x1 = max(0, x2 - crop_size) 129 | if y2 - y1 < crop_size: 130 | y1 = max(0, y2 - crop_size) 131 | 132 | # Crop the image 133 | cropped_face = image[y1:y2, x1:x2] 134 | 135 | # Return the cropped face and the crop bounding box 136 | return cropped_face, (x1, y1, x2 - x1, y2 - y1) 137 | 138 | @staticmethod 139 | def draw_landmarks(image_size: Union[int, Tuple[int, int]], 140 | landmarks_df: pd.DataFrame, 141 | transparency: float = 0.5, 142 | color: Tuple[int, int, int] = (0, 255, 0), 143 | radius: int = 2, 144 | label: bool = False) -> Union[np.ndarray, None]: 145 | """ 146 | Draw facial landmarks on a transparent image. 147 | 148 | Args: 149 | image_size: Target image size (int or tuple of width, height) 150 | landmarks_df: DataFrame containing landmark coordinates 151 | transparency: Alpha channel value (0.0-1.0) 152 | color: RGB color tuple for landmarks 153 | radius: Radius of landmark points 154 | label: Whether to show landmark IDs 155 | 156 | Returns: 157 | numpy array: RGBA image with drawn landmarks 158 | """ 159 | if landmarks_df is None or landmarks_df.empty: 160 | return None 161 | 162 | # Create transparent image 163 | if isinstance(image_size, tuple): 164 | width, height = image_size 165 | else: 166 | width = height = image_size 167 | 168 | # Create RGBA image (alpha channel for transparency) 169 | image = np.zeros((height, width, 4), dtype=np.uint8) 170 | 171 | # Set alpha channel based on transparency 172 | alpha = int(transparency * 255) 173 | 174 | # Draw landmarks 175 | for _, landmark in landmarks_df.iterrows(): 176 | x = int(landmark['x']) 177 | y = int(landmark['y']) 178 | index = int(landmark['index']) 179 | 180 | # Ensure coordinates are within image bounds 181 | if 0 <= x < width and 0 <= y < height: 182 | # Draw filled circle 183 | cv2.circle( 184 | image, 185 | (x, y), 186 | radius, 187 | (*color, alpha), 188 | -1, # Filled circle 189 | cv2.LINE_AA 190 | ) 191 | 192 | # Add landmark ID if requested 193 | if label: 194 | # Text parameters 195 | font = cv2.FONT_HERSHEY_SIMPLEX 196 | font_scale = 0.4 197 | text_thickness = 1 198 | 199 | # Add small offset to text position 200 | text_x = x + radius + 5 201 | text_y = y + radius 202 | 203 | # Draw text with white color and black outline for better visibility 204 | text = str(index) 205 | 206 | # Draw text outline 207 | cv2.putText( 208 | image, 209 | text, 210 | (text_x, text_y), 211 | font, 212 | font_scale, 213 | (0, 0, 0, alpha), # Black outline 214 | text_thickness + 1, 215 | cv2.LINE_AA 216 | ) 217 | 218 | # Draw text 219 | cv2.putText( 220 | image, 221 | text, 222 | (text_x, text_y), 223 | font, 224 | font_scale, 225 | (*color, alpha), # Main color 226 | text_thickness, 227 | cv2.LINE_AA 228 | ) 229 | 230 | return image 231 | 232 | @staticmethod 233 | def draw_dynamic_histogram(composite: np.ndarray, 234 | input_black: int, input_white: int, 235 | gamma: float) -> np.ndarray: 236 | """Creates histogram visualization with dynamic control lines""" 237 | # Convert to grayscale for histogram calculation 238 | gray = cv2.cvtColor(composite, cv2.COLOR_RGB2GRAY) if len(composite.shape) == 3 else composite 239 | 240 | # Calculate histogram 241 | hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) 242 | hist_img = np.zeros((200, 256, 3), dtype=np.uint8) 243 | cv2.normalize(hist, hist, 0, hist_img.shape[0], cv2.NORM_MINMAX) 244 | 245 | # Draw histogram baseline 246 | for i in range(1, 256): 247 | cv2.line(hist_img, 248 | (i - 1, 200 - int(hist[i - 1])), 249 | (i, 200 - int(hist[i])), 250 | (128, 128, 128), 1) 251 | 252 | # Calculate gamma position 253 | gamma_pos = int(255 * (0.5 ** (1 / gamma))) 254 | 255 | # Draw control lines 256 | cv2.line(hist_img, (input_black, 0), (input_black, 200), (255, 0, 0), 2) # Blue - black level 257 | cv2.line(hist_img, (input_white, 0), (input_white, 200), (0, 255, 0), 2) # Green - white level 258 | cv2.line(hist_img, (gamma_pos, 0), (gamma_pos, 200), (255, 255, 255), 2) # White - gamma 259 | 260 | # Add legend 261 | cv2.putText(hist_img, f"Black: {input_black}", (10, 20), 262 | cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1) 263 | cv2.putText(hist_img, f"White: {input_white}", (10, 40), 264 | cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) 265 | cv2.putText(hist_img, f"Gamma: {gamma:.2f}", (10, 60), 266 | cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) 267 | 268 | return hist_img 269 | 270 | # Image Converters 271 | @staticmethod 272 | def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: 273 | """Convert PyTorch tensor to numpy array. 274 | 275 | Args: 276 | tensor: Input tensor in format (B,H,W,C) or (B,C,H,W), values 0-1 277 | 278 | Returns: 279 | np.ndarray: RGB image (H,W,3), uint8 values 0-255 280 | """ 281 | if tensor is None: 282 | return None 283 | 284 | # Handle different tensor formats 285 | img = tensor.detach().cpu().numpy() 286 | if len(img.shape) == 4: 287 | img = img[0] # Remove batch dimension 288 | 289 | # Handle channel dimension if in (C,H,W) format 290 | if img.shape[0] in [1, 3, 4]: 291 | img = img.transpose(1, 2, 0) 292 | 293 | # Handle single channel 294 | if len(img.shape) == 2 or img.shape[2] == 1: 295 | img = np.stack((img,) * 3, axis=-1) 296 | 297 | # Handle RGBA 298 | elif img.shape[2] == 4: 299 | img = img[..., :3] 300 | 301 | # Convert to uint8 if float 302 | if img.dtype == np.float32 or img.dtype == np.float64: 303 | img = (img * 255).clip(0, 255).astype(np.uint8) 304 | 305 | return img 306 | 307 | @staticmethod 308 | def numpy_to_tensor(array: np.ndarray) -> torch.Tensor: 309 | """Convert numpy array to PyTorch tensor. 310 | 311 | Args: 312 | array: RGB image (H,W,3), uint8 or float32 313 | 314 | Returns: 315 | torch.Tensor: (1,H,W,3) tensor, values 0-1 316 | """ 317 | if array is None: 318 | return None 319 | 320 | # Convert to float32 if uint8 321 | if array.dtype == np.uint8: 322 | array = array.astype(np.float32) / 255.0 323 | 324 | # Add batch dimension 325 | tensor = torch.from_numpy(array).unsqueeze(0) 326 | return tensor 327 | 328 | @staticmethod 329 | def pil_to_numpy(image: Image.Image) -> np.ndarray: 330 | """Convert PIL image to numpy array. 331 | 332 | Args: 333 | image: PIL Image 334 | 335 | Returns: 336 | np.ndarray: RGB image (H,W,3), uint8 values 0-255 337 | """ 338 | if image is None: 339 | return None 340 | 341 | # Convert to RGB if needed 342 | if image.mode != 'RGB': 343 | image = image.convert('RGB') 344 | 345 | return np.array(image) 346 | 347 | @staticmethod 348 | def numpy_to_pil(array: np.ndarray) -> Image.Image: 349 | """Convert numpy array to PIL image. 350 | 351 | Args: 352 | array: RGB image (H,W,3), uint8 or float32 353 | 354 | Returns: 355 | PIL.Image: RGB PIL image 356 | """ 357 | if array is None: 358 | return None 359 | 360 | # Convert to uint8 if float 361 | if array.dtype == np.float32 or array.dtype == np.float64: 362 | array = (array * 255).clip(0, 255).astype(np.uint8) 363 | 364 | return Image.fromarray(array) 365 | 366 | @staticmethod 367 | def tensor_to_pil(tensor: torch.Tensor) -> Image.Image: 368 | """Convert PyTorch tensor to PIL image. 369 | 370 | Args: 371 | tensor: Input tensor in format (B,H,W,C), values 0-1 372 | 373 | Returns: 374 | PIL.Image: RGB PIL image 375 | """ 376 | return ImageProcessor.numpy_to_pil(ImageProcessor.tensor_to_numpy(tensor)) 377 | 378 | @staticmethod 379 | def pil_to_tensor(image: Image.Image) -> torch.Tensor: 380 | """Convert PIL image to PyTorch tensor. 381 | 382 | Args: 383 | image: PIL Image 384 | 385 | Returns: 386 | torch.Tensor: (1,H,W,3) tensor, values 0-1 387 | """ 388 | return ImageProcessor.numpy_to_tensor(ImageProcessor.pil_to_numpy(image)) 389 | 390 | @staticmethod 391 | def convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray: 392 | """Universal converter to numpy array. 393 | 394 | Args: 395 | image: Input in any supported format 396 | 397 | Returns: 398 | np.ndarray: RGB image (H,W,3), uint8 values 0-255 399 | """ 400 | if image is None: 401 | return None 402 | 403 | if torch.is_tensor(image): 404 | return ImageProcessor.tensor_to_numpy(image) 405 | elif isinstance(image, Image.Image): 406 | return ImageProcessor.pil_to_numpy(image) 407 | elif isinstance(image, np.ndarray): 408 | # Handle grayscale 409 | if len(image.shape) == 2: 410 | image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) 411 | # Handle RGBA 412 | elif image.shape[2] == 4: 413 | image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) 414 | # Convert to uint8 if float 415 | if image.dtype == np.float32 or image.dtype == np.float64: 416 | image = (image * 255).clip(0, 255).astype(np.uint8) 417 | return image 418 | else: 419 | raise ValueError(f"Unsupported image type: {type(image)}") 420 | 421 | # Mask conversion methods 422 | @staticmethod 423 | def convert_mask_to_numpy(mask: Union[torch.Tensor, np.ndarray]) -> np.ndarray: 424 | """Convert mask to numpy array. 425 | 426 | Args: 427 | mask: Input mask as tensor (1,H,W) or array 428 | 429 | Returns: 430 | np.ndarray: Mask as (H,W) uint8 array, values 0-255 431 | """ 432 | if mask is None: 433 | return None 434 | 435 | if torch.is_tensor(mask): 436 | mask = mask.detach().cpu().numpy() 437 | if len(mask.shape) == 3: 438 | mask = mask[0] # Remove batch dimension 439 | 440 | if mask.dtype == np.float32 or mask.dtype == np.float64: 441 | mask = (mask * 255).clip(0, 255).astype(np.uint8) 442 | 443 | return mask 444 | 445 | @staticmethod 446 | def convert_mask_to_tensor(mask: np.ndarray) -> torch.Tensor: 447 | """Convert numpy mask to tensor. 448 | 449 | Args: 450 | mask: Input mask as (H,W) array 451 | 452 | Returns: 453 | torch.Tensor: Mask as (1,H,W) tensor, values 0-1 454 | """ 455 | if mask is None: 456 | return None 457 | 458 | if mask.dtype == np.uint8: 459 | mask = mask.astype(np.float32) / 255.0 460 | 461 | return torch.from_numpy(mask).unsqueeze(0) 462 | 463 | # Loaders 464 | @staticmethod 465 | def load_image_from_path(image_path): 466 | """Load and convert image from file path to tensor format.""" 467 | try: 468 | image = Image.open(image_path) 469 | if image.mode != 'RGB': 470 | image = image.convert('RGB') 471 | return ImageProcessor.pil_to_tensor(image) 472 | except Exception as e: 473 | print(f"Error loading image from {image_path}: {str(e)}") 474 | return None 475 | 476 | -------------------------------------------------------------------------------- /core/lm_mapping.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import numpy as np 3 | from typing import Dict, List, Tuple, Optional 4 | 5 | 6 | class LandmarkMappings: 7 | """ 8 | Class for mapping and transforming facial landmarks between MediaPipe and Dlib coordinate spaces. 9 | Supports interpolated control points derived from Dlib landmark pairs. 10 | """ 11 | 12 | # Base landmark pairs mapping MediaPipe to Dlib indices 13 | # Format: (mediapipe_idx, dlib_idx) 14 | LANDMARKS_PAIRS = { 15 | "left_eye": [ 16 | (133, 40), # Left eye outer corner 17 | (33, 37) # Left eye inner corner 18 | ], 19 | "right_eye": [ 20 | (362, 43), # Right eye outer corner 21 | (263, 46) # Right eye inner corner 22 | ], 23 | "mouth": [ 24 | (0, 52), # Upper lip top 25 | (17, 58), # Bottom lip bottom 26 | (61, 49), # Left mouth corner 27 | (291, 55), # Right mouth corner 28 | (13, 63), # Upper lip bottom 29 | (14, 67), # Bottom lip top 30 | (78, 61), # Left mouth inner corner 31 | (308, 65) # Right mouth inner corner 32 | ], 33 | "nose_edge": [ 34 | (6, 28), # Nose bridge top 35 | (4, 31) # Nose tip 36 | ], 37 | "face_oval": [ 38 | (152, 9), # Left jawline 39 | (127, 1), # Chin left 40 | (356, 17), # Right jawline 41 | (172, 5), # Mid jawline left 42 | (397, 13) # Mid jawline right 43 | ] 44 | } 45 | 46 | # Virtual points interpolated from Dlib landmark pairs 47 | # Format: mediapipe_idx: (dlib_idx1, dlib_idx2) 48 | VIRTUAL_PAIRS = { 49 | "9": (22, 23) 50 | } 51 | 52 | @classmethod 53 | def interpolate_point(cls, dlib_landmarks: pd.DataFrame, dlib_idx1: int, dlib_idx2: int) -> Optional[ 54 | Tuple[float, float]]: 55 | """ 56 | Interpolate a point between two Dlib landmarks. 57 | 58 | Args: 59 | dlib_landmarks: DataFrame with Dlib landmarks 60 | dlib_idx1: First Dlib landmark index 61 | dlib_idx2: Second Dlib landmark index 62 | 63 | Returns: 64 | Tuple[float, float]: Interpolated (x, y) coordinates or None if error 65 | """ 66 | try: 67 | # Get the two Dlib points 68 | point1 = dlib_landmarks[dlib_landmarks['index'] == dlib_idx1] 69 | point2 = dlib_landmarks[dlib_landmarks['index'] == dlib_idx2] 70 | 71 | if point1.empty or point2.empty: 72 | print(f"Warning: Could not find Dlib landmarks {dlib_idx1} and/or {dlib_idx2}") 73 | return None 74 | 75 | # Calculate midpoint 76 | x = (point1['x'].iloc[0] + point2['x'].iloc[0]) / 2 77 | y = (point1['y'].iloc[0] + point2['y'].iloc[0]) / 2 78 | 79 | return float(x), float(y) 80 | 81 | except Exception as e: 82 | print(f"Error interpolating point between Dlib landmarks {dlib_idx1} and {dlib_idx2}: {str(e)}") 83 | return None 84 | 85 | @classmethod 86 | def get_control_points(cls, dlib_landmarks: pd.DataFrame) -> Optional[pd.DataFrame]: 87 | """ 88 | Generate control points DataFrame based on Dlib landmarks. 89 | Includes both direct mappings and interpolated points. 90 | 91 | Args: 92 | dlib_landmarks: DataFrame containing Dlib landmark coordinates with columns [x, y, index] 93 | 94 | Returns: 95 | pd.DataFrame: DataFrame containing mapped MediaPipe landmark positions with columns [x, y, index] 96 | or None if input data is invalid 97 | """ 98 | try: 99 | # Validate input DataFrame 100 | required_columns = {'x', 'y', 'index'} 101 | if not all(col in dlib_landmarks.columns for col in required_columns): 102 | print("Error: Input DataFrame missing required columns [x, y, index]") 103 | return None 104 | 105 | # Prepare data structures for control points 106 | control_points = { 107 | 'x': [], 108 | 'y': [], 109 | 'index': [] 110 | } 111 | 112 | # Process regular landmark pairs 113 | print("Processing direct landmark mappings...") 114 | for feature_group, pairs in cls.LANDMARKS_PAIRS.items(): 115 | for mp_idx, dlib_idx in pairs: 116 | dlib_point = dlib_landmarks[dlib_landmarks['index'] == dlib_idx] 117 | 118 | if dlib_point.empty: 119 | print(f"Warning: Dlib landmark {dlib_idx} not found for MediaPipe index {mp_idx}") 120 | continue 121 | 122 | control_points['x'].append(float(dlib_point['x'].iloc[0])) 123 | control_points['y'].append(float(dlib_point['y'].iloc[0])) 124 | control_points['index'].append(mp_idx) 125 | 126 | # Process interpolated virtual points 127 | print("Processing interpolated control points...") 128 | for mp_idx, (dlib_idx1, dlib_idx2) in cls.VIRTUAL_PAIRS.items(): 129 | interpolated = cls.interpolate_point(dlib_landmarks, dlib_idx1, dlib_idx2) 130 | if interpolated is not None: 131 | x, y = interpolated 132 | control_points['x'].append(x) 133 | control_points['y'].append(y) 134 | control_points['index'].append(int(mp_idx)) 135 | print(f"Added interpolated point for MediaPipe {mp_idx} at ({x:.1f}, {y:.1f})") 136 | 137 | # Create DataFrame with all control points 138 | control_df = pd.DataFrame(control_points) 139 | 140 | # Sort by MediaPipe landmark index for consistency 141 | control_df = control_df.sort_values('index').reset_index(drop=True) 142 | 143 | print(f"Successfully mapped {len(control_df)} control points " 144 | f"({len(control_df) - len(cls.VIRTUAL_PAIRS)} regular, " 145 | f"{len(cls.VIRTUAL_PAIRS)} interpolated)") 146 | 147 | return control_df 148 | 149 | except Exception as e: 150 | print(f"Error generating control points: {str(e)}") 151 | return None 152 | 153 | # TODO: filter_trackers() 154 | def filter_trackers(self, distance): 155 | """ 156 | return list of base mediapipe landmarks ids as a list of trackers to use for cv2.goodFeaturesToTrack 157 | based on given distance 158 | """ 159 | trackers = [] 160 | return trackers -------------------------------------------------------------------------------- /core/resources/model_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Optional, Dict, List 3 | 4 | import dlib 5 | import mediapipe as mp 6 | from mediapipe.tasks import python as mp_python 7 | from mediapipe.tasks.python import vision 8 | import numpy as np 9 | import requests 10 | from tqdm import tqdm 11 | 12 | 13 | class ModelDownload: 14 | """Class for managing model downloads and verification""" 15 | 16 | MODEL_URLS = { 17 | # Mediapipe 468 landmark model 18 | "canonical_face_model": "https://raw.githubusercontent.com/google-ai-edge/mediapipe/master/mediapipe/modules/face_geometry/data/canonical_face_model.obj", 19 | 20 | # Dlib 68 landmark model 21 | "shape_predictor_68_face_landmarks": "https://huggingface.co/spaces/asdasdasdasd/Face-forgery-detection/resolve/ccfc24642e0210d4d885bc7b3dbc9a68ed948ad6/shape_predictor_68_face_landmarks.dat", 22 | # MediaPipe face landmarker .task model 23 | "face_landmarker": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" 24 | } 25 | 26 | def __init__(self): 27 | """Initialize model downloader""" 28 | self.chunk_size = 8192 29 | self.models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") 30 | 31 | def download_model(self, model_name: str, target_dir: Optional[str] = None) -> Optional[str]: 32 | """ 33 | Download model from predefined URL. 34 | 35 | Args: 36 | model_name: Name of the model to download 37 | target_dir: Optional target directory (defaults to models/{model_name}) 38 | 39 | Returns: 40 | str: Path to downloaded model or None if failed 41 | """ 42 | try: 43 | if model_name not in self.MODEL_URLS: 44 | print(f"Unknown model: {model_name}") 45 | return None 46 | 47 | # Set target directory 48 | if target_dir is None: 49 | target_dir = os.path.join(self.models_dir, os.path.splitext(model_name)[0]) 50 | 51 | # Create target directory if it doesn't exist 52 | os.makedirs(target_dir, exist_ok=True) 53 | 54 | # Set target file path 55 | url = self.MODEL_URLS[model_name] 56 | if not url: # Skip if URL is empty 57 | print(f"No URL defined for model: {model_name}") 58 | return None 59 | 60 | file_name = os.path.basename(url.split('?')[0]) # Remove URL parameters if any 61 | target_path = os.path.join(target_dir, file_name) 62 | 63 | # Check if file already exists 64 | if os.path.exists(target_path): 65 | print(f"Model {model_name} already exists at: {target_path}") 66 | return target_path 67 | 68 | print(f"Downloading {model_name} from: {url}") 69 | print(f"Target path: {target_path}") 70 | 71 | # Download file with progress bar 72 | response = requests.get(url, stream=True) 73 | total_size = int(response.headers.get('content-length', 0)) 74 | 75 | with open(target_path, 'wb') as f: 76 | with tqdm(total=total_size, unit='B', unit_scale=True, desc=model_name) as pbar: 77 | for data in response.iter_content(chunk_size=self.chunk_size): 78 | f.write(data) 79 | pbar.update(len(data)) 80 | 81 | print(f"Successfully downloaded {model_name} to: {target_path}") 82 | return target_path 83 | 84 | except Exception as e: 85 | print(f"Error downloading model {model_name}: {str(e)}") 86 | return None 87 | 88 | def get_model_path(self, model_name: str, target_dir: Optional[str] = None) -> Optional[str]: 89 | """ 90 | Get path to model, download if not exists. 91 | 92 | Args: 93 | model_name: Name of the model 94 | target_dir: Optional target directory 95 | 96 | Returns: 97 | str: Path to model or None if failed 98 | """ 99 | if target_dir is None: 100 | target_dir = os.path.join(self.models_dir, os.path.splitext(model_name)[0]) 101 | 102 | # Check if model exists 103 | if os.path.exists(target_dir): 104 | expected_file = os.path.join(target_dir, os.path.basename(self.MODEL_URLS[model_name].split('?')[0])) 105 | if os.path.exists(expected_file): 106 | return expected_file 107 | 108 | # Download if not exists 109 | return self.download_model(model_name, target_dir) 110 | 111 | 112 | class ModelOBJ: 113 | """Class for loading and parsing OBJ files with landmark transformation support""" 114 | 115 | _instance: Optional['ModelOBJ'] = None 116 | _vertices: List[List[float]] = [] 117 | _uvs: List[List[float]] = [] 118 | _faces: List[List[int]] = [] 119 | _vertex_to_uv: Dict[int, int] = {} 120 | _landmarks: Optional[np.ndarray] = None 121 | _num_landmarks: int = 468 122 | _model_loaded: bool = False 123 | 124 | def __new__(cls): 125 | """Implement singleton pattern""" 126 | if cls._instance is None: 127 | cls._instance = super(ModelOBJ, cls).__new__(cls) 128 | return cls._instance 129 | 130 | def __init__(self): 131 | """Initialize ModelOBJ and load face model""" 132 | if not self._model_loaded: 133 | self.downloader = ModelDownload() 134 | self._load_obj() 135 | 136 | def _load_obj(self) -> None: 137 | """Load and parse MediaPipe face model OBJ file""" 138 | try: 139 | # Get model path, download if needed 140 | model_path = self.downloader.get_model_path("canonical_face_model") 141 | if not model_path: 142 | raise FileNotFoundError("Failed to get canonical face model") 143 | 144 | print(f"Loading 3D model from: {model_path}") 145 | 146 | with open(model_path, 'r') as f: 147 | for line in f: 148 | if line.startswith('v '): 149 | # Parse vertex coordinates 150 | parts = line.strip().split() 151 | self._vertices.append([float(p) for p in parts[1:4]]) 152 | 153 | elif line.startswith('vt '): 154 | # Parse UV coordinates (flip V coordinate) 155 | parts = line.strip().split() 156 | u = float(parts[1]) 157 | v = 1.0 - float(parts[2]) # Flip V coordinate 158 | self._uvs.append([u, v]) 159 | 160 | elif line.startswith('f '): 161 | # Parse face indices 162 | parts = line.strip().split()[1:] 163 | face_verts = [] 164 | for part in parts: 165 | indices = part.split('/') 166 | if len(indices) >= 2: 167 | vertex_idx = int(indices[0]) - 1 # OBJ indices are 1-based 168 | uv_idx = int(indices[1]) - 1 169 | self._vertex_to_uv[vertex_idx] = uv_idx 170 | face_verts.append(vertex_idx) 171 | if len(face_verts) == 3: # Only store triangular faces 172 | self._faces.append(face_verts) 173 | 174 | # Initialize landmarks 175 | self._landmarks = self._initialize_landmarks() 176 | self._model_loaded = True 177 | 178 | print(f"Successfully loaded model data:") 179 | print(f"- Vertices: {len(self._vertices)}") 180 | print(f"- UVs: {len(self._uvs)}") 181 | print(f"- Faces: {len(self._faces)}") 182 | print(f"- Vertex to UV mappings: {len(self._vertex_to_uv)}") 183 | 184 | except Exception as e: 185 | print(f"Error loading OBJ file: {str(e)}") 186 | raise 187 | 188 | def _initialize_landmarks(self) -> np.ndarray: 189 | """ 190 | Initialize landmarks array from UV coordinates 191 | 192 | Returns: 193 | numpy.ndarray: Array of landmark coordinates in UV space 194 | """ 195 | landmarks = np.zeros((self._num_landmarks, 2), dtype=np.float32) 196 | uvs = np.array(self._uvs, dtype=np.float32) 197 | 198 | for vertex_idx, uv_idx in self._vertex_to_uv.items(): 199 | if vertex_idx < self._num_landmarks: 200 | landmarks[vertex_idx] = uvs[uv_idx] 201 | 202 | return landmarks 203 | 204 | def get_faces(self) -> np.ndarray: 205 | """ 206 | Get face indices as numpy array 207 | 208 | Returns: 209 | numpy.ndarray: Array of face indices 210 | """ 211 | return np.array(self._faces, dtype=np.int32) 212 | 213 | def get_transformed_landmarks(self, x_scale: float = 1.0, y_translation: float = 0.0) -> np.ndarray: 214 | """ 215 | Get landmarks with applied transformations 216 | 217 | Args: 218 | x_scale: Horizontal scaling factor (0.5 to 1.0) 219 | y_translation: Vertical translation (-0.5 to 0.5) 220 | 221 | Returns: 222 | numpy.ndarray: Transformed landmarks 223 | """ 224 | if self._landmarks is None: 225 | raise RuntimeError("Landmarks not initialized") 226 | 227 | transformed = self._landmarks.copy() 228 | 229 | # Calculate center point 230 | center_x = (self._landmarks[:, 0].min() + self._landmarks[:, 0].max()) / 2 231 | 232 | # Calculate horizontal distance from center (normalized to 0-1) 233 | dx = np.abs(transformed[:, 0] - center_x) 234 | 235 | # Avoid division by zero 236 | if np.abs(center_x) < 1e-6: # Use small epsilon value 237 | influence = np.zeros_like(dx) 238 | else: 239 | influence = np.clip(dx / center_x, 0, 1) 240 | 241 | # Calculate scaled positions with horizontal influence 242 | scale_factor = 1.0 + (x_scale - 1.0) * influence 243 | transformed[:, 0] = center_x + (transformed[:, 0] - center_x) * scale_factor 244 | 245 | # Apply vertical translation 246 | transformed[:, 1] += y_translation 247 | 248 | return transformed 249 | 250 | @property 251 | def num_landmarks(self) -> int: 252 | """ 253 | Get the number of landmarks 254 | 255 | Returns: 256 | int: Number of landmarks 257 | """ 258 | return self._num_landmarks 259 | 260 | def _initialize_landmarks(self) -> np.ndarray: 261 | """ 262 | Initialize landmarks array from UV coordinates 263 | 264 | Returns: 265 | numpy.ndarray: Array of landmark coordinates in UV space 266 | """ 267 | landmarks = np.zeros((self._num_landmarks, 2), dtype=np.float32) 268 | uvs = np.array(self._uvs, dtype=np.float32) 269 | 270 | for vertex_idx, uv_idx in self._vertex_to_uv.items(): 271 | if vertex_idx < self._num_landmarks: 272 | landmarks[vertex_idx] = uvs[uv_idx] 273 | 274 | return landmarks 275 | 276 | def get_faces(self) -> np.ndarray: 277 | """ 278 | Get face indices as numpy array 279 | 280 | Returns: 281 | numpy.ndarray: Array of face indices 282 | """ 283 | return np.array(self._faces, dtype=np.int32) 284 | 285 | def get_transformed_landmarks(self, x_scale: float = 1.0, y_translation: float = 0.0) -> np.ndarray: 286 | """ 287 | Get landmarks with applied transformations 288 | 289 | Args: 290 | x_scale: Horizontal scaling factor (0.5 to 1.0) 291 | y_translation: Vertical translation (-0.5 to 0.5) 292 | 293 | Returns: 294 | numpy.ndarray: Transformed landmarks 295 | """ 296 | if self._landmarks is None: 297 | raise RuntimeError("Landmarks not initialized") 298 | 299 | transformed = self._landmarks.copy() 300 | 301 | # Calculate center point 302 | center_x = (self._landmarks[:, 0].min() + self._landmarks[:, 0].max()) / 2 303 | 304 | # Calculate horizontal distance from center (normalized to 0-1) 305 | dx = np.abs(transformed[:, 0] - center_x) 306 | 307 | # Avoid division by zero 308 | if np.abs(center_x) < 1e-6: # Use small epsilon value 309 | influence = np.zeros_like(dx) 310 | else: 311 | influence = np.clip(dx / center_x, 0, 1) 312 | 313 | # Calculate scaled positions with horizontal influence 314 | scale_factor = 1.0 + (x_scale - 1.0) * influence 315 | transformed[:, 0] = center_x + (transformed[:, 0] - center_x) * scale_factor 316 | 317 | # Apply vertical translation 318 | transformed[:, 1] += y_translation 319 | 320 | return transformed 321 | 322 | 323 | class ModelDlib: 324 | """Class for loading and managing Dlib face detection models using singleton pattern""" 325 | 326 | _instance: Optional['ModelDlib'] = None 327 | _face_detector: Optional[dlib.fhog_object_detector] = None 328 | _shape_predictor: Optional[dlib.shape_predictor] = None 329 | _model_loaded: bool = False 330 | 331 | def __new__(cls): 332 | """Implement singleton pattern""" 333 | if cls._instance is None: 334 | cls._instance = super(ModelDlib, cls).__new__(cls) 335 | return cls._instance 336 | 337 | def __init__(self): 338 | """Initialize Dlib model loader""" 339 | if not self._model_loaded: 340 | self.downloader = ModelDownload() 341 | self._initialize_models() 342 | 343 | def _initialize_models(self) -> None: 344 | """Initialize Dlib face detector and shape predictor""" 345 | try: 346 | print("Initializing Dlib face detector...") 347 | self._face_detector = dlib.get_frontal_face_detector() 348 | 349 | # Get model path, download if needed 350 | model_path = self.downloader.get_model_path("shape_predictor_68_face_landmarks") 351 | if not model_path: 352 | raise FileNotFoundError("Failed to get shape predictor model") 353 | 354 | print(f"Loading Dlib shape predictor from: {model_path}") 355 | self._shape_predictor = dlib.shape_predictor(model_path) 356 | self._model_loaded = True 357 | print("Dlib models initialized successfully") 358 | 359 | except Exception as e: 360 | print(f"Error initializing Dlib models: {str(e)}") 361 | raise 362 | 363 | @property 364 | def face_detector(self) -> dlib.fhog_object_detector: 365 | """Get Dlib face detector instance""" 366 | if not self._face_detector: 367 | self._initialize_models() 368 | return self._face_detector 369 | 370 | @property 371 | def shape_predictor(self) -> dlib.shape_predictor: 372 | """Get Dlib shape predictor instance""" 373 | if not self._shape_predictor: 374 | self._initialize_models() 375 | return self._shape_predictor 376 | 377 | 378 | class _FaceLandmarkerWrapper: 379 | """Wrapper to provide a FaceMesh-like interface for MediaPipe Tasks.""" 380 | 381 | def __init__(self, landmarker: vision.FaceLandmarker) -> None: 382 | self._landmarker = landmarker 383 | 384 | def process(self, image: np.ndarray): 385 | mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image) 386 | return self._landmarker.detect(mp_image) 387 | 388 | def close(self) -> None: 389 | self._landmarker.close() 390 | 391 | 392 | class ModelMediaPipe: 393 | """Class for loading and managing MediaPipe face mesh models using singleton pattern""" 394 | 395 | _instance: Optional['ModelMediaPipe'] = None 396 | _face_mesh: Optional[_FaceLandmarkerWrapper] = None 397 | _model_loaded: bool = False 398 | 399 | def __new__(cls): 400 | if cls._instance is None: 401 | cls._instance = super(ModelMediaPipe, cls).__new__(cls) 402 | return cls._instance 403 | 404 | def __init__(self): 405 | if not self._model_loaded: 406 | self._initialize_models() 407 | 408 | def _initialize_models(self) -> None: 409 | try: 410 | print("Initializing MediaPipe Face Mesh...") 411 | 412 | downloader = ModelDownload() 413 | model_path = downloader.get_model_path("face_landmarker") 414 | if not model_path: 415 | raise FileNotFoundError("Failed to get face landmarker model") 416 | 417 | base_options = mp_python.BaseOptions(model_asset_path=model_path) 418 | options = vision.FaceLandmarkerOptions( 419 | base_options=base_options, 420 | num_faces=1, 421 | output_face_blendshapes=False, 422 | output_facial_transformation_matrixes=False, 423 | ) 424 | landmarker = vision.FaceLandmarker.create_from_options(options) 425 | self._face_mesh = _FaceLandmarkerWrapper(landmarker) 426 | 427 | self._model_loaded = True 428 | print("MediaPipe Face Mesh initialized successfully") 429 | 430 | except Exception as e: 431 | print(f"Error initializing MediaPipe models: {str(e)}") 432 | raise 433 | 434 | @property 435 | def face_mesh(self) -> _FaceLandmarkerWrapper: 436 | if not self._face_mesh: 437 | self._initialize_models() 438 | return self._face_mesh 439 | 440 | def __del__(self): 441 | if self._face_mesh: 442 | self._face_mesh.close() 443 | self._face_mesh = None 444 | self._model_loaded = False 445 | -------------------------------------------------------------------------------- /nodes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SykkoAtHome/ComfyUI_FaceProcessor/7687b344c15f8c5c0d8983d60e98447a20d01757/nodes/__init__.py -------------------------------------------------------------------------------- /nodes/face_fit_and_restore.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | 5 | from ..core.face_detector import FaceDetector 6 | from ..core.image_processor import ImageProcessor 7 | 8 | 9 | class FaceFitAndRestore: 10 | """ComfyUI node for processing face images in Fit or Restore mode.""" 11 | 12 | def __init__(self): 13 | self.face_detector = FaceDetector() 14 | self.image_processor = ImageProcessor() 15 | 16 | @classmethod 17 | def INPUT_TYPES(cls): 18 | return { 19 | "required": { 20 | "workflow": (["image", "image_sequence"], { 21 | "default": "image" 22 | }), 23 | "mode": (["Fit", "Restore"], { 24 | "default": "Fit" 25 | }), 26 | "output_mode": (["current_frame", "batch_sequence"], { 27 | "default": "current_frame" 28 | }), 29 | "padding_percent": ("FLOAT", { 30 | "default": 0.0, 31 | "min": 0.0, 32 | "max": 1.0, 33 | "step": 0.05 34 | }), 35 | "bbox_size": (["512", "1024", "2048"], { 36 | "default": "1024" 37 | }), 38 | }, 39 | "optional": { 40 | "image": ("IMAGE",), 41 | "fp_pipe": ("DICT", { 42 | "default": None 43 | }), 44 | "image_sequence": ("DICT", { 45 | "default": None 46 | }), 47 | } 48 | } 49 | 50 | RETURN_TYPES = ("IMAGE", "DICT", "MASK", "INT") 51 | RETURN_NAMES = ("image", "fp_pipe", "mask", "bbox_size") 52 | FUNCTION = "process_image" 53 | CATEGORY = "Face Processor" 54 | 55 | def process_image(self, mode, workflow, output_mode, padding_percent=0.0, bbox_size="1024", 56 | image=None, fp_pipe=None, image_sequence=None): 57 | """Process images in either Fit or Restore mode with unified fp_pipe structure.""" 58 | 59 | # Validate inputs 60 | if image is None and (workflow == "image" or mode == "Restore"): 61 | print("Error: Image is required for image workflow and Restore mode") 62 | return None, fp_pipe, None, int(bbox_size) 63 | 64 | if workflow == "image_sequence" and mode == "Fit" and (not image_sequence or "frames" not in image_sequence): 65 | print("Error: Valid image sequence data is required for Fit mode") 66 | return None, fp_pipe, None, int(bbox_size) 67 | 68 | if mode == "Restore" and (not fp_pipe or "frames" not in fp_pipe): 69 | print("Error: Valid fp_pipe with frames is required for Restore mode") 70 | return None, fp_pipe, None, int(bbox_size) 71 | 72 | # Initialize fp_pipe if None 73 | if fp_pipe is None: 74 | fp_pipe = { 75 | "workflow": workflow, 76 | "current_frame": 0, 77 | "padding_percent": padding_percent, 78 | "target_lm": {}, 79 | "frames": {} 80 | } 81 | 82 | try: 83 | if workflow == "image": 84 | # Handle potential batch input 85 | if torch.is_tensor(image) and len(image.shape) == 4 and image.shape[0] > 1: 86 | processed_frames = [] 87 | processed_masks = [] 88 | batch_size = image.shape[0] 89 | for idx in range(batch_size): 90 | result = self._process_frame( 91 | mode=mode, 92 | image=image[idx:idx + 1], # keep batch dimension 93 | frame_number=idx, 94 | padding_percent=padding_percent, 95 | bbox_size=bbox_size, 96 | fp_pipe=fp_pipe 97 | ) 98 | processed_frames.append(result[0]) 99 | if result[2] is not None: 100 | processed_masks.append(result[2]) 101 | 102 | batched_frames = torch.cat(processed_frames, dim=0) 103 | batched_masks = torch.cat(processed_masks, dim=0) if processed_masks else None 104 | return batched_frames, fp_pipe, batched_masks, int(bbox_size) 105 | 106 | else: 107 | return self._process_frame( 108 | mode=mode, 109 | image=image, 110 | frame_number=0, 111 | padding_percent=padding_percent, 112 | bbox_size=bbox_size, 113 | fp_pipe=fp_pipe 114 | ) 115 | 116 | else: # image_sequence 117 | current_frame = (fp_pipe.get("current_frame", 0) if mode == "Restore" 118 | else image_sequence.get("current_frame", 0)) 119 | fp_pipe["current_frame"] = current_frame 120 | 121 | if mode == "Restore": 122 | total_frames = len(fp_pipe["frames"]) 123 | frames_to_process = { 124 | idx: None for idx in range(total_frames) 125 | } 126 | else: # Fit mode 127 | total_frames = len(image_sequence["frames"]) 128 | frames_to_process = image_sequence["frames"] 129 | 130 | print(f"Processing sequence of {total_frames} frames...") 131 | results = {} 132 | current_frame_result = None 133 | 134 | # Process each frame 135 | for frame_idx in range(total_frames): 136 | print(f"Processing frame {frame_idx + 1}/{total_frames}") 137 | 138 | if mode == "Fit": 139 | frame_path = frames_to_process[frame_idx] 140 | frame_image = self.image_processor.load_image_from_path(frame_path) 141 | if frame_image is None: 142 | print(f"Warning: Could not load frame {frame_idx}") 143 | continue 144 | else: # Restore 145 | # Handle potential batch input 146 | if len(image.shape) == 4 and image.shape[0] > 1: # We have a batch 147 | if frame_idx < image.shape[0]: 148 | frame_image = image[frame_idx:frame_idx + 1] # Keep the batch dimension 149 | else: 150 | print(f"Warning: Frame index {frame_idx} exceeds batch size {image.shape[0]}") 151 | continue 152 | else: # Single image 153 | frame_image = image 154 | frame_path = None 155 | 156 | # Process frame 157 | result = self._process_frame( 158 | mode=mode, 159 | image=frame_image, 160 | frame_number=frame_idx, 161 | padding_percent=padding_percent, 162 | bbox_size=bbox_size, 163 | fp_pipe=fp_pipe, 164 | original_path=frame_path if mode == "Fit" else None 165 | ) 166 | 167 | # Store results for both current frame and batch processing 168 | results[frame_idx] = result 169 | if frame_idx == current_frame: 170 | current_frame_result = result 171 | 172 | # Return based on output mode 173 | if output_mode == "current_frame": 174 | if current_frame_result is None: 175 | print(f"Warning: Frame {current_frame} not found in sequence") 176 | return None, fp_pipe, None, int(bbox_size) 177 | return current_frame_result 178 | else: # batch_sequence mode 179 | # Collect all processed frames 180 | all_frames = [] 181 | all_masks = [] 182 | 183 | for idx in range(total_frames): 184 | frame_result = results.get(idx) 185 | if frame_result is not None: 186 | processed_image = frame_result[0] 187 | if processed_image is not None: 188 | all_frames.append(processed_image) 189 | if frame_result[2] is not None: # If mask exists 190 | all_masks.append(frame_result[2]) 191 | 192 | if not all_frames: 193 | print("Warning: No frames were successfully processed") 194 | return None, fp_pipe, None, int(bbox_size) 195 | 196 | # Stack all frames into a single batch tensor 197 | batched_frames = torch.cat(all_frames, dim=0) 198 | 199 | # Stack masks if they exist 200 | batched_masks = torch.cat(all_masks, dim=0) if all_masks else None 201 | 202 | return batched_frames, fp_pipe, batched_masks, int(bbox_size) 203 | 204 | except Exception as e: 205 | print(f"Error in process_image: {str(e)}") 206 | import traceback 207 | traceback.print_exc() 208 | return None, fp_pipe, None, int(bbox_size) 209 | 210 | def _process_frame(self, mode, image, frame_number, padding_percent, bbox_size, fp_pipe, original_path=None): 211 | """Process a single frame in either Fit or Restore mode.""" 212 | try: 213 | frame_key = f"frame_{frame_number}" 214 | 215 | if mode == "Fit": 216 | # Process frame in Fit mode 217 | result_image, frame_settings, result_mask = self._fit( 218 | image=image, 219 | padding_percent=padding_percent, 220 | bbox_size=bbox_size 221 | ) 222 | 223 | if frame_settings: 224 | # Create standardized frame data structure 225 | frame_data = { 226 | "bbox_size": int(bbox_size), 227 | "crop_bbox": frame_settings["crop_bbox"], 228 | "original_image_shape": frame_settings["original_image_shape"], 229 | "rotation_angle": frame_settings["rotation_angle"], 230 | "detected_lm": {} # Will be populated by face_wrapper 231 | } 232 | 233 | # Add original path for sequence processing 234 | if original_path: 235 | frame_data["original_image_path"] = original_path 236 | 237 | # Store frame data 238 | fp_pipe["frames"][frame_key] = frame_data 239 | 240 | else: # Restore mode 241 | frame_settings = fp_pipe["frames"].get(frame_key) 242 | if frame_settings is None: 243 | print(f"Error: No settings found for frame {frame_number}") 244 | return None, fp_pipe, None, int(bbox_size) 245 | 246 | result = self._restore(image, frame_settings) 247 | if result is None: 248 | return None, fp_pipe, None, int(bbox_size) 249 | 250 | result_image, result_mask = result[0], result[2] 251 | 252 | return result_image, fp_pipe, result_mask, int(bbox_size) 253 | 254 | except Exception as e: 255 | print(f"Error processing frame {frame_number}: {str(e)}") 256 | import traceback 257 | traceback.print_exc() 258 | return None, fp_pipe, None, int(bbox_size) 259 | 260 | def _fit(self, image, padding_percent, bbox_size): 261 | """Fit mode: Crop and process the face.""" 262 | # Convert image to numpy array 263 | image_np = self.image_processor.convert_to_numpy(image) 264 | if image_np is None: 265 | return image, {}, self._create_empty_mask(image) 266 | 267 | # Detect facial landmarks 268 | landmarks_df = self.face_detector.detect_landmarks_mp(image_np) 269 | if landmarks_df is None: 270 | print("No face detected, returning original image") 271 | return image, {}, self._create_empty_mask(image) 272 | 273 | # Calculate rotation and rotate image 274 | rotation_angle = self.image_processor.calculate_rotation_angle(landmarks_df) 275 | rotated_image, updated_landmarks = self.image_processor.rotate_image(image_np, landmarks_df) 276 | if rotated_image is None: 277 | print("Failed to rotate image, returning original") 278 | return image, {}, self._create_empty_mask(image) 279 | 280 | # Crop face region 281 | cropped_face, crop_bbox = self.image_processor.crop_face_to_square( 282 | rotated_image, updated_landmarks, padding_percent 283 | ) 284 | if cropped_face is None: 285 | print("Failed to crop face, returning original image") 286 | return image, {}, self._create_empty_mask(image) 287 | 288 | # Resize to target size 289 | target_size = int(bbox_size) 290 | final_image = self.image_processor.resize_image(cropped_face, target_size) 291 | if final_image is None: 292 | print("Failed to resize image, returning original") 293 | return image, {}, self._create_empty_mask(image) 294 | 295 | # Convert to tensor 296 | final_image = self.image_processor.numpy_to_tensor(final_image) 297 | 298 | # Create frame settings 299 | frame_settings = { 300 | "original_image_shape": image_np.shape, 301 | "rotation_angle": rotation_angle, 302 | "crop_bbox": crop_bbox, 303 | "bbox_size": target_size 304 | } 305 | 306 | # Create mask 307 | mask = self._create_mask( 308 | image_shape=frame_settings["original_image_shape"], 309 | crop_bbox=frame_settings["crop_bbox"], 310 | rotation_angle=frame_settings["rotation_angle"] 311 | ) 312 | 313 | return final_image, frame_settings, mask 314 | 315 | def _restore(self, image, frame_settings): 316 | """Restore mode: Restore the face back to original position.""" 317 | try: 318 | # Validate frame settings 319 | required_keys = ["original_image_shape", "rotation_angle", "crop_bbox"] 320 | if not all(key in frame_settings for key in required_keys): 321 | print("Error: Missing required frame settings") 322 | return None 323 | 324 | # Convert input image to numpy 325 | processed_face_np = self.image_processor.convert_to_numpy(image) 326 | if processed_face_np is None: 327 | return None 328 | 329 | # Get settings 330 | original_shape = frame_settings["original_image_shape"] 331 | rotation_angle = frame_settings["rotation_angle"] 332 | x1, y1, w, h = frame_settings["crop_bbox"] 333 | 334 | # Create output image 335 | restored_image = np.zeros(original_shape, dtype=np.uint8) 336 | 337 | # Resize face to original crop size 338 | resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4) 339 | 340 | # Place face in position 341 | restored_image[y1:y1 + h, x1:x1 + w] = resized_face 342 | 343 | # Rotate back if needed 344 | if rotation_angle != 0: 345 | height, width = original_shape[:2] 346 | center = (width // 2, height // 2) 347 | rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) 348 | restored_image = cv2.warpAffine( 349 | restored_image, 350 | rotation_matrix, 351 | (width, height), 352 | flags=cv2.INTER_LANCZOS4 353 | ) 354 | 355 | # Convert to tensor 356 | restored_image = self.image_processor.numpy_to_tensor(restored_image) 357 | 358 | # Create mask 359 | mask = self._create_mask(original_shape, (x1, y1, w, h), rotation_angle) 360 | 361 | return restored_image, frame_settings, mask 362 | 363 | except Exception as e: 364 | print(f"Error in restore: {str(e)}") 365 | return None 366 | 367 | def _create_mask(self, image_shape, crop_bbox, rotation_angle): 368 | """Create a mask for the face area.""" 369 | mask = np.zeros(image_shape[:2], dtype=np.float32) 370 | x1, y1, w, h = crop_bbox 371 | 372 | # Create face area mask 373 | mask[y1:y1 + h, x1:x1 + w] = 1.0 374 | 375 | # Rotate mask if needed 376 | if rotation_angle != 0: 377 | height, width = image_shape[:2] 378 | center = (width // 2, height // 2) 379 | rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) 380 | mask = cv2.warpAffine( 381 | mask, 382 | rotation_matrix, 383 | (width, height), 384 | flags=cv2.INTER_LINEAR 385 | ) 386 | 387 | return self.image_processor.convert_mask_to_tensor(mask) 388 | 389 | def _create_empty_mask(self, image): 390 | """Create an empty mask matching the image dimensions.""" 391 | if torch.is_tensor(image): 392 | shape = image.shape[1:3] 393 | else: 394 | shape = image.shape[:2] 395 | return torch.zeros((1, *shape), dtype=torch.float32) 396 | -------------------------------------------------------------------------------- /nodes/face_tracker.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | from ..core.face_detector import FaceDetector 4 | from ..core.image_processor import ImageProcessor 5 | 6 | 7 | class FaceTracker: 8 | """ComfyUI node for tracking facial features across image sequences.""" 9 | 10 | # Optical flow parameters optimized for face tracking 11 | OPTICAL_FLOW_PARAMS = { 12 | "pyramid_scale": 0.5, # Each pyramid level is half the size of the previous 13 | "levels": 5, # Number of pyramid levels 14 | "window_size": 21, # Size of the search window at each pyramid level 15 | "iterations": 3, # Number of iterations the algorithm does at each pyramid level 16 | "poly_n": 7, # Size of the pixel neighborhood used for polynomial expansion 17 | "poly_sigma": 1.5, # Standard deviation of the Gaussian used to smooth derivatives 18 | "flags": 0 # Optional flags (can be cv2.OPTFLOW_USE_INITIAL_FLOW or cv2.OPTFLOW_FARNEBACK_GAUSSIAN) 19 | } 20 | 21 | def __init__(self): 22 | self.face_detector = FaceDetector() 23 | self.image_processor = ImageProcessor() 24 | 25 | @classmethod 26 | def INPUT_TYPES(cls): 27 | return { 28 | "required": { 29 | "image": ("IMAGE",), 30 | "fp_pipe": ("DICT",), 31 | "debug": ("BOOLEAN", { 32 | "default": False 33 | }), 34 | "use_motion_vectors": ("BOOLEAN", { 35 | "default": False 36 | }), 37 | "reset_points_interval": ("INT", { 38 | "default": 10, 39 | "min": 3, 40 | "max": 50, 41 | "step": 1 42 | }), 43 | "tracker_region_size": (["32", "48", "64", "96", "128"], { 44 | "default": "64" 45 | }), 46 | "proxy_scale": ("FLOAT", { 47 | "default": 1.0, 48 | "min": 0.2, 49 | "max": 1.0, 50 | "step": 0.05 51 | }), 52 | "show_detection": ("BOOLEAN", { 53 | "default": False 54 | }), 55 | "show_region": ("BOOLEAN", { 56 | "default": False 57 | }), 58 | } 59 | } 60 | 61 | RETURN_TYPES = ("IMAGE", "DICT") 62 | RETURN_NAMES = ("image", "fp_pipe") 63 | FUNCTION = "track_face" 64 | CATEGORY = "Face Processor" 65 | 66 | def track_face(self, image, fp_pipe, debug, use_motion_vectors, reset_points_interval, 67 | tracker_region_size, proxy_scale, show_detection, show_region): 68 | input_image_np = self.image_processor.convert_to_numpy(image) 69 | try: 70 | # Convert input image for processing 71 | image_np = input_image_np 72 | if image_np is None: 73 | print("Error: Failed to convert input image") 74 | return image, fp_pipe 75 | 76 | # Apply proxy scale if needed 77 | if proxy_scale < 1.0: 78 | h, w = image_np.shape[:2] 79 | new_h, new_w = int(h * proxy_scale), int(w * proxy_scale) 80 | image_np = cv2.resize(image_np, (new_w, new_h)) 81 | 82 | # Detect facial landmarks 83 | # TODO: Handle No face detected. Fill with NaN values and keep going 84 | landmarks_df = self.face_detector.detect_landmarks_mp(image_np) 85 | if landmarks_df is None: 86 | print("No face detected") 87 | return image, fp_pipe 88 | 89 | # If using proxy scale, rescale landmarks back to original size 90 | if proxy_scale < 1.0: 91 | landmarks_df['x'] = landmarks_df['x'] / proxy_scale 92 | landmarks_df['y'] = landmarks_df['y'] / proxy_scale 93 | 94 | # Update tracking data in fp_pipe 95 | current_frame = fp_pipe.get("current_frame", 0) 96 | frame_key = f"frame_{current_frame}" 97 | 98 | # Initialize tracking data structure if needed 99 | if "tracking_data" not in fp_pipe: 100 | fp_pipe["tracking_data"] = {} 101 | 102 | if frame_key not in fp_pipe["tracking_data"]: 103 | fp_pipe["tracking_data"][frame_key] = {} 104 | 105 | # Store MediaPipe detection results 106 | fp_pipe["tracking_data"][frame_key]["mediapipe"] = { 107 | "landmarks": landmarks_df.to_dict('records'), 108 | "timestamp": current_frame, 109 | "confidence": 1.0 # MediaPipe detection confidence 110 | } 111 | 112 | # Debug visualization 113 | if debug and show_detection: 114 | output_image = self._draw_debug_visualization(input_image_np, landmarks_df, show_region, 115 | int(tracker_region_size)) 116 | return self.image_processor.numpy_to_tensor(output_image), fp_pipe 117 | 118 | return image, fp_pipe 119 | 120 | except Exception as e: 121 | print(f"Error in track_face: {str(e)}") 122 | import traceback 123 | traceback.print_exc() 124 | return image, fp_pipe 125 | 126 | def _draw_debug_visualization(self, image, landmarks_df, show_region, region_size): 127 | """Draw debug visualization including landmarks and optionally tracking regions.""" 128 | debug_image = image.copy() 129 | 130 | # Draw landmarks 131 | for _, landmark in landmarks_df.iterrows(): 132 | x, y = int(landmark['x']), int(landmark['y']) 133 | cv2.circle(debug_image, (x, y), 2, (0, 255, 0), -1) 134 | 135 | # Draw tracking region if requested 136 | if show_region: 137 | half_size = region_size // 2 138 | x1 = max(0, x - half_size) 139 | y1 = max(0, y - half_size) 140 | x2 = min(image.shape[1], x + half_size) 141 | y2 = min(image.shape[0], y + half_size) 142 | cv2.rectangle(debug_image, (x1, y1), (x2, y2), (255, 0, 0), 1) 143 | 144 | return debug_image -------------------------------------------------------------------------------- /nodes/face_wrapper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import torch 4 | 5 | from ..core.base_mesh import MediapipeBaseLandmarks 6 | from ..core.cpu_deformer import CPUDeformer 7 | from ..core.face_detector import FaceDetector 8 | from ..core.gpu_deformer import GPUDeformer 9 | from ..core.image_processor import ImageProcessor 10 | 11 | 12 | class FaceWrapper: 13 | """ComfyUI node for detecting facial landmarks with optional visualization and warping.""" 14 | 15 | def __init__(self): 16 | self.face_detector = FaceDetector() 17 | self.image_processor = ImageProcessor() 18 | 19 | @classmethod 20 | def INPUT_TYPES(cls): 21 | return { 22 | "required": { 23 | "image": ("IMAGE",), 24 | "mode": (["Debug", "Un-Wrap", "Wrap"], {"default": "Debug"}), 25 | "device": (["CPU", "CUDA"], {"default": "CUDA"}), 26 | "show_detection": ("BOOLEAN", {"default": False}), 27 | "show_target": ("BOOLEAN", {"default": False}), 28 | "refiner": (["None", "Dlib"], {"default": "None"}), 29 | "landmark_size": ("INT", {"default": 4, "min": 1, "max": 10, "step": 1}), 30 | "show_labels": ("BOOLEAN", {"default": False}), 31 | "x_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 1.0, "step": 0.01}), 32 | "y_transform": ("FLOAT", {"default": 0.0, "min": -0.5, "max": 0.5, "step": 0.01}) 33 | }, 34 | "optional": { 35 | "fp_pipe": ("DICT", {"default": None}), 36 | "mask": ("MASK", {"default": None}) 37 | } 38 | } 39 | 40 | RETURN_TYPES = ("IMAGE", "DICT", "MASK") 41 | RETURN_NAMES = ("image", "fp_pipe", "mask") 42 | FUNCTION = "detect_face" 43 | CATEGORY = "Face Processor" 44 | 45 | def detect_face(self, image, mode, device, show_detection, show_target, refiner, landmark_size, 46 | show_labels, x_scale, y_transform, fp_pipe=None, mask=None): 47 | # Convert input image to numpy with proper RGB format 48 | image_np = self.image_processor.convert_to_numpy(image) 49 | height, width = image_np.shape[:2] 50 | 51 | # Convert mask if provided 52 | mask_np = None 53 | if mask is not None: 54 | mask_np = self.image_processor.convert_mask_to_numpy(mask) 55 | 56 | if mode == "Wrap": 57 | return self._wrap_mode(image_np, None, width, height, 58 | device, x_scale, y_transform, fp_pipe, mask_np) 59 | 60 | # Detect facial landmarks 61 | landmarks_df = self.face_detector.detect_landmarks(image_np, refiner=(None if refiner == "None" else refiner)) 62 | 63 | if landmarks_df is None: 64 | print("No face detected") 65 | empty_mask = torch.zeros((1, height, width), dtype=torch.float32) if mask is not None else None 66 | return image, fp_pipe or {}, empty_mask 67 | 68 | # Handle different modes 69 | if mode == "Debug": 70 | return self._debug_mode(image_np, landmarks_df, width, height, 71 | show_detection, show_target, landmark_size, 72 | show_labels, x_scale, y_transform, fp_pipe, mask_np) 73 | elif mode == "Un-Wrap": 74 | return self._unwrap_mode(image_np, landmarks_df, width, height, 75 | device, x_scale, y_transform, fp_pipe, mask_np) 76 | 77 | 78 | def _debug_mode(self, image_np, landmarks_df, width, height, show_detection, 79 | show_target, landmark_size, show_labels, x_scale, y_transform, 80 | fp_pipe, mask_np): 81 | result_image = image_np.astype(np.float32) / 255.0 82 | overlays = [] 83 | 84 | if show_detection: 85 | det_overlay = self.image_processor.draw_landmarks( 86 | (width, height), landmarks_df, 87 | transparency=0.4, color=(0, 255, 0), 88 | radius=landmark_size, label=show_labels 89 | ) 90 | if det_overlay is not None: 91 | overlays.append(det_overlay) 92 | 93 | # Get base landmarks 94 | base_landmarks = MediapipeBaseLandmarks.get_base_landmarks( 95 | (width, height), x_scale=x_scale, y_translation=y_transform 96 | ) 97 | 98 | if show_target: 99 | base_df = pd.DataFrame({ 100 | 'x': base_landmarks[:, 0], 101 | 'y': base_landmarks[:, 1], 102 | 'z': np.zeros(len(base_landmarks)), 103 | 'index': range(len(base_landmarks)) 104 | }) 105 | target_overlay = self.image_processor.draw_landmarks( 106 | (width, height), base_df, 107 | transparency=0.4, color=(255, 0, 0), 108 | radius=landmark_size, label=show_labels 109 | ) 110 | if target_overlay is not None: 111 | overlays.append(target_overlay) 112 | 113 | for overlay in overlays: 114 | overlay = overlay.astype(np.float32) / 255.0 115 | alpha = overlay[:, :, 3:] 116 | rgb = overlay[:, :, :3] 117 | result_image = result_image * (1 - alpha) + rgb * alpha 118 | 119 | output_image = self.image_processor.numpy_to_tensor(result_image) 120 | output_mask = self.image_processor.convert_mask_to_tensor(mask_np) 121 | 122 | # Prepare and update landmarks data 123 | landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) 124 | updated_pipe = self._update_pipe(fp_pipe, landmarks_data) 125 | 126 | return output_image, updated_pipe, output_mask 127 | 128 | def _unwrap_mode(self, image_np, landmarks_df, width, height, device, 129 | x_scale, y_transform, fp_pipe, mask_np): 130 | # Get base landmarks 131 | base_landmarks = MediapipeBaseLandmarks.get_base_landmarks( 132 | (width, height), x_scale=x_scale, y_translation=y_transform 133 | ) 134 | source_landmarks = landmarks_df.iloc[:468][['x', 'y']].values.astype(np.float32) 135 | 136 | # Process image 137 | pil_image = self.image_processor.numpy_to_pil(image_np) 138 | warped_image = self._apply_warping(pil_image, source_landmarks, base_landmarks, device) 139 | 140 | # Process mask if provided 141 | warped_mask = None 142 | if mask_np is not None: 143 | pil_mask = self.image_processor.numpy_to_pil(mask_np) 144 | warped_mask = self._apply_warping(pil_mask, source_landmarks, base_landmarks, device) 145 | warped_mask = self.image_processor.convert_mask_to_tensor(np.array(warped_mask)) 146 | 147 | output_image = self.image_processor.pil_to_tensor(warped_image) 148 | 149 | # Prepare and update landmarks data 150 | landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks) 151 | updated_pipe = self._update_pipe(fp_pipe, landmarks_data) 152 | 153 | return output_image, updated_pipe, warped_mask 154 | 155 | def _wrap_mode(self, image_np, landmarks_df, width, height, device, 156 | x_scale, y_transform, fp_pipe, mask_np): 157 | if not fp_pipe or 'target_lm' not in fp_pipe: 158 | print("No landmarks found in face processor pipe") 159 | output_image = self.image_processor.numpy_to_tensor(image_np) 160 | output_mask = self.image_processor.convert_mask_to_tensor(mask_np) 161 | return output_image, {}, output_mask 162 | 163 | current_frame = fp_pipe.get("current_frame", 0) 164 | frame_key = f"frame_{current_frame}" 165 | 166 | # Check if we have required data 167 | if frame_key not in fp_pipe["frames"] or "detected_lm" not in fp_pipe["frames"][frame_key]: 168 | print(f"No detected landmarks found for frame {current_frame}") 169 | output_image = self.image_processor.numpy_to_tensor(image_np) 170 | output_mask = self.image_processor.convert_mask_to_tensor(mask_np) 171 | return output_image, fp_pipe, output_mask 172 | 173 | # Get source landmarks from target_lm at root level 174 | source_x = fp_pipe['target_lm']['x'] 175 | source_y = fp_pipe['target_lm']['y'] 176 | source_landmarks = np.column_stack((source_x, source_y))[:468] 177 | 178 | # Get target landmarks from detected_lm in current frame 179 | frame_data = fp_pipe["frames"][frame_key] 180 | detected_x = frame_data['detected_lm']['x'] 181 | detected_y = frame_data['detected_lm']['y'] 182 | target_landmarks = np.column_stack((detected_x, detected_y))[:468] 183 | 184 | # Process image 185 | pil_image = self.image_processor.numpy_to_pil(image_np) 186 | warped_image = self._apply_warping(pil_image, source_landmarks, target_landmarks, device) 187 | 188 | # Process mask if provided 189 | warped_mask = None 190 | if mask_np is not None: 191 | pil_mask = self.image_processor.numpy_to_pil(mask_np) 192 | warped_mask = self._apply_warping(pil_mask, source_landmarks, target_landmarks, device) 193 | warped_mask = self.image_processor.convert_mask_to_tensor(np.array(warped_mask)) 194 | 195 | output_image = self.image_processor.pil_to_tensor(warped_image) 196 | 197 | return output_image, fp_pipe, warped_mask 198 | 199 | def _apply_warping(self, image, source_landmarks, target_landmarks, device): 200 | """Apply warping to image using selected device""" 201 | if device == "CUDA" and torch.cuda.is_available(): 202 | return GPUDeformer.warp_face(image, source_landmarks, target_landmarks) 203 | else: 204 | return CPUDeformer.warp_face(image, source_landmarks, target_landmarks) 205 | 206 | def _prepare_landmarks_data(self, detected_df, target_lm): 207 | """ 208 | Prepare landmarks data in the correct format for fp_pipe 209 | 210 | Args: 211 | detected_df: DataFrame with detected landmarks 212 | target_lm: Target landmarks array 213 | 214 | Returns: 215 | dict: Landmarks data in the format compatible with fp_pipe 216 | """ 217 | # Prepare detected landmarks data 218 | detected_data = { 219 | 'x': detected_df['x'].tolist(), 220 | 'y': detected_df['y'].tolist(), 221 | 'indices': detected_df['index'].tolist() 222 | } 223 | 224 | # Prepare target landmarks data 225 | target_data = { 226 | 'x': target_lm[:, 0].tolist(), 227 | 'y': target_lm[:, 1].tolist(), 228 | 'indices': list(range(len(target_lm))) 229 | } 230 | 231 | return detected_data, target_data 232 | 233 | def _update_pipe(self, pipe, landmarks_data): 234 | """ 235 | Update fp_pipe with new landmarks data 236 | 237 | Args: 238 | pipe: Existing fp_pipe dictionary or None 239 | landmarks_data: Tuple of (detected_landmarks, target_landmarks) 240 | 241 | Returns: 242 | dict: Updated fp_pipe structure 243 | """ 244 | if pipe is None: 245 | pipe = { 246 | "workflow": "image", 247 | "current_frame": 0, 248 | "frames": {} 249 | } 250 | 251 | detected_lm, target_lm = landmarks_data 252 | 253 | # Update target landmarks at the root level 254 | pipe["target_lm"] = target_lm 255 | 256 | # Update detected landmarks for the current frame 257 | current_frame = pipe.get("current_frame", 0) 258 | frame_key = f"frame_{current_frame}" 259 | 260 | # Ensure the frame exists in the structure 261 | if frame_key not in pipe["frames"]: 262 | pipe["frames"][frame_key] = {} 263 | 264 | # Update detected landmarks for the current frame 265 | pipe["frames"][frame_key]["detected_lm"] = detected_lm 266 | 267 | return pipe -------------------------------------------------------------------------------- /nodes/image_feeder.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import List, Dict, Tuple 3 | 4 | import torch 5 | from PIL import Image 6 | 7 | from ..core.image_processor import ImageProcessor 8 | 9 | 10 | class ImageFeeder: 11 | """ComfyUI node for feeding images from a directory.""" 12 | 13 | def __init__(self): 14 | self.current_dir = None 15 | self.image_files = [] 16 | self.image_processor = ImageProcessor() 17 | 18 | @classmethod 19 | def INPUT_TYPES(cls): 20 | return { 21 | "required": { 22 | "directory": ("STRING", { 23 | "default": "", 24 | "multiline": False 25 | }), 26 | "frame_number": ("INT", { 27 | "default": 0, 28 | "min": 0, 29 | "step": 1 30 | }), 31 | } 32 | } 33 | 34 | RETURN_TYPES = ("IMAGE", "DICT") 35 | RETURN_NAMES = ("image", "image_sequence") 36 | FUNCTION = "feed_images" 37 | CATEGORY = "Face Processor/Image" 38 | 39 | def _load_image(self, image_path: str) -> torch.Tensor: 40 | """ 41 | Load an image from path and convert it to ComfyUI format 42 | 43 | Args: 44 | image_path: Path to the image file 45 | 46 | Returns: 47 | torch.Tensor: Image in ComfyUI format (B,H,W,C) normalized to 0-1 range 48 | """ 49 | try: 50 | image = Image.open(image_path) 51 | 52 | # Convert to RGB if necessary 53 | if image.mode != 'RGB': 54 | image = image.convert('RGB') 55 | 56 | # Use ImageProcessor to convert to tensor 57 | return self.image_processor.pil_to_tensor(image) 58 | 59 | except Exception as e: 60 | print(f"Error loading image {image_path}: {str(e)}") 61 | return None 62 | 63 | def _scan_directory(self, directory: str) -> List[str]: 64 | """ 65 | Scan directory for image files 66 | 67 | Args: 68 | directory: Path to directory to scan 69 | 70 | Returns: 71 | List of image file paths 72 | """ 73 | # List of supported image extensions 74 | valid_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.webp'} 75 | 76 | image_files = [] 77 | 78 | try: 79 | # Get absolute path 80 | abs_dir = os.path.abspath(directory) 81 | 82 | if not os.path.exists(abs_dir): 83 | print(f"Directory not found: {abs_dir}") 84 | return image_files 85 | 86 | # Scan directory for image files 87 | print(f"Scanning directory: {abs_dir}") 88 | 89 | for file in os.listdir(abs_dir): 90 | # Check file extension 91 | if os.path.splitext(file)[1].lower() in valid_extensions: 92 | full_path = os.path.join(abs_dir, file) 93 | image_files.append(full_path) 94 | 95 | # Sort files for consistent ordering 96 | image_files.sort() 97 | 98 | print(f"Found {len(image_files)} image files") 99 | 100 | except Exception as e: 101 | print(f"Error scanning directory {directory}: {str(e)}") 102 | 103 | return image_files 104 | 105 | def feed_images(self, directory: str, frame_number: int) -> Tuple[torch.Tensor, Dict]: 106 | """ 107 | Main processing function 108 | 109 | Args: 110 | directory: Path to images directory 111 | frame_number: Index of frame to return 112 | 113 | Returns: 114 | Tuple containing: 115 | - Selected frame as tensor 116 | - Dictionary with image_sequence 117 | """ 118 | # Check if directory changed 119 | if directory != self.current_dir: 120 | print(f"New directory detected, scanning: {directory}") 121 | self.image_files = self._scan_directory(directory) 122 | self.current_dir = directory 123 | 124 | # Prepare return data with frame to file mapping 125 | image_sequence = { 126 | "frames": { 127 | idx: file_path 128 | for idx, file_path in enumerate(self.image_files) 129 | }, 130 | "current_frame": frame_number 131 | } 132 | 133 | # Handle empty directory case 134 | if not self.image_files: 135 | print("No images found in directory") 136 | # Return empty image and data 137 | empty_image = torch.zeros((1, 64, 64, 3)) 138 | return empty_image, image_sequence 139 | 140 | # Validate frame number 141 | max_frame = len(self.image_files) - 1 142 | if frame_number > max_frame: 143 | print(f"Warning: Requested frame {frame_number} is out of range. Using last available frame ({max_frame})") 144 | frame_number = max_frame 145 | elif frame_number < 0: 146 | print(f"Warning: Requested frame {frame_number} is negative. Using first frame (0)") 147 | frame_number = 0 148 | 149 | # Load selected frame 150 | selected_frame = self._load_image(self.image_files[frame_number]) 151 | 152 | if selected_frame is None: 153 | print(f"Failed to load frame {frame_number}") 154 | empty_image = torch.zeros((1, 64, 64, 3)) 155 | return empty_image, image_sequence 156 | 157 | return selected_frame, image_sequence -------------------------------------------------------------------------------- /nodes/image_filters.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import cv2 4 | from torch import Tensor 5 | from ..core.image_processor import ImageProcessor 6 | 7 | 8 | class HighPassFilter: 9 | """ComfyUI node implementing AE-style high-pass filter with dynamic histogram visualization""" 10 | 11 | @classmethod 12 | def INPUT_TYPES(cls): 13 | return { 14 | "required": { 15 | "image": ("IMAGE",), 16 | "blur_radius": ("INT", { 17 | "default": 2, 18 | "min": 1, 19 | "max": 20, 20 | "step": 1 21 | }), 22 | "blur_iterations": ("INT", { 23 | "default": 6, 24 | "min": 1, 25 | "max": 20, 26 | "step": 1 27 | }), 28 | "blend_opacity": ("FLOAT", { 29 | "default": 0.5, 30 | "min": 0.45, 31 | "max": 0.55, 32 | "step": 0.01 33 | }), 34 | "input_black": ("INT", { 35 | "default": 117, 36 | "min": 0, 37 | "max": 255, 38 | "step": 1 39 | }), 40 | "input_white": ("INT", { 41 | "default": 137, 42 | "min": 0, 43 | "max": 255, 44 | "step": 1 45 | }), 46 | "gamma": ("FLOAT", { 47 | "default": 1.0, 48 | "min": 0.0, 49 | "max": 2.0, 50 | "step": 0.01 51 | }), 52 | "show_histogram": ("BOOLEAN", { 53 | "default": False 54 | }), 55 | } 56 | } 57 | 58 | RETURN_TYPES = ("IMAGE", "DICT") 59 | RETURN_NAMES = ("image", "HPF Settings") 60 | FUNCTION = "apply_hpf" 61 | CATEGORY = "Face Processor/Tools" 62 | 63 | def apply_hpf(self, image: torch.Tensor, blur_radius: int, blur_iterations: int, 64 | blend_opacity: float, input_black: int, input_white: int, 65 | gamma: float, show_histogram: bool) -> tuple[Tensor, dict]: 66 | """Main processing function implementing the High Pass Filter""" 67 | 68 | # Convert input tensor to numpy array 69 | np_img = ImageProcessor.tensor_to_numpy(image) 70 | 71 | # 1. Apply box blur with multiple iterations 72 | kernel_size = 2 * blur_radius + 1 73 | kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size ** 2) 74 | blurred = np_img.copy() 75 | for _ in range(blur_iterations): 76 | blurred = cv2.filter2D(blurred, -1, kernel, borderType=cv2.BORDER_REPLICATE) 77 | 78 | # 2. Invert blurred image 79 | inverted = cv2.bitwise_not(blurred) 80 | 81 | # 3. Blend original with inverted using specified opacity 82 | composite = cv2.addWeighted(np_img, 1 - blend_opacity, inverted, blend_opacity, 0) 83 | 84 | # 4. Apply levels adjustment 85 | in_black = np.clip(input_black, 0, 255) 86 | in_white = np.clip(input_white, in_black + 1, 255) 87 | levels = np.clip((composite.astype(np.float32) - in_black) / (in_white - in_black), 0, 1) 88 | levels = np.power(levels, 1.0 / gamma) if gamma > 0 else levels 89 | result = (levels * 255).astype(np.uint8) 90 | 91 | # Add histogram visualization if enabled 92 | if show_histogram: 93 | hist_img = ImageProcessor.draw_dynamic_histogram(composite, in_black, in_white, gamma) 94 | target_width = result.shape[1] 95 | hist_img = cv2.resize(hist_img, (target_width, 200)) 96 | result = np.vstack([result, hist_img]) 97 | 98 | # Convert back to tensor format 99 | result = ImageProcessor.numpy_to_tensor(result) 100 | 101 | # Prepare settings dictionary 102 | settings = { 103 | "high_pass_filter": { 104 | "blur_radius": blur_radius, 105 | "blur_iterations": blur_iterations, 106 | "blend_opacity": blend_opacity, 107 | "input_black": input_black, 108 | "input_white": input_white, 109 | "gamma": gamma, 110 | "show_histogram": False 111 | } 112 | } 113 | 114 | return result, settings -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Face Processor for ComfyUI 2 | 3 | A custom node collection for ComfyUI that provides advanced face detection, alignment, and transformation capabilities using MediaPipe Face Mesh. 4 | 5 | ## Features 6 | 7 | - **Face Detection & Landmark Extraction**: Uses MediaPipe Face Mesh to detect and extract 468 facial landmarks 8 | - **Face Alignment**: Automatic face alignment based on eye positions 9 | - **Face Transformation**: 10 | - Warping between source and target face landmarks 11 | - Scale and translation controls for face shape adjustment 12 | - CPU and CUDA-accelerated processing options 13 | - **Debug Visualization**: 14 | - Visual landmark overlay with customizable parameters 15 | - Support for both detected and target landmark visualization 16 | - Optional landmark labels 17 | 18 | ## Installation 19 | 20 | 1. Clone this repository into your ComfyUI's `custom_nodes` directory: 21 | ```bash 22 | cd ComfyUI/custom_nodes 23 | git clone https://github.com/SykkoAtHome/ComfyUI_FaceProcessor.git face_processor 24 | ``` 25 | 26 | 2. Install required dependencies: 27 | ```bash 28 | pip install mediapipe opencv-python numpy pandas pillow torch 29 | ``` 30 | 31 | For CUDA acceleration: 32 | 1. Install [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit) 33 | 2. Install [CuPy](https://docs.cupy.dev/en/stable/install.html): 34 | ```bash 35 | pip install cupy-cuda12x # Replace with your CUDA version 36 | ``` 37 | 38 | ## Nodes 39 | 40 | ### FaceWrapper 41 | Main node for face detection and transformation operations. 42 | 43 | #### Inputs: 44 | - `image`: Input image (ComfyUI IMAGE type) 45 | - `mode`: Operating mode 46 | - `Debug`: Visualization of detected landmarks 47 | - `Un-Wrap`: Transform face to normalized position 48 | - `Wrap`: Transform normalized face back to original position 49 | - `device`: Processing device (`CPU` or `CUDA`) 50 | - `show_detection`: Toggle detected landmarks visualization 51 | - `show_target`: Toggle target landmarks visualization 52 | - `landmark_size`: Size of landmark points in visualization 53 | - `show_labels`: Toggle landmark index labels 54 | - `x_scale`: Horizontal scaling factor (0.5 to 1.0) 55 | - `y_transform`: Vertical translation (-0.5 to 0.5) 56 | - `fp_pipe`: Optional settings dictionary 57 | 58 | #### Outputs: 59 | - `image`: Processed image 60 | - `fp_pipe`: Updated settings dictionary 61 | 62 | ### FaceFitAndRestore 63 | Node for face cropping and restoration operations. 64 | 65 | #### Inputs: 66 | - `mode`: Operating mode 67 | - `Fit`: Crop and align face 68 | - `Restore`: Place processed face back in original image 69 | - `image`: Input image 70 | - `padding_percent`: Additional padding around face (0.0 to 1.0) 71 | - `bbox_size`: Output size for cropped face (512, 1024, or 2048) 72 | - `fp_pipe`: Required for Restore mode 73 | 74 | #### Outputs: 75 | - `image`: Processed image 76 | - `fp_pipe`: Updated settings dictionary 77 | - `mask`: Mask indicating face region 78 | 79 | ## Technical Details 80 | 81 | ### Core Components 82 | 83 | #### Face Detection 84 | - Uses MediaPipe Face Mesh for robust face detection and landmark extraction 85 | - Provides 468 facial landmarks with 3D coordinates 86 | - Supports various input formats (PIL Image, numpy array, torch tensor) 87 | 88 | #### Image Processing 89 | - Automatic face rotation based on eye positions 90 | - Aspect ratio-preserving resizing 91 | - Support for square cropping with configurable padding 92 | - Boundary triangulation for complete face warping 93 | 94 | #### Face Warping 95 | - Triangle-based warping using predefined mesh topology from MediaPipe Face Mesh 96 | - CPU implementation using pure Python/NumPy 97 | - CUDA-accelerated GPU implementation using CuPy 98 | - Handles both forward and inverse warping 99 | 100 | ### Performance Considerations 101 | 102 | - GPU acceleration requires CUDA toolkit and CuPy 103 | - CPU fallback available for all operations 104 | - Progressive feedback during long operations 105 | - Memory-efficient processing for large images 106 | 107 | ## Example Usage 108 | 109 | Basic face detection and visualization: 110 | ```python 111 | face_wrapper = FaceWrapper() 112 | result_image, settings = face_wrapper.detect_face( 113 | image=input_image, 114 | mode="Debug", 115 | device="CPU", 116 | show_detection=True, 117 | show_target=False, 118 | landmark_size=4, 119 | show_labels=True, 120 | x_scale=1.0, 121 | y_transform=0.0 122 | ) 123 | ``` 124 | 125 | Face normalization workflow: 126 | 1. Detect and normalize face: 127 | ```python 128 | # Unwrap face to normalized position 129 | normalized_face, settings = face_wrapper.detect_face( 130 | image=input_image, 131 | mode="Un-Wrap", 132 | device="CUDA", 133 | x_scale=1.0, 134 | y_transform=0.0 135 | ) 136 | ``` 137 | 138 | 2. Process normalized face with your preferred method 139 | 140 | 3. Restore face to original position: 141 | 142 | ```python 143 | # Wrap processed face back 144 | final_image, _ = face_wrapper.detect_face( 145 | image=processed_face, 146 | mode="Wrap", 147 | device="CUDA", 148 | fp_pipe=settings 149 | ) 150 | ``` 151 | 152 | ## License 153 | 154 | MIT License 155 | 156 | Copyright (c) 2024 157 | 158 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 159 | 160 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 161 | 162 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 163 | 164 | ## Acknowledgments 165 | 166 | - MediaPipe Face Mesh for facial landmark detection 167 | - ComfyUI project for the node system framework -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | mediapipe>=0.10.0 2 | dlib>=19.24.0 3 | pandas>=2.0.0 4 | #cupy-cuda11x>=12.0.0 # for CUDA 11.x 5 | cupy-cuda12x>=12.0.0 # for CUDA 12.x (uncomment appropriate version) 6 | -------------------------------------------------------------------------------- /workflow/FaceProcessor_basic.json: -------------------------------------------------------------------------------- 1 | {"last_node_id":121,"last_link_id":204,"nodes":[{"id":7,"type":"CLIPTextEncode","pos":[413,389],"size":[425.27801513671875,180.6060791015625],"flags":{},"order":4,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":5,"localized_name":"clip"}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","links":[6],"slot_index":0,"localized_name":"CONDITIONING"}],"properties":{"Node name for S&R":"CLIPTextEncode"},"widgets_values":["text, watermark"]},{"id":6,"type":"CLIPTextEncode","pos":[415,186],"size":[422.84503173828125,164.31304931640625],"flags":{},"order":3,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":3,"localized_name":"clip"}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","links":[4],"slot_index":0,"localized_name":"CONDITIONING"}],"properties":{"Node name for S&R":"CLIPTextEncode"},"widgets_values":["beautiful scenery nature glass bottle landscape, , purple galaxy bottle,"]},{"id":5,"type":"EmptyLatentImage","pos":[473,609],"size":[315,106],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"LATENT","type":"LATENT","links":[2],"slot_index":0,"localized_name":"LATENT"}],"properties":{"Node name for S&R":"EmptyLatentImage"},"widgets_values":[512,512,1]},{"id":3,"type":"KSampler","pos":[863,186],"size":[315,262],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","type":"MODEL","link":1,"localized_name":"model"},{"name":"positive","type":"CONDITIONING","link":4,"localized_name":"positive"},{"name":"negative","type":"CONDITIONING","link":6,"localized_name":"negative"},{"name":"latent_image","type":"LATENT","link":2,"localized_name":"latent_image"}],"outputs":[{"name":"LATENT","type":"LATENT","links":[7],"slot_index":0,"localized_name":"LATENT"}],"properties":{"Node name for S&R":"KSampler"},"widgets_values":[437615434401321,"randomize",20,8,"euler","normal",1]},{"id":8,"type":"VAEDecode","pos":[1209,188],"size":[210,46],"flags":{},"order":9,"mode":0,"inputs":[{"name":"samples","type":"LATENT","link":7,"localized_name":"samples"},{"name":"vae","type":"VAE","link":8,"localized_name":"vae"}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[9],"slot_index":0,"localized_name":"IMAGE"}],"properties":{"Node name for S&R":"VAEDecode"},"widgets_values":[]},{"id":9,"type":"SaveImage","pos":[1451,189],"size":[210,58],"flags":{},"order":12,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":9,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"SaveImage"},"widgets_values":["ComfyUI"]},{"id":4,"type":"CheckpointLoaderSimple","pos":[26,474],"size":[315,98],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"MODEL","type":"MODEL","links":[1],"slot_index":0,"localized_name":"MODEL"},{"name":"CLIP","type":"CLIP","links":[3,5],"slot_index":1,"localized_name":"CLIP"},{"name":"VAE","type":"VAE","links":[8],"slot_index":2,"localized_name":"VAE"}],"properties":{"Node name for S&R":"CheckpointLoaderSimple"},"widgets_values":["v1-5-pruned-emaonly.ckpt"]},{"id":62,"type":"PreviewImage","pos":[1750.8646240234375,2628.38037109375],"size":[991.3868408203125,580.22119140625],"flags":{},"order":17,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":93,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":64,"type":"PreviewImage","pos":[55.24797058105469,2738.969970703125],"size":[457.8550109863281,481.67730712890625],"flags":{},"order":10,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":97,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[],"color":"#232","bgcolor":"#353"},{"id":118,"type":"PreviewImage","pos":[-702.2293090820312,2737.677001953125],"size":[457.8550109863281,481.67730712890625],"flags":{},"order":7,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":199,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":61,"type":"FaceFitAndRestore","pos":[1748.9951171875,2412.573486328125],"size":[468.5999755859375,166],"flags":{},"order":16,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":101,"localized_name":"image"},{"name":"processor_settings","type":"DICT","link":100,"shape":7,"localized_name":"processor_settings"}],"outputs":[{"name":"image","type":"IMAGE","links":[93,137],"slot_index":0,"localized_name":"image"},{"name":"processor_settings","type":"DICT","links":null,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","links":[148],"slot_index":2,"localized_name":"mask"},{"name":"bbox_size","type":"INT","links":null,"localized_name":"bbox_size"}],"properties":{"Node name for S&R":"FaceFitAndRestore"},"widgets_values":["Restore",0.15,"1024"],"color":"#432","bgcolor":"#653"},{"id":95,"type":"InvertMask","pos":[2284.236328125,2480.36181640625],"size":[210,26],"flags":{"collapsed":false},"order":18,"mode":0,"inputs":[{"name":"mask","type":"MASK","link":148,"localized_name":"mask"}],"outputs":[{"name":"MASK","type":"MASK","links":[150],"slot_index":0,"localized_name":"MASK"}],"properties":{"Node name for S&R":"InvertMask"},"widgets_values":[]},{"id":96,"type":"GrowMask","pos":[2541.130615234375,2481.22509765625],"size":[315,82],"flags":{"collapsed":false},"order":19,"mode":0,"inputs":[{"name":"mask","type":"MASK","link":150,"localized_name":"mask"}],"outputs":[{"name":"MASK","type":"MASK","links":[151],"slot_index":0,"localized_name":"MASK"}],"properties":{"Node name for S&R":"GrowMask"},"widgets_values":[3,true]},{"id":117,"type":"LoadImage","pos":[-1319.765625,2192.198486328125],"size":[315,314],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[197,198],"slot_index":0,"localized_name":"IMAGE"},{"name":"MASK","type":"MASK","links":null,"localized_name":"MASK"}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["1-F-2546-3-300x341.jpg","image"]},{"id":59,"type":"FaceFitAndRestore","pos":[-711.7913818359375,2426.533447265625],"size":[468.5999755859375,166],"flags":{},"order":5,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":197,"localized_name":"image"},{"name":"processor_settings","type":"DICT","link":null,"shape":7,"localized_name":"processor_settings"}],"outputs":[{"name":"image","type":"IMAGE","links":[170,199],"slot_index":0,"localized_name":"image"},{"name":"processor_settings","type":"DICT","links":[95],"slot_index":1,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","links":null,"localized_name":"mask"},{"name":"bbox_size","type":"INT","links":[],"slot_index":3,"localized_name":"bbox_size"}],"properties":{"Node name for S&R":"FaceFitAndRestore"},"widgets_values":["Fit",0.25,"2048"],"color":"#432","bgcolor":"#653"},{"id":91,"type":"ImageCompositeMasked","pos":[2671.207275390625,2187.55126953125],"size":[315,146],"flags":{"collapsed":false},"order":20,"mode":0,"inputs":[{"name":"destination","type":"IMAGE","link":137,"localized_name":"destination"},{"name":"source","type":"IMAGE","link":198,"localized_name":"source"},{"name":"mask","type":"MASK","link":151,"shape":7,"localized_name":"mask"}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[182],"slot_index":0,"localized_name":"IMAGE"}],"properties":{"Node name for S&R":"ImageCompositeMasked"},"widgets_values":[0,0,false],"color":"#223","bgcolor":"#335"},{"id":111,"type":"PreviewImage","pos":[3081.21044921875,2186.16748046875],"size":[991.3868408203125,580.22119140625],"flags":{},"order":21,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":182,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[],"color":"#223","bgcolor":"#335"},{"id":63,"type":"FaceWrapper","pos":[42.43104934692383,2422.59619140625],"size":[468.5999755859375,266],"flags":{},"order":8,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":170,"localized_name":"image"},{"name":"processor_settings","type":"DICT","link":95,"shape":7,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","link":null,"shape":7,"localized_name":"mask"}],"outputs":[{"name":"image","type":"IMAGE","links":[97,202],"slot_index":0,"localized_name":"image"},{"name":"processor_settings","type":"DICT","links":[99],"slot_index":1,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","links":null,"localized_name":"mask"}],"properties":{"Node name for S&R":"FaceWrapper"},"widgets_values":["Un-Wrap","CUDA",false,true,4,true,0.75,-0.03],"color":"#232","bgcolor":"#353"},{"id":120,"type":"Reroute","pos":[970.80712890625,2570.6123046875],"size":[75,26],"flags":{},"order":13,"mode":0,"inputs":[{"name":"","type":"*","link":203}],"outputs":[{"name":"","type":"IMAGE","links":[204],"slot_index":0}],"properties":{"showOutputText":false,"horizontal":false}},{"id":119,"type":"Reroute","pos":[603.6091918945312,2571.401611328125],"size":[75,26],"flags":{},"order":11,"mode":0,"inputs":[{"name":"","type":"*","link":202}],"outputs":[{"name":"","type":"IMAGE","links":[203],"slot_index":0}],"properties":{"showOutputText":false,"horizontal":false}},{"id":65,"type":"FaceWrapper","pos":[1119.014404296875,2424.967041015625],"size":[468.5999755859375,266],"flags":{},"order":14,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":204,"localized_name":"image"},{"name":"processor_settings","type":"DICT","link":99,"shape":7,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","link":null,"shape":7,"localized_name":"mask"}],"outputs":[{"name":"image","type":"IMAGE","links":[101,102],"slot_index":0,"localized_name":"image"},{"name":"processor_settings","type":"DICT","links":[100],"slot_index":1,"localized_name":"processor_settings"},{"name":"mask","type":"MASK","links":null,"localized_name":"mask"}],"properties":{"Node name for S&R":"FaceWrapper"},"widgets_values":["Wrap","CUDA",false,false,4,false,0.85,-0.03],"color":"#232","bgcolor":"#353"},{"id":66,"type":"PreviewImage","pos":[1129.1875,2753.408447265625],"size":[457.8550109863281,481.67730712890625],"flags":{},"order":15,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":102,"localized_name":"images"}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[],"color":"#232","bgcolor":"#353"}],"links":[[1,4,0,3,0,"MODEL"],[2,5,0,3,3,"LATENT"],[3,4,1,6,0,"CLIP"],[4,6,0,3,1,"CONDITIONING"],[5,4,1,7,0,"CLIP"],[6,7,0,3,2,"CONDITIONING"],[7,3,0,8,0,"LATENT"],[8,4,2,8,1,"VAE"],[9,8,0,9,0,"IMAGE"],[93,61,0,62,0,"IMAGE"],[95,59,1,63,1,"DICT"],[97,63,0,64,0,"IMAGE"],[99,63,1,65,1,"DICT"],[100,65,1,61,1,"DICT"],[101,65,0,61,0,"IMAGE"],[102,65,0,66,0,"IMAGE"],[137,61,0,91,0,"IMAGE"],[148,61,2,95,0,"MASK"],[150,95,0,96,0,"MASK"],[151,96,0,91,2,"MASK"],[170,59,0,63,0,"IMAGE"],[182,91,0,111,0,"IMAGE"],[197,117,0,59,0,"IMAGE"],[198,117,0,91,1,"IMAGE"],[199,59,0,118,0,"IMAGE"],[202,63,0,119,0,"*"],[203,119,0,120,0,"*"],[204,120,0,65,0,"IMAGE"]],"groups":[{"id":2,"title":"Do your Stuff Here!","bounding":[565.589599609375,2513.28759765625,506.4100036621094,836.4593505859375],"color":"#3f789e","font_size":22,"flags":{}}],"config":{},"extra":{"ds":{"scale":0.6303940863128641,"offset":[1641.7318147015405,-1791.7634118005708]},"node_versions":{"comfy-core":"0.3.12","ComfyUI_FaceProcessor":"817a44be8fe6074cd55650a503b68f307e15f0c4"},"ue_links":[]},"version":0.4} --------------------------------------------------------------------------------