nodes/SykkoAtHome-ComfyUI_FaceProcessor
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

dlib implementation for landmark refinement

Browse Source 
with kernel thin_plate_spline
This commit is contained in:
Rafał Tłuchowski
2025-01-28 20:01:13 +01:00
parent 42031431fa
commit 48aaca39ea
8 changed files with 212 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
-5
View File
@@ -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__))
core/__init__.py
View File
core/cpu_deformer.py
+1 -1
View File
@@ -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:
"""
core/face_detector.py
+164 -15
View File
@@ -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."""
core/gpu_deformer.py
+1 -1
View File
@@ -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:
core/lm_mapping.py
+33
View File
@@ -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)
]
}
nodes/face_fit_and_restore.py
+3 -3
View File
@@ -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))
nodes/face_wrapper.py
+10 -8
View File
@@ -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