feat: support face swap
This commit is contained in:
Jeffrey Wu
@@ -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.
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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'] = {}
|
||||
+43
-3
@@ -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'
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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 ()
|
||||
|
||||
@@ -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,)
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
+112
-3
@@ -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']
|
||||
|
||||
+52
-4
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user