real_finals / roop /face_analyser.py
VanNguyen1214's picture
Upload 58 files
8dff9a2 verified
import threading
from typing import Any
import insightface
import numpy as np
from PIL import Image
import roop.globals
from roop.typing import Frame
FACE_ANALYSER = None
THREAD_LOCK = threading.Lock()
def get_face_analyser() -> Any:
global FACE_ANALYSER
with THREAD_LOCK:
if FACE_ANALYSER is None:
FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=roop.globals.execution_providers)
FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))
return FACE_ANALYSER
def get_precise_face_mask(frame: Frame) -> Any:
"""
Get precise face mask using advanced segmentation (same as detect_face_and_forehead_no_hair).
Returns both InsightFace detection and precise mask.
"""
try:
# Import the precise detection function
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from segmentation import detect_face_and_forehead_no_hair
# Convert frame to PIL Image
if isinstance(frame, np.ndarray):
pil_image = Image.fromarray(frame)
else:
pil_image = frame
# Get precise face mask (clean skin only)
precise_mask = detect_face_and_forehead_no_hair(pil_image)
# Also get InsightFace detection for face swapping compatibility
insightface_faces = get_face_analyser().get(frame)
return {
'precise_mask': precise_mask,
'insightface_faces': insightface_faces,
'has_face': precise_mask.sum() > 0 and len(insightface_faces) > 0
}
except Exception as e:
print(f"Precise face detection failed: {e}")
# Fallback to regular InsightFace
insightface_faces = get_face_analyser().get(frame)
return {
'precise_mask': None,
'insightface_faces': insightface_faces,
'has_face': len(insightface_faces) > 0
}
def get_one_face(frame: Frame) -> Any:
"""
Get one face with enhanced precision detection.
"""
# Get precise detection info
face_info = get_precise_face_mask(frame)
if face_info['has_face'] and face_info['insightface_faces']:
try:
# Select face (leftmost) for compatibility
selected_face = min(face_info['insightface_faces'], key=lambda x: x.bbox[0])
# Add precise mask info to face object
if face_info['precise_mask'] is not None:
selected_face.precise_mask = face_info['precise_mask']
print(f"✅ Enhanced face detection: {face_info['precise_mask'].sum()} precise pixels")
return selected_face
except (ValueError, IndexError):
return None
# Fallback to original method
face = get_face_analyser().get(frame)
try:
selected_face = min(face, key=lambda x: x.bbox[0])
return selected_face
except ValueError:
return None
def get_many_faces(frame: Frame) -> Any:
"""
Get many faces with enhanced precision detection.
"""
# Get precise detection info
face_info = get_precise_face_mask(frame)
if face_info['has_face'] and face_info['insightface_faces']:
faces = face_info['insightface_faces']
# Add precise mask info to all face objects
if face_info['precise_mask'] is not None:
for face in faces:
face.precise_mask = face_info['precise_mask']
print(f"✅ Enhanced multi-face detection: {len(faces)} faces with precise masks")
return faces
# Fallback to original method
try:
return get_face_analyser().get(frame)
except IndexError:
return None
def has_precise_face_mask(face_obj) -> bool:
"""
Check if face object has precise mask attached.
"""
return hasattr(face_obj, 'precise_mask') and face_obj.precise_mask is not None