From 48aaca39ea4cef83c616652a85df148331e84e2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafa=C5=82=20T=C5=82uchowski?= Date: Tue, 28 Jan 2025 20:01:13 +0100 Subject: [PATCH] dlib implementation for landmark refinement with kernel thin_plate_spline --- __init__.py | 5 - core/__init__.py | 0 core/cpu_deformer.py | 2 +- core/face_detector.py | 179 +++++++++++++++++++++++++++++++--- core/gpu_deformer.py | 2 +- core/lm_mapping.py | 33 +++++++ nodes/face_fit_and_restore.py | 6 +- nodes/face_wrapper.py | 18 ++-- 8 files changed, 212 insertions(+), 33 deletions(-) create mode 100644 core/__init__.py create mode 100644 core/lm_mapping.py diff --git a/__init__.py b/__init__.py index 6c63b7f..aea3370 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,10 @@ -# __init__.py import os -import sys from .nodes.image_feeder import ImageFeeder from .nodes.image_filters import HighPassFilter from .nodes.face_wrapper import FaceWrapper from .nodes.face_fit_and_restore import FaceFitAndRestore -# Add current directory to system path -sys.path.append(os.path.dirname(os.path.realpath(__file__))) - # Get the path to the current directory NODE_PATH = os.path.dirname(os.path.realpath(__file__)) diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/core/cpu_deformer.py b/core/cpu_deformer.py index 0f0ecb1..ef8e79b 100644 --- a/core/cpu_deformer.py +++ b/core/cpu_deformer.py @@ -1,6 +1,6 @@ import numpy as np from PIL import Image -from ..core.base_mesh import MediapipeBaseLandmarks +from core.base_mesh import MediapipeBaseLandmarks class CPUDeformer: """ diff --git a/core/face_detector.py b/core/face_detector.py index f131ed7..7defb5c 100644 --- a/core/face_detector.py +++ b/core/face_detector.py @@ -3,13 +3,20 @@ import numpy as np import pandas as pd import cv2 import torch +import os from PIL import Image +from pandas import DataFrame from typing import Union, Optional +from core.lm_mapping import LandmarkMappings +import dlib class FaceDetector: + _model_loaded = False # Class variable to track if model was loaded + def __init__(self): - """Initialize the FaceDetector with MediaPipe Face Mesh.""" + """Initialize the FaceDetector with MediaPipe Face Mesh and Dlib.""" + # MediaPipe initialization self.mp_face_mesh = mp.solutions.face_mesh self.face_mesh = self.mp_face_mesh.FaceMesh( static_image_mode=True, @@ -18,15 +25,48 @@ class FaceDetector: min_detection_confidence=0.5 ) - def detect_landmarks(self, image: Union[torch.Tensor, np.ndarray, Image.Image], method: str = "mediapipe"): - if method == "mediapipe": - return self.detect_landmarks_mp(image) - else: - return self.detect_landmarks_dlib(image) - pass + # Dlib initialization + self.face_detector = dlib.get_frontal_face_detector() + + # Get path to the model file + current_dir = os.path.dirname(os.path.realpath(__file__)) + model_path = os.path.join(current_dir, "resources", "shape_predictor_68_face_landmarks.dat") + + if not os.path.exists(model_path): + print(f"Dlib landmarks model not found at: {model_path}") + print("Please ensure the model file is placed in the core/resources directory") + raise FileNotFoundError(f"Missing model file at: {model_path}") + + if not FaceDetector._model_loaded: + print(f"Loading dlib model from: {model_path}") + FaceDetector._model_loaded = True + + self.shape_predictor = dlib.shape_predictor(model_path) + + def detect_landmarks(self, image: Union[torch.Tensor, np.ndarray, Image.Image], + refine: bool = False) -> Union[DataFrame, None, str]: + mp_landmarks = self.detect_landmarks_mp(image) + if mp_landmarks is None: + return None + + if refine: + print("Detect with dlib") + dlib_landmarks = self.detect_landmarks_dlib(image) + if dlib_landmarks is not None: + return self.landmarks_interpolation(mp_landmarks, dlib_landmarks, LandmarkMappings.landmarks_pairs) + + return mp_landmarks def detect_landmarks_mp(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: - """Detect facial landmarks in the given image.""" + """ + Detect facial landmarks in the given image using MediaPipe. + + Args: + image: Input image in various formats + + Returns: + Optional[pd.DataFrame]: DataFrame containing landmark coordinates (x, y) and indices or None if no face detected + """ image_np = self._convert_to_numpy(image) if image_np is None: return None @@ -37,27 +77,136 @@ class FaceDetector: return None face_landmarks = results.multi_face_landmarks[0] - landmarks_data = {'x': [], 'y': [], 'z': [], 'index': []} + landmarks_data = {'x': [], 'y': [], 'index': []} image_height, image_width = image_np.shape[:2] for idx, landmark in enumerate(face_landmarks.landmark): x = landmark.x * image_width y = landmark.y * image_height - z = landmark.z landmarks_data['x'].append(x) landmarks_data['y'].append(y) - landmarks_data['z'].append(z) landmarks_data['index'].append(idx) return pd.DataFrame(landmarks_data) - def detect_landmarks_dlib(self, image: Union[torch.Tensor, np.ndarray, Image.Image]): + def detect_landmarks_dlib(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: """ - PlaceHolder - returns pd.dataframe with dlib landmarks + Detect facial landmarks using dlib's 68 point predictor. + + Args: + image: Input image in various formats + + Returns: + Optional[pd.DataFrame]: DataFrame containing landmark coordinates (x, y) and indices or None if no face detected """ - pass + try: + # Convert image to numpy format suitable for dlib + image_np = self._convert_to_numpy(image) + if image_np is None: + print("Failed to convert image for dlib processing") + return None + + # Convert to grayscale for dlib + gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY) + + # Detect faces + faces = self.face_detector(gray) + if not faces: + print("No face detected by dlib") + return None + + # Get first face + face = faces[0] + + # Detect landmarks + shape = self.shape_predictor(gray, face) + + # Create landmark data structure + landmarks_data = { + 'x': [], + 'y': [], + 'index': [] + } + + # Extract landmarks (adding 1 to index to match dlib's 1-based indexing) + for i in range(68): + point = shape.part(i) + landmarks_data['x'].append(float(point.x)) + landmarks_data['y'].append(float(point.y)) + landmarks_data['index'].append(i + 1) # Adding 1 to match dlib's indexing + + return pd.DataFrame(landmarks_data) + + except Exception as e: + print(f"Error in dlib landmark detection: {str(e)}") + return None + + def landmarks_interpolation(self, mp_landmarks: pd.DataFrame, dlib_landmarks: pd.DataFrame, + mapping: dict) -> pd.DataFrame: + """ + Interpolates all MediaPipe landmarks based on dlib reference points using RBF interpolation. + + Args: + mp_landmarks: DataFrame with columns [x, y, index] + dlib_landmarks: DataFrame with columns [x, y, index] + mapping: Dictionary of corresponding landmark pairs between MediaPipe and dlib + + Returns: + pd.DataFrame: DataFrame with interpolated landmark positions + """ + from scipy.interpolate import RBFInterpolator + # Validate input DataFrames + required_columns = ['x', 'y', 'index'] + if not all(col in mp_landmarks.columns for col in required_columns): + print("Error: MediaPipe landmarks missing required columns") + return mp_landmarks + + if not all(col in dlib_landmarks.columns for col in required_columns): + print("Error: Dlib landmarks missing required columns") + return mp_landmarks + + try: + # Collect control points + control_src = [] # MediaPipe points + control_dst = [] # Dlib points + + for feature, pairs in mapping.items(): + for mp_idx, dlib_idx in pairs: + mp_point = mp_landmarks[mp_landmarks['index'] == mp_idx].iloc[0] + dlib_point = dlib_landmarks[dlib_landmarks['index'] == dlib_idx].iloc[0] + + control_src.append([mp_point['x'], mp_point['y']]) + control_dst.append([dlib_point['x'], dlib_point['y']]) + + control_src = np.array(control_src) + control_dst = np.array(control_dst) + + print(f"Using {len(control_src)} control points for RBF interpolation") + + # Create and fit RBF interpolation for both x and y coordinates + kernel = 'thin_plate_spline' + rbf_x = RBFInterpolator(control_src, control_dst[:, 0], kernel=kernel) + rbf_y = RBFInterpolator(control_src, control_dst[:, 1], kernel=kernel) + + # Transform all MediaPipe points + all_points = mp_landmarks[['x', 'y']].values + + # Apply transformation + transformed_x = rbf_x(all_points) + transformed_y = rbf_y(all_points) + + # Create result DataFrame + result = mp_landmarks.copy() + result['x'] = transformed_x + result['y'] = transformed_y + + print("Successfully interpolated all landmarks using RBF") + return result + + except Exception as e: + print(f"Error during landmark interpolation: {str(e)}") + return mp_landmarks def _convert_to_numpy(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[np.ndarray]: """Convert different image types to numpy array.""" diff --git a/core/gpu_deformer.py b/core/gpu_deformer.py index c2c71b2..627a684 100644 --- a/core/gpu_deformer.py +++ b/core/gpu_deformer.py @@ -1,7 +1,7 @@ import numpy as np from PIL import Image import torch -from ..core.base_mesh import MediapipeBaseLandmarks +from core.base_mesh import MediapipeBaseLandmarks class GPUDeformer: diff --git a/core/lm_mapping.py b/core/lm_mapping.py new file mode 100644 index 0000000..f8937fe --- /dev/null +++ b/core/lm_mapping.py @@ -0,0 +1,33 @@ +class LandmarkMappings: + # Each tuple represents (mediapipe_point_id, dlib_point_id) + landmarks_pairs = { + "left_eye": [ + (133, 40), + (33, 37) + ], + "right_eye": [ + (362, 43), + (263, 46) + ], + "mouth": [ + (0, 52), + (17, 58), + (61, 49), + (291, 55), + (13, 63), + (14, 67), + (78, 61), + (308, 65) + ], + "node_edge": [ + (6, 28), + (4, 31) + ], + "oval": [ + (152, 9), + (127, 1), + (356, 17), + (172, 5), + (397, 13) + ] + } diff --git a/nodes/face_fit_and_restore.py b/nodes/face_fit_and_restore.py index c9bc963..9320de0 100644 --- a/nodes/face_fit_and_restore.py +++ b/nodes/face_fit_and_restore.py @@ -4,7 +4,7 @@ import cv2 import pandas as pd from typing import Union, Optional, Tuple -from ..core.face_detector import FaceDetector +from core.face_detector import FaceDetector class FaceFitAndRestore: @@ -62,8 +62,8 @@ class FaceFitAndRestore: if image_np is None: return (image, {}, self._create_empty_mask(image), int(bbox_size)) - # Detect facial landmarks - landmarks_df = self.face_detector.detect_landmarks(image_np, "mediapipe") + # Detect facial landmarks using mediapipe + landmarks_df = self.face_detector.detect_landmarks_mp(image_np) if landmarks_df is None: print("No face detected, returning original image") return (image, {}, self._create_empty_mask(image), int(bbox_size)) diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 38363d3..9a04683 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -4,11 +4,11 @@ import pandas as pd from PIL import Image import cv2 -from ..core.face_detector import FaceDetector -from ..core.image_processor import ImageProcessor -from ..core.base_mesh import MediapipeBaseLandmarks -from ..core.cpu_deformer import CPUDeformer -from ..core.gpu_deformer import GPUDeformer +from core.face_detector import FaceDetector +from core.image_processor import ImageProcessor +from core.base_mesh import MediapipeBaseLandmarks +from core.cpu_deformer import CPUDeformer +from core.gpu_deformer import GPUDeformer class FaceWrapper: @@ -27,6 +27,7 @@ class FaceWrapper: "device": (["CPU", "CUDA"], {"default": "CPU"}), "show_detection": ("BOOLEAN", {"default": False}), "show_target": ("BOOLEAN", {"default": False}), + "dlib_refine": ("BOOLEAN", {"default": False}), "landmark_size": ("INT", {"default": 4, "min": 1, "max": 10, "step": 1}), "show_labels": ("BOOLEAN", {"default": False}), "x_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 1.0, "step": 0.01}), @@ -43,7 +44,7 @@ class FaceWrapper: FUNCTION = "detect_face" CATEGORY = "Face Processor" - def detect_face(self, image, mode, device, show_detection, show_target, landmark_size, + def detect_face(self, image, mode, device, show_detection, show_target, dlib_refine, landmark_size, show_labels, x_scale, y_transform, processor_settings=None, mask=None): # Convert input image to numpy with proper RGB format image_np = self._convert_to_numpy(image) @@ -58,8 +59,9 @@ class FaceWrapper: return self._wrap_mode(image_np, None, width, height, device, x_scale, y_transform, processor_settings, mask_np) - # Detect facial landmarks - landmarks_df = self.face_detector.detect_landmarks(image_np, "mediapipe") + # Detect facial landmarks using mediapipe + landmarks_df = self.face_detector.detect_landmarks(image_np, refine=bool(dlib_refine)) + if landmarks_df is None: print("No face detected") empty_mask = torch.zeros((1, height, width), dtype=torch.float32) if mask is not None else None