diff --git a/core/resources/model_loader.py b/core/resources/model_loader.py new file mode 100644 index 0000000..dc5860e --- /dev/null +++ b/core/resources/model_loader.py @@ -0,0 +1,237 @@ +import os +from typing import Optional, Dict, List + +import dlib +import numpy as np +from insightface.app import FaceAnalysis + + +class ModelOBJ: + """Class for loading and parsing OBJ files with landmark transformation support""" + + def __init__(self): + """Initialize ModelOBJ and load face model""" + self.vertices: List[List[float]] = [] + self.uvs: List[List[float]] = [] + self.faces: List[List[int]] = [] + self.vertex_to_uv: Dict[int, int] = {} + self.landmarks: Optional[np.ndarray] = None + self.num_landmarks: int = 468 + + # Load the model during initialization + self._load_obj() + + 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') + + print(f"Loading 3D model from: {obj_path}") + + if not os.path.exists(obj_path): + raise FileNotFoundError(f"OBJ file not found at: {obj_path}") + + try: + with open(obj_path, 'r') as f: + for line in f: + if line.startswith('v '): + # Parse vertex coordinates + parts = line.strip().split() + self.vertices.append([float(p) for p in parts[1:4]]) + + elif line.startswith('vt '): + # Parse UV coordinates (flip V coordinate) + parts = line.strip().split() + u = float(parts[1]) + v = 1.0 - float(parts[2]) # Flip V coordinate + self.uvs.append([u, v]) + + elif line.startswith('f '): + # Parse face indices + parts = line.strip().split()[1:] + face_verts = [] + for part in parts: + indices = part.split('/') + if len(indices) >= 2: + vertex_idx = int(indices[0]) - 1 # OBJ indices are 1-based + uv_idx = int(indices[1]) - 1 + self.vertex_to_uv[vertex_idx] = uv_idx + face_verts.append(vertex_idx) + if len(face_verts) == 3: # Only store triangular faces + self.faces.append(face_verts) + + # Initialize landmarks + self.landmarks = self._initialize_landmarks() + + print(f"Successfully loaded model data:") + print(f"- Vertices: {len(self.vertices)}") + print(f"- UVs: {len(self.uvs)}") + print(f"- Faces: {len(self.faces)}") + print(f"- Vertex to UV mappings: {len(self.vertex_to_uv)}") + + except Exception as e: + print(f"Error loading OBJ file: {str(e)}") + raise + + 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) + 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""" + + _instance: Optional['ModelDlib'] = None + _face_detector: Optional[dlib.fhog_object_detector] = None + _shape_predictor: Optional[dlib.shape_predictor] = None + _model_loaded: bool = False + + def __new__(cls): + """Implement singleton pattern""" + if cls._instance is None: + cls._instance = super(ModelDlib, cls).__new__(cls) + return cls._instance + + def __init__(self): + """Initialize Dlib model loader""" + if not self._model_loaded: + self._initialize_models() + + def _initialize_models(self) -> None: + """Initialize Dlib face detector and shape predictor""" + try: + # Initialize face detector + print("Initializing Dlib face detector...") + 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, "models", "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/models directory") + raise FileNotFoundError(f"Missing model file at: {model_path}") + + # Load shape predictor + print(f"Loading Dlib shape predictor from: {model_path}") + self._shape_predictor = dlib.shape_predictor(model_path) + self._model_loaded = True + print("Dlib models initialized successfully") + + except Exception as e: + print(f"Error initializing Dlib models: {str(e)}") + raise + + @property + def face_detector(self) -> dlib.fhog_object_detector: + """Get Dlib face detector instance""" + if not self._face_detector: + self._initialize_models() + return self._face_detector + + @property + def shape_predictor(self) -> dlib.shape_predictor: + """Get Dlib shape predictor instance""" + if not self._shape_predictor: + self._initialize_models() + 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 + + 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...") + + # Initialize FaceAnalysis with required modules + self._app = FaceAnalysis( + allowed_modules=['detection', 'landmark_2d_106'], + providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] + ) + + # 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 \ No newline at end of file