diff --git a/README.md b/README.md index aad9db1..da8fdbb 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,24 @@ # Faceless Node for ComfyUI -> Morden face toolkit for ComfyUI. +> Next generation face toolkit for ComfyUI. + +## Installation + +### Install custom node + +```bash +git clone https://github.com/jeffy5/comfyui-faceless-node +pip install -r requirements.txt +``` + +### Download model + +All models is same as [facefusion](https://github.com/facefusion/facefusion). + +Download what your needs models from [facefusion assets](https://github.com/facefusion/facefusion-assets) + +Put all into directory `CmofyUI/models/faceless`. + +## Thanks and Credits + +Thanks to [Facefusion](https://github.com/facefusion/facefusion). This project is based on facefusion to implenment a special version for ComfyUI. diff --git a/faceless/execution.py b/faceless/execution.py new file mode 100644 index 0000000..1a8b962 --- /dev/null +++ b/faceless/execution.py @@ -0,0 +1,81 @@ +from functools import lru_cache +import subprocess +import xml.etree.ElementTree as ElementTree +from typing import List, Any + +from .typing import ValueAndUnit, ExecutionDevice + +def apply_execution_provider_options(execution_providers: List[str]) -> List[Any]: + execution_providers_with_options : List[Any] = [] + + for execution_provider in execution_providers: + if execution_provider == 'CUDAExecutionProvider': + execution_providers_with_options.append((execution_provider, + { + 'cudnn_conv_algo_search': 'EXHAUSTIVE' if use_exhaustive() else 'DEFAULT' + })) + else: + execution_providers_with_options.append(execution_provider) + return execution_providers_with_options + + +def use_exhaustive() -> bool: + execution_devices = detect_static_execution_devices() + product_names = ('GeForce GTX 1630', 'GeForce GTX 1650', 'GeForce GTX 1660') + + return any(execution_device.get('product').get('name').startswith(product_names) for execution_device in execution_devices) + + +def run_nvidia_smi() -> subprocess.Popen[bytes]: + commands = [ 'nvidia-smi', '--query', '--xml-format' ] + return subprocess.Popen(commands, stdout = subprocess.PIPE) + +@lru_cache(maxsize = None) +def detect_static_execution_devices() -> List[ExecutionDevice]: + return detect_execution_devices() + +def detect_execution_devices() -> List[ExecutionDevice]: + execution_devices : List[ExecutionDevice] = [] + try: + output, _ = run_nvidia_smi().communicate() + root_element = ElementTree.fromstring(output) + except Exception: + root_element = ElementTree.Element('xml') + + for gpu_element in root_element.findall('gpu'): + execution_devices.append( + { + 'driver_version': root_element.find('driver_version').text, + 'framework': + { + 'name': 'CUDA', + 'version': root_element.find('cuda_version').text, + }, + 'product': + { + 'vendor': 'NVIDIA', + 'name': gpu_element.find('product_name').text.replace('NVIDIA ', ''), + 'architecture': gpu_element.find('product_architecture').text, + }, + 'video_memory': + { + 'total': create_value_and_unit(gpu_element.find('fb_memory_usage/total').text), + 'free': create_value_and_unit(gpu_element.find('fb_memory_usage/free').text) + }, + 'utilization': + { + 'gpu': create_value_and_unit(gpu_element.find('utilization/gpu_util').text), + 'memory': create_value_and_unit(gpu_element.find('utilization/memory_util').text) + } + }) + return execution_devices + +def create_value_and_unit(text : str) -> ValueAndUnit: + value, unit = text.split() + value_and_unit : ValueAndUnit =\ + { + 'value': value, + 'unit': unit + } + + return value_and_unit diff --git a/faceless/face_helper.py b/faceless/face_helper.py new file mode 100644 index 0000000..515bef5 --- /dev/null +++ b/faceless/face_helper.py @@ -0,0 +1,140 @@ +from functools import lru_cache +from typing import Any, Tuple, List + +import numpy +import cv2 +from cv2.typing import Size + +from .typing import BoundingBox, FaceLandmark5, FaceLandmark68, VisionFrame, WarpTemplateSet, WarpTemplate, Matrix, Translation, FaceAnalyserGender, FaceAnalyserAge, Mask + +WARP_TEMPLATES : WarpTemplateSet =\ +{ + 'arcface_112_v1': numpy.array( + [ + [ 0.35473214, 0.45658929 ], + [ 0.64526786, 0.45658929 ], + [ 0.50000000, 0.61154464 ], + [ 0.37913393, 0.77687500 ], + [ 0.62086607, 0.77687500 ] + ]), + 'arcface_112_v2': numpy.array( + [ + [ 0.34191607, 0.46157411 ], + [ 0.65653393, 0.45983393 ], + [ 0.50022500, 0.64050536 ], + [ 0.37097589, 0.82469196 ], + [ 0.63151696, 0.82325089 ] + ]), + 'arcface_128_v2': numpy.array( + [ + [ 0.36167656, 0.40387734 ], + [ 0.63696719, 0.40235469 ], + [ 0.50019687, 0.56044219 ], + [ 0.38710391, 0.72160547 ], + [ 0.61507734, 0.72034453 ] + ]), + 'ffhq_512': numpy.array( + [ + [ 0.37691676, 0.46864664 ], + [ 0.62285697, 0.46912813 ], + [ 0.50123859, 0.61331904 ], + [ 0.39308822, 0.72541100 ], + [ 0.61150205, 0.72490465 ] + ]) +} + +@lru_cache(maxsize = None) +def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> numpy.ndarray[Any, Any]: + y, x = numpy.mgrid[:stride_height, :stride_width][::-1] + anchors = numpy.stack((y, x), axis = -1) + anchors = (anchors * feature_stride).reshape((-1, 2)) + anchors = numpy.stack([ anchors ] * anchor_total, axis = 1).reshape((-1, 2)) + return anchors + +def distance_to_bounding_box(points : numpy.ndarray[Any, Any], distance : numpy.ndarray[Any, Any]) -> BoundingBox: + x1 = points[:, 0] - distance[:, 0] + y1 = points[:, 1] - distance[:, 1] + x2 = points[:, 0] + distance[:, 2] + y2 = points[:, 1] + distance[:, 3] + bounding_box = numpy.column_stack([ x1, y1, x2, y2 ]) + return bounding_box + +def distance_to_face_landmark_5(points : numpy.ndarray[Any, Any], distance : numpy.ndarray[Any, Any]) -> FaceLandmark5: + x = points[:, 0::2] + distance[:, 0::2] + y = points[:, 1::2] + distance[:, 1::2] + face_landmark_5 = numpy.stack((x, y), axis = -1) + return face_landmark_5 + +def estimate_matrix_by_face_landmark_5(face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Matrix: + normed_warp_template = WARP_TEMPLATES[warp_template] * crop_size + affine_matrix = cv2.estimateAffinePartial2D(face_landmark_5, normed_warp_template, method = cv2.RANSAC, ransacReprojThreshold = 100)[0] + return affine_matrix + +def warp_face_by_face_landmark_5(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Tuple[VisionFrame, Matrix]: + affine_matrix = estimate_matrix_by_face_landmark_5(face_landmark_5, warp_template, crop_size) + crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size, borderMode = cv2.BORDER_REPLICATE, flags = cv2.INTER_AREA) + return crop_vision_frame, affine_matrix + +def warp_face_by_translation(temp_vision_frame : VisionFrame, translation : Translation, scale : float, crop_size : Size) -> Tuple[VisionFrame, Matrix]: + affine_matrix = numpy.array([ [ scale, 0, translation[0] ], [ 0, scale, translation[1] ] ]) + crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size) + return crop_vision_frame, affine_matrix + +def categorize_age(age : int) -> FaceAnalyserAge: + if age < 13: + return 'child' + elif age < 19: + return 'teen' + elif age < 60: + return 'adult' + return 'senior' + +def categorize_gender(gender : int) -> FaceAnalyserGender: + if gender == 0: + return 'female' + return 'male' + +def apply_nms(bounding_box_list : List[BoundingBox], iou_threshold : float) -> List[int]: + keep_indices = [] + dimension_list = numpy.reshape(bounding_box_list, (-1, 4)) + x1 = dimension_list[:, 0] + y1 = dimension_list[:, 1] + x2 = dimension_list[:, 2] + y2 = dimension_list[:, 3] + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + indices = numpy.arange(len(bounding_box_list)) + while indices.size > 0: + index = indices[0] + remain_indices = indices[1:] + keep_indices.append(index) + xx1 = numpy.maximum(x1[index], x1[remain_indices]) + yy1 = numpy.maximum(y1[index], y1[remain_indices]) + xx2 = numpy.minimum(x2[index], x2[remain_indices]) + yy2 = numpy.minimum(y2[index], y2[remain_indices]) + width = numpy.maximum(0, xx2 - xx1 + 1) + height = numpy.maximum(0, yy2 - yy1 + 1) + iou = width * height / (areas[index] + areas[remain_indices] - width * height) + indices = indices[numpy.where(iou <= iou_threshold)[0] + 1] + return keep_indices + +def convert_face_landmark_68_to_5(face_landmark_68 : FaceLandmark68) -> FaceLandmark5: + face_landmark_5 = numpy.array( + [ + numpy.mean(face_landmark_68[36:42], axis = 0), + numpy.mean(face_landmark_68[42:48], axis = 0), + face_landmark_68[30], + face_landmark_68[48], + face_landmark_68[54] + ]) + return face_landmark_5 + +def paste_back(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, crop_mask : Mask, affine_matrix : Matrix) -> VisionFrame: + inverse_matrix = cv2.invertAffineTransform(affine_matrix) + temp_size = temp_vision_frame.shape[:2][::-1] + inverse_mask = cv2.warpAffine(crop_mask, inverse_matrix, temp_size).clip(0, 1) + inverse_vision_frame = cv2.warpAffine(crop_vision_frame, inverse_matrix, temp_size, borderMode = cv2.BORDER_REPLICATE) + paste_vision_frame = temp_vision_frame.copy() + paste_vision_frame[:, :, 0] = inverse_mask * inverse_vision_frame[:, :, 0] + (1 - inverse_mask) * temp_vision_frame[:, :, 0] + paste_vision_frame[:, :, 1] = inverse_mask * inverse_vision_frame[:, :, 1] + (1 - inverse_mask) * temp_vision_frame[:, :, 1] + paste_vision_frame[:, :, 2] = inverse_mask * inverse_vision_frame[:, :, 2] + (1 - inverse_mask) * temp_vision_frame[:, :, 2] + return paste_vision_frame diff --git a/faceless/face_masker.py b/faceless/face_masker.py new file mode 100644 index 0000000..40a5542 --- /dev/null +++ b/faceless/face_masker.py @@ -0,0 +1,129 @@ +from typing import Any, Dict, List +from cv2.typing import Size +from functools import lru_cache +import threading +import cv2 +import numpy +import onnxruntime + +from .typing import FaceLandmark68, VisionFrame, Mask, Padding, FaceMaskRegion, ModelSet +from .execution import apply_execution_provider_options +from .filesystem import resolve_relative_path + +# TODO load from options +execution_providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider'] + +FACE_OCCLUDER = None +FACE_PARSER = None +THREAD_LOCK : threading.Lock = threading.Lock() +MODELS : ModelSet =\ +{ + 'face_occluder': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_occluder.onnx', + 'path': resolve_relative_path('../../../models/faceless/face_occluder.onnx') + }, + 'face_parser': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_parser.onnx', + 'path': resolve_relative_path('../../../models/faceless/face_parser.onnx') + } +} +FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\ +{ + 'skin': 1, + 'left-eyebrow': 2, + 'right-eyebrow': 3, + 'left-eye': 4, + 'right-eye': 5, + 'glasses': 6, + 'nose': 10, + 'mouth': 11, + 'upper-lip': 12, + 'lower-lip': 13 +} + + +def get_face_occluder() -> Any: + global FACE_OCCLUDER + + with THREAD_LOCK: + if FACE_OCCLUDER is None: + model_path = MODELS['face_occluder']['path'] + FACE_OCCLUDER = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(execution_providers)) + return FACE_OCCLUDER + + +def get_face_parser() -> Any: + global FACE_PARSER + + with THREAD_LOCK: + if FACE_PARSER is None: + model_path = MODELS['face_parser']['path'] + FACE_PARSER = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(execution_providers)) + return FACE_PARSER + + +def clear_face_occluder() -> None: + global FACE_OCCLUDER + + FACE_OCCLUDER = None + + +def clear_face_parser() -> None: + global FACE_PARSER + + FACE_PARSER = None + + +@lru_cache(maxsize = None) +def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_padding : Padding) -> Mask: + blur_amount = int(crop_size[0] * 0.5 * face_mask_blur) + blur_area = max(blur_amount // 2, 1) + box_mask : Mask = numpy.ones(crop_size, numpy.float32) + box_mask[:max(blur_area, int(crop_size[1] * face_mask_padding[0] / 100)), :] = 0 + box_mask[-max(blur_area, int(crop_size[1] * face_mask_padding[2] / 100)):, :] = 0 + box_mask[:, :max(blur_area, int(crop_size[0] * face_mask_padding[3] / 100))] = 0 + box_mask[:, -max(blur_area, int(crop_size[0] * face_mask_padding[1] / 100)):] = 0 + if blur_amount > 0: + box_mask = cv2.GaussianBlur(box_mask, (0, 0), blur_amount * 0.25) + return box_mask + + +def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask: + face_occluder = get_face_occluder() + prepare_vision_frame = cv2.resize(crop_vision_frame, face_occluder.get_inputs()[0].shape[1:3][::-1]) + prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0).astype(numpy.float32) / 255 + prepare_vision_frame = prepare_vision_frame.transpose(0, 1, 2, 3) + occlusion_mask : Mask = face_occluder.run(None, + { + face_occluder.get_inputs()[0].name: prepare_vision_frame + })[0][0] + occlusion_mask = occlusion_mask.transpose(0, 1, 2).clip(0, 1).astype(numpy.float32) + occlusion_mask = cv2.resize(occlusion_mask, crop_vision_frame.shape[:2][::-1]) + occlusion_mask = (cv2.GaussianBlur(occlusion_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 + return occlusion_mask + + +def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask: + face_parser = get_face_parser() + prepare_vision_frame = cv2.flip(cv2.resize(crop_vision_frame, (512, 512)), 1) + prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0).astype(numpy.float32)[:, :, ::-1] / 127.5 - 1 + prepare_vision_frame = prepare_vision_frame.transpose(0, 3, 1, 2) + region_mask : Mask = face_parser.run(None, + { + face_parser.get_inputs()[0].name: prepare_vision_frame + })[0][0] + region_mask = numpy.isin(region_mask.argmax(0), [ FACE_MASK_REGIONS[region] for region in face_mask_regions ]) + region_mask = cv2.resize(region_mask.astype(numpy.float32), crop_vision_frame.shape[:2][::-1]) + region_mask = (cv2.GaussianBlur(region_mask.clip(0, 1), (0, 0), 5).clip(0.5, 1) - 0.5) * 2 + return region_mask + + +def create_mouth_mask(face_landmark_68 : FaceLandmark68) -> Mask: + convex_hull = cv2.convexHull(face_landmark_68[numpy.r_[3:14, 31:36]].astype(numpy.int32)) + mouth_mask : Mask = numpy.zeros((512, 512)).astype(numpy.float32) + mouth_mask = cv2.fillConvexPoly(mouth_mask, convex_hull, 1.0) + mouth_mask = cv2.erode(mouth_mask.clip(0, 1), numpy.ones((21, 3))) + mouth_mask = cv2.GaussianBlur(mouth_mask, (0, 0), sigmaX = 1, sigmaY = 15) + return mouth_mask diff --git a/faceless/face_store.py b/faceless/face_store.py new file mode 100644 index 0000000..7dfca7f --- /dev/null +++ b/faceless/face_store.py @@ -0,0 +1,48 @@ +from typing import Optional, List +import hashlib +import numpy + +from .typing import VisionFrame, Face, FaceStore, FaceSet + +FACE_STORE: FaceStore =\ +{ + 'static_faces': {}, + 'reference_faces': {} +} + + +def get_static_faces(vision_frame : VisionFrame) -> Optional[List[Face]]: + frame_hash = create_frame_hash(vision_frame) + if frame_hash in FACE_STORE['static_faces']: + return FACE_STORE['static_faces'][frame_hash] + return None + + +def set_static_faces(vision_frame : VisionFrame, faces : List[Face]) -> None: + frame_hash = create_frame_hash(vision_frame) + if frame_hash: + FACE_STORE['static_faces'][frame_hash] = faces + + +def clear_static_faces() -> None: + FACE_STORE['static_faces'] = {} + + +def create_frame_hash(vision_frame : VisionFrame) -> Optional[str]: + return hashlib.sha1(vision_frame.tobytes()).hexdigest() if numpy.any(vision_frame) else None + + +def get_reference_faces() -> Optional[FaceSet]: + if FACE_STORE['reference_faces']: + return FACE_STORE['reference_faces'] + return None + + +def append_reference_face(name : str, face : Face) -> None: + if name not in FACE_STORE['reference_faces']: + FACE_STORE['reference_faces'][name] = [] + FACE_STORE['reference_faces'][name].append(face) + + +def clear_reference_faces() -> None: + FACE_STORE['reference_faces'] = {} diff --git a/faceless/ffmpeg.py b/faceless/ffmpeg.py index e5972f7..4050886 100644 --- a/faceless/ffmpeg.py +++ b/faceless/ffmpeg.py @@ -2,9 +2,9 @@ from typing import List, Optional import subprocess -from .vision import pack_resolution +from .vision import pack_resolution, restrict_video_fps -from .typing import Fps, FrameFormat, Resolution +from .typing import Fps, FrameFormat, OutputVideoEncoder, OutputVideoPreset, Resolution from .filesystem import get_temp_frames_pattern def run_ffmpeg(args : List[str]): @@ -12,7 +12,18 @@ def run_ffmpeg(args : List[str]): commands.extend(args) process = subprocess.Popen(commands, stderr = subprocess.PIPE, stdout = subprocess.PIPE) # TODO Make timeout be a option? - return process.wait(timeout = 300) == 0 + code = process.wait(timeout = 300) + + if code != 0: + msg = f"code: {code}" + stderr = "stderr: " + stdout = "stdout: " + if process.stderr is not None: + stderr += process.stderr.readline().decode('utf-8') + if process.stdout is not None: + stdout += process.stdout.readline().decode('utf-8') + print(', '.join([msg, stderr, stdout])) + return code == 0 def extract_frames(video_path: str, frames_path: str, video_resolution : Resolution, video_fps : Fps, trim_frame_start : Optional[int] = None, trim_frame_end: Optional[int] = None, frame_format: FrameFormat = 'png') -> bool: # TODO Get frame image format from options. @@ -31,3 +42,32 @@ def extract_frames(video_path: str, frames_path: str, video_resolution : Resolut commands.extend([ '-vf', 'scale=' + format_resolution + ',fps=' + str(video_fps) ]) commands.extend([ '-vsync', '0', temp_frames_pattern ]) return run_ffmpeg(commands) + +def merge_video(target_path: str, output_path: str, video_resolution: Resolution, video_fps: Fps, output_video_encoder: OutputVideoEncoder = 'libx264', output_video_quality: int = 80, output_video_preset: OutputVideoPreset = 'veryfast', frame_format: FrameFormat = 'png') -> bool: + temp_video_fps = restrict_video_fps(target_path, video_fps) + temp_frames_pattern = get_temp_frames_pattern(target_path, '%04d', frame_format) + commands = [ '-hwaccel', 'auto', '-s', pack_resolution(video_resolution), '-r', str(temp_video_fps), '-i', temp_frames_pattern, '-c:v', output_video_encoder ] + + if output_video_encoder in [ 'libx264', 'libx265' ]: + output_video_compression = round(51 - (output_video_quality * 0.51)) + commands.extend([ '-crf', str(output_video_compression), '-preset', output_video_preset ]) + if output_video_encoder in [ 'libvpx-vp9' ]: + output_video_compression = round(63 - (output_video_quality * 0.63)) + commands.extend([ '-crf', str(output_video_compression) ]) + if output_video_encoder in [ 'h264_nvenc', 'hevc_nvenc' ]: + output_video_compression = round(51 - (output_video_quality * 0.51)) + commands.extend([ '-cq', str(output_video_compression), '-preset', output_video_preset ]) + if output_video_encoder in [ 'h264_amf', 'hevc_amf' ]: + output_video_compression = round(51 - (output_video_quality * 0.51)) + commands.extend([ '-qp_i', str(output_video_compression), '-qp_p', str(output_video_compression), '-quality', map_amf_preset(output_video_preset) ]) + commands.extend([ '-vf', 'framerate=fps=' + str(video_fps), '-pix_fmt', 'yuv420p', '-colorspace', 'bt709', '-y', output_path ]) + return run_ffmpeg(commands) + +def map_amf_preset(output_video_preset : OutputVideoPreset) -> str: + if output_video_preset in [ 'ultrafast', 'superfast', 'veryfast' ]: + return 'speed' + if output_video_preset in [ 'faster', 'fast', 'medium' ]: + return 'balanced' + if output_video_preset in [ 'slow', 'slower', 'veryslow' ]: + return 'quality' + return 'balanced' diff --git a/faceless/filesystem.py b/faceless/filesystem.py index 40bb6ac..3001af5 100644 --- a/faceless/filesystem.py +++ b/faceless/filesystem.py @@ -16,4 +16,5 @@ def is_video(video_path : str) -> bool: def get_temp_frames_pattern(target_path : str, temp_frame_prefix : str, format: FrameFormat) -> str: return os.path.join(target_path, temp_frame_prefix + '.' + format) - +def resolve_relative_path(path : str) -> str: + return os.path.abspath(os.path.join(os.path.dirname(__file__), path)) diff --git a/faceless/nodes/nodes_face_swap.py b/faceless/nodes/nodes_face_swap.py index 533cc7f..4fe2c85 100644 --- a/faceless/nodes/nodes_face_swap.py +++ b/faceless/nodes/nodes_face_swap.py @@ -1,9 +1,23 @@ +import os +import shutil +import time +from PIL import Image, ImageOps, ImageSequence + +import numpy as np +import torch + +import folder_paths + +from ..vision import is_image +from ..processors.face_swapper import process_images + class NodesFaceSwap: @classmethod def INPUT_TYPES(cls): return { "required": { - "images": ("IMAGE",), + "source_image": ("IMAGE",), + "target_images": ("IMAGE",), }, } @@ -13,6 +27,31 @@ class NodesFaceSwap: RETURN_NAMES = ("IMAGE",) FUNCTION = "swapFace" - def swapFace(self, images): - print("swap face") - return (images,) + def swapFace(self, source_image, target_images): + now = f"{int(time.time())}" + output_path = os.path.join(folder_paths.get_temp_directory(), "faceless/swapped_frames", now) + if os.path.exists(output_path): + shutil.rmtree(output_path) + os.makedirs(output_path) + + process_images(source_image[0], target_images, output_path) + + images = [] + for file in sorted(os.listdir(output_path)): + file_path = os.path.join(output_path, file) + if not is_image(file_path): + continue + img = Image.open(file_path) + for i in ImageSequence.Iterator(img): + i = ImageOps.exif_transpose(i) + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + images.append(image) + if len(images) > 1: + output_image = torch.cat(images, dim=0) + else: + output_image = images[0] + return (output_image,) diff --git a/faceless/nodes/nodes_load_frames.py b/faceless/nodes/nodes_load_frames.py index a2c17e8..a524c42 100644 --- a/faceless/nodes/nodes_load_frames.py +++ b/faceless/nodes/nodes_load_frames.py @@ -1,4 +1,5 @@ import os + from PIL import Image, ImageOps, ImageSequence import torch import numpy as np @@ -7,7 +8,6 @@ from ..filesystem import is_image from ..typing import FacelessVideo class NodesLoadFrames: - @classmethod def INPUT_TYPES(cls): return { @@ -21,7 +21,7 @@ class NodesLoadFrames: FUNCTION = "load_frames" def load_frames(self, video: FacelessVideo): - frames_path = video['frames_path'] + frames_path = video['output_path'] images = [] for file in sorted(os.listdir(frames_path)): @@ -37,7 +37,6 @@ class NodesLoadFrames: image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] images.append(image) - images.append(image) if len(images) > 1: output_image = torch.cat(images, dim=0) else: diff --git a/faceless/nodes/nodes_load_video.py b/faceless/nodes/nodes_load_video.py index 958192b..e3ffcd9 100644 --- a/faceless/nodes/nodes_load_video.py +++ b/faceless/nodes/nodes_load_video.py @@ -3,7 +3,7 @@ import folder_paths import os import shutil -from ..filesystem import is_image, is_video +from ..filesystem import is_video from ..ffmpeg import extract_frames from ..vision import detect_video_fps, detect_video_resolution from ..typing import FacelessVideo @@ -20,11 +20,13 @@ class NodesLoadVideo: "trim_frame_start": ("INT", { "default": -1, "min": -1, + "max": 999999, "display": "number", }), "trim_frame_end": ("INT", { "default": -1, "min": -1, + "max": 999999, "display": "number", }), }, @@ -44,7 +46,7 @@ class NodesLoadVideo: def process(self, video, trim_frame_start: int, trim_frame_end: int): video_path = folder_paths.get_annotated_filepath(video) video_name, _ = os.path.splitext(os.path.basename(video_path)) - frames_path = os.path.join(folder_paths.get_temp_directory(), "faceless_frames", video_name) + frames_path = os.path.join(folder_paths.get_temp_directory(), "faceless/frames", video_name) print("frames path: " + frames_path) # Remove all cached frames @@ -57,7 +59,6 @@ class NodesLoadVideo: if video_resolution is None or video_fps is None: raise Exception("Failed to detect video resolution and fps") - # TODO Get trim start, trim end and frame format from options. if trim_frame_start == -1: final_trim_frame_start = None else: @@ -71,7 +72,9 @@ class NodesLoadVideo: raise Exception("Failed to extract frames") faceless_video: FacelessVideo = { - "video_path": video_path, - "frames_path": frames_path + 'video_path': video_path, + 'output_path': frames_path, + 'resolution': video_resolution, + 'fps': video_fps, } return (faceless_video,) diff --git a/faceless/nodes/nodes_save_video.py b/faceless/nodes/nodes_save_video.py index 68fc972..13a765a 100644 --- a/faceless/nodes/nodes_save_video.py +++ b/faceless/nodes/nodes_save_video.py @@ -1,3 +1,12 @@ +import os +import time + +import folder_paths + +from ..ffmpeg import merge_video + +from ..typing import FacelessVideo + class NodesSaveVideo: @classmethod @@ -13,6 +22,21 @@ class NodesSaveVideo: FUNCTION = "save_video" OUTPUT_NODE = True - def save_video(self): - pass + def save_video(self, video: FacelessVideo): + video_path = video.get("video_path") + frames_path = video.get("output_path") + now = int(time.time()) + output_path = os.path.join(folder_paths.get_output_directory(), "faceless") + if not os.path.exists(output_path): + os.makedirs(output_path) + output_filepath = os.path.join(output_path, f"{now}_" + os.path.basename(video_path)) + + resolution = video.get("resolution") + fps = video.get("fps") + + if not merge_video(frames_path, output_filepath, resolution, fps): + raise Exception("Failed to merge video") + + # TODO Restore audio + return () diff --git a/faceless/nodes/nodes_video_face_swap.py b/faceless/nodes/nodes_video_face_swap.py index b2164ef..2f4541d 100644 --- a/faceless/nodes/nodes_video_face_swap.py +++ b/faceless/nodes/nodes_video_face_swap.py @@ -1,9 +1,18 @@ +import os +import shutil + +import folder_paths + +from ..processors.face_swapper import process_frames +from ..typing import FacelessVideo + class NodesVideoFaceSwap: @classmethod def INPUT_TYPES(cls): return { "required": { - "video": ("FACELESS_VIDEO",), + "source_image": ("IMAGE",), + "target_video": ("FACELESS_VIDEO",), }, } @@ -13,6 +22,22 @@ class NodesVideoFaceSwap: RETURN_NAMES = ("video",) FUNCTION = "swapVideoFace" - def swapVideoFace(self, video): - print("video face swap" + video.video_path + " " + video.frames_path) - return (video,) + def swapVideoFace(self, source_image, target_video: FacelessVideo): + video_path = target_video.get("video_path") + video_name, _ = os.path.splitext(os.path.basename(video_path)) + + frames_path = os.path.join(folder_paths.get_temp_directory(), "faceless/frames", video_name) + output_path = os.path.join(folder_paths.get_temp_directory(), "faceless/swapped_frames", video_name) + + if os.path.exists(output_path): + shutil.rmtree(output_path) + os.makedirs(output_path) + + # TODO Check if has face on source image + + print("source image", len(source_image)) + # Fetch source image or change process_frames argument. + process_frames(source_image[0], frames_path, output_path) + + target_video['output_path'] = output_path + return (target_video,) diff --git a/faceless/processors/face_analyser.py b/faceless/processors/face_analyser.py new file mode 100644 index 0000000..814b814 --- /dev/null +++ b/faceless/processors/face_analyser.py @@ -0,0 +1,510 @@ +from typing import List, Optional, Tuple, Any + +import numpy +import cv2 +import threading +import onnxruntime +import traceback + +from ..face_store import get_static_faces, set_static_faces +from ..face_helper import create_static_anchors, distance_to_bounding_box, distance_to_face_landmark_5, warp_face_by_face_landmark_5, warp_face_by_translation, estimate_matrix_by_face_landmark_5, categorize_age, categorize_gender, apply_nms, convert_face_landmark_68_to_5 +from ..execution import apply_execution_provider_options +from ..vision import unpack_resolution, resize_frame_resolution +from ..filesystem import resolve_relative_path +from ..typing import FaceLandmark68, FaceLandmarkSet, FaceScoreSet, FaceRecognizerModel, VisionFrame, Face, FaceDetectorModel, BoundingBox, FaceLandmark5, Score, ModelSet, FaceAnalyserOrder, FaceAnalyserAge, FaceAnalyserGender, Embedding + +THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore() +THREAD_LOCK : threading.Lock = threading.Lock() + +FACE_ANALYSER = None + +# TODO load from options +face_detector_model: FaceDetectorModel = "yoloface" +face_recognizer_model: Optional[FaceRecognizerModel] = 'arcface_inswapper' +face_detector_score = 0.5 +face_landmarker_score = 0.5 +face_detector_size = '640x640' +execution_providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider'] +face_analyser_order: FaceAnalyserOrder = 'left-right' +face_analyser_age: Optional[FaceAnalyserAge] = None +face_analyser_gender: Optional[FaceAnalyserGender] = None + +MODELS : ModelSet =\ +{ + 'face_detector_retinaface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/retinaface_10g.onnx', + 'path': resolve_relative_path('../../../models/faceless/retinaface_10g.onnx') + }, + 'face_detector_scrfd': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/scrfd_2.5g.onnx', + 'path': resolve_relative_path('../../../models/faceless/scrfd_2.5g.onnx') + }, + 'face_detector_yoloface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/yoloface_8n.onnx', + 'path': resolve_relative_path('../../../models/faceless/yoloface_8n.onnx') + }, + 'face_detector_yunet': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/yunet_2023mar.onnx', + 'path': resolve_relative_path('../../../models/faceless/yunet_2023mar.onnx') + }, + 'face_recognizer_arcface_blendswap': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', + 'path': resolve_relative_path('../../../models/faceless/arcface_w600k_r50.onnx') + }, + 'face_recognizer_arcface_inswapper': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', + 'path': resolve_relative_path('../../../models/faceless/arcface_w600k_r50.onnx') + }, + 'face_recognizer_arcface_simswap': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_simswap.onnx', + 'path': resolve_relative_path('../../../models/faceless/arcface_simswap.onnx') + }, + 'face_recognizer_arcface_uniface': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx', + 'path': resolve_relative_path('../../../models/faceless/arcface_w600k_r50.onnx') + }, + 'face_landmarker_68': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/2dfan4.onnx', + 'path': resolve_relative_path('../../../models/faceless/2dfan4.onnx') + }, + 'face_landmarker_68_5': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_landmarker_68_5.onnx', + 'path': resolve_relative_path('../../../models/faceless/face_landmarker_68_5.onnx') + }, + 'gender_age': + { + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gender_age.onnx', + 'path': resolve_relative_path('../../../models/faceless/gender_age.onnx') + } +} + +def get_face_analyser() -> Any: + global FACE_ANALYSER + + face_detectors = {} + face_landmarkers = {} + face_recognizer = None + + with THREAD_LOCK: + if FACE_ANALYSER is None: + if face_detector_model in [ 'many', 'retinaface' ]: + face_detectors['retinaface'] = onnxruntime.InferenceSession(MODELS['face_detector_retinaface']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_detector_model in [ 'many', 'scrfd' ]: + face_detectors['scrfd'] = onnxruntime.InferenceSession(MODELS['face_detector_scrfd']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_detector_model in [ 'many', 'yoloface' ]: + face_detectors['yoloface'] = onnxruntime.InferenceSession(MODELS['face_detector_yoloface']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_detector_model in [ 'yunet' ]: + face_detectors['yunet'] = cv2.FaceDetectorYN.create(MODELS['face_detector_yunet']['path'], '', (0, 0)) + if face_recognizer_model == 'arcface_blendswap': + face_recognizer = onnxruntime.InferenceSession(MODELS['face_recognizer_arcface_blendswap']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_recognizer_model == 'arcface_inswapper': + face_recognizer = onnxruntime.InferenceSession(MODELS['face_recognizer_arcface_inswapper']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_recognizer_model == 'arcface_simswap': + face_recognizer = onnxruntime.InferenceSession(MODELS['face_recognizer_arcface_simswap']['path'], providers = apply_execution_provider_options(execution_providers)) + if face_recognizer_model == 'arcface_uniface': + face_recognizer = onnxruntime.InferenceSession(MODELS['face_recognizer_arcface_uniface']['path'], providers = apply_execution_provider_options(execution_providers)) + face_landmarkers['68'] = onnxruntime.InferenceSession(MODELS['face_landmarker_68']['path'], providers = apply_execution_provider_options(execution_providers)) + face_landmarkers['68_5'] = onnxruntime.InferenceSession(MODELS['face_landmarker_68_5']['path'], providers = apply_execution_provider_options(execution_providers)) + gender_age = onnxruntime.InferenceSession(MODELS['gender_age']['path'], providers = apply_execution_provider_options(execution_providers)) + FACE_ANALYSER =\ + { + 'face_detectors': face_detectors, + 'face_recognizer': face_recognizer, + 'face_landmarkers': face_landmarkers, + 'gender_age': gender_age + } + return FACE_ANALYSER + +def detect_with_retinaface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: + face_detector = get_face_analyser().get('face_detectors').get('retinaface') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + feature_strides = [ 8, 16, 32 ] + feature_map_channel = 3 + anchor_total = 2 + bounding_box_list = [] + face_landmark_5_list = [] + score_list = [] + + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + with THREAD_SEMAPHORE: + detections = face_detector.run(None, + { + face_detector.get_inputs()[0].name: detect_vision_frame + }) + for index, feature_stride in enumerate(feature_strides): + keep_indices = numpy.where(detections[index] >= face_detector_score)[0] + if keep_indices.any(): + stride_height = face_detector_height // feature_stride + stride_width = face_detector_width // feature_stride + anchors = create_static_anchors(feature_stride, anchor_total, stride_height, stride_width) + bounding_box_raw = detections[index + feature_map_channel] * feature_stride + face_landmark_5_raw = detections[index + feature_map_channel * 2] * feature_stride + for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: + bounding_box_list.append(numpy.array( + [ + bounding_box[0] * ratio_width, + bounding_box[1] * ratio_height, + bounding_box[2] * ratio_width, + bounding_box[3] * ratio_height + ])) + for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: + face_landmark_5_list.append(face_landmark_5 * [ ratio_width, ratio_height ]) + for score in detections[index][keep_indices]: + score_list.append(score[0]) + return bounding_box_list, face_landmark_5_list, score_list + + +def detect_with_scrfd(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: + face_detector = get_face_analyser().get('face_detectors').get('scrfd') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + feature_strides = [ 8, 16, 32 ] + feature_map_channel = 3 + anchor_total = 2 + bounding_box_list = [] + face_landmark_5_list = [] + score_list = [] + + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + with THREAD_SEMAPHORE: + detections = face_detector.run(None, + { + face_detector.get_inputs()[0].name: detect_vision_frame + }) + for index, feature_stride in enumerate(feature_strides): + keep_indices = numpy.where(detections[index] >= face_detector_score)[0] + if keep_indices.any(): + stride_height = face_detector_height // feature_stride + stride_width = face_detector_width // feature_stride + anchors = create_static_anchors(feature_stride, anchor_total, stride_height, stride_width) + bounding_box_raw = detections[index + feature_map_channel] * feature_stride + face_landmark_5_raw = detections[index + feature_map_channel * 2] * feature_stride + for bounding_box in distance_to_bounding_box(anchors, bounding_box_raw)[keep_indices]: + bounding_box_list.append(numpy.array( + [ + bounding_box[0] * ratio_width, + bounding_box[1] * ratio_height, + bounding_box[2] * ratio_width, + bounding_box[3] * ratio_height + ])) + for face_landmark_5 in distance_to_face_landmark_5(anchors, face_landmark_5_raw)[keep_indices]: + face_landmark_5_list.append(face_landmark_5 * [ ratio_width, ratio_height ]) + for score in detections[index][keep_indices]: + score_list.append(score[0]) + return bounding_box_list, face_landmark_5_list, score_list + + +def detect_with_yoloface(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: + face_detector = get_face_analyser().get('face_detectors').get('yoloface') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + bounding_box_list = [] + face_landmark_5_list = [] + score_list = [] + + detect_vision_frame = prepare_detect_frame(temp_vision_frame, face_detector_size) + with THREAD_SEMAPHORE: + detections = face_detector.run(None, + { + face_detector.get_inputs()[0].name: detect_vision_frame + }) + detections = numpy.squeeze(detections).T + bounding_box_raw, score_raw, face_landmark_5_raw = numpy.split(detections, [ 4, 5 ], axis = 1) + keep_indices = numpy.where(score_raw > face_detector_score)[0] + if keep_indices.any(): + bounding_box_raw, face_landmark_5_raw, score_raw = bounding_box_raw[keep_indices], face_landmark_5_raw[keep_indices], score_raw[keep_indices] + for bounding_box in bounding_box_raw: + bounding_box_list.append(numpy.array( + [ + (bounding_box[0] - bounding_box[2] / 2) * ratio_width, + (bounding_box[1] - bounding_box[3] / 2) * ratio_height, + (bounding_box[0] + bounding_box[2] / 2) * ratio_width, + (bounding_box[1] + bounding_box[3] / 2) * ratio_height + ])) + face_landmark_5_raw[:, 0::3] = (face_landmark_5_raw[:, 0::3]) * ratio_width + face_landmark_5_raw[:, 1::3] = (face_landmark_5_raw[:, 1::3]) * ratio_height + for face_landmark_5 in face_landmark_5_raw: + face_landmark_5_list.append(numpy.array(face_landmark_5.reshape(-1, 3)[:, :2])) + score_list = score_raw.ravel().tolist() + return bounding_box_list, face_landmark_5_list, score_list + + +def detect_with_yunet(vision_frame : VisionFrame, face_detector_size : str) -> Tuple[List[BoundingBox], List[FaceLandmark5], List[Score]]: + face_detector = get_face_analyser().get('face_detectors').get('yunet') + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + temp_vision_frame = resize_frame_resolution(vision_frame, (face_detector_width, face_detector_height)) + ratio_height = vision_frame.shape[0] / temp_vision_frame.shape[0] + ratio_width = vision_frame.shape[1] / temp_vision_frame.shape[1] + bounding_box_list = [] + face_landmark_5_list = [] + score_list = [] + + face_detector.setInputSize((temp_vision_frame.shape[1], temp_vision_frame.shape[0])) + face_detector.setScoreThreshold(face_detector_score) + with THREAD_SEMAPHORE: + _, detections = face_detector.detect(temp_vision_frame) + if numpy.any(detections): + for detection in detections: + bounding_box_list.append(numpy.array( + [ + detection[0] * ratio_width, + detection[1] * ratio_height, + (detection[0] + detection[2]) * ratio_width, + (detection[1] + detection[3]) * ratio_height + ])) + face_landmark_5_list.append(detection[4:14].reshape((5, 2)) * [ ratio_width, ratio_height ]) + score_list.append(detection[14]) + return bounding_box_list, face_landmark_5_list, score_list + + +def prepare_detect_frame(temp_vision_frame : VisionFrame, face_detector_size : str) -> VisionFrame: + face_detector_width, face_detector_height = unpack_resolution(face_detector_size) + detect_vision_frame = numpy.zeros((face_detector_height, face_detector_width, 3)) + detect_vision_frame[:temp_vision_frame.shape[0], :temp_vision_frame.shape[1], :] = temp_vision_frame + detect_vision_frame = (detect_vision_frame - 127.5) / 128.0 + detect_vision_frame = numpy.expand_dims(detect_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32) + return detect_vision_frame + + +def create_faces(vision_frame : VisionFrame, bounding_box_list : List[BoundingBox], face_landmark_5_list : List[FaceLandmark5], score_list : List[Score]) -> List[Face]: + faces = [] + if face_detector_score > 0: + sort_indices = numpy.argsort(-numpy.array(score_list)) + bounding_box_list = [ bounding_box_list[index] for index in sort_indices ] + face_landmark_5_list = [face_landmark_5_list[index] for index in sort_indices] + score_list = [ score_list[index] for index in sort_indices ] + iou_threshold = 0.1 if face_detector_model == 'many' else 0.4 + keep_indices = apply_nms(bounding_box_list, iou_threshold) + for index in keep_indices: + bounding_box = bounding_box_list[index] + face_landmark_5_68 = face_landmark_5_list[index] + face_landmark_68_5 = expand_face_landmark_68_from_5(face_landmark_5_68) + face_landmark_68 = face_landmark_68_5 + face_landmark_68_score = 0.0 + if face_landmarker_score > 0: + face_landmark_68, face_landmark_68_score = detect_face_landmark_68(vision_frame, bounding_box) + if face_landmark_68_score > face_landmarker_score: + face_landmark_5_68 = convert_face_landmark_68_to_5(face_landmark_68) + landmarks : FaceLandmarkSet =\ + { + '5': face_landmark_5_list[index], + '5/68': face_landmark_5_68, + '68': face_landmark_68, + '68/5': face_landmark_68_5 + } + scores : FaceScoreSet = \ + { + 'detector': score_list[index], + 'landmarker': face_landmark_68_score + } + embedding, normed_embedding = calc_embedding(vision_frame, landmarks.get('5/68')) + gender, age = detect_gender_age(vision_frame, bounding_box) + faces.append(Face( + bounding_box = bounding_box, + landmarks = landmarks, + scores = scores, + embedding = embedding, + normed_embedding = normed_embedding, + gender = gender, + age = age + )) + return faces + + +def calc_embedding(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5) -> Tuple[Embedding, Embedding]: + face_recognizer = get_face_analyser().get('face_recognizer') + crop_vision_frame, _ = warp_face_by_face_landmark_5(temp_vision_frame, face_landmark_5, 'arcface_112_v2', (112, 112)) + crop_vision_frame = crop_vision_frame / 127.5 - 1 + crop_vision_frame = crop_vision_frame[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + embedding = face_recognizer.run(None, + { + face_recognizer.get_inputs()[0].name: crop_vision_frame + })[0] + embedding = embedding.ravel() + normed_embedding = embedding / numpy.linalg.norm(embedding) + return embedding, normed_embedding + + +def detect_face_landmark_68(temp_vision_frame : VisionFrame, bounding_box : BoundingBox) -> Tuple[FaceLandmark68, Score]: + face_landmarker = get_face_analyser().get('face_landmarkers').get('68') + scale = 195 / numpy.subtract(bounding_box[2:], bounding_box[:2]).max() + translation = (256 - numpy.add(bounding_box[2:], bounding_box[:2]) * scale) * 0.5 + crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, (256, 256)) + crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_RGB2Lab) + if numpy.mean(crop_vision_frame[:, :, 0]) < 30: + crop_vision_frame[:, :, 0] = cv2.createCLAHE(clipLimit = 2).apply(crop_vision_frame[:, :, 0]) + crop_vision_frame = cv2.cvtColor(crop_vision_frame, cv2.COLOR_Lab2RGB) + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1).astype(numpy.float32) / 255.0 + face_landmark_68, face_heatmap = face_landmarker.run(None, + { + face_landmarker.get_inputs()[0].name: [ crop_vision_frame ] + }) + face_landmark_68 = face_landmark_68[:, :, :2][0] / 64 + face_landmark_68 = face_landmark_68.reshape(1, -1, 2) * 256 + face_landmark_68 = cv2.transform(face_landmark_68, cv2.invertAffineTransform(affine_matrix)) + face_landmark_68 = face_landmark_68.reshape(-1, 2) + face_landmark_68_score = numpy.amax(face_heatmap, axis = (2, 3)) + face_landmark_68_score = numpy.mean(face_landmark_68_score) + return face_landmark_68, face_landmark_68_score + + +def expand_face_landmark_68_from_5(face_landmark_5 : FaceLandmark5) -> FaceLandmark68: + face_landmarker = get_face_analyser().get('face_landmarkers').get('68_5') + affine_matrix = estimate_matrix_by_face_landmark_5(face_landmark_5, 'ffhq_512', (1, 1)) + face_landmark_5 = cv2.transform(face_landmark_5.reshape(1, -1, 2), affine_matrix).reshape(-1, 2) + face_landmark_68_5 = face_landmarker.run(None, + { + face_landmarker.get_inputs()[0].name: [ face_landmark_5 ] + })[0][0] + face_landmark_68_5 = cv2.transform(face_landmark_68_5.reshape(1, -1, 2), cv2.invertAffineTransform(affine_matrix)).reshape(-1, 2) + return face_landmark_68_5 + + +def detect_gender_age(temp_vision_frame : VisionFrame, bounding_box : BoundingBox) -> Tuple[int, int]: + gender_age = get_face_analyser().get('gender_age') + bounding_box = bounding_box.reshape(2, -1) + scale = 64 / numpy.subtract(*bounding_box[::-1]).max() + translation = 48 - bounding_box.sum(axis = 0) * scale * 0.5 + crop_vision_frame, affine_matrix = warp_face_by_translation(temp_vision_frame, translation, scale, (96, 96)) + crop_vision_frame = crop_vision_frame[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0) + prediction = gender_age.run(None, + { + gender_age.get_inputs()[0].name: crop_vision_frame + })[0][0] + gender = int(numpy.argmax(prediction[:2])) + age = int(numpy.round(prediction[2] * 100)) + return gender, age + +def get_average_face(vision_frames : List[VisionFrame], position : int = 0) -> Optional[Face]: + average_face = None + faces = [] + embedding_list = [] + normed_embedding_list = [] + + for vision_frame in vision_frames: + face = get_one_face(vision_frame, position) + if face: + faces.append(face) + embedding_list.append(face.embedding) + normed_embedding_list.append(face.normed_embedding) + if faces: + first_face = faces[0] + average_face = Face( + bounding_box = first_face.bounding_box, + landmarks = first_face.landmarks, + scores = first_face.scores, + embedding = numpy.mean(embedding_list, axis = 0), + normed_embedding = numpy.mean(normed_embedding_list, axis = 0), + gender = first_face.gender, + age = first_face.age + ) + return average_face + +def get_one_face(vision_frame : VisionFrame, position : int = 0) -> Optional[Face]: + many_faces = get_many_faces(vision_frame) + if many_faces: + try: + return many_faces[position] + except IndexError: + return many_faces[-1] + return None + +def get_many_faces(vision_frame : VisionFrame) -> List[Face]: + faces = [] + # try: + faces_cache = get_static_faces(vision_frame) + if faces_cache: + faces = faces_cache + else: + bounding_box_list = [] + face_landmark_5_list = [] + score_list = [] + + if face_detector_model in [ 'many', 'retinaface']: + bounding_box_list_retinaface, face_landmark_5_list_retinaface, score_list_retinaface = detect_with_retinaface(vision_frame, face_detector_size) + bounding_box_list.extend(bounding_box_list_retinaface) + face_landmark_5_list.extend(face_landmark_5_list_retinaface) + score_list.extend(score_list_retinaface) + if face_detector_model in [ 'many', 'scrfd' ]: + bounding_box_list_scrfd, face_landmark_5_list_scrfd, score_list_scrfd = detect_with_scrfd(vision_frame, face_detector_size) + bounding_box_list.extend(bounding_box_list_scrfd) + face_landmark_5_list.extend(face_landmark_5_list_scrfd) + score_list.extend(score_list_scrfd) + if face_detector_model in [ 'many', 'yoloface' ]: + bounding_box_list_yoloface, face_landmark_5_list_yoloface, score_list_yoloface = detect_with_yoloface(vision_frame, face_detector_size) + bounding_box_list.extend(bounding_box_list_yoloface) + face_landmark_5_list.extend(face_landmark_5_list_yoloface) + score_list.extend(score_list_yoloface) + if face_detector_model in [ 'yunet' ]: + bounding_box_list_yunet, face_landmark_5_list_yunet, score_list_yunet = detect_with_yunet(vision_frame, face_detector_size) + bounding_box_list.extend(bounding_box_list_yunet) + face_landmark_5_list.extend(face_landmark_5_list_yunet) + score_list.extend(score_list_yunet) + + if bounding_box_list and face_landmark_5_list and score_list: + faces = create_faces(vision_frame, bounding_box_list, face_landmark_5_list, score_list) + if faces: + set_static_faces(vision_frame, faces) + if face_analyser_order: + faces = sort_by_order(faces, face_analyser_order) + if face_analyser_age: + faces = filter_by_age(faces, face_analyser_age) + if face_analyser_gender: + faces = filter_by_gender(faces, face_analyser_gender) + # except (AttributeError, ValueError) as e: + # print("error", e) + return faces + +def sort_by_order(faces : List[Face], order : FaceAnalyserOrder) -> List[Face]: + if order == 'left-right': + return sorted(faces, key = lambda face: face.bounding_box[0]) + if order == 'right-left': + return sorted(faces, key = lambda face: face.bounding_box[0], reverse = True) + if order == 'top-bottom': + return sorted(faces, key = lambda face: face.bounding_box[1]) + if order == 'bottom-top': + return sorted(faces, key = lambda face: face.bounding_box[1], reverse = True) + if order == 'small-large': + return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1])) + if order == 'large-small': + return sorted(faces, key = lambda face: (face.bounding_box[2] - face.bounding_box[0]) * (face.bounding_box[3] - face.bounding_box[1]), reverse = True) + if order == 'best-worst': + return sorted(faces, key = lambda face: face.scores.get('detector'), reverse = True) + if order == 'worst-best': + return sorted(faces, key = lambda face: face.scores.get('detector')) + return faces + + +def filter_by_age(faces : List[Face], age : FaceAnalyserAge) -> List[Face]: + filter_faces = [] + for face in faces: + if categorize_age(face.age) == age: + filter_faces.append(face) + return filter_faces + + +def filter_by_gender(faces : List[Face], gender : FaceAnalyserGender) -> List[Face]: + filter_faces = [] + for face in faces: + if categorize_gender(face.gender) == gender: + filter_faces.append(face) + return filter_faces diff --git a/faceless/processors/face_swapper.py b/faceless/processors/face_swapper.py new file mode 100644 index 0000000..17d1d26 --- /dev/null +++ b/faceless/processors/face_swapper.py @@ -0,0 +1,261 @@ +import os +import threading +from typing import Optional, Literal, Any + +import numpy +import onnx +from onnx import numpy_helper +import onnxruntime + +from ..processors.face_analyser import get_average_face, get_many_faces, get_one_face +from ..face_helper import warp_face_by_face_landmark_5, paste_back +from ..face_masker import create_static_box_mask, create_occlusion_mask, create_region_mask +from ..execution import apply_execution_provider_options +from ..typing import Face, VisionFrame, FaceSelectorMode, ModelSet, OptionsWithModel, Embedding +from ..vision import read_image, write_image, tensor_to_vision_frame +from ..filesystem import resolve_relative_path + +# TODO load from options +face_selector_mode: FaceSelectorMode = 'many' +face_mask_blur = 0.3 +face_mask_padding = (0, 0, 0, 0) +face_mask_regions = [] +face_mask_types = ['box'] +face_swapper_model = 'inswapper_128_fp16' + +execution_providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider'] + +THREAD_LOCK : threading.Lock = threading.Lock() +MODEL_INITIALIZER = None +FRAME_PROCESSOR = None + +model_template = 'arcface_128_v2' +model_size = (128, 128) + +MODELS : ModelSet =\ +{ + 'blendswap_256': + { + 'type': 'blendswap', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/blendswap_256.onnx', + 'path': resolve_relative_path('../../../models/faceless/blendswap_256.onnx'), + 'template': 'ffhq_512', + 'size': (256, 256), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'inswapper_128': + { + 'type': 'inswapper', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx', + 'path': resolve_relative_path('../../../models/faceless/inswapper_128.onnx'), + 'template': 'arcface_128_v2', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'inswapper_128_fp16': + { + 'type': 'inswapper', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128_fp16.onnx', + 'path': resolve_relative_path('../../../models/faceless/inswapper_128_fp16.onnx'), + 'template': 'arcface_128_v2', + 'size': (128, 128), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'simswap_256': + { + 'type': 'simswap', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_256.onnx', + 'path': resolve_relative_path('../../../models/faceless/simswap_256.onnx'), + 'template': 'arcface_112_v1', + 'size': (256, 256), + 'mean': [ 0.485, 0.456, 0.406 ], + 'standard_deviation': [ 0.229, 0.224, 0.225 ] + }, + 'simswap_512_unofficial': + { + 'type': 'simswap', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_512_unofficial.onnx', + 'path': resolve_relative_path('../../../models/faceless/simswap_512_unofficial.onnx'), + 'template': 'arcface_112_v1', + 'size': (512, 512), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + }, + 'uniface_256': + { + 'type': 'uniface', + 'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/uniface_256.onnx', + 'path': resolve_relative_path('../../../models/faceless/uniface_256.onnx'), + 'template': 'ffhq_512', + 'size': (256, 256), + 'mean': [ 0.0, 0.0, 0.0 ], + 'standard_deviation': [ 1.0, 1.0, 1.0 ] + } +} +OPTIONS : Optional[OptionsWithModel] = None + +def swap_face(source_face: Face, target_face: Face, source_vision_frame, target_vision_frame: VisionFrame) -> VisionFrame: + crop_vision_frame, affine_matrix = warp_face_by_face_landmark_5(target_vision_frame, target_face.landmarks.get('5/68'), model_template, model_size) + crop_mask_list = [] + + if 'box' in face_mask_types: + box_mask = create_static_box_mask(crop_vision_frame.shape[:2][::-1], face_mask_blur, face_mask_padding) + crop_mask_list.append(box_mask) + if 'occlusion' in face_mask_types: + occlusion_mask = create_occlusion_mask(crop_vision_frame) + crop_mask_list.append(occlusion_mask) + crop_vision_frame = prepare_crop_frame(crop_vision_frame) + crop_vision_frame = apply_swap(source_face, source_vision_frame, crop_vision_frame) + crop_vision_frame = normalize_crop_frame(crop_vision_frame) + if 'region' in face_mask_types: + region_mask = create_region_mask(crop_vision_frame, face_mask_regions) + crop_mask_list.append(region_mask) + crop_mask = numpy.minimum.reduce(crop_mask_list).clip(0, 1) + target_vision_frame = paste_back(target_vision_frame, crop_vision_frame, crop_mask, affine_matrix) + return target_vision_frame + +def process_frame(source_face: Face, source_vision_frame: VisionFrame, target_vision_frame: VisionFrame) -> Optional[VisionFrame]: + if face_selector_mode == 'many': + target_faces = get_many_faces(target_vision_frame) + for target_face in target_faces: + target_vision_frame = swap_face(source_face, target_face, source_vision_frame, target_vision_frame) + if face_selector_mode == 'one': + target_face = get_one_face(target_vision_frame) + if target_face: + target_vision_frame = swap_face(source_face, target_face, source_vision_frame, target_vision_frame) + return target_vision_frame + +def process_images(source_image, target_images, output_frames_path): + source_frame = tensor_to_vision_frame(source_image) + if source_frame is None: + raise Exception('cannot read source image') + source_face = get_average_face([source_frame]) + if source_face is None: + raise Exception('cannot find source face') + + count = len(target_images) + for (index, target_image) in enumerate(target_images): + print(f"progress: {index + 1}/{count}") + filename = f"{index + 1}".ljust(4, "0") + ".png" + output_filepath = os.path.join(output_frames_path, filename) + + target_vision_frame = tensor_to_vision_frame(target_image) + if target_vision_frame is None: + raise Exception("invalid target image") + output_vision_frame = process_frame(source_face, source_frame, target_vision_frame) + if output_vision_frame is None: + raise Exception("process frame failed") + write_image(output_filepath, output_vision_frame) + +def process_frames(source_image, target_frames_path: str, output_frames_path: str): + source_frame = tensor_to_vision_frame(source_image) + if source_frame is None: + raise Exception("cannot read source image") + source_face = get_average_face([source_frame]) + if source_face is None: + raise Exception("cannot find source face") + + frames_filenames = os.listdir(target_frames_path) + count = len(frames_filenames) + for index, frame_filename in enumerate(sorted(frames_filenames)): + print(f"progress: {index + 1}/{count}") + frame_filepath = os.path.join(target_frames_path, frame_filename) + output_filepath = os.path.join(output_frames_path, frame_filename) + + target_vision_frame = read_image(frame_filepath) + if target_vision_frame is None: + raise Exception("invalid target image") + output_vision_frame = process_frame(source_face, source_frame, target_vision_frame) + if output_vision_frame is None: + raise Exception("process frame failed") + write_image(output_filepath, output_vision_frame) + +def apply_swap(source_face : Face, source_vision_frame: VisionFrame, crop_vision_frame : VisionFrame) -> VisionFrame: + frame_processor = get_frame_processor() + model_type = get_options('model').get('type') + frame_processor_inputs = {} + + for frame_processor_input in frame_processor.get_inputs(): + if frame_processor_input.name == 'source': + if model_type == 'blendswap' or model_type == 'uniface': + frame_processor_inputs[frame_processor_input.name] = prepare_source_frame(source_face, source_vision_frame) + else: + frame_processor_inputs[frame_processor_input.name] = prepare_source_embedding(source_face) + if frame_processor_input.name == 'target': + frame_processor_inputs[frame_processor_input.name] = crop_vision_frame + crop_vision_frame = frame_processor.run(None, frame_processor_inputs)[0][0] + return crop_vision_frame + +def prepare_source_frame(source_face : Face, source_vision_frame: VisionFrame) -> VisionFrame: + model_type = get_options('model').get('type') + if model_type == 'blendswap': + source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmarks.get('5/68'), 'arcface_112_v2', (112, 112)) + if model_type == 'uniface': + source_vision_frame, _ = warp_face_by_face_landmark_5(source_vision_frame, source_face.landmarks.get('5/68'), 'ffhq_512', (256, 256)) + source_vision_frame = source_vision_frame[:, :, ::-1] / 255.0 + source_vision_frame = source_vision_frame.transpose(2, 0, 1) + source_vision_frame = numpy.expand_dims(source_vision_frame, axis = 0).astype(numpy.float32) + return source_vision_frame + + +def prepare_source_embedding(source_face : Face) -> Embedding: + model_type = get_options('model').get('type') + if model_type == 'inswapper': + model_initializer = get_model_initializer() + source_embedding = source_face.embedding.reshape((1, -1)) + source_embedding = numpy.dot(source_embedding, model_initializer) / numpy.linalg.norm(source_embedding) + else: + source_embedding = source_face.normed_embedding.reshape(1, -1) + return source_embedding + + +def prepare_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + model_mean = get_options('model').get('mean') + model_standard_deviation = get_options('model').get('standard_deviation') + crop_vision_frame = crop_vision_frame[:, :, ::-1] / 255.0 + crop_vision_frame = (crop_vision_frame - model_mean) / model_standard_deviation + crop_vision_frame = crop_vision_frame.transpose(2, 0, 1) + crop_vision_frame = numpy.expand_dims(crop_vision_frame, axis = 0).astype(numpy.float32) + return crop_vision_frame + + +def normalize_crop_frame(crop_vision_frame : VisionFrame) -> VisionFrame: + crop_vision_frame = crop_vision_frame.transpose(1, 2, 0) + crop_vision_frame = (crop_vision_frame * 255.0).round() + crop_vision_frame = crop_vision_frame[:, :, ::-1] + return crop_vision_frame + + +def get_options(key : Literal['model']) -> Any: + global OPTIONS + + if OPTIONS is None: + OPTIONS =\ + { + 'model': MODELS[face_swapper_model] + } + return OPTIONS.get(key) + + +def get_model_initializer() -> Any: + global MODEL_INITIALIZER + + with THREAD_LOCK: + if MODEL_INITIALIZER is None: + model_path = get_options('model').get('path') + model = onnx.load(model_path) + MODEL_INITIALIZER = numpy_helper.to_array(model.graph.initializer[-1]) + return MODEL_INITIALIZER + + +def get_frame_processor() -> Any: + global FRAME_PROCESSOR + + with THREAD_LOCK: + if FRAME_PROCESSOR is None: + model_path = get_options('model').get('path') + FRAME_PROCESSOR = onnxruntime.InferenceSession(model_path, providers = apply_execution_provider_options(execution_providers)) + return FRAME_PROCESSOR diff --git a/faceless/typing.py b/faceless/typing.py index 9adaac0..af33d2f 100644 --- a/faceless/typing.py +++ b/faceless/typing.py @@ -1,13 +1,122 @@ -from typing import Literal, Tuple, TypedDict +from typing import Literal, Tuple, TypedDict, Any, Dict, List +from collections import namedtuple + +import numpy + + +BoundingBox = numpy.ndarray[Any, Any] +FaceLandmark5 = numpy.ndarray[Any, Any] +FaceLandmark68 = numpy.ndarray[Any, Any] +FaceLandmarkSet = TypedDict('FaceLandmarkSet', +{ + '5' : FaceLandmark5, # type: ignore[valid-type] + '5/68' : FaceLandmark5, # type: ignore[valid-type] + '68' : FaceLandmark68, # type: ignore[valid-type] + '68/5' : FaceLandmark68 # type: ignore[valid-type] +}) +Score = float +FaceScoreSet = TypedDict('FaceScoreSet', +{ + 'detector' : Score, + 'landmarker' : Score +}) + +WarpTemplate = Literal['arcface_112_v1', 'arcface_112_v2', 'arcface_128_v2', 'ffhq_512'] +WarpTemplateSet = Dict[WarpTemplate, numpy.ndarray[Any, Any]] # Vision -Resolution = Tuple[int, int] Fps = float +Padding = Tuple[int, int, int, int] +Resolution = Tuple[int, int] + +VisionFrame = numpy.ndarray[Any, Any] +Mask = numpy.ndarray[Any, Any] +Matrix = numpy.ndarray[Any, Any] +Translation = numpy.ndarray[Any, Any] FrameFormat = Literal['jpg', 'png', 'bmp'] FacelessVideo = TypedDict('FacelessVideo', { 'video_path': str, - 'frames_path': str, + 'output_path': str, + 'resolution': Resolution, + 'fps': Fps, }) + +Embedding = numpy.ndarray[Any, Any] +Face = namedtuple('Face', +[ + 'bounding_box', + 'landmarks', + 'scores', + 'embedding', + 'normed_embedding', + 'gender', + 'age' +]) + +# Face Store +FaceSet = Dict[str, List[Face]] +FaceStore = TypedDict('FaceStore', +{ + 'static_faces' : FaceSet, + 'reference_faces': FaceSet +}) + + +FaceDetectorModel = Literal['many', 'retinaface', 'scrfd', 'yoloface', 'yunet'] +FaceRecognizerModel = Literal['arcface_blendswap', 'arcface_inswapper', 'arcface_simswap', 'arcface_uniface'] + +ModelValue = Dict[str, Any] +ModelSet = Dict[str, ModelValue] +OptionsWithModel = TypedDict('OptionsWithModel', +{ + 'model' : ModelValue +}) + +ValueAndUnit = TypedDict('ValueAndUnit', +{ + 'value' : str, + 'unit' : str +}) +ExecutionDeviceFramework = TypedDict('ExecutionDeviceFramework', +{ + 'name' : str, + 'version' : str +}) +ExecutionDeviceProduct = TypedDict('ExecutionDeviceProduct', +{ + 'vendor' : str, + 'name' : str, + 'architecture' : str, +}) +ExecutionDeviceVideoMemory = TypedDict('ExecutionDeviceVideoMemory', +{ + 'total' : ValueAndUnit, + 'free' : ValueAndUnit +}) +ExecutionDeviceUtilization = TypedDict('ExecutionDeviceUtilization', +{ + 'gpu' : ValueAndUnit, + 'memory' : ValueAndUnit +}) +ExecutionDevice = TypedDict('ExecutionDevice', +{ + 'driver_version' : str, + 'framework' : ExecutionDeviceFramework, + 'product' : ExecutionDeviceProduct, + 'video_memory' : ExecutionDeviceVideoMemory, + 'utilization' : ExecutionDeviceUtilization +}) + + +FaceAnalyserOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best'] +FaceAnalyserAge = Literal['child', 'teen', 'adult', 'senior'] +FaceAnalyserGender = Literal['female', 'male'] +FaceSelectorMode = Literal['many', 'one', 'reference'] + +FaceMaskRegion = Literal['skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip'] + +OutputVideoEncoder = Literal['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc', 'h264_amf', 'hevc_amf'] +OutputVideoPreset = Literal['ultrafast', 'superfast', 'veryfast', 'faster', 'fast', 'medium', 'slow', 'slower', 'veryslow'] diff --git a/faceless/vision.py b/faceless/vision.py index 87af18e..1dc120b 100644 --- a/faceless/vision.py +++ b/faceless/vision.py @@ -1,11 +1,14 @@ from typing import Optional, Tuple +from functools import lru_cache import cv2 +import numpy as np +from PIL import Image -from .filesystem import is_video -from .typing import Resolution +from .filesystem import is_video, is_image +from .typing import Resolution, VisionFrame, Fps -def detect_video_fps(video_path : str) -> Optional[float]: +def detect_video_fps(video_path : str) -> Optional[Fps]: if is_video(video_path): video_capture = cv2.VideoCapture(video_path) if video_capture.isOpened(): @@ -13,6 +16,15 @@ def detect_video_fps(video_path : str) -> Optional[float]: video_capture.release() return video_fps +def restrict_video_fps(video_path : str, fps : Fps) -> Fps: + if is_video(video_path): + video_fps = detect_video_fps(video_path) + if video_fps is None: + return fps + if video_fps < fps: + return video_fps + return fps + def detect_video_resolution(video_path : str) -> Optional[Resolution]: if is_video(video_path): video_capture = cv2.VideoCapture(video_path) @@ -31,7 +43,43 @@ def normalize_resolution(resolution : Tuple[float, float]) -> Resolution: return normalize_width, normalize_height return 0, 0 - def pack_resolution(resolution : Resolution) -> str: width, height = normalize_resolution(resolution) return str(width) + 'x' + str(height) + +def unpack_resolution(resolution : str) -> Resolution: + width, height = map(int, resolution.split('x')) + return width, height + +@lru_cache(maxsize = 128) +def read_static_image(image_path : str) -> Optional[VisionFrame]: + return read_image(image_path) + +def read_image(image_path : str) -> Optional[VisionFrame]: + if is_image(image_path): + return cv2.imread(image_path) + return None + +def tensor_to_vision_frame(image_tensor) -> Optional[VisionFrame]: + i = 255. * image_tensor.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # cv2_image = np.transpose(np.array(img), (1, 2, 0)) + return cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB) + +def write_image(image_path : str, vision_frame : VisionFrame) -> bool: + if image_path: + return cv2.imwrite(image_path, vision_frame) + return False + +def resize_frame_resolution(vision_frame : VisionFrame, max_resolution : Resolution) -> VisionFrame: + height, width = vision_frame.shape[:2] + max_width, max_height = max_resolution + + if height > max_height or width > max_width: + scale = min(max_height / height, max_width / width) + new_width = int(width * scale) + new_height = int(height * scale) + return cv2.resize(vision_frame, (new_width, new_height)) + return vision_frame +