Facefusion_Master / face_selector.py
andro1241's picture
Upload 55 files
773f76d verified
Raw
History Blame Contribute Delete
5.72 kB
from typing import List
import numpy
import facefusion.choices
from facefusion import state_manager
from facefusion.common_helper import get_first, get_middle
from facefusion.face_creator import get_one_face, get_static_faces
from facefusion.face_tracker import track_faces
from facefusion.types import Face, FaceSelectorOrder, Gender, Race, Score, VisionFrame
def select_faces(reference_vision_frame : VisionFrame, source_vision_frames : List[VisionFrame], target_vision_frames : List[VisionFrame]) -> List[Face]:
source_faces = get_static_faces(source_vision_frames)
if state_manager.get_item('face_tracker_score') > 0:
target_faces = track_faces(target_vision_frames, state_manager.get_item('face_tracker_score'))
else:
target_faces = get_static_faces([ get_middle(target_vision_frames) ])
if state_manager.get_item('face_selector_mode') == 'many':
return sort_and_filter_faces(source_faces, target_faces)
if state_manager.get_item('face_selector_mode') == 'one':
target_face = get_one_face(sort_and_filter_faces(source_faces, target_faces))
if target_face:
return [ target_face ]
if state_manager.get_item('face_selector_mode') == 'reference':
reference_faces = get_static_faces([ reference_vision_frame ])
reference_faces = sort_and_filter_faces(source_faces, reference_faces)
reference_face = get_one_face(reference_faces, state_manager.get_item('reference_face_position'))
if reference_face:
match_faces = find_match_faces([ reference_face ], target_faces, state_manager.get_item('reference_face_distance'))
return match_faces
return []
def find_match_faces(reference_faces : List[Face], target_faces : List[Face], face_distance : float) -> List[Face]:
match_faces : List[Face] = []
for reference_face in reference_faces:
if reference_face:
for index, target_face in enumerate(target_faces):
if compare_faces(target_face, reference_face, face_distance):
match_faces.append(target_faces[index])
return match_faces
def compare_faces(face : Face, reference_face : Face, face_distance : float) -> bool:
current_face_distance = calculate_face_distance(face, reference_face)
current_face_distance = float(numpy.interp(current_face_distance, [ 0, 2 ], [ 0, 1 ]))
return current_face_distance < face_distance
def calculate_face_distance(face : Face, reference_face : Face) -> float:
if hasattr(face, 'embedding_norm') and hasattr(reference_face, 'embedding_norm'):
return 1 - numpy.dot(face.embedding_norm, reference_face.embedding_norm)
return 0
def sort_and_filter_faces(source_faces : List[Face], target_faces : List[Face]) -> List[Face]:
if target_faces:
if state_manager.get_item('face_selector_order'):
target_faces = sort_faces_by_order(target_faces, state_manager.get_item('face_selector_order'))
face_selector_gender = state_manager.get_item('face_selector_gender')
face_selector_race = state_manager.get_item('face_selector_race')
if source_faces and face_selector_gender == 'auto' or face_selector_race == 'auto':
source_face = get_first(sort_faces_by_order(source_faces, 'large-small'))
if source_face:
if face_selector_gender == 'auto':
face_selector_gender = source_face.gender
if face_selector_race == 'auto':
face_selector_race = source_face.race
if face_selector_gender in facefusion.choices.genders:
target_faces = filter_faces_by_gender(target_faces, face_selector_gender)
if face_selector_race in facefusion.choices.races:
target_faces = filter_faces_by_race(target_faces, face_selector_race)
if state_manager.get_item('face_selector_age_start') or state_manager.get_item('face_selector_age_end'):
target_faces = filter_faces_by_age(target_faces, state_manager.get_item('face_selector_age_start'), state_manager.get_item('face_selector_age_end'))
return target_faces
def sort_faces_by_order(faces : List[Face], order : FaceSelectorOrder) -> List[Face]:
if order == 'left-right':
return sorted(faces, key = get_bounding_box_left)
if order == 'right-left':
return sorted(faces, key = get_bounding_box_left, reverse = True)
if order == 'top-bottom':
return sorted(faces, key = get_bounding_box_top)
if order == 'bottom-top':
return sorted(faces, key = get_bounding_box_top, reverse = True)
if order == 'small-large':
return sorted(faces, key = get_bounding_box_area)
if order == 'large-small':
return sorted(faces, key = get_bounding_box_area, reverse = True)
if order == 'best-worst':
return sorted(faces, key = get_face_detector_score, reverse = True)
if order == 'worst-best':
return sorted(faces, key = get_face_detector_score)
return faces
def get_bounding_box_left(face : Face) -> float:
return face.bounding_box[0]
def get_bounding_box_top(face : Face) -> float:
return face.bounding_box[1]
def get_bounding_box_area(face : Face) -> float:
return (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1])
def get_face_detector_score(face : Face) -> Score:
return face.score_set.get('detector')
def filter_faces_by_gender(faces : List[Face], gender : Gender) -> List[Face]:
filter_faces = []
for face in faces:
if face.gender == gender:
filter_faces.append(face)
return filter_faces
def filter_faces_by_age(faces : List[Face], face_selector_age_start : int, face_selector_age_end : int) -> List[Face]:
filter_faces = []
age = range(face_selector_age_start, face_selector_age_end)
for face in faces:
if set(face.age) & set(age):
filter_faces.append(face)
return filter_faces
def filter_faces_by_race(faces : List[Face], race : Race) -> List[Face]:
filter_faces = []
for face in faces:
if face.race == race:
filter_faces.append(face)
return filter_faces