237 lines
8.7 KiB
Python
237 lines
8.7 KiB
Python
import mediapipe as mp
|
|
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 and Dlib."""
|
|
# MediaPipe initialization
|
|
self.mp_face_mesh = mp.solutions.face_mesh
|
|
self.face_mesh = self.mp_face_mesh.FaceMesh(
|
|
static_image_mode=True,
|
|
max_num_faces=1,
|
|
refine_landmarks=True,
|
|
min_detection_confidence=0.5
|
|
)
|
|
|
|
# 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 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
|
|
|
|
results = self.face_mesh.process(image_np)
|
|
if not results.multi_face_landmarks:
|
|
print("No face detected in the image")
|
|
return None
|
|
|
|
face_landmarks = results.multi_face_landmarks[0]
|
|
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
|
|
|
|
landmarks_data['x'].append(x)
|
|
landmarks_data['y'].append(y)
|
|
landmarks_data['index'].append(idx)
|
|
|
|
return pd.DataFrame(landmarks_data)
|
|
|
|
def detect_landmarks_dlib(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]:
|
|
"""
|
|
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
|
|
"""
|
|
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."""
|
|
if torch.is_tensor(image):
|
|
image = image.detach().cpu().numpy()
|
|
if len(image.shape) == 4:
|
|
image = image[0]
|
|
image = (image * 255).astype(np.uint8)
|
|
elif isinstance(image, Image.Image):
|
|
image = np.array(image)
|
|
elif isinstance(image, np.ndarray):
|
|
if image.dtype == np.float32 and image.max() <= 1.0:
|
|
image = (image * 255).astype(np.uint8)
|
|
if len(image.shape) == 4:
|
|
image = image[0]
|
|
|
|
if len(image.shape) == 2:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
|
elif image.shape[2] == 4:
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
|
|
|
|
return image
|
|
|
|
def __del__(self):
|
|
"""Clean up resources."""
|
|
if hasattr(self, 'face_mesh'):
|
|
self.face_mesh.close()
|