diff --git a/core/face_detector.py b/core/face_detector.py index dd22771..baf6d7e 100644 --- a/core/face_detector.py +++ b/core/face_detector.py @@ -9,7 +9,7 @@ from PIL import Image from pandas import DataFrame from core.lm_mapping import LandmarkMappings -from core.resources.model_loader import ModelDlib, ModelInsightFace, ModelMediaPipe +from core.resources.model_loader import ModelDlib, ModelMediaPipe class FaceDetector: @@ -46,11 +46,6 @@ class FaceDetector: if dlib_landmarks is not None: return self.landmarks_interpolation(mp_landmarks, dlib_landmarks) - if refiner.lower() == 'insightface': - insight_landmarks = self.detect_landmarks_insight(image) - if insight_landmarks is not None: - return self.landmarks_interpolation(mp_landmarks, insight_landmarks) - return mp_landmarks def __del__(self): @@ -157,50 +152,6 @@ class FaceDetector: print(f"Error in dlib landmark detection: {str(e)}") return None - def detect_landmarks_insight(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: - """ - Detect facial landmarks using InsightFace. - - Args: - image: Input image in various formats - - Returns: - Optional[pd.DataFrame]: DataFrame containing landmark coordinates (x, y) and indices - """ - try: - # Initialize InsightFace model if not already done - if self.insight_model is None: - self.insight_model = ModelInsightFace() - - # Convert image to numpy format - image_np = self._convert_to_numpy(image) - if image_np is None: - return None - - # Detect faces - faces = self.insight_model.app.get(image_np) - if not faces: - print("No face detected by InsightFace") - return None - - # Get landmarks from first face - face = faces[0] - landmarks = face.landmark_2d_106 - - # Create DataFrame - landmarks_data = { - 'x': landmarks[:, 0], - 'y': landmarks[:, 1], - 'index': range(len(landmarks)) - } - - print(f"Detected {len(landmarks)} landmarks using InsightFace") - return pd.DataFrame(landmarks_data) - - except Exception as e: - print(f"Error in InsightFace landmark detection: {str(e)}") - return None - def landmarks_interpolation(self, mp_landmarks: pd.DataFrame, dlib_landmarks: pd.DataFrame) -> pd.DataFrame: """ diff --git a/core/resources/model_loader.py b/core/resources/model_loader.py index bef937d..9571238 100644 --- a/core/resources/model_loader.py +++ b/core/resources/model_loader.py @@ -1,13 +1,11 @@ import os +from typing import Optional, Dict, List + +import dlib +import mediapipe as mp import numpy as np import requests -import hashlib from tqdm import tqdm -from typing import Optional, Dict, List -import dlib -import insightface -from insightface.app import FaceAnalysis -import mediapipe as mp class ModelDownload: @@ -245,6 +243,78 @@ class ModelOBJ: return transformed + @property + def num_landmarks(self) -> int: + """ + Get the number of landmarks + + Returns: + int: Number of landmarks + """ + return self._num_landmarks + + def _initialize_landmarks(self) -> np.ndarray: + """ + Initialize landmarks array from UV coordinates + + Returns: + numpy.ndarray: Array of landmark coordinates in UV space + """ + landmarks = np.zeros((self._num_landmarks, 2), dtype=np.float32) + uvs = np.array(self._uvs, dtype=np.float32) + + for vertex_idx, uv_idx in self._vertex_to_uv.items(): + if vertex_idx < self._num_landmarks: + landmarks[vertex_idx] = uvs[uv_idx] + + return landmarks + + def get_faces(self) -> np.ndarray: + """ + Get face indices as numpy array + + Returns: + numpy.ndarray: Array of face indices + """ + return np.array(self._faces, dtype=np.int32) + + def get_transformed_landmarks(self, x_scale: float = 1.0, y_translation: float = 0.0) -> np.ndarray: + """ + Get landmarks with applied transformations + + Args: + x_scale: Horizontal scaling factor (0.5 to 1.0) + y_translation: Vertical translation (-0.5 to 0.5) + + Returns: + numpy.ndarray: Transformed landmarks + """ + if self._landmarks is None: + raise RuntimeError("Landmarks not initialized") + + transformed = self._landmarks.copy() + + # Calculate center point + center_x = (self._landmarks[:, 0].min() + self._landmarks[:, 0].max()) / 2 + + # Calculate horizontal distance from center (normalized to 0-1) + dx = np.abs(transformed[:, 0] - center_x) + + # Avoid division by zero + if np.abs(center_x) < 1e-6: # Use small epsilon value + influence = np.zeros_like(dx) + else: + influence = np.clip(dx / center_x, 0, 1) + + # Calculate scaled positions with horizontal influence + scale_factor = 1.0 + (x_scale - 1.0) * influence + transformed[:, 0] = center_x + (transformed[:, 0] - center_x) * scale_factor + + # Apply vertical translation + transformed[:, 1] += y_translation + + return transformed + class ModelDlib: """Class for loading and managing Dlib face detection models using singleton pattern""" @@ -301,65 +371,6 @@ class ModelDlib: return self._shape_predictor -class ModelInsightFace: - """Class for loading and managing InsightFace models using singleton pattern""" - - _instance: Optional['ModelInsightFace'] = None - _app: Optional[FaceAnalysis] = None - _model_loaded: bool = False - _models_dir = None - - def __new__(cls): - """Implement singleton pattern""" - if cls._instance is None: - cls._instance = super(ModelInsightFace, cls).__new__(cls) - return cls._instance - - def __init__(self): - """Initialize InsightFace model loader""" - if not self._model_loaded: - self._initialize_models() - - def _initialize_models(self) -> None: - """Initialize InsightFace models""" - try: - print("Initializing InsightFace detector...") - - # Set base models directory (without 'models' as InsightFace adds it) - current_dir = os.path.dirname(os.path.realpath(__file__)) - base_dir = os.path.dirname(os.path.dirname(current_dir)) # Go up to custom_nodes dir - self._models_dir = os.path.join(base_dir, "models", "insightface") - os.makedirs(self._models_dir, exist_ok=True) - - # Set model root directory for InsightFace - os.environ['INSIGHTFACE_HOME'] = self._models_dir - - # Initialize FaceAnalysis - self._app = FaceAnalysis( - allowed_modules=['detection', 'landmark_2d_106'], - providers=['CUDAExecutionProvider', 'CPUExecutionProvider'], - root=self._models_dir - ) - - # Prepare the model - self._app.prepare(ctx_id=0, det_size=(640, 640)) - - self._model_loaded = True - print("InsightFace models initialized successfully") - - except Exception as e: - print(f"Error initializing InsightFace models: {str(e)}") - raise - - - @property - def app(self) -> FaceAnalysis: - """Get InsightFace FaceAnalysis instance""" - if not self._app: - self._initialize_models() - return self._app - - class ModelMediaPipe: """Class for loading and managing MediaPipe face mesh models using singleton pattern""" diff --git a/nodes/face_wrapper.py b/nodes/face_wrapper.py index 551eff3..0771fb9 100644 --- a/nodes/face_wrapper.py +++ b/nodes/face_wrapper.py @@ -24,10 +24,10 @@ class FaceWrapper: "required": { "image": ("IMAGE",), "mode": (["Debug", "Un-Wrap", "Wrap"], {"default": "Debug"}), - "device": (["CPU", "CUDA"], {"default": "CPU"}), + "device": (["CPU", "CUDA"], {"default": "CUDA"}), "show_detection": ("BOOLEAN", {"default": False}), "show_target": ("BOOLEAN", {"default": False}), - "refiner": (["None", "Dlib", "InsightFace"], {"default": "None"}), + "refiner": (["None", "Dlib"], {"default": "None"}), "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}),