445 lines
16 KiB
Python
445 lines
16 KiB
Python
import os
|
|
from typing import Optional, Dict, List
|
|
|
|
import dlib
|
|
import mediapipe as mp
|
|
from mediapipe.tasks import python as mp_python
|
|
from mediapipe.tasks.python import vision
|
|
import numpy as np
|
|
import requests
|
|
from tqdm import tqdm
|
|
|
|
|
|
class ModelDownload:
|
|
"""Class for managing model downloads and verification"""
|
|
|
|
MODEL_URLS = {
|
|
# Mediapipe 468 landmark model
|
|
"canonical_face_model": "https://raw.githubusercontent.com/google-ai-edge/mediapipe/master/mediapipe/modules/face_geometry/data/canonical_face_model.obj",
|
|
|
|
# Dlib 68 landmark model
|
|
"shape_predictor_68_face_landmarks": "https://huggingface.co/spaces/asdasdasdasd/Face-forgery-detection/resolve/ccfc24642e0210d4d885bc7b3dbc9a68ed948ad6/shape_predictor_68_face_landmarks.dat",
|
|
# MediaPipe face landmarker .task model
|
|
"face_landmarker": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
|
|
}
|
|
|
|
def __init__(self):
|
|
"""Initialize model downloader"""
|
|
self.chunk_size = 8192
|
|
self.models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
|
|
|
|
def download_model(self, model_name: str, target_dir: Optional[str] = None) -> Optional[str]:
|
|
"""
|
|
Download model from predefined URL.
|
|
|
|
Args:
|
|
model_name: Name of the model to download
|
|
target_dir: Optional target directory (defaults to models/{model_name})
|
|
|
|
Returns:
|
|
str: Path to downloaded model or None if failed
|
|
"""
|
|
try:
|
|
if model_name not in self.MODEL_URLS:
|
|
print(f"Unknown model: {model_name}")
|
|
return None
|
|
|
|
# Set target directory
|
|
if target_dir is None:
|
|
target_dir = os.path.join(self.models_dir, os.path.splitext(model_name)[0])
|
|
|
|
# Create target directory if it doesn't exist
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
# Set target file path
|
|
url = self.MODEL_URLS[model_name]
|
|
if not url: # Skip if URL is empty
|
|
print(f"No URL defined for model: {model_name}")
|
|
return None
|
|
|
|
file_name = os.path.basename(url.split('?')[0]) # Remove URL parameters if any
|
|
target_path = os.path.join(target_dir, file_name)
|
|
|
|
# Check if file already exists
|
|
if os.path.exists(target_path):
|
|
print(f"Model {model_name} already exists at: {target_path}")
|
|
return target_path
|
|
|
|
print(f"Downloading {model_name} from: {url}")
|
|
print(f"Target path: {target_path}")
|
|
|
|
# Download file with progress bar
|
|
response = requests.get(url, stream=True)
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
|
|
with open(target_path, 'wb') as f:
|
|
with tqdm(total=total_size, unit='B', unit_scale=True, desc=model_name) as pbar:
|
|
for data in response.iter_content(chunk_size=self.chunk_size):
|
|
f.write(data)
|
|
pbar.update(len(data))
|
|
|
|
print(f"Successfully downloaded {model_name} to: {target_path}")
|
|
return target_path
|
|
|
|
except Exception as e:
|
|
print(f"Error downloading model {model_name}: {str(e)}")
|
|
return None
|
|
|
|
def get_model_path(self, model_name: str, target_dir: Optional[str] = None) -> Optional[str]:
|
|
"""
|
|
Get path to model, download if not exists.
|
|
|
|
Args:
|
|
model_name: Name of the model
|
|
target_dir: Optional target directory
|
|
|
|
Returns:
|
|
str: Path to model or None if failed
|
|
"""
|
|
if target_dir is None:
|
|
target_dir = os.path.join(self.models_dir, os.path.splitext(model_name)[0])
|
|
|
|
# Check if model exists
|
|
if os.path.exists(target_dir):
|
|
expected_file = os.path.join(target_dir, os.path.basename(self.MODEL_URLS[model_name].split('?')[0]))
|
|
if os.path.exists(expected_file):
|
|
return expected_file
|
|
|
|
# Download if not exists
|
|
return self.download_model(model_name, target_dir)
|
|
|
|
|
|
class ModelOBJ:
|
|
"""Class for loading and parsing OBJ files with landmark transformation support"""
|
|
|
|
_instance: Optional['ModelOBJ'] = None
|
|
_vertices: List[List[float]] = []
|
|
_uvs: List[List[float]] = []
|
|
_faces: List[List[int]] = []
|
|
_vertex_to_uv: Dict[int, int] = {}
|
|
_landmarks: Optional[np.ndarray] = None
|
|
_num_landmarks: int = 468
|
|
_model_loaded: bool = False
|
|
|
|
def __new__(cls):
|
|
"""Implement singleton pattern"""
|
|
if cls._instance is None:
|
|
cls._instance = super(ModelOBJ, cls).__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
"""Initialize ModelOBJ and load face model"""
|
|
if not self._model_loaded:
|
|
self.downloader = ModelDownload()
|
|
self._load_obj()
|
|
|
|
def _load_obj(self) -> None:
|
|
"""Load and parse MediaPipe face model OBJ file"""
|
|
try:
|
|
# Get model path, download if needed
|
|
model_path = self.downloader.get_model_path("canonical_face_model")
|
|
if not model_path:
|
|
raise FileNotFoundError("Failed to get canonical face model")
|
|
|
|
print(f"Loading 3D model from: {model_path}")
|
|
|
|
with open(model_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()
|
|
self._model_loaded = True
|
|
|
|
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)
|
|
|
|
# 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
|
|
|
|
@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"""
|
|
|
|
_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.downloader = ModelDownload()
|
|
self._initialize_models()
|
|
|
|
def _initialize_models(self) -> None:
|
|
"""Initialize Dlib face detector and shape predictor"""
|
|
try:
|
|
print("Initializing Dlib face detector...")
|
|
self._face_detector = dlib.get_frontal_face_detector()
|
|
|
|
# Get model path, download if needed
|
|
model_path = self.downloader.get_model_path("shape_predictor_68_face_landmarks")
|
|
if not model_path:
|
|
raise FileNotFoundError("Failed to get shape predictor model")
|
|
|
|
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 _FaceLandmarkerWrapper:
|
|
"""Wrapper to provide a FaceMesh-like interface for MediaPipe Tasks."""
|
|
|
|
def __init__(self, landmarker: vision.FaceLandmarker) -> None:
|
|
self._landmarker = landmarker
|
|
|
|
def process(self, image: np.ndarray):
|
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
|
return self._landmarker.detect(mp_image)
|
|
|
|
def close(self) -> None:
|
|
self._landmarker.close()
|
|
|
|
|
|
class ModelMediaPipe:
|
|
"""Class for loading and managing MediaPipe face mesh models using singleton pattern"""
|
|
|
|
_instance: Optional['ModelMediaPipe'] = None
|
|
_face_mesh: Optional[_FaceLandmarkerWrapper] = None
|
|
_model_loaded: bool = False
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super(ModelMediaPipe, cls).__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
if not self._model_loaded:
|
|
self._initialize_models()
|
|
|
|
def _initialize_models(self) -> None:
|
|
try:
|
|
print("Initializing MediaPipe Face Mesh...")
|
|
|
|
downloader = ModelDownload()
|
|
model_path = downloader.get_model_path("face_landmarker")
|
|
if not model_path:
|
|
raise FileNotFoundError("Failed to get face landmarker model")
|
|
|
|
base_options = mp_python.BaseOptions(model_asset_path=model_path)
|
|
options = vision.FaceLandmarkerOptions(
|
|
base_options=base_options,
|
|
num_faces=1,
|
|
output_face_blendshapes=False,
|
|
output_facial_transformation_matrixes=False,
|
|
)
|
|
landmarker = vision.FaceLandmarker.create_from_options(options)
|
|
self._face_mesh = _FaceLandmarkerWrapper(landmarker)
|
|
|
|
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) -> _FaceLandmarkerWrapper:
|
|
if not self._face_mesh:
|
|
self._initialize_models()
|
|
return self._face_mesh
|
|
|
|
def __del__(self):
|
|
if self._face_mesh:
|
|
self._face_mesh.close()
|
|
self._face_mesh = None
|
|
self._model_loaded = False
|