diff --git a/core/face_detector.py b/core/face_detector.py index f7c5a9e..dd22771 100644 --- a/core/face_detector.py +++ b/core/face_detector.py @@ -9,22 +9,14 @@ from PIL import Image from pandas import DataFrame from core.lm_mapping import LandmarkMappings -from core.resources.model_loader import ModelDlib, ModelInsightFace +from core.resources.model_loader import ModelDlib, ModelInsightFace, ModelMediaPipe class FaceDetector: def __init__(self): """Initialize the FaceDetector with MediaPipe Face Mesh.""" - # 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 - ) - # Models will be initialized on demand + self.mediapipe_model = None self.dlib_model = None self.insight_model = None @@ -61,6 +53,11 @@ class FaceDetector: return mp_landmarks + def __del__(self): + """Clean up resources""" + if hasattr(self, 'mediapipe_model') and self.mediapipe_model: + self.mediapipe_model.face_mesh.close() + def detect_landmarks_mp(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]: """ @@ -72,29 +69,38 @@ class FaceDetector: 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: + try: + # Initialize MediaPipe model if not already done + if self.mediapipe_model is None: + self.mediapipe_model = ModelMediaPipe() + + image_np = self._convert_to_numpy(image) + if image_np is None: + return None + + results = self.mediapipe_model.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) + + except Exception as e: + print(f"Error in MediaPipe landmark detection: {str(e)}") 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. diff --git a/core/resources/model_loader.py b/core/resources/model_loader.py index dc5860e..792438f 100644 --- a/core/resources/model_loader.py +++ b/core/resources/model_loader.py @@ -4,8 +4,71 @@ from typing import Optional, Dict, List import dlib import numpy as np from insightface.app import FaceAnalysis +import mediapipe as mp + +class ModelDownload: + """ + placeholder for model download + """ + MODEL_URLS = { + "canonical_face_model": "https://github.com/google-ai-edge/mediapipe/blob/master/mediapipe/modules/face_geometry/data/canonical_face_model.obj", # Mediapipe 468 obj model + "shape_predictor_68_face_landmarks": "https://huggingface.co/spaces/asdasdasdasd/Face-forgery-detection/resolve/ccfc24642e0210d4d885bc7b3dbc9a68ed948ad6/shape_predictor_68_face_landmarks.dat" # dlib model + } + def download_model(self, model_name): + pass + + +class ModelMediaPipe: + """Class for loading and managing MediaPipe face mesh models using singleton pattern""" + + _instance: Optional['ModelMediaPipe'] = None + _face_mesh: Optional[mp.solutions.face_mesh.FaceMesh] = None + _model_loaded: bool = False + + def __new__(cls): + """Implement singleton pattern""" + if cls._instance is None: + cls._instance = super(ModelMediaPipe, cls).__new__(cls) + return cls._instance + + def __init__(self): + """Initialize MediaPipe model loader""" + if not self._model_loaded: + self._initialize_models() + + def _initialize_models(self) -> None: + """Initialize MediaPipe face mesh model""" + try: + print("Initializing MediaPipe Face Mesh...") + + self._face_mesh = mp.solutions.face_mesh.FaceMesh( + static_image_mode=True, + max_num_faces=1, + refine_landmarks=True, + min_detection_confidence=0.5 + ) + + self._model_loaded = True + print("MediaPipe Face Mesh initialized successfully") + + except Exception as e: + print(f"Error initializing MediaPipe models: {str(e)}") + raise + + @property + def face_mesh(self) -> mp.solutions.face_mesh.FaceMesh: + """Get MediaPipe Face Mesh instance""" + if not self._face_mesh: + self._initialize_models() + return self._face_mesh + + def __del__(self): + """Clean up resources""" + if self._face_mesh: + self._face_mesh.close() + class ModelOBJ: """Class for loading and parsing OBJ files with landmark transformation support""" @@ -24,7 +87,7 @@ class ModelOBJ: def _load_obj(self) -> None: """Load and parse MediaPipe face model OBJ file""" current_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) - obj_path = os.path.join(current_dir, 'resources', 'models', 'mediapipe_landmarks_468.obj') + obj_path = os.path.join(current_dir, 'resources', 'models', 'canonical_face_model.obj') print(f"Loading 3D model from: {obj_path}") @@ -189,7 +252,6 @@ class ModelDlib: self._initialize_models() return self._shape_predictor - class ModelInsightFace: """Class for loading and managing InsightFace models using singleton pattern""" diff --git a/core/resources/models/mediapipe_landmarks_468.obj b/core/resources/models/canonical_face_model.obj similarity index 100% rename from core/resources/models/mediapipe_landmarks_468.obj rename to core/resources/models/canonical_face_model.obj