diff --git a/facefusion.ini b/facefusion.ini index 17f3a2e2..a336b0c6 100644 --- a/facefusion.ini +++ b/facefusion.ini @@ -35,6 +35,9 @@ reference_face_position = reference_face_distance = reference_frame_number = +[face_tracker] +face_tracker_score = + [face_masker] face_occluder_model = face_parser_model = @@ -52,6 +55,9 @@ trim_frame_start = trim_frame_end = temp_frame_format = +[frame_distribution] +target_frame_amount = + [output_creation] output_image_quality = output_image_scale = diff --git a/facefusion/args_helper.py b/facefusion/args_helper.py index f89055aa..963050f0 100644 --- a/facefusion/args_helper.py +++ b/facefusion/args_helper.py @@ -36,6 +36,7 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('reference_face_position', args.get('reference_face_position')) apply_state_item('reference_face_distance', args.get('reference_face_distance')) apply_state_item('reference_frame_number', args.get('reference_frame_number')) + apply_state_item('face_tracker_score', args.get('face_tracker_score')) apply_state_item('face_occluder_model', args.get('face_occluder_model')) apply_state_item('face_parser_model', args.get('face_parser_model')) apply_state_item('face_mask_types', args.get('face_mask_types')) @@ -47,6 +48,7 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('trim_frame_start', args.get('trim_frame_start')) apply_state_item('trim_frame_end', args.get('trim_frame_end')) apply_state_item('temp_frame_format', args.get('temp_frame_format')) + apply_state_item('target_frame_amount', args.get('target_frame_amount')) apply_state_item('output_image_quality', args.get('output_image_quality')) apply_state_item('output_image_scale', args.get('output_image_scale')) apply_state_item('output_audio_encoder', args.get('output_audio_encoder')) diff --git a/facefusion/choices.py b/facefusion/choices.py index bde7dcbc..8eb25318 100755 --- a/facefusion/choices.py +++ b/facefusion/choices.py @@ -2,7 +2,7 @@ import logging from typing import List, Sequence, get_args from facefusion.common_helper import create_float_range, create_int_range -from facefusion.types import Angle, ApiSecurityStrategy, AudioEncoder, AudioFormat, AudioSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorMode, FaceSelectorOrder, Gender, ImageEncoder, ImageFormat, ImageSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoSet, VoiceExtractorModel, WorkFlow +from facefusion.types import Angle, ApiSecurityStrategy, AudioEncoder, AudioFormat, AudioSet, BenchmarkMode, BenchmarkResolution, BenchmarkSet, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskArea, FaceMaskAreaSet, FaceMaskRegion, FaceMaskRegionSet, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, FaceSelectorRace, Gender, ImageEncoder, ImageFormat, ImageSet, JobStatus, LogLevel, LogLevelSet, Race, Score, TempFrameFormat, VideoEncoder, VideoFormat, VideoMemoryStrategy, VideoPreset, VideoSet, VoiceExtractorModel, WorkFlow face_detector_set : FaceDetectorSet =\ { @@ -16,8 +16,10 @@ face_detector_models : List[FaceDetectorModel] = list(get_args(FaceDetectorModel face_landmarker_models : List[FaceLandmarkerModel] = list(get_args(FaceLandmarkerModel)) face_selector_modes : List[FaceSelectorMode] = list(get_args(FaceSelectorMode)) face_selector_orders : List[FaceSelectorOrder] = list(get_args(FaceSelectorOrder)) -face_selector_genders : List[Gender] = list(get_args(Gender)) -face_selector_races : List[Race] = list(get_args(Race)) +genders : List[Gender] = list(get_args(Gender)) +races : List[Race] = list(get_args(Race)) +face_selector_genders : List[FaceSelectorGender] = list(get_args(FaceSelectorGender)) +face_selector_races : List[FaceSelectorRace] = list(get_args(FaceSelectorRace)) face_occluder_models : List[FaceOccluderModel] = list(get_args(FaceOccluderModel)) face_parser_models : List[FaceParserModel] = list(get_args(FaceParserModel)) face_mask_types : List[FaceMaskType] = list(get_args(FaceMaskType)) @@ -159,6 +161,8 @@ face_mask_blur_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) face_mask_padding_range : Sequence[int] = create_int_range(0, 100, 1) face_selector_age_range : Sequence[int] = create_int_range(0, 100, 1) reference_face_distance_range : Sequence[float] = create_float_range(0.0, 1.0, 0.05) +face_tracker_score_range : Sequence[Score] = create_float_range(0.0, 0.5, 0.05) +target_frame_amount_range : Sequence[int] = create_int_range(0, 10, 1) output_image_quality_range : Sequence[int] = create_int_range(0, 100, 1) output_image_scale_range : Sequence[float] = create_float_range(0.25, 8.0, 0.25) output_audio_quality_range : Sequence[int] = create_int_range(0, 100, 1) diff --git a/facefusion/common_helper.py b/facefusion/common_helper.py index 3587784d..05e4c4e9 100644 --- a/facefusion/common_helper.py +++ b/facefusion/common_helper.py @@ -78,6 +78,12 @@ def get_first(__list__ : Any) -> Any: return None +def get_middle(__list__ : Any) -> Any: + if isinstance(__list__, Sequence) and __list__: + return __list__[len(__list__) // 2] + return None + + def get_last(__list__ : Any) -> Any: if isinstance(__list__, Reversible): return next(reversed(__list__), None) diff --git a/facefusion/face_analyser.py b/facefusion/face_creator.py similarity index 70% rename from facefusion/face_analyser.py rename to facefusion/face_creator.py index 7e264ccd..fb4a8a81 100644 --- a/facefusion/face_analyser.py +++ b/facefusion/face_creator.py @@ -3,10 +3,10 @@ from typing import List, Optional import numpy from facefusion import face_store, state_manager -from facefusion.common_helper import get_first +from facefusion.common_helper import get_first, get_middle from facefusion.face_classifier import classify_face from facefusion.face_detector import detect_faces, detect_faces_by_angle -from facefusion.face_helper import apply_nms, convert_to_face_landmark_5, estimate_face_angle, get_nms_threshold +from facefusion.face_helper import apply_nms, average_points, convert_to_face_landmark_5, estimate_face_angle, get_nms_threshold from facefusion.face_landmarker import detect_face_landmark, estimate_face_landmark_68_5 from facefusion.face_recognizer import calculate_face_embedding from facefusion.types import BoundingBox, Face, FaceLandmark5, FaceLandmarkSet, FaceScoreSet, Score, VisionFrame @@ -46,7 +46,9 @@ def create_faces(vision_frame : VisionFrame, bounding_boxes : List[BoundingBox], } face_embedding, face_embedding_norm = calculate_face_embedding(vision_frame, face_landmark_set.get('5/68')) gender, age, race = classify_face(vision_frame, face_landmark_set.get('5/68')) + faces.append(Face( + origin = 'detect', bounding_box = bounding_box, score_set = face_score_set, landmark_set = face_landmark_set, @@ -67,48 +69,6 @@ def get_one_face(faces : List[Face], position : int = 0) -> Optional[Face]: return None -def get_average_face(faces : List[Face]) -> Optional[Face]: - face_embeddings = [] - face_embeddings_norm = [] - - if faces: - first_face = get_first(faces) - - for face in faces: - face_embeddings.append(face.embedding) - face_embeddings_norm.append(face.embedding_norm) - - return Face( - bounding_box = first_face.bounding_box, - score_set = first_face.score_set, - landmark_set = first_face.landmark_set, - angle = first_face.angle, - embedding = numpy.mean(face_embeddings, axis = 0), - embedding_norm = numpy.mean(face_embeddings_norm, axis = 0), - gender = first_face.gender, - age = first_face.age, - race = first_face.race - ) - return None - - -def get_static_faces(vision_frames : List[VisionFrame]) -> List[Face]: - many_faces : List[Face] = [] - - for vision_frame in vision_frames: - faces = face_store.get_faces(vision_frame) - - if not faces: - faces = get_many_faces([ vision_frame ]) - - if faces: - face_store.set_faces(vision_frame, faces) - - many_faces.extend(faces) - - return many_faces - - def get_many_faces(vision_frames : List[VisionFrame]) -> List[Face]: many_faces : List[Face] = [] @@ -136,6 +96,100 @@ def get_many_faces(vision_frames : List[VisionFrame]) -> List[Face]: return many_faces +def get_static_faces(vision_frames : List[VisionFrame]) -> List[Face]: + many_faces : List[Face] = [] + + for vision_frame in vision_frames: + faces = face_store.get_faces(vision_frame) + + if not faces: + with face_store.resolve_lock(vision_frame): + faces = face_store.get_faces(vision_frame) + + if not faces: + faces = get_many_faces([ vision_frame ]) + + if faces: + face_store.set_faces(vision_frame, faces) + + many_faces.extend(faces) + + return many_faces + + +def refill_faces(faces : List[Optional[Face]]) -> List[Face]: + fill_faces = [] + anchor_index_previous = -1 + + for index, face in enumerate(faces): + if face: + for gap_index in range(anchor_index_previous + 1, index): + average_factor = (gap_index - anchor_index_previous) / (index - anchor_index_previous) + average_face = average_face_geometry([faces[anchor_index_previous], face], average_factor) + fill_faces.append(average_face) + + fill_faces.append(face) + anchor_index_previous = index + + return fill_faces + + +def average_face_geometry(faces : List[Face], average_factor : float) -> Face: + face_first = get_first(faces) + face_middle = get_middle(faces) + face_anchor = face_middle + + if average_factor < 0.5: + face_anchor = face_first + + landmark_set : FaceLandmarkSet =\ + { + '5': average_points(face_first.landmark_set.get('5'), face_middle.landmark_set.get('5'), average_factor), + '5/68': average_points(face_first.landmark_set.get('5/68'), face_middle.landmark_set.get('5/68'), average_factor), + '68': average_points(face_first.landmark_set.get('68'), face_middle.landmark_set.get('68'), average_factor), + '68/5': average_points(face_first.landmark_set.get('68/5'), face_middle.landmark_set.get('68/5'), average_factor) + } + + return Face( + origin = 'refill', + bounding_box = average_points(face_first.bounding_box, face_middle.bounding_box, average_factor), + score_set = face_anchor.score_set, + landmark_set = landmark_set, + angle = estimate_face_angle(landmark_set.get('68/5')), + embedding = face_anchor.embedding, + embedding_norm = face_anchor.embedding_norm, + gender = face_anchor.gender, + age = face_anchor.age, + race = face_anchor.race + ) + + +def average_face_identity(faces : List[Face]) -> Optional[Face]: + face_embeddings = [] + face_embeddings_norm = [] + + if faces: + first_face = get_first(faces) + + for face in faces: + face_embeddings.append(face.embedding) + face_embeddings_norm.append(face.embedding_norm) + + return Face( + origin = first_face.origin, + bounding_box = first_face.bounding_box, + score_set = first_face.score_set, + landmark_set = first_face.landmark_set, + angle = first_face.angle, + embedding = numpy.mean(face_embeddings, axis = 0), + embedding_norm = numpy.mean(face_embeddings_norm, axis = 0), + gender = first_face.gender, + age = first_face.age, + race = first_face.race + ) + return None + + def scale_face(target_face : Face, target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> Face: scale_x = temp_vision_frame.shape[1] / target_vision_frame.shape[1] scale_y = temp_vision_frame.shape[0] / target_vision_frame.shape[0] diff --git a/facefusion/face_helper.py b/facefusion/face_helper.py index 2dfe222a..5a5de311 100644 --- a/facefusion/face_helper.py +++ b/facefusion/face_helper.py @@ -254,3 +254,23 @@ def merge_matrix(temp_matrices : List[Matrix]) -> Matrix: matrix = numpy.dot(temp_matrix, matrix) return matrix[:2, :] + + +def calculate_bounding_box_overlap(bounding_box_a : BoundingBox, bounding_box_b : BoundingBox) -> float: + intersection_x1 = max(bounding_box_a[0], bounding_box_b[0]) + intersection_y1 = max(bounding_box_a[1], bounding_box_b[1]) + intersection_x2 = min(bounding_box_a[2], bounding_box_b[2]) + intersection_y2 = min(bounding_box_a[3], bounding_box_b[3]) + intersection = max(0, intersection_x2 - intersection_x1) * max(0, intersection_y2 - intersection_y1) + bounding_box_area = (bounding_box_a[2] - bounding_box_a[0]) * (bounding_box_a[3] - bounding_box_a[1]) + reference_bounding_box_area = (bounding_box_b[2] - bounding_box_b[0]) * (bounding_box_b[3] - bounding_box_b[1]) + union = bounding_box_area + reference_bounding_box_area - intersection + + if union > 0: + return intersection / union + + return 0.0 + + +def average_points(points_previous : Points, points_next : Points, average_factor : float) -> Points: + return points_previous * (1 - average_factor) + points_next * average_factor diff --git a/facefusion/face_selector.py b/facefusion/face_selector.py index 6b6545d3..a20761cc 100644 --- a/facefusion/face_selector.py +++ b/facefusion/face_selector.py @@ -2,26 +2,35 @@ from typing import List import numpy +import facefusion.choices from facefusion import state_manager -from facefusion.face_analyser import get_many_faces, get_one_face, get_static_faces +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, target_vision_frame : VisionFrame) -> List[Face]: - target_faces = get_many_faces([ target_vision_frame ]) +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(target_faces) + 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(target_faces)) + 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(reference_faces) + 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 @@ -53,17 +62,33 @@ def calculate_face_distance(face : Face, reference_face : Face) -> float: return 0 -def sort_and_filter_faces(faces : List[Face]) -> List[Face]: - if faces: +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'): - faces = sort_faces_by_order(faces, state_manager.get_item('face_selector_order')) - if state_manager.get_item('face_selector_gender'): - faces = filter_faces_by_gender(faces, state_manager.get_item('face_selector_gender')) - if state_manager.get_item('face_selector_race'): - faces = filter_faces_by_race(faces, state_manager.get_item('face_selector_race')) + 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'): - faces = filter_faces_by_age(faces, state_manager.get_item('face_selector_age_start'), state_manager.get_item('face_selector_age_end')) - return faces + 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]: diff --git a/facefusion/face_store.py b/facefusion/face_store.py index 8e627b0a..022210a6 100644 --- a/facefusion/face_store.py +++ b/facefusion/face_store.py @@ -1,3 +1,4 @@ +import threading from typing import List, Optional import numpy @@ -11,14 +12,30 @@ FACE_STORE : FaceStore = {} def get_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: if numpy.any(vision_frame): vision_hash = create_hash(vision_frame.tobytes()) - return FACE_STORE.get(vision_hash) + + if FACE_STORE.get(vision_hash): + return FACE_STORE.get(vision_hash).get('faces') + return None def set_faces(vision_frame : VisionFrame, faces : List[Face]) -> None: if numpy.any(vision_frame): vision_hash = create_hash(vision_frame.tobytes()) - FACE_STORE[vision_hash] = faces + FACE_STORE.setdefault(vision_hash, + { + 'lock': threading.Lock() + })['faces'] = faces + + +def resolve_lock(vision_frame : VisionFrame) -> threading.Lock: + if numpy.any(vision_frame): + vision_hash = create_hash(vision_frame.tobytes()) + return FACE_STORE.setdefault(vision_hash, + { + 'lock': threading.Lock() + }).get('lock') + return threading.Lock() def clear_faces() -> None: diff --git a/facefusion/face_tracker.py b/facefusion/face_tracker.py new file mode 100644 index 00000000..c14afec3 --- /dev/null +++ b/facefusion/face_tracker.py @@ -0,0 +1,61 @@ +from typing import List + +from facefusion.common_helper import get_first, get_last +from facefusion.face_creator import get_static_faces, refill_faces +from facefusion.face_helper import calculate_bounding_box_overlap +from facefusion.types import Face, FaceTrack, Score, VisionFrame + + +def track_faces(vision_frames : List[VisionFrame], score : Score) -> List[Face]: + target_index = len(vision_frames) // 2 + face_tracks = create_face_tracks(vision_frames, score) + temp_faces = [] + + for face_track in face_tracks: + track_indices = sorted(face_track) + track_index_first = get_first(track_indices) + track_index_last = get_last(track_indices) + track_range = range(track_index_first, track_index_last + 1) + + if target_index in track_range: + fill_faces = [] + + for index in track_range: + fill_faces.append(face_track.get(index)) + + temp_faces.append(refill_faces(fill_faces)[target_index - track_index_first]) + + return temp_faces + + +def create_face_tracks(vision_frames : List[VisionFrame], score : Score) -> List[FaceTrack]: + face_tracks : List[FaceTrack] = [] + + for frame_index, vision_frame in enumerate(vision_frames): + for face in get_static_faces([ vision_frame ]): + face_track = select_face_track(face_tracks, face, score) + + if face_track: + face_track[frame_index] = face + else: + face_tracks.append( + { + frame_index : face + }) + + return face_tracks + + +def select_face_track(face_tracks : List[FaceTrack], face : Face, score : Score) -> FaceTrack: + select_track : FaceTrack = {} + select_score = score + + for face_track in face_tracks: + track_face = face_track.get(get_last(face_track)) + track_score = calculate_bounding_box_overlap(face.bounding_box, track_face.bounding_box) + + if track_score > select_score: + select_score = track_score + select_track = face_track + + return select_track diff --git a/facefusion/ffmpeg.py b/facefusion/ffmpeg.py index 886c0fd7..26d5a784 100644 --- a/facefusion/ffmpeg.py +++ b/facefusion/ffmpeg.py @@ -132,6 +132,7 @@ def extract_frames(target_path : str, output_path : str, temp_video_resolution : ffmpeg_builder.enforce_pixel_format('rgb24'), ffmpeg_builder.select_frame_range(trim_frame_start, trim_frame_end, temp_video_fps), ffmpeg_builder.prevent_frame_drop(), + ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_output(temp_frames_pattern) ) @@ -207,6 +208,7 @@ def restore_audio(target_path : str, output_path : str, trim_frame_start : int, temp_video_path = get_temp_file_path(state_manager.get_temp_path(), output_path) temp_video_format = cast(VideoFormat, get_file_format(output_path)) temp_video_duration = detect_video_duration(temp_video_path) + output_video_format = cast(VideoFormat, get_file_format(output_path)) output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) commands = ffmpeg_builder.chain( @@ -220,6 +222,7 @@ def restore_audio(target_path : str, output_path : str, trim_frame_start : int, ffmpeg_builder.select_media_stream('0:v:0'), ffmpeg_builder.select_media_stream('1:a:0'), ffmpeg_builder.set_video_duration(temp_video_duration), + ffmpeg_builder.set_faststart(output_video_format), ffmpeg_builder.force_output(output_path) ) return run_ffmpeg(commands).returncode == 0 @@ -232,6 +235,7 @@ def replace_audio(audio_path : str, output_path : str) -> bool: temp_video_path = get_temp_file_path(state_manager.get_temp_path(), output_path) temp_video_format = cast(VideoFormat, get_file_format(output_path)) temp_video_duration = detect_video_duration(temp_video_path) + output_video_format = cast(VideoFormat, get_file_format(output_path)) output_audio_encoder = fix_audio_encoder(temp_video_format, output_audio_encoder) commands = ffmpeg_builder.chain( @@ -242,6 +246,7 @@ def replace_audio(audio_path : str, output_path : str) -> bool: ffmpeg_builder.set_audio_quality(output_audio_encoder, output_audio_quality), ffmpeg_builder.set_audio_volume(output_audio_volume), ffmpeg_builder.set_video_duration(temp_video_duration), + ffmpeg_builder.set_faststart(output_video_format), ffmpeg_builder.force_output(output_path) ) return run_ffmpeg(commands).returncode == 0 @@ -259,9 +264,11 @@ def merge_video(target_path : str, output_path : str, temp_video_fps : Fps, outp output_video_encoder = fix_video_encoder(temp_video_format, output_video_encoder) commands = ffmpeg_builder.chain( ffmpeg_builder.set_input_fps(temp_video_fps), + ffmpeg_builder.set_start_number(trim_frame_start), ffmpeg_builder.set_input(temp_frames_pattern), ffmpeg_builder.set_media_resolution(pack_resolution(output_video_resolution)), ffmpeg_builder.set_video_encoder(output_video_encoder), + ffmpeg_builder.set_video_tag(output_video_encoder, temp_video_format), ffmpeg_builder.set_video_quality(output_video_encoder, output_video_quality), ffmpeg_builder.set_video_preset(output_video_encoder, output_video_preset), ffmpeg_builder.concat( @@ -288,11 +295,13 @@ def concat_video(output_path : str, temp_output_paths : List[str]) -> bool: concat_video_file.close() output_path = os.path.abspath(output_path) + output_video_format = cast(VideoFormat, get_file_format(output_path)) commands = ffmpeg_builder.chain( ffmpeg_builder.unsafe_concat(), ffmpeg_builder.set_input(concat_video_file.name), ffmpeg_builder.copy_video_encoder(), ffmpeg_builder.copy_audio_encoder(), + ffmpeg_builder.set_faststart(output_video_format), ffmpeg_builder.force_output(output_path) ) process = run_ffmpeg(commands) diff --git a/facefusion/ffmpeg_builder.py b/facefusion/ffmpeg_builder.py index 17471d9e..3c20ae01 100644 --- a/facefusion/ffmpeg_builder.py +++ b/facefusion/ffmpeg_builder.py @@ -5,7 +5,7 @@ from typing import List, Optional import numpy from facefusion.filesystem import get_file_format -from facefusion.types import AudioEncoder, Command, CommandSet, Duration, Fps, SampleRate, StreamMode, VideoEncoder, VideoPreset +from facefusion.types import AudioEncoder, Command, CommandSet, Duration, Fps, SampleRate, StreamMode, VideoEncoder, VideoFormat, VideoPreset def run(commands : List[Command]) -> List[Command]: @@ -51,6 +51,10 @@ def set_input_fps(input_fps : Fps) -> List[Command]: return [ '-r', str(input_fps) ] +def set_start_number(frame_number : int) -> List[Command]: + return [ '-start_number', str(frame_number) ] + + def set_output(output_path : str) -> List[Command]: return [ output_path ] @@ -207,6 +211,18 @@ def copy_video_encoder() -> List[Command]: return set_video_encoder('copy') +def set_faststart(video_format : VideoFormat) -> List[Command]: + if video_format in [ 'm4v', 'mov', 'mp4' ]: + return [ '-movflags', '+faststart' ] + return [] + + +def set_video_tag(video_encoder : VideoEncoder, video_format : VideoFormat) -> List[Command]: + if video_format in [ 'm4v', 'mov', 'mp4' ] and video_encoder in [ 'libx265', 'hevc_nvenc', 'hevc_amf', 'hevc_qsv', 'hevc_videotoolbox' ]: + return [ '-tag:v', 'hvc1' ] + return [] + + def set_video_quality(video_encoder : VideoEncoder, video_quality : int) -> List[Command]: if video_encoder in [ 'libx264', 'libx264rgb', 'libx265' ]: video_compression = numpy.round(numpy.interp(video_quality, [ 0, 100 ], [ 51, 0 ])).astype(int).item() diff --git a/facefusion/locales.py b/facefusion/locales.py index 19af72d8..7931c20c 100644 --- a/facefusion/locales.py +++ b/facefusion/locales.py @@ -129,6 +129,7 @@ LOCALES : Locales =\ 'reference_face_position': 'specify the position used to create the reference face', 'reference_face_distance': 'specify the similarity between the reference face and target face', 'reference_frame_number': 'specify the frame used to create the reference face', + 'face_tracker_score': 'specify the overlap score used to match the tracked faces', 'face_occluder_model': 'choose the model responsible for the occlusion mask', 'face_parser_model': 'choose the model responsible for the region mask', 'face_mask_types': 'mix and match different face mask types (choices: {choices})', @@ -140,6 +141,7 @@ LOCALES : Locales =\ 'trim_frame_start': 'specify the starting frame of the target video', 'trim_frame_end': 'specify the ending frame of the target video', 'temp_frame_format': 'specify the temporary resources format', + 'target_frame_amount': 'specify the amount of target frames forwarded to the processor', 'output_image_quality': 'specify the image quality which translates to the image compression', 'output_image_scale': 'specify the image scale based on the target image', 'output_audio_encoder': 'specify the encoder used for the audio', @@ -228,6 +230,7 @@ LOCALES : Locales =\ 'face_selector_mode_dropdown': 'FACE SELECTOR MODE', 'face_selector_order_dropdown': 'FACE SELECTOR ORDER', 'face_selector_race_dropdown': 'FACE SELECTOR RACE', + 'face_tracker_score_slider': 'FACE TRACKER SCORE', 'face_occluder_model_dropdown': 'FACE OCCLUDER MODEL', 'face_parser_model_dropdown': 'FACE PARSER MODEL', 'voice_extractor_model_dropdown': 'VOICE EXTRACTOR MODEL', diff --git a/facefusion/processors/core.py b/facefusion/processors/core.py index 09d45e5e..fe395445 100644 --- a/facefusion/processors/core.py +++ b/facefusion/processors/core.py @@ -12,6 +12,7 @@ PROCESSORS_METHODS =\ 'clear_inference_pool', 'register_args', 'apply_args', + 'get_common_modules', 'pre_check', 'pre_process', 'post_process', diff --git a/facefusion/processors/modules/age_modifier/core.py b/facefusion/processors/modules/age_modifier/core.py index afc095ca..f5f023d9 100755 --- a/facefusion/processors/modules/age_modifier/core.py +++ b/facefusion/processors/modules/age_modifier/core.py @@ -1,5 +1,7 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType +from typing import List import cv2 import numpy @@ -8,10 +10,9 @@ import facefusion.capability_store import facefusion.choices import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import create_int_metavar, is_macos +from facefusion.common_helper import create_int_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url -from facefusion.execution import has_execution_provider -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import merge_matrix, paste_back, scale_face_landmark_5, warp_face_by_face_landmark_5 from facefusion.face_masker import create_box_mask, create_occlusion_mask from facefusion.face_selector import select_faces @@ -22,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import match_frame_color, read_static_image, read_static_video_frame +from facefusion.vision import match_frame_color, read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -150,10 +151,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('age_modifier_direction', args.get('age_modifier_direction')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -171,16 +180,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def modify_age(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -245,9 +253,6 @@ def forward(crop_vision_frame : VisionFrame, extend_vision_frame : VisionFrame, age_modifier = get_inference_pool().get('age_modifier') age_modifier_inputs = {} - if is_macos() and has_execution_provider('coreml'): - age_modifier.set_providers([ facefusion.choices.execution_provider_set.get('cpu') ]) - for age_modifier_input in age_modifier.get_inputs(): if age_modifier_input.name == 'target': age_modifier_inputs[age_modifier_input.name] = crop_vision_frame @@ -294,10 +299,13 @@ def normalize_extend_frame(extend_vision_frame : VisionFrame) -> VisionFrame: def process_frame(inputs : AgeModifierInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/age_modifier/types.py b/facefusion/processors/modules/age_modifier/types.py index 3a204abc..5f968131 100644 --- a/facefusion/processors/modules/age_modifier/types.py +++ b/facefusion/processors/modules/age_modifier/types.py @@ -1,4 +1,4 @@ -from typing import Any, Literal, TypeAlias, TypedDict +from typing import Any, List, Literal, TypeAlias, TypedDict from numpy.typing import NDArray @@ -7,7 +7,8 @@ from facefusion.types import Mask, VisionFrame AgeModifierInputs = TypedDict('AgeModifierInputs', { 'reference_vision_frame' : VisionFrame, - 'target_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/background_remover/core.py b/facefusion/processors/modules/background_remover/core.py index 1958d1e8..735d19f3 100644 --- a/facefusion/processors/modules/background_remover/core.py +++ b/facefusion/processors/modules/background_remover/core.py @@ -1,11 +1,13 @@ from argparse import ArgumentParser from functools import lru_cache, partial +from types import ModuleType from typing import List, Tuple import cv2 import numpy import facefusion.capability_store +import facefusion.choices import facefusion.jobs.job_manager from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager from facefusion.common_helper import is_macos, is_windows @@ -19,8 +21,8 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.sanitizer import sanitize_int_range from facefusion.thread_helper import thread_semaphore -from facefusion.types import Args, DownloadScope, ExecutionProvider, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_frame +from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -477,12 +479,13 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_execution_providers() -> List[ExecutionProvider]: +def resolve_inference_providers() -> List[InferenceProvider]: model_type = get_model_options().get('type') if is_macos() and has_execution_provider('coreml') or is_windows() and has_execution_provider('directml') and model_type == 'corridor_key': - return [ 'cpu' ] - return state_manager.get_item('execution_providers') + return [ facefusion.choices.execution_provider_set.get('cpu') ] + + return [] def get_model_options() -> ModelOptions: @@ -526,10 +529,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('background_remover_despill_color', normalize_color(args.get('background_remover_despill_color'))) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -547,11 +558,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def remove_background(temp_vision_frame : VisionFrame) -> Tuple[VisionFrame, Mask]: diff --git a/facefusion/processors/modules/background_remover/types.py b/facefusion/processors/modules/background_remover/types.py index e2c5290f..6e19dc0f 100644 --- a/facefusion/processors/modules/background_remover/types.py +++ b/facefusion/processors/modules/background_remover/types.py @@ -1,10 +1,10 @@ -from typing import Literal, TypedDict +from typing import List, Literal, TypedDict from facefusion.types import Mask, VisionFrame BackgroundRemoverInputs = TypedDict('BackgroundRemoverInputs', { - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/deep_swapper/core.py b/facefusion/processors/modules/deep_swapper/core.py index a10a9290..1de063df 100755 --- a/facefusion/processors/modules/deep_swapper/core.py +++ b/facefusion/processors/modules/deep_swapper/core.py @@ -1,6 +1,7 @@ from argparse import ArgumentParser from functools import lru_cache -from typing import Tuple +from types import ModuleType +from typing import List, Tuple import cv2 import numpy @@ -9,9 +10,9 @@ from cv2.typing import Size import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import create_int_metavar +from facefusion.common_helper import create_int_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask from facefusion.face_selector import select_faces @@ -22,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, Mask, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import conditional_match_frame_color, read_static_image, read_static_video_frame +from facefusion.vision import conditional_match_frame_color, read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -302,10 +303,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('deep_swapper_morph', args.get('deep_swapper_morph')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + if model_hash_set and model_source_set: return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) return True @@ -325,16 +334,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def swap_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -425,10 +433,13 @@ def prepare_crop_mask(crop_source_mask : Mask, crop_target_mask : Mask) -> Mask: def process_frame(inputs : DeepSwapperInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/deep_swapper/types.py b/facefusion/processors/modules/deep_swapper/types.py index b404fef8..41e4f979 100644 --- a/facefusion/processors/modules/deep_swapper/types.py +++ b/facefusion/processors/modules/deep_swapper/types.py @@ -1,4 +1,4 @@ -from typing import Any, TypeAlias, TypedDict +from typing import Any, List, TypeAlias, TypedDict from numpy.typing import NDArray @@ -7,7 +7,8 @@ from facefusion.types import Mask, VisionFrame DeepSwapperInputs = TypedDict('DeepSwapperInputs', { 'reference_vision_frame' : VisionFrame, - 'target_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/expression_restorer/core.py b/facefusion/processors/modules/expression_restorer/core.py index 04fd0421..6a6808e9 100755 --- a/facefusion/processors/modules/expression_restorer/core.py +++ b/facefusion/processors/modules/expression_restorer/core.py @@ -1,6 +1,7 @@ from argparse import ArgumentParser from functools import lru_cache -from typing import Tuple +from types import ModuleType +from typing import List, Tuple import cv2 import numpy @@ -8,9 +9,9 @@ import numpy import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import create_int_metavar +from facefusion.common_helper import create_int_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 from facefusion.face_masker import create_box_mask, create_occlusion_mask from facefusion.face_selector import select_faces @@ -22,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, LivePortraitExpression, from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -134,10 +135,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('expression_restorer_areas', args.get('expression_restorer_areas')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -158,16 +167,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def restore_expression(target_face : Face, target_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -278,10 +286,13 @@ def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: def process_frame(inputs : ExpressionRestorerInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/expression_restorer/types.py b/facefusion/processors/modules/expression_restorer/types.py index 1cee5a52..e0dec6e3 100644 --- a/facefusion/processors/modules/expression_restorer/types.py +++ b/facefusion/processors/modules/expression_restorer/types.py @@ -6,7 +6,7 @@ ExpressionRestorerInputs = TypedDict('ExpressionRestorerInputs', { 'reference_vision_frame' : VisionFrame, 'source_vision_frames' : List[VisionFrame], - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/face_debugger/core.py b/facefusion/processors/modules/face_debugger/core.py index fc05e7b8..8e14ced1 100755 --- a/facefusion/processors/modules/face_debugger/core.py +++ b/facefusion/processors/modules/face_debugger/core.py @@ -1,4 +1,6 @@ from argparse import ArgumentParser +from types import ModuleType +from typing import List import cv2 import numpy @@ -6,7 +8,8 @@ import numpy import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, logger, state_manager, translator, video_manager -from facefusion.face_analyser import scale_face +from facefusion.common_helper import get_middle +from facefusion.face_creator import scale_face from facefusion.face_helper import warp_face_by_face_landmark_5 from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask from facefusion.face_selector import select_faces @@ -16,7 +19,7 @@ from facefusion.processors.modules.face_debugger.types import FaceDebuggerInputs from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.types import Args, Face, InferencePool, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame def get_inference_pool() -> InferencePool: @@ -49,7 +52,14 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('face_debugger_items', args.get('face_debugger_items')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False return True @@ -67,14 +77,12 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def debug_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -103,21 +111,22 @@ def debug_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFra def draw_bounding_box(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) - box_color = 0, 0, 255 - border_color = 100, 100, 255 bounding_box = target_face.bounding_box.astype(numpy.int32) x1, y1, x2, y2 = bounding_box + box_color = 0, 0, 255 + border_scale = calculate_scale(temp_vision_frame) + border_color = 100, 100, 255 - cv2.rectangle(temp_vision_frame, (x1, y1), (x2, y2), box_color, 2) + cv2.rectangle(temp_vision_frame, (x1, y1), (x2, y2), box_color, border_scale) if target_face.angle == 0: - cv2.line(temp_vision_frame, (x1, y1), (x2, y1), border_color, 3) + cv2.line(temp_vision_frame, (x1, y1), (x2, y1), border_color, border_scale + 1) if target_face.angle == 180: - cv2.line(temp_vision_frame, (x1, y2), (x2, y2), border_color, 3) + cv2.line(temp_vision_frame, (x1, y2), (x2, y2), border_color, border_scale + 1) if target_face.angle == 90: - cv2.line(temp_vision_frame, (x2, y1), (x2, y2), border_color, 3) + cv2.line(temp_vision_frame, (x2, y1), (x2, y2), border_color, border_scale + 1) if target_face.angle == 270: - cv2.line(temp_vision_frame, (x1, y1), (x1, y2), border_color, 3) + cv2.line(temp_vision_frame, (x1, y1), (x1, y2), border_color, border_scale + 1) return temp_vision_frame @@ -131,11 +140,15 @@ def draw_face_mask(target_face : Face, temp_vision_frame : VisionFrame) -> Visio crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5_68, 'arcface_128', (512, 512)) inverse_matrix = cv2.invertAffineTransform(affine_matrix) temp_size = temp_vision_frame.shape[:2][::-1] + mask_scale = calculate_scale(temp_vision_frame) mask_color = 0, 255, 0 if numpy.array_equal(face_landmark_5, face_landmark_5_68): mask_color = 255, 255, 0 + if target_face.origin == 'refill': + mask_color = 0, 165, 255 + if 'box' in state_manager.get_item('face_mask_types'): box_mask = create_box_mask(crop_vision_frame, 0, state_manager.get_item('face_mask_padding')) crop_masks.append(box_mask) @@ -158,7 +171,7 @@ def draw_face_mask(target_face : Face, temp_vision_frame : VisionFrame) -> Visio inverse_vision_frame = cv2.warpAffine(crop_mask, inverse_matrix, temp_size) inverse_vision_frame = cv2.threshold(inverse_vision_frame, 100, 255, cv2.THRESH_BINARY)[1] inverse_contours, _ = cv2.findContours(inverse_vision_frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE) - cv2.drawContours(temp_vision_frame, inverse_contours, -1, mask_color, 2) + cv2.drawContours(temp_vision_frame, inverse_contours, -1, mask_color, mask_scale) return temp_vision_frame @@ -166,13 +179,17 @@ def draw_face_mask(target_face : Face, temp_vision_frame : VisionFrame) -> Visio def draw_face_landmark_5(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) face_landmark_5 = target_face.landmark_set.get('5') + point_scale = calculate_scale(temp_vision_frame) point_color = 0, 0, 255 + if target_face.origin == 'refill': + point_color = 0, 165, 255 + if numpy.any(face_landmark_5): face_landmark_5 = face_landmark_5.astype(numpy.int32) for point in face_landmark_5: - cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + cv2.circle(temp_vision_frame, tuple(point), point_scale, point_color, -1) return temp_vision_frame @@ -181,16 +198,20 @@ def draw_face_landmark_5_68(target_face : Face, temp_vision_frame : VisionFrame) temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) face_landmark_5 = target_face.landmark_set.get('5') face_landmark_5_68 = target_face.landmark_set.get('5/68') + point_scale = calculate_scale(temp_vision_frame) point_color = 0, 255, 0 if numpy.array_equal(face_landmark_5, face_landmark_5_68): point_color = 255, 255, 0 + if target_face.origin == 'refill': + point_color = 0, 165, 255 + if numpy.any(face_landmark_5_68): face_landmark_5_68 = face_landmark_5_68.astype(numpy.int32) for point in face_landmark_5_68: - cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + cv2.circle(temp_vision_frame, tuple(point), point_scale, point_color, -1) return temp_vision_frame @@ -199,16 +220,20 @@ def draw_face_landmark_68(target_face : Face, temp_vision_frame : VisionFrame) - temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) face_landmark_68 = target_face.landmark_set.get('68') face_landmark_68_5 = target_face.landmark_set.get('68/5') + point_scale = calculate_scale(temp_vision_frame) point_color = 0, 255, 0 if numpy.array_equal(face_landmark_68, face_landmark_68_5): point_color = 255, 255, 0 + if target_face.origin == 'refill': + point_color = 0, 165, 255 + if numpy.any(face_landmark_68): face_landmark_68 = face_landmark_68.astype(numpy.int32) for point in face_landmark_68: - cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + cv2.circle(temp_vision_frame, tuple(point), point_scale, point_color, -1) return temp_vision_frame @@ -216,23 +241,36 @@ def draw_face_landmark_68(target_face : Face, temp_vision_frame : VisionFrame) - def draw_face_landmark_68_5(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: temp_vision_frame = numpy.ascontiguousarray(temp_vision_frame) face_landmark_68_5 = target_face.landmark_set.get('68/5') + point_scale = calculate_scale(temp_vision_frame) point_color = 255, 255, 0 + if target_face.origin == 'refill': + point_color = 0, 165, 255 + if numpy.any(face_landmark_68_5): face_landmark_68_5 = face_landmark_68_5.astype(numpy.int32) for point in face_landmark_68_5: - cv2.circle(temp_vision_frame, tuple(point), 3, point_color, -1) + cv2.circle(temp_vision_frame, tuple(point), point_scale, point_color, -1) return temp_vision_frame +def calculate_scale(temp_vision_frame : VisionFrame) -> int: + frame_height, _ = temp_vision_frame.shape[:2] + frame_scale = round(frame_height / 270) + return max(1, min(10, frame_scale)) + + def process_frame(inputs : FaceDebuggerInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: @@ -240,5 +278,3 @@ def process_frame(inputs : FaceDebuggerInputs) -> ProcessorOutputs: temp_vision_frame = debug_face(target_face, temp_vision_frame) return temp_vision_frame, temp_vision_mask - - diff --git a/facefusion/processors/modules/face_debugger/types.py b/facefusion/processors/modules/face_debugger/types.py index a51539b6..afb7e5b3 100644 --- a/facefusion/processors/modules/face_debugger/types.py +++ b/facefusion/processors/modules/face_debugger/types.py @@ -1,11 +1,12 @@ -from typing import Literal, TypedDict +from typing import List, Literal, TypedDict from facefusion.types import Mask, VisionFrame FaceDebuggerInputs = TypedDict('FaceDebuggerInputs', { 'reference_vision_frame' : VisionFrame, - 'target_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/face_editor/core.py b/facefusion/processors/modules/face_editor/core.py index 9f9174a1..78d96e75 100755 --- a/facefusion/processors/modules/face_editor/core.py +++ b/facefusion/processors/modules/face_editor/core.py @@ -1,6 +1,7 @@ from argparse import ArgumentParser from functools import lru_cache -from typing import Tuple +from types import ModuleType +from typing import List, Tuple import cv2 import numpy @@ -8,9 +9,9 @@ import numpy import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import create_float_metavar +from facefusion.common_helper import create_float_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import paste_back, scale_face_landmark_5, warp_face_by_face_landmark_5 from facefusion.face_masker import create_box_mask from facefusion.face_selector import select_faces @@ -22,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, LivePortraitExpression, from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore, thread_semaphore from facefusion.types import Args, DownloadScope, Face, FaceLandmark68, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -272,10 +273,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('face_editor_head_roll', args.get('face_editor_head_roll')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -293,16 +302,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def edit_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -591,10 +599,13 @@ def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: def process_frame(inputs : FaceEditorInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/face_editor/types.py b/facefusion/processors/modules/face_editor/types.py index 6e246466..2df890ff 100644 --- a/facefusion/processors/modules/face_editor/types.py +++ b/facefusion/processors/modules/face_editor/types.py @@ -1,11 +1,12 @@ -from typing import Literal, TypedDict +from typing import List, Literal, TypedDict from facefusion.types import Mask, VisionFrame FaceEditorInputs = TypedDict('FaceEditorInputs', { 'reference_vision_frame' : VisionFrame, - 'target_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/face_enhancer/core.py b/facefusion/processors/modules/face_enhancer/core.py index bc40538d..d73a0f3d 100755 --- a/facefusion/processors/modules/face_enhancer/core.py +++ b/facefusion/processors/modules/face_enhancer/core.py @@ -1,14 +1,16 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType +from typing import List import numpy import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import create_float_metavar, create_int_metavar +from facefusion.common_helper import create_float_metavar, create_int_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 from facefusion.face_masker import create_box_mask, create_occlusion_mask from facefusion.face_selector import select_faces @@ -19,7 +21,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore from facefusion.types import Args, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, read_static_image, read_static_video_frame +from facefusion.vision import blend_frame, read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -327,10 +329,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('face_enhancer_weight', args.get('face_enhancer_weight')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -348,16 +358,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def enhance_face(target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -434,10 +443,13 @@ def blend_paste_frame(temp_vision_frame : VisionFrame, paste_vision_frame : Visi def process_frame(inputs : FaceEnhancerInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') - target_vision_frame = inputs.get('target_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/face_enhancer/types.py b/facefusion/processors/modules/face_enhancer/types.py index 484a98c6..2104dbb8 100644 --- a/facefusion/processors/modules/face_enhancer/types.py +++ b/facefusion/processors/modules/face_enhancer/types.py @@ -1,4 +1,4 @@ -from typing import Any, Literal, TypeAlias, TypedDict +from typing import Any, List, Literal, TypeAlias, TypedDict from numpy.typing import NDArray @@ -7,7 +7,8 @@ from facefusion.types import Mask, VisionFrame FaceEnhancerInputs = TypedDict('FaceEnhancerInputs', { 'reference_vision_frame' : VisionFrame, - 'target_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/face_swapper/core.py b/facefusion/processors/modules/face_swapper/core.py index 368c1a63..d7fa11a6 100755 --- a/facefusion/processors/modules/face_swapper/core.py +++ b/facefusion/processors/modules/face_swapper/core.py @@ -1,5 +1,6 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType from typing import List, Optional, Tuple import cv2 @@ -9,10 +10,10 @@ import facefusion.capability_store import facefusion.choices import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager -from facefusion.common_helper import get_first, is_macos +from facefusion.common_helper import get_first, get_middle, is_macos from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url from facefusion.execution import has_execution_provider -from facefusion.face_analyser import get_average_face, get_one_face, get_static_faces, scale_face +from facefusion.face_creator import average_face_identity, get_one_face, get_static_faces, scale_face from facefusion.face_helper import paste_back, warp_face_by_face_landmark_5 from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask, create_region_mask from facefusion.face_selector import select_faces, sort_faces_by_order @@ -24,8 +25,8 @@ from facefusion.processors.pixel_boost import explode_pixel_boost, implode_pixel from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore -from facefusion.types import Args, DownloadScope, Embedding, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_images, read_static_video_frame, unpack_resolution +from facefusion.types import Args, DownloadScope, Embedding, Face, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import read_static_image, read_static_images, read_static_video_chunk, read_static_video_frame, unpack_resolution @lru_cache() @@ -246,6 +247,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/hyperswap_1a_256.onnx') } }, + 'precision': 'fp16', 'type': 'hyperswap', 'template': 'arcface_128', 'size': (256, 256), @@ -276,6 +278,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/hyperswap_1b_256.onnx') } }, + 'precision': 'fp16', 'type': 'hyperswap', 'template': 'arcface_128', 'size': (256, 256), @@ -306,6 +309,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/hyperswap_1c_256.onnx') } }, + 'precision': 'fp16', 'type': 'hyperswap', 'template': 'arcface_128', 'size': (256, 256), @@ -366,6 +370,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/inswapper_128_fp16.onnx') } }, + 'precision': 'fp16', 'type': 'inswapper', 'template': 'arcface_128', 'size': (128, 128), @@ -486,28 +491,38 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: def get_inference_pool() -> InferencePool: - model_names = [ get_model_name() ] + model_names = [ state_manager.get_item('face_swapper_model') ] model_source_set = get_model_options().get('sources') return inference_manager.get_inference_pool(__name__, model_names, model_source_set) def clear_inference_pool() -> None: - model_names = [ get_model_name() ] + model_names = [ state_manager.get_item('face_swapper_model') ] inference_manager.clear_inference_pool(__name__, model_names) +def resolve_inference_providers() -> List[InferenceProvider]: + model_precision = get_model_options().get('precision') + model_type = get_model_options().get('type') + + if is_macos() and has_execution_provider('coreml'): + if model_type in [ 'ghost', 'uniface' ] or model_precision == 'fp16': + return\ + [ + (facefusion.choices.execution_provider_set.get('coreml'), + { + 'ModelFormat': 'MLProgram', + 'SpecializationStrategy': 'FastPrediction' + }) + ] + + return [] + + def get_model_options() -> ModelOptions: - model_name = get_model_name() - return create_static_model_set('full').get(model_name) - - -def get_model_name() -> str: model_name = state_manager.get_item('face_swapper_model') - - if is_macos() and has_execution_provider('coreml') and model_name == 'inswapper_128_fp16': - return 'inswapper_128' - return model_name + return create_static_model_set('full').get(model_name) def register_args(program : ArgumentParser) -> None: @@ -552,10 +567,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('face_swapper_weight', args.get('face_swapper_weight')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -587,20 +610,19 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: get_static_model_initializer.cache_clear() clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() -def swap_face(source_face : Face, target_face : Face, temp_vision_frame : VisionFrame) -> VisionFrame: +def swap_face(source_face : Face, target_face : Face, source_vision_frame : VisionFrame, temp_vision_frame : VisionFrame) -> VisionFrame: model_template = get_model_options().get('template') model_size = get_model_options().get('size') pixel_boost_size = unpack_resolution(state_manager.get_item('face_swapper_pixel_boost')) @@ -620,7 +642,7 @@ def swap_face(source_face : Face, target_face : Face, temp_vision_frame : Vision pixel_boost_vision_frames = implode_pixel_boost(crop_vision_frame, pixel_boost_total, model_size) for pixel_boost_vision_frame in pixel_boost_vision_frames: pixel_boost_vision_frame = prepare_crop_frame(pixel_boost_vision_frame) - pixel_boost_vision_frame = forward_swap_face(source_face, target_face, pixel_boost_vision_frame) + pixel_boost_vision_frame = forward_swap_face(source_face, target_face, source_vision_frame, pixel_boost_vision_frame) pixel_boost_vision_frame = normalize_crop_frame(pixel_boost_vision_frame) temp_vision_frames.append(pixel_boost_vision_frame) crop_vision_frame = explode_pixel_boost(temp_vision_frames, pixel_boost_total, model_size, pixel_boost_size) @@ -639,18 +661,15 @@ def swap_face(source_face : Face, target_face : Face, temp_vision_frame : Vision return paste_vision_frame -def forward_swap_face(source_face : Face, target_face : Face, crop_vision_frame : VisionFrame) -> VisionFrame: +def forward_swap_face(source_face : Face, target_face : Face, source_vision_frame : VisionFrame, crop_vision_frame : VisionFrame) -> VisionFrame: face_swapper = get_inference_pool().get('face_swapper') model_type = get_model_options().get('type') face_swapper_inputs = {} - if is_macos() and has_execution_provider('coreml') and model_type in [ 'ghost', 'uniface' ]: - face_swapper.set_providers([ facefusion.choices.execution_provider_set.get('cpu') ]) - for face_swapper_input in face_swapper.get_inputs(): if face_swapper_input.name == 'source': if model_type in [ 'blendswap', 'uniface' ]: - face_swapper_inputs[face_swapper_input.name] = prepare_source_frame(source_face) + face_swapper_inputs[face_swapper_input.name] = prepare_source_frame(source_face, source_vision_frame) else: source_embedding = prepare_source_embedding(source_face) source_embedding = balance_source_embedding(source_embedding, target_face.embedding) @@ -676,9 +695,8 @@ def forward_convert_embedding(face_embedding : Embedding) -> Embedding: return face_embedding -def prepare_source_frame(source_face : Face) -> VisionFrame: +def prepare_source_frame(source_face : Face, source_vision_frame : VisionFrame) -> VisionFrame: model_type = get_model_options().get('type') - source_vision_frame = read_static_image(get_first(state_manager.get_item('source_paths'))) if model_type == 'blendswap': source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmark_set.get('5/68'), 'arcface_112_v2', (112, 112)) @@ -776,21 +794,25 @@ def extract_source_face(source_vision_frames : List[VisionFrame]) -> Optional[Fa if temp_faces: source_faces.append(get_first(temp_faces)) - return get_average_face(source_faces) + return average_face_identity(source_faces) def process_frame(inputs : FaceSwapperInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') source_vision_frames = inputs.get('source_vision_frames') - target_vision_frame = inputs.get('target_vision_frame') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') + + target_vision_frame = get_middle(target_vision_frames) source_face = extract_source_face(source_vision_frames) - target_faces = select_faces(reference_vision_frame, target_vision_frame) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if source_face and target_faces: + source_vision_frame = get_first(source_vision_frames) + for target_face in target_faces: target_face = scale_face(target_face, target_vision_frame, temp_vision_frame) - temp_vision_frame = swap_face(source_face, target_face, temp_vision_frame) + temp_vision_frame = swap_face(source_face, target_face, source_vision_frame, temp_vision_frame) return temp_vision_frame, temp_vision_mask diff --git a/facefusion/processors/modules/face_swapper/types.py b/facefusion/processors/modules/face_swapper/types.py index 4afc3cf5..addda8de 100644 --- a/facefusion/processors/modules/face_swapper/types.py +++ b/facefusion/processors/modules/face_swapper/types.py @@ -6,7 +6,7 @@ FaceSwapperInputs = TypedDict('FaceSwapperInputs', { 'reference_vision_frame' : VisionFrame, 'source_vision_frames' : List[VisionFrame], - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/frame_colorizer/core.py b/facefusion/processors/modules/frame_colorizer/core.py index 75966067..cdef0e1c 100644 --- a/facefusion/processors/modules/frame_colorizer/core.py +++ b/facefusion/processors/modules/frame_colorizer/core.py @@ -1,11 +1,13 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType from typing import List import cv2 import numpy import facefusion.capability_store +import facefusion.choices import facefusion.jobs.job_manager from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager from facefusion.common_helper import create_int_metavar, is_macos @@ -17,8 +19,8 @@ from facefusion.processors.modules.frame_colorizer.types import FrameColorizerIn from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import thread_semaphore -from facefusion.types import Args, DownloadScope, ExecutionProvider, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, read_static_image, read_static_video_frame, unpack_resolution +from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import blend_frame, read_static_image, read_static_video_chunk, read_static_video_frame, unpack_resolution @lru_cache() @@ -170,10 +172,11 @@ def clear_inference_pool() -> None: inference_manager.clear_inference_pool(__name__, model_names) -def resolve_execution_providers() -> List[ExecutionProvider]: +def resolve_inference_providers() -> List[InferenceProvider]: if is_macos() and has_execution_provider('coreml'): - return [ 'cpu' ] - return state_manager.get_item('execution_providers') + return [ facefusion.choices.execution_provider_set.get('cpu') ] + + return [] def get_model_options() -> ModelOptions: @@ -218,10 +221,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('frame_colorizer_size', args.get('frame_colorizer_size')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -239,11 +250,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def colorize_frame(temp_vision_frame : VisionFrame) -> VisionFrame: diff --git a/facefusion/processors/modules/frame_colorizer/types.py b/facefusion/processors/modules/frame_colorizer/types.py index b0e152f5..8670646d 100644 --- a/facefusion/processors/modules/frame_colorizer/types.py +++ b/facefusion/processors/modules/frame_colorizer/types.py @@ -1,10 +1,10 @@ -from typing import Literal, TypedDict +from typing import List, Literal, TypedDict from facefusion.types import Mask, VisionFrame FrameColorizerInputs = TypedDict('FrameColorizerInputs', { - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/frame_enhancer/core.py b/facefusion/processors/modules/frame_enhancer/core.py index 2adc0ffc..d0d07b79 100644 --- a/facefusion/processors/modules/frame_enhancer/core.py +++ b/facefusion/processors/modules/frame_enhancer/core.py @@ -1,10 +1,13 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType +from typing import List import cv2 import numpy import facefusion.capability_store +import facefusion.choices import facefusion.jobs.job_manager from facefusion import config, content_analyser, inference_manager, logger, state_manager, translator, video_manager from facefusion.common_helper import create_int_metavar, is_macos @@ -16,8 +19,8 @@ from facefusion.processors.modules.frame_enhancer.types import FrameEnhancerInpu from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore -from facefusion.types import Args, DownloadScope, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import blend_frame, create_tile_frames, merge_tile_frames, read_static_image, read_static_video_frame +from facefusion.types import Args, DownloadScope, InferencePool, InferenceProvider, ModelOptions, ModelSet, ProcessMode, VisionFrame +from facefusion.vision import blend_frame, create_tile_frames, merge_tile_frames, read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -156,6 +159,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/real_esrgan_x2_fp16.onnx') } }, + 'precision': 'fp16', 'size': (256, 16, 8), 'scale': 2 }, @@ -210,6 +214,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/real_esrgan_x4_fp16.onnx') } }, + 'precision': 'fp16', 'size': (256, 16, 8), 'scale': 4 }, @@ -264,6 +269,7 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: 'path': resolve_relative_path('../.assets/models/real_esrgan_x8_fp16.onnx') } }, + 'precision': 'fp16', 'size': (256, 16, 8), 'scale': 8 }, @@ -541,35 +547,38 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet: def get_inference_pool() -> InferencePool: - model_names = [ get_frame_enhancer_model() ] + model_names = [ state_manager.get_item('frame_enhancer_model') ] model_source_set = get_model_options().get('sources') return inference_manager.get_inference_pool(__name__, model_names, model_source_set) def clear_inference_pool() -> None: - model_names = [ get_frame_enhancer_model() ] + model_names = [ state_manager.get_item('frame_enhancer_model') ] inference_manager.clear_inference_pool(__name__, model_names) +def resolve_inference_providers() -> List[InferenceProvider]: + model_precision = get_model_options().get('precision') + + if is_macos() and has_execution_provider('coreml') and model_precision == 'fp16': + return\ + [ + (facefusion.choices.execution_provider_set.get('coreml'), + { + 'ModelFormat': 'MLProgram', + 'SpecializationStrategy': 'FastPrediction' + }) + ] + + return [] + + def get_model_options() -> ModelOptions: - model_name = get_frame_enhancer_model() + model_name = state_manager.get_item('frame_enhancer_model') return create_static_model_set('full').get(model_name) -def get_frame_enhancer_model() -> str: - frame_enhancer_model = state_manager.get_item('frame_enhancer_model') - - if is_macos() and has_execution_provider('coreml'): - if frame_enhancer_model == 'real_esrgan_x2_fp16': - return 'real_esrgan_x2' - if frame_enhancer_model == 'real_esrgan_x4_fp16': - return 'real_esrgan_x4' - if frame_enhancer_model == 'real_esrgan_x8_fp16': - return 'real_esrgan_x8' - return frame_enhancer_model - - def register_args(program : ArgumentParser) -> None: group_processors = find_argument_group(program, 'processors') if group_processors: @@ -599,10 +608,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('frame_enhancer_blend', args.get('frame_enhancer_blend')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -620,11 +637,15 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def enhance_frame(temp_vision_frame : VisionFrame) -> VisionFrame: diff --git a/facefusion/processors/modules/frame_enhancer/types.py b/facefusion/processors/modules/frame_enhancer/types.py index d9786a39..d593b5fb 100644 --- a/facefusion/processors/modules/frame_enhancer/types.py +++ b/facefusion/processors/modules/frame_enhancer/types.py @@ -1,10 +1,10 @@ -from typing import Literal, TypedDict +from typing import List, Literal, TypedDict from facefusion.types import Mask, VisionFrame FrameEnhancerInputs = TypedDict('FrameEnhancerInputs', { - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/processors/modules/lip_syncer/core.py b/facefusion/processors/modules/lip_syncer/core.py index 9291afa1..227f4b71 100755 --- a/facefusion/processors/modules/lip_syncer/core.py +++ b/facefusion/processors/modules/lip_syncer/core.py @@ -1,5 +1,7 @@ from argparse import ArgumentParser from functools import lru_cache +from types import ModuleType +from typing import List import cv2 import numpy @@ -8,9 +10,9 @@ import facefusion.capability_store import facefusion.jobs.job_manager from facefusion import config, content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, inference_manager, logger, state_manager, translator, video_manager, voice_extractor from facefusion.audio import read_static_voice -from facefusion.common_helper import create_float_metavar +from facefusion.common_helper import create_float_metavar, get_middle from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url -from facefusion.face_analyser import scale_face +from facefusion.face_creator import scale_face from facefusion.face_helper import create_bounding_box, paste_back, warp_face_by_bounding_box, warp_face_by_face_landmark_5 from facefusion.face_masker import create_area_mask, create_box_mask, create_occlusion_mask from facefusion.face_selector import select_faces @@ -21,7 +23,7 @@ from facefusion.processors.types import ApplyStateItem, ProcessorOutputs from facefusion.program_helper import find_argument_group from facefusion.thread_helper import conditional_thread_semaphore from facefusion.types import Args, AudioFrame, DownloadScope, Face, InferencePool, ModelOptions, ModelSet, ProcessMode, VisionFrame -from facefusion.vision import read_static_image, read_static_video_frame +from facefusion.vision import read_static_image, read_static_video_chunk, read_static_video_frame @lru_cache() @@ -158,10 +160,18 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None: apply_state_item('lip_syncer_weight', args.get('lip_syncer_weight')) +def get_common_modules() -> List[ModuleType]: + return [ content_analyser, face_classifier, face_detector, face_landmarker, face_masker, face_recognizer, voice_extractor ] + + def pre_check() -> bool: model_hash_set = get_model_options().get('hashes') model_source_set = get_model_options().get('sources') + for common_module in get_common_modules(): + if not common_module.pre_check(): + return False + return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) @@ -175,18 +185,16 @@ def pre_process(mode : ProcessMode) -> bool: def post_process() -> None: read_static_image.cache_clear() read_static_video_frame.cache_clear() + read_static_video_chunk.cache_clear() read_static_voice.cache_clear() video_manager.clear_video_pool() + if state_manager.get_item('video_memory_strategy') in [ 'strict', 'moderate' ]: clear_inference_pool() + if state_manager.get_item('video_memory_strategy') == 'strict': - content_analyser.clear_inference_pool() - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_masker.clear_inference_pool() - face_recognizer.clear_inference_pool() - voice_extractor.clear_inference_pool() + for common_module in get_common_modules(): + common_module.clear_inference_pool() def sync_lip(target_face : Face, source_voice_frame : AudioFrame, temp_vision_frame : VisionFrame) -> VisionFrame: @@ -298,11 +306,14 @@ def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: def process_frame(inputs : LipSyncerInputs) -> ProcessorOutputs: reference_vision_frame = inputs.get('reference_vision_frame') + source_vision_frames = inputs.get('source_vision_frames') source_voice_frame = inputs.get('source_voice_frame') - target_vision_frame = inputs.get('target_vision_frame') + target_vision_frames = inputs.get('target_vision_frames') temp_vision_frame = inputs.get('temp_vision_frame') temp_vision_mask = inputs.get('temp_vision_mask') - target_faces = select_faces(reference_vision_frame, target_vision_frame) + + target_vision_frame = get_middle(target_vision_frames) + target_faces = select_faces(reference_vision_frame, source_vision_frames, target_vision_frames) if target_faces: for target_face in target_faces: diff --git a/facefusion/processors/modules/lip_syncer/types.py b/facefusion/processors/modules/lip_syncer/types.py index 32861f6a..d13b6a9a 100644 --- a/facefusion/processors/modules/lip_syncer/types.py +++ b/facefusion/processors/modules/lip_syncer/types.py @@ -1,4 +1,4 @@ -from typing import Any, Literal, TypeAlias, TypedDict +from typing import Any, List, Literal, TypeAlias, TypedDict from numpy.typing import NDArray @@ -7,8 +7,9 @@ from facefusion.types import AudioFrame, Mask, VisionFrame LipSyncerInputs = TypedDict('LipSyncerInputs', { 'reference_vision_frame' : VisionFrame, + 'source_vision_frames' : List[VisionFrame], 'source_voice_frame' : AudioFrame, - 'target_vision_frame' : VisionFrame, + 'target_vision_frames' : List[VisionFrame], 'temp_vision_frame' : VisionFrame, 'temp_vision_mask' : Mask }) diff --git a/facefusion/program.py b/facefusion/program.py index 964fda79..d11198fb 100755 --- a/facefusion/program.py +++ b/facefusion/program.py @@ -421,6 +421,27 @@ def create_face_selector_program() -> ArgumentParser: return program +def create_face_tracker_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_face_tracker = program.add_argument_group('face tracker') + + capability_store.register_capability_set( + [ + group_face_tracker.add_argument( + '--face-tracker-score', + help = translator.get('help.face_tracker_score'), + type = float, + default = config.get_float_value('face_tracker', 'face_tracker_score', '0.0'), + choices = facefusion.choices.face_tracker_score_range, + metavar = create_float_metavar(facefusion.choices.face_tracker_score_range) + ) + ], + scopes = [ 'api', 'cli' ] + ) + + return program + + def create_face_masker_program() -> ArgumentParser: program = ArgumentParser(add_help = False) group_face_masker = program.add_argument_group('face masker') @@ -575,6 +596,27 @@ def create_frame_extraction_program() -> ArgumentParser: return program +def create_frame_distribution_program() -> ArgumentParser: + program = ArgumentParser(add_help = False) + group_frame_distribution = program.add_argument_group('frame distribution') + + capability_store.register_capability_set( + [ + group_frame_distribution.add_argument( + '--target-frame-amount', + help = translator.get('help.target_frame_amount'), + type = int, + default = config.get_int_value('frame_distribution', 'target_frame_amount', '5'), + choices = facefusion.choices.target_frame_amount_range, + metavar = create_int_metavar(facefusion.choices.target_frame_amount_range) + ) + ], + scopes = [ 'api', 'cli' ] + ) + + return program + + def create_output_creation_program() -> ArgumentParser: program = ArgumentParser(add_help = False) available_encoder_set = get_available_encoder_set() @@ -985,9 +1027,11 @@ def collect_step_program() -> ArgumentParser: create_face_detector_program(), create_face_landmarker_program(), create_face_selector_program(), + create_face_tracker_program(), create_face_masker_program(), create_voice_extractor_program(), create_frame_extraction_program(), + create_frame_distribution_program(), create_output_creation_program(), create_processors_program() ], diff --git a/facefusion/streamer.py b/facefusion/streamer.py index 97357dce..6cfccdd5 100644 --- a/facefusion/streamer.py +++ b/facefusion/streamer.py @@ -60,7 +60,7 @@ def process_frame(stream_audio_frame : AudioFrame, stream_vision_frame : VisionF 'source_vision_frames': source_vision_frames, 'source_audio_frame': stream_audio_frame, 'source_voice_frame': source_voice_frame, - 'target_vision_frame': stream_vision_frame, + 'target_vision_frames': [ stream_vision_frame ], 'temp_vision_frame': temp_vision_frame, 'temp_vision_mask': temp_vision_mask }) diff --git a/facefusion/types.py b/facefusion/types.py index 60bf7477..f06d1b0c 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -1,6 +1,7 @@ import ctypes from collections import namedtuple from datetime import datetime +from threading import Lock from typing import Any, Callable, Dict, List, Literal, NotRequired, Optional, Tuple, TypeAlias, TypedDict, Union import cv2 @@ -31,22 +32,34 @@ FaceScoreSet = TypedDict('FaceScoreSet', 'landmarker' : Score }) Embedding : TypeAlias = NDArray[numpy.float64] -Gender = Literal['female', 'male'] + Age : TypeAlias = range +Gender = Literal['female', 'male'] Race = Literal['white', 'black', 'latino', 'asian', 'indian', 'arabic'] + +FaceSelectorGender = Literal['auto', 'female', 'male'] +FaceSelectorRace = Literal['auto', 'white', 'black', 'latino', 'asian', 'indian', 'arabic'] + Face = namedtuple('Face', [ + 'origin', 'bounding_box', 'score_set', 'landmark_set', 'angle', 'embedding', 'embedding_norm', - 'gender', 'age', + 'gender', 'race' ]) -FaceStore : TypeAlias = Dict[str, List[Face]] +FaceSet = TypedDict('FaceSet', +{ + 'lock': Lock, + 'faces': NotRequired[List[Face]] +}) +FaceStore : TypeAlias = Dict[str, FaceSet] +FaceTrack : TypeAlias = Dict[int, Face] Language = Literal['en'] Locales : TypeAlias = Dict[Language, Dict[str, Any]] @@ -185,6 +198,8 @@ ImageFormat = Literal['bmp', 'jpeg', 'png', 'tiff', 'webp'] VideoFormat = Literal['avi', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mxf', 'webm', 'wmv'] TempFrameFormat = Literal['bmp', 'jpeg', 'png', 'tiff'] +FrameSet : TypeAlias = Dict[int, str] + AudioEncoder = Literal['flac', 'aac', 'libmp3lame', 'libopus', 'libvorbis', 'pcm_s16le', 'pcm_s32le'] ImageEncoder = Literal['bmp', 'mjpeg', 'png', 'tiff', 'libwebp'] VideoEncoder = Literal['libx264', 'libx264rgb', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf', 'h264_qsv', 'hevc_qsv', 'h264_videotoolbox', 'hevc_videotoolbox', 'rawvideo'] @@ -491,6 +506,7 @@ StateKey = Literal\ 'reference_face_position', 'reference_face_distance', 'reference_frame_number', + 'face_tracker_score', 'face_occluder_model', 'face_parser_model', 'face_mask_types', @@ -503,6 +519,7 @@ StateKey = Literal\ 'trim_frame_end', 'temp_frame_format', 'keep_temp', + 'target_frame_amount', 'output_image_quality', 'output_image_scale', 'output_audio_encoder', @@ -556,13 +573,14 @@ State = TypedDict('State', 'face_landmarker_score' : Score, 'face_selector_mode' : FaceSelectorMode, 'face_selector_order' : FaceSelectorOrder, - 'face_selector_race' : Race, - 'face_selector_gender' : Gender, + 'face_selector_race' : FaceSelectorRace, + 'face_selector_gender' : FaceSelectorGender, 'face_selector_age_start' : int, 'face_selector_age_end' : int, 'reference_face_position' : int, 'reference_face_distance' : float, 'reference_frame_number' : int, + 'face_tracker_score' : Score, 'face_occluder_model' : FaceOccluderModel, 'face_parser_model' : FaceParserModel, 'face_mask_types' : List[FaceMaskType], @@ -575,6 +593,7 @@ State = TypedDict('State', 'trim_frame_end' : int, 'temp_frame_format' : TempFrameFormat, 'keep_temp' : bool, + 'target_frame_amount' : int, 'output_image_quality' : int, 'output_image_scale' : Scale, 'output_audio_encoder' : AudioEncoder, diff --git a/facefusion/vision.py b/facefusion/vision.py index 42291f72..a60290c5 100644 --- a/facefusion/vision.py +++ b/facefusion/vision.py @@ -1,6 +1,6 @@ import math from functools import lru_cache -from typing import List, Optional, Tuple +from typing import Dict, List, Optional, Tuple import cv2 import numpy @@ -9,7 +9,7 @@ from cv2.typing import Size from facefusion.common_helper import is_windows from facefusion.filesystem import get_file_extension, is_image, is_video from facefusion.media_helper import restrict_trim_frame -from facefusion.thread_helper import thread_semaphore +from facefusion.thread_helper import thread_lock, thread_semaphore from facefusion.types import ColorMode, Duration, Fps, Mask, Orientation, Resolution, Scale, VisionFrame from facefusion.video_manager import get_video_capture @@ -82,9 +82,10 @@ def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[Visio if video_capture and video_capture.isOpened(): frame_total = video_capture.get(cv2.CAP_PROP_FRAME_COUNT) + frame_position = min(frame_total, frame_number) with thread_semaphore(): - video_capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1)) + video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_position) has_vision_frame, vision_frame = video_capture.read() if has_vision_frame: @@ -93,6 +94,51 @@ def read_video_frame(video_path : str, frame_number : int = 0) -> Optional[Visio return None +@lru_cache(maxsize = 2) +def read_static_video_chunk(video_path : str, chunk_number : int, chunk_size : int) -> Dict[int, VisionFrame]: + return read_video_chunk(video_path, chunk_number, chunk_size) + + +def read_video_chunk(video_path : str, chunk_number : int, chunk_size : int) -> Dict[int, VisionFrame]: + video_frame_chunk = {} + + if is_video(video_path) and chunk_number > -1: + video_capture = get_video_capture(video_path) + + if video_capture and video_capture.isOpened(): + frame_total = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) + frame_position = chunk_number * chunk_size + + with thread_semaphore(): + video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_position) + + for frame_number in range(frame_position, min(frame_position + chunk_size, frame_total)): + has_vision_frame, vision_frame = video_capture.read() + + if has_vision_frame: + video_frame_chunk[frame_number] = vision_frame + + return video_frame_chunk + + +def select_video_frames(video_path : str, frame_number : int = 0, frame_offset : int = 5) -> List[VisionFrame]: + vision_frames = [] + chunk_size = (frame_offset * 2 + 1) * 4 + + if is_video(video_path): + with thread_lock(): + for current_number in range(frame_number - frame_offset, frame_number + frame_offset + 1): + video_frame_chunk = read_static_video_chunk(video_path, current_number // chunk_size, chunk_size) + vision_frame = create_empty_vision_frame() + + if current_number in video_frame_chunk: + vision_frame = video_frame_chunk.get(current_number) + + vision_frames.append(vision_frame) + + return vision_frames + + def count_video_frame_total(video_path : str) -> int: if is_video(video_path): video_capture = get_video_capture(video_path) @@ -290,6 +336,10 @@ def blend_vision_frames(source_vision_frame : VisionFrame, target_vision_frame : return blend_vision_frame +def create_empty_vision_frame() -> VisionFrame: + return numpy.zeros((1, 1, 3)).astype(numpy.uint8) + + def create_tile_frames(vision_frame : VisionFrame, size : Size) -> Tuple[List[VisionFrame], int, int]: tile_width = size[0] - 2 * size[2] pad_size_top = size[1] + size[2] diff --git a/facefusion/workflows/core.py b/facefusion/workflows/core.py index c7f6603e..fad5d486 100644 --- a/facefusion/workflows/core.py +++ b/facefusion/workflows/core.py @@ -1,4 +1,5 @@ from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import List import numpy from tqdm import tqdm @@ -10,7 +11,7 @@ from facefusion.filesystem import filter_audio_paths from facefusion.processors.core import get_processors_modules from facefusion.temp_helper import clear_temp_directory, create_temp_directory, resolve_temp_frame_paths from facefusion.types import AudioFrame, ErrorCode, VisionFrame -from facefusion.vision import conditional_merge_vision_mask, extract_vision_mask, read_static_image, read_static_images, read_static_video_frame, restrict_video_fps, write_image +from facefusion.vision import conditional_merge_vision_mask, extract_vision_mask, read_static_image, read_static_images, read_static_video_frame, restrict_video_fps, select_video_frames, write_image def is_process_stopping() -> bool: @@ -76,13 +77,19 @@ def conditional_get_reference_vision_frame() -> VisionFrame: return read_static_image(state_manager.get_item('target_path')) +def conditional_get_target_vision_frames(temp_frame_path : str, frame_number : int) -> List[VisionFrame]: + if state_manager.get_item('workflow') in [ 'image-to-video', 'image-to-video:frames' ]: + return select_video_frames(state_manager.get_item('target_path'), frame_number, state_manager.get_item('target_frame_amount')) + return [ read_static_image(temp_frame_path) ] + + def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: reference_vision_frame = conditional_get_reference_vision_frame() source_vision_frames = read_static_images(state_manager.get_item('source_paths')) - target_vision_frame = read_static_image(temp_frame_path, 'rgba') + target_vision_frames = conditional_get_target_vision_frames(temp_frame_path, frame_number) source_audio_frame = conditional_get_source_audio_frame(frame_number) source_voice_frame = conditional_get_source_voice_frame(frame_number) - temp_vision_frame = target_vision_frame.copy() + temp_vision_frame = read_static_image(temp_frame_path, 'rgba').copy() temp_vision_mask = extract_vision_mask(temp_vision_frame) for processor_module in get_processors_modules(state_manager.get_item('processors')): @@ -92,7 +99,7 @@ def process_temp_frame(temp_frame_path : str, frame_number : int) -> bool: 'source_vision_frames': source_vision_frames, 'source_audio_frame': source_audio_frame, 'source_voice_frame': source_voice_frame, - 'target_vision_frame': target_vision_frame[:, :, :3], + 'target_vision_frames': target_vision_frames, 'temp_vision_frame': temp_vision_frame[:, :, :3], 'temp_vision_mask': temp_vision_mask }) diff --git a/tests/test_common_helper.py b/tests/test_common_helper.py index 12e78853..b886a84e 100644 --- a/tests/test_common_helper.py +++ b/tests/test_common_helper.py @@ -1,4 +1,4 @@ -from facefusion.common_helper import calculate_float_step, calculate_int_step, create_float_metavar, create_float_range, create_int_metavar, create_int_range +from facefusion.common_helper import calculate_float_step, calculate_int_step, create_float_metavar, create_float_range, create_int_metavar, create_int_range, get_middle def test_create_int_metavar() -> None: @@ -25,3 +25,8 @@ def test_calc_int_step() -> None: def test_calc_float_step() -> None: assert calculate_float_step([ 0.1, 0.2 ]) == 0.1 + + +def test_get_middle() -> None: + assert get_middle([ 1, 2, 3, 4, 5 ]) == 3 + assert get_middle([ 1 ]) == 1 diff --git a/tests/test_face_analyser.py b/tests/test_face_analyser.py deleted file mode 100644 index 4c8eb54d..00000000 --- a/tests/test_face_analyser.py +++ /dev/null @@ -1,80 +0,0 @@ -import subprocess - -import pytest - -from facefusion import face_classifier, face_detector, face_landmarker, face_recognizer, state_manager -from facefusion.download import conditional_download -from facefusion.face_analyser import get_many_faces -from facefusion.face_store import clear_faces -from facefusion.vision import read_static_image -from .assert_helper import get_test_example_file, get_test_examples_directory - - -@pytest.fixture(scope = 'module', autouse = True) -def before_all() -> None: - state_manager.init_item('execution_device_ids', [ 0 ]) - state_manager.init_item('execution_providers', [ 'cpu' ]) - state_manager.init_item('download_providers', [ 'github' ]) - state_manager.init_item('face_detector_angles', [ 0 ]) - state_manager.init_item('face_detector_model', 'many') - state_manager.init_item('face_detector_score', 0.5) - state_manager.init_item('face_landmarker_model', 'many') - state_manager.init_item('face_landmarker_score', 0.5) - - face_classifier.pre_check() - face_landmarker.pre_check() - face_recognizer.pre_check() - - conditional_download(get_test_examples_directory(), - [ - 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' - ]) - - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.8:ih*0.8', get_test_example_file('source-80crop.jpg') ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.7:ih*0.7', get_test_example_file('source-70crop.jpg') ]) - subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.6:ih*0.6', get_test_example_file('source-60crop.jpg') ]) - - -@pytest.fixture(autouse = True) -def before_each() -> None: - face_classifier.clear_inference_pool() - face_detector.clear_inference_pool() - face_landmarker.clear_inference_pool() - face_recognizer.clear_inference_pool() - clear_faces() - - -@pytest.mark.parametrize('face_detector_model, face_detector_size', -[ - ('retinaface', '320x320'), - ('scrfd', '320x320'), - ('yolo_face', '640x640'), - ('yunet', '640x640') -]) -def test_get_one_face(face_detector_model : str, face_detector_size : str) -> None: - state_manager.init_item('face_detector_model', face_detector_model) - state_manager.init_item('face_detector_size', face_detector_size) - state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) - face_detector.pre_check() - - source_paths =\ - [ - get_test_example_file('source.jpg'), - get_test_example_file('source-80crop.jpg'), - get_test_example_file('source-70crop.jpg'), - get_test_example_file('source-60crop.jpg') - ] - - for source_path in source_paths: - source_frame = read_static_image(source_path) - many_faces = get_many_faces([ source_frame ]) - - assert len(many_faces) == 1 - - -def test_get_many_faces() -> None: - source_path = get_test_example_file('source.jpg') - source_frame = read_static_image(source_path) - many_faces = get_many_faces([ source_frame, source_frame, source_frame ]) - - assert len(many_faces) == 3 diff --git a/tests/test_face_creator.py b/tests/test_face_creator.py new file mode 100644 index 00000000..da7495a2 --- /dev/null +++ b/tests/test_face_creator.py @@ -0,0 +1,105 @@ +import subprocess + +import numpy +import pytest + +from facefusion import face_classifier, face_detector, face_landmarker, face_recognizer, state_manager +from facefusion.download import conditional_download +from facefusion.face_creator import average_face_geometry, get_many_faces, get_one_face, refill_faces +from facefusion.face_store import clear_faces +from facefusion.vision import read_static_image +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.8:ih*0.8', get_test_example_file('source-80crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.7:ih*0.7', get_test_example_file('source-70crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.6:ih*0.6', get_test_example_file('source-60crop.jpg') ]) + + state_manager.init_item('execution_device_ids', [ 0 ]) + state_manager.init_item('execution_providers', [ 'cpu' ]) + state_manager.init_item('download_providers', [ 'github' ]) + state_manager.init_item('face_detector_angles', [ 0 ]) + state_manager.init_item('face_detector_model', 'many') + state_manager.init_item('face_detector_size', '640x640') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + state_manager.init_item('face_detector_score', 0.5) + state_manager.init_item('face_landmarker_model', 'many') + state_manager.init_item('face_landmarker_score', 0.5) + + face_classifier.pre_check() + face_detector.pre_check() + face_landmarker.pre_check() + face_recognizer.pre_check() + + +@pytest.fixture(autouse = True) +def before_each() -> None: + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_recognizer.clear_inference_pool() + clear_faces() + + +def test_get_one_face() -> None: + source_vision_frame = read_static_image(get_test_example_file('source.jpg')) + face = get_one_face(get_many_faces([ source_vision_frame ])) + + assert face.bounding_box.size == 4 + + +def test_get_many_faces() -> None: + source_path = get_test_example_file('source.jpg') + source_vision_frame = read_static_image(source_path) + many_faces = get_many_faces([ source_vision_frame, source_vision_frame, source_vision_frame ]) + + assert len(many_faces) == 3 + + +def test_refill_faces() -> None: + source_vision_frame = read_static_image(get_test_example_file('source.jpg')) + face = get_one_face(get_many_faces([ source_vision_frame ])) + face_first = face._replace(bounding_box = numpy.array([ 0, 0, 10, 10 ])) + face_middle = face._replace(bounding_box = numpy.array([ 40, 40, 50, 50 ])) + face_last = face._replace(bounding_box = numpy.array([ 80, 80, 90, 90 ])) + + fill_faces = refill_faces([ face_first, None, face_last ]) + + assert fill_faces[0].bounding_box.tolist() == [ 0.0, 0.0, 10.0, 10.0 ] + assert fill_faces[1].bounding_box.tolist() == [ 40.0, 40.0, 50.0, 50.0 ] + assert fill_faces[2].bounding_box.tolist() == [ 80.0, 80.0, 90.0, 90.0 ] + + fill_faces = refill_faces([ face_first, None, None, None, face_last ]) + + assert fill_faces[0].bounding_box.tolist() == [ 0.0, 0.0, 10.0, 10.0 ] + assert fill_faces[1].bounding_box.tolist() == [ 20.0, 20.0, 30.0, 30.0 ] + assert fill_faces[2].bounding_box.tolist() == [ 40.0, 40.0, 50.0, 50.0 ] + assert fill_faces[3].bounding_box.tolist() == [ 60.0, 60.0, 70.0, 70.0 ] + assert fill_faces[4].bounding_box.tolist() == [ 80.0, 80.0, 90.0, 90.0 ] + + fill_faces = refill_faces([ face_first, None, face_middle, None, face_last ]) + + assert fill_faces[0].bounding_box.tolist() == [ 0.0, 0.0, 10.0, 10.0 ] + assert fill_faces[1].bounding_box.tolist() == [ 20.0, 20.0, 30.0, 30.0 ] + assert fill_faces[2].bounding_box.tolist() == [ 40.0, 40.0, 50.0, 50.0 ] + assert fill_faces[3].bounding_box.tolist() == [ 60.0, 60.0, 70.0, 70.0 ] + assert fill_faces[4].bounding_box.tolist() == [ 80.0, 80.0, 90.0, 90.0 ] + + +def test_average_face_geometry() -> None: + source_vision_frame = read_static_image(get_test_example_file('source.jpg')) + face_previous = get_one_face(get_many_faces([ source_vision_frame ])) + face_next = get_one_face(get_many_faces([ source_vision_frame ])) + face_previous = face_previous._replace(bounding_box = numpy.array([ 0, 0, 10, 10 ])) + face_next = face_next._replace(bounding_box = numpy.array([ 80, 80, 90, 90 ])) + + assert average_face_geometry([face_previous, face_next], 0.5).bounding_box.tolist() == [40.0, 40.0, 50.0, 50.0] + assert average_face_geometry([face_previous, face_next], 0.5).angle == face_next.angle + assert average_face_geometry([face_previous, face_next], 0.5).embedding is face_next.embedding + assert average_face_geometry([face_previous, face_next], 0.25).embedding is face_previous.embedding diff --git a/tests/test_face_detector.py b/tests/test_face_detector.py new file mode 100644 index 00000000..800e147f --- /dev/null +++ b/tests/test_face_detector.py @@ -0,0 +1,112 @@ +import subprocess + +import pytest + +from facefusion import face_detector, state_manager +from facefusion.download import conditional_download +from facefusion.face_detector import detect_with_retinaface, detect_with_scrfd, detect_with_yolo_face, detect_with_yunet +from facefusion.face_helper import apply_nms, get_nms_threshold +from facefusion.vision import read_static_image +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' + ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.8:ih*0.8', get_test_example_file('source-80crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.7:ih*0.7', get_test_example_file('source-70crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.6:ih*0.6', get_test_example_file('source-60crop.jpg') ]) + + state_manager.init_item('execution_device_ids', [ 0 ]) + state_manager.init_item('execution_providers', [ 'cpu' ]) + state_manager.init_item('download_providers', [ 'github' ]) + state_manager.init_item('face_detector_angles', [ 0 ]) + state_manager.init_item('face_detector_model', 'many') + state_manager.init_item('face_detector_score', 0.5) + + face_detector.pre_check() + + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/source.jpg' + ]) + + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.8:ih*0.8', get_test_example_file('source-80crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.7:ih*0.7', get_test_example_file('source-70crop.jpg') ]) + subprocess.run([ 'ffmpeg', '-i', get_test_example_file('source.jpg'), '-vf', 'crop=iw*0.6:ih*0.6', get_test_example_file('source-60crop.jpg') ]) + + +@pytest.fixture(autouse = True) +def before_each() -> None: + face_detector.clear_inference_pool() + + +def test_detect_with_retinaface() -> None: + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + bounding_boxes, face_scores, face_landmarks_5 = detect_with_retinaface(source_frame, '320x320') + keep_indices = apply_nms(bounding_boxes, face_scores, 0.5, get_nms_threshold('retinaface', [ 0 ])) + + assert len(keep_indices) == 1 + + +def test_detect_with_scrfd() -> None: + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + bounding_boxes, face_scores, face_landmarks_5 = detect_with_scrfd(source_frame, '320x320') + keep_indices = apply_nms(bounding_boxes, face_scores, 0.5, get_nms_threshold('scrfd', [ 0 ])) + + assert len(keep_indices) == 1 + + +def test_detect_with_yolo_face() -> None: + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + bounding_boxes, face_scores, face_landmarks_5 = detect_with_yolo_face(source_frame, '640x640') + keep_indices = apply_nms(bounding_boxes, face_scores, 0.5, get_nms_threshold('yolo_face', [ 0 ])) + + assert len(keep_indices) == 1 + + +def test_detect_with_yunet() -> None: + source_paths =\ + [ + get_test_example_file('source.jpg'), + get_test_example_file('source-80crop.jpg'), + get_test_example_file('source-70crop.jpg'), + get_test_example_file('source-60crop.jpg') + ] + + for source_path in source_paths: + source_frame = read_static_image(source_path) + bounding_boxes, face_scores, face_landmarks_5 = detect_with_yunet(source_frame, '640x640') + keep_indices = apply_nms(bounding_boxes, face_scores, 0.5, get_nms_threshold('yunet', [ 0 ])) + + assert len(keep_indices) == 1 diff --git a/tests/test_face_tracker.py b/tests/test_face_tracker.py new file mode 100644 index 00000000..0c655e6d --- /dev/null +++ b/tests/test_face_tracker.py @@ -0,0 +1,102 @@ +import numpy +import pytest + +from facefusion import face_classifier, face_detector, face_landmarker, face_recognizer, state_manager +from facefusion.common_helper import get_first, get_last +from facefusion.download import conditional_download +from facefusion.face_creator import get_many_faces, get_one_face +from facefusion.face_store import clear_faces +from facefusion.face_tracker import create_face_tracks, select_face_track, track_faces +from facefusion.vision import read_static_video_chunk, read_static_video_frame +from .assert_helper import get_test_example_file, get_test_examples_directory + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + conditional_download(get_test_examples_directory(), + [ + 'https://github.com/facefusion/facefusion-assets/releases/download/examples-3.0.0/target-240p.mp4' + ]) + + state_manager.init_item('execution_device_ids', [ 0 ]) + state_manager.init_item('execution_providers', [ 'cpu' ]) + state_manager.init_item('download_providers', [ 'github' ]) + state_manager.init_item('face_detector_angles', [ 0 ]) + state_manager.init_item('face_detector_model', 'yolo_face') + state_manager.init_item('face_detector_size', '640x640') + state_manager.init_item('face_detector_margin', (0, 0, 0, 0)) + state_manager.init_item('face_detector_score', 0.5) + state_manager.init_item('face_landmarker_model', 'many') + state_manager.init_item('face_landmarker_score', 0.5) + state_manager.init_item('face_tracker_score', 0.3) + + face_classifier.pre_check() + face_detector.pre_check() + face_landmarker.pre_check() + face_recognizer.pre_check() + + +@pytest.fixture(autouse = True) +def before_each() -> None: + face_classifier.clear_inference_pool() + face_detector.clear_inference_pool() + face_landmarker.clear_inference_pool() + face_recognizer.clear_inference_pool() + clear_faces() + + +def test_track_faces() -> None: + target_path = get_test_example_file('target-240p.mp4') + video_frame_chunk = read_static_video_chunk(target_path, 0, 7) + target_vision_frames = [ video_frame_chunk.get(frame_number) for frame_number in sorted(video_frame_chunk) ] + empty_vision_frame = numpy.zeros_like(get_first(target_vision_frames)) + + target_vision_frames[2] = empty_vision_frame + target_vision_frames[3] = empty_vision_frame + target_vision_frames[4] = empty_vision_frame + target_vision_frames[5] = empty_vision_frame + + assert len(track_faces(target_vision_frames, 0.3)) == 1 + + target_vision_frames = [ video_frame_chunk.get(frame_number) for frame_number in sorted(video_frame_chunk)[:5] ] + target_vision_frames[0] = empty_vision_frame + target_vision_frames[1] = empty_vision_frame + target_vision_frames[2] = empty_vision_frame + + assert len(track_faces(target_vision_frames, 0.3)) == 0 + + +def test_create_face_tracks() -> None: + target_vision_frame = read_static_video_frame(get_test_example_file('target-240p.mp4'), 0) + multi_face_vision_frame = numpy.hstack([ target_vision_frame, target_vision_frame ]) + + face_tracks = create_face_tracks([ target_vision_frame, target_vision_frame ], 0.3) + + assert len(face_tracks) == 1 + assert sorted(get_first(face_tracks)) == [ 0, 1 ] + + face_tracks = create_face_tracks([ multi_face_vision_frame, multi_face_vision_frame ], 0.3) + + assert len(face_tracks) == 2 + assert sorted(get_first(face_tracks)) == [ 0, 1 ] + assert sorted(get_last(face_tracks)) == [ 0, 1 ] + + assert len(create_face_tracks([ target_vision_frame, target_vision_frame ], 1.0)) == 2 + + +def test_select_face_track() -> None: + target_vision_frame = read_static_video_frame(get_test_example_file('target-240p.mp4'), 0) + face = get_one_face(get_many_faces([ target_vision_frame ])) + face_overlap = face._replace(bounding_box = numpy.array([ 12, 12, 52, 52 ])) + face_distant = face._replace(bounding_box = numpy.array([ 200, 200, 240, 240 ])) + face_track_overlap =\ + { + 0 : face._replace(bounding_box = numpy.array([ 10, 10, 50, 50 ])) + } + face_track_distant =\ + { + 0 : face._replace(bounding_box = numpy.array([ 100, 100, 140, 140 ])) + } + + assert select_face_track([ face_track_overlap, face_track_distant ], face_overlap, 0.3) is face_track_overlap + assert select_face_track([ face_track_overlap, face_track_distant ], face_distant, 0.3) == {} diff --git a/tests/test_ffmpeg_builder.py b/tests/test_ffmpeg_builder.py index 9805a5f1..43cd9104 100644 --- a/tests/test_ffmpeg_builder.py +++ b/tests/test_ffmpeg_builder.py @@ -1,7 +1,7 @@ from shutil import which from facefusion import ffmpeg_builder -from facefusion.ffmpeg_builder import capture_video, chain, concat, enforce_pixel_format, keep_video_alpha, run, select_frame_range, set_audio_quality, set_audio_sample_size, set_stream_mode, set_stream_quality, set_video_encoder, set_video_fps, set_video_quality +from facefusion.ffmpeg_builder import capture_video, chain, concat, enforce_pixel_format, keep_video_alpha, run, select_frame_range, set_audio_quality, set_audio_sample_size, set_faststart, set_stream_mode, set_stream_quality, set_video_encoder, set_video_fps, set_video_quality, set_video_tag def test_run() -> None: @@ -71,6 +71,24 @@ def test_set_audio_quality() -> None: assert set_audio_quality('flac', 100) == [] +def test_set_faststart() -> None: + assert set_faststart('m4v') == [ '-movflags', '+faststart' ] + assert set_faststart('mov') == [ '-movflags', '+faststart' ] + assert set_faststart('mp4') == [ '-movflags', '+faststart' ] + assert set_faststart('mkv') == [] + assert set_faststart('webm') == [] + + +def test_set_video_tag() -> None: + assert set_video_tag('libx265', 'm4v') == [ '-tag:v', 'hvc1' ] + assert set_video_tag('hevc_nvenc', 'mov') == [ '-tag:v', 'hvc1' ] + assert set_video_tag('hevc_videotoolbox', 'mp4') == [ '-tag:v', 'hvc1' ] + assert set_video_tag('libx265', 'mkv') == [] + assert set_video_tag('libx265', 'webm') == [] + assert set_video_tag('libx264', 'mp4') == [] + assert set_video_tag('h264_nvenc', 'mp4') == [] + + def test_set_video_quality() -> None: assert set_video_quality('libx264', 0) == [ '-crf', '51' ] assert set_video_quality('libx264', 50) == [ '-crf', '26' ] diff --git a/tests/test_vision.py b/tests/test_vision.py index c85d4a04..ec4b8e77 100644 --- a/tests/test_vision.py +++ b/tests/test_vision.py @@ -1,9 +1,10 @@ import subprocess +import numpy import pytest from facefusion.download import conditional_download -from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_frame, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, unpack_resolution, write_image +from facefusion.vision import calculate_histogram_difference, count_video_frame_total, detect_image_resolution, detect_video_duration, detect_video_fps, detect_video_resolution, match_frame_color, normalize_resolution, pack_resolution, predict_video_frame_total, read_image, read_video_chunk, read_video_frame, restrict_image_resolution, restrict_trim_video_frame, restrict_video_fps, restrict_video_resolution, scale_resolution, select_video_frames, unpack_resolution, write_image from .assert_helper import get_test_example_file, get_test_examples_directory, get_test_output_path, prepare_test_output_directory @@ -62,10 +63,27 @@ def test_restrict_image_resolution() -> None: def test_read_video_frame() -> None: - assert hasattr(read_video_frame(get_test_example_file('target-240p-25fps.mp4')), '__array_interface__') + target_path = get_test_example_file('target-240p-25fps.mp4') + + assert read_video_frame(target_path).shape == (226, 426, 3) + assert numpy.array_equal(read_video_frame(target_path, 49), select_video_frames(target_path, 49, 5)[5]) + assert numpy.array_equal(read_video_frame(target_path, 50), select_video_frames(target_path, 50, 5)[5]) + assert numpy.array_equal(read_video_frame(target_path, 51), select_video_frames(target_path, 51, 5)[5]) assert read_video_frame('invalid') is None +def test_read_video_chunk() -> None: + assert len(read_video_chunk(get_test_example_file('target-240p-25fps.mp4'), 1, 40)) == 40 + assert read_video_chunk('invalid', 1, 40) == {} + + +def test_select_video_frames() -> None: + assert len(select_video_frames(get_test_example_file('target-240p-25fps.mp4'), 50, 5)) == 11 + assert len(select_video_frames(get_test_example_file('target-240p-25fps.mp4'), 1, 5)) == 11 + assert len(select_video_frames(get_test_example_file('target-240p-25fps.mp4'), 269, 5)) == 11 + assert select_video_frames('invalid', 50, 5) == [] + + def test_count_video_frame_total() -> None: assert count_video_frame_total(get_test_example_file('target-240p-25fps.mp4')) == 270 assert count_video_frame_total(get_test_example_file('target-240p-30fps.mp4')) == 324