Merge pull request #566 from pereiraroland26/main
Added support for multiple faces
This commit is contained in:
Kenneth Estanislao
@@ -0,0 +1,32 @@
|
||||
import numpy as np
|
||||
from sklearn.cluster import KMeans
|
||||
from sklearn.metrics import silhouette_score
|
||||
from typing import Any
|
||||
|
||||
|
||||
def find_cluster_centroids(embeddings, max_k=10) -> Any:
|
||||
inertia = []
|
||||
cluster_centroids = []
|
||||
K = range(1, max_k+1)
|
||||
|
||||
for k in K:
|
||||
kmeans = KMeans(n_clusters=k, random_state=0)
|
||||
kmeans.fit(embeddings)
|
||||
inertia.append(kmeans.inertia_)
|
||||
cluster_centroids.append({"k": k, "centroids": kmeans.cluster_centers_})
|
||||
|
||||
diffs = [inertia[i] - inertia[i+1] for i in range(len(inertia)-1)]
|
||||
optimal_centroids = cluster_centroids[diffs.index(max(diffs)) + 1]['centroids']
|
||||
|
||||
return optimal_centroids
|
||||
|
||||
def find_closest_centroid(centroids: list, normed_face_embedding) -> list:
|
||||
try:
|
||||
centroids = np.array(centroids)
|
||||
normed_face_embedding = np.array(normed_face_embedding)
|
||||
similarities = np.dot(centroids, normed_face_embedding)
|
||||
closest_centroid_index = np.argmax(similarities)
|
||||
|
||||
return closest_centroid_index, centroids[closest_centroid_index]
|
||||
except ValueError:
|
||||
return None
|
||||
+99
-190
@@ -1,17 +1,16 @@
|
||||
import os
|
||||
import sys
|
||||
# single thread doubles cuda performance - needs to be set before torch import
|
||||
if any(arg.startswith('--execution-provider') for arg in sys.argv):
|
||||
os.environ['OMP_NUM_THREADS'] = '1'
|
||||
# reduce tensorflow log level
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
|
||||
import warnings
|
||||
from typing import List
|
||||
import platform
|
||||
import signal
|
||||
import shutil
|
||||
import argparse
|
||||
from typing import List
|
||||
|
||||
# Set environment variables for CUDA performance and TensorFlow logging
|
||||
if any(arg.startswith('--execution-provider') for arg in sys.argv):
|
||||
os.environ['OMP_NUM_THREADS'] = '1'
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
|
||||
|
||||
import torch
|
||||
import onnxruntime
|
||||
import tensorflow
|
||||
@@ -20,73 +19,38 @@ import modules.globals
|
||||
import modules.metadata
|
||||
import modules.ui as ui
|
||||
from modules.processors.frame.core import get_frame_processors_modules
|
||||
from modules.utilities import (
|
||||
has_image_extension,
|
||||
is_image,
|
||||
is_video,
|
||||
detect_fps,
|
||||
create_video,
|
||||
extract_frames,
|
||||
get_temp_frame_paths,
|
||||
restore_audio,
|
||||
create_temp,
|
||||
move_temp,
|
||||
clean_temp,
|
||||
normalize_output_path
|
||||
)
|
||||
from modules.utilities import has_image_extension, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clean_temp, normalize_output_path
|
||||
|
||||
if 'ROCMExecutionProvider' in modules.globals.execution_providers:
|
||||
del torch
|
||||
|
||||
# Filter warnings
|
||||
warnings.filterwarnings('ignore', category=FutureWarning, module='insightface')
|
||||
warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')
|
||||
|
||||
# Cross-platform resource management
|
||||
if platform.system() == 'Darwin' and 'ROCMExecutionProvider' in modules.globals.execution_providers:
|
||||
del torch
|
||||
|
||||
|
||||
def parse_args() -> None:
|
||||
signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
|
||||
program = argparse.ArgumentParser()
|
||||
program.add_argument('-s', '--source', help='Select a source image', dest='source_path')
|
||||
program.add_argument('-t', '--target', help='Select a target image or video', dest='target_path')
|
||||
program.add_argument('-o', '--output', help='Select output file or directory', dest='output_path')
|
||||
program.add_argument('--frame-processor', help='Pipeline of frame processors', dest='frame_processor',
|
||||
default=['face_swapper'], choices=['face_swapper', 'face_enhancer', 'super_resolution'],
|
||||
nargs='+')
|
||||
program.add_argument('--keep-fps', help='Keep original fps', dest='keep_fps', action='store_true', default=False)
|
||||
program.add_argument('--keep-audio', help='Keep original audio', dest='keep_audio', action='store_true',
|
||||
default=True)
|
||||
program.add_argument('--keep-frames', help='Keep temporary frames', dest='keep_frames', action='store_true',
|
||||
default=False)
|
||||
program.add_argument('--many-faces', help='Process every face', dest='many_faces', action='store_true',
|
||||
default=False)
|
||||
program.add_argument('--video-encoder', help='Adjust output video encoder', dest='video_encoder', default='libx264',
|
||||
choices=['libx264', 'libx265', 'libvpx-vp9'])
|
||||
program.add_argument('--video-quality', help='Adjust output video quality', dest='video_quality', type=int,
|
||||
default=18,
|
||||
choices=range(52), metavar='[0-51]')
|
||||
program.add_argument('--live-mirror', help='The live camera display as you see it in the front-facing camera frame',
|
||||
dest='live_mirror', action='store_true', default=False)
|
||||
program.add_argument('--live-resizable', help='The live camera frame is resizable',
|
||||
dest='live_resizable', action='store_true', default=False)
|
||||
program.add_argument('--max-memory', help='Maximum amount of RAM in GB', dest='max_memory', type=int,
|
||||
default=suggest_max_memory())
|
||||
program.add_argument('--execution-provider', help='Execution provider', dest='execution_provider', default=['cpu'],
|
||||
choices=suggest_execution_providers(), nargs='+')
|
||||
program.add_argument('--execution-threads', help='Number of execution threads', dest='execution_threads', type=int,
|
||||
default=suggest_execution_threads())
|
||||
program.add_argument('--headless', help='Run in headless mode', dest='headless', default=False, action='store_true')
|
||||
program.add_argument('--enhancer-upscale-factor',
|
||||
help='Sets the upscale factor for the enhancer. Only applies if `face_enhancer` is set as a frame-processor',
|
||||
dest='enhancer_upscale_factor', type=int, default=1)
|
||||
program.add_argument('--source-image-scaling-factor', help='Set the upscale factor for source images',
|
||||
dest='source_image_scaling_factor', default=2, type=int)
|
||||
program.add_argument('-r', '--super-resolution-scale-factor', dest='super_resolution_scale_factor',
|
||||
help='Set the upscale factor for super resolution', default=4, choices=[2, 3, 4], type=int)
|
||||
program.add_argument('-v', '--version', action='version',
|
||||
version=f'{modules.metadata.name} {modules.metadata.version}')
|
||||
program.add_argument('-s', '--source', help='select an source image', dest='source_path')
|
||||
program.add_argument('-t', '--target', help='select an target image or video', dest='target_path')
|
||||
program.add_argument('-o', '--output', help='select output file or directory', dest='output_path')
|
||||
program.add_argument('--frame-processor', help='pipeline of frame processors', dest='frame_processor', default=['face_swapper'], choices=['face_swapper', 'face_enhancer'], nargs='+')
|
||||
program.add_argument('--keep-fps', help='keep original fps', dest='keep_fps', action='store_true', default=False)
|
||||
program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True)
|
||||
program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False)
|
||||
program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False)
|
||||
program.add_argument('--nsfw-filter', help='filter the NSFW image or video', dest='nsfw_filter', action='store_true', default=False)
|
||||
program.add_argument('--map-faces', help='map source target faces', dest='map_faces', action='store_true', default=False)
|
||||
program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
|
||||
program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]')
|
||||
program.add_argument('--live-mirror', help='The live camera display as you see it in the front-facing camera frame', dest='live_mirror', action='store_true', default=False)
|
||||
program.add_argument('--live-resizable', help='The live camera frame is resizable', dest='live_resizable', action='store_true', default=False)
|
||||
program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
|
||||
program.add_argument('--execution-provider', help='execution provider', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+')
|
||||
program.add_argument('--execution-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
|
||||
program.add_argument('-v', '--version', action='version', version=f'{modules.metadata.name} {modules.metadata.version}')
|
||||
|
||||
# Register deprecated args
|
||||
# register deprecated args
|
||||
program.add_argument('-f', '--face', help=argparse.SUPPRESS, dest='source_path_deprecated')
|
||||
program.add_argument('--cpu-cores', help=argparse.SUPPRESS, dest='cpu_cores_deprecated', type=int)
|
||||
program.add_argument('--gpu-vendor', help=argparse.SUPPRESS, dest='gpu_vendor_deprecated')
|
||||
@@ -96,14 +60,15 @@ def parse_args() -> None:
|
||||
|
||||
modules.globals.source_path = args.source_path
|
||||
modules.globals.target_path = args.target_path
|
||||
modules.globals.output_path = normalize_output_path(modules.globals.source_path, modules.globals.target_path,
|
||||
args.output_path)
|
||||
modules.globals.output_path = normalize_output_path(modules.globals.source_path, modules.globals.target_path, args.output_path)
|
||||
modules.globals.frame_processors = args.frame_processor
|
||||
modules.globals.headless = args.source_path or args.target_path or args.output_path
|
||||
modules.globals.keep_fps = args.keep_fps
|
||||
modules.globals.keep_audio = args.keep_audio
|
||||
modules.globals.keep_frames = args.keep_frames
|
||||
modules.globals.many_faces = args.many_faces
|
||||
modules.globals.nsfw_filter = args.nsfw_filter
|
||||
modules.globals.map_faces = args.map_faces
|
||||
modules.globals.video_encoder = args.video_encoder
|
||||
modules.globals.video_quality = args.video_quality
|
||||
modules.globals.live_mirror = args.live_mirror
|
||||
@@ -111,26 +76,18 @@ def parse_args() -> None:
|
||||
modules.globals.max_memory = args.max_memory
|
||||
modules.globals.execution_providers = decode_execution_providers(args.execution_provider)
|
||||
modules.globals.execution_threads = args.execution_threads
|
||||
modules.globals.headless = args.headless
|
||||
modules.globals.enhancer_upscale_factor = args.enhancer_upscale_factor
|
||||
modules.globals.source_image_scaling_factor = args.source_image_scaling_factor
|
||||
modules.globals.sr_scale_factor = args.super_resolution_scale_factor
|
||||
# Handle face enhancer tumbler
|
||||
modules.globals.fp_ui['face_enhancer'] = 'face_enhancer' in args.frame_processor
|
||||
|
||||
modules.globals.nsfw = False
|
||||
#for ENHANCER tumbler:
|
||||
if 'face_enhancer' in args.frame_processor:
|
||||
modules.globals.fp_ui['face_enhancer'] = True
|
||||
else:
|
||||
modules.globals.fp_ui['face_enhancer'] = False
|
||||
|
||||
# Handle deprecated arguments
|
||||
handle_deprecated_args(args)
|
||||
|
||||
|
||||
def handle_deprecated_args(args) -> None:
|
||||
"""Handle deprecated arguments by translating them to the new format."""
|
||||
# translate deprecated args
|
||||
if args.source_path_deprecated:
|
||||
print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m')
|
||||
modules.globals.source_path = args.source_path_deprecated
|
||||
modules.globals.output_path = normalize_output_path(args.source_path_deprecated, modules.globals.target_path,
|
||||
args.output_path)
|
||||
modules.globals.output_path = normalize_output_path(args.source_path_deprecated, modules.globals.target_path, args.output_path)
|
||||
if args.cpu_cores_deprecated:
|
||||
print('\033[33mArgument --cpu-cores is deprecated. Use --execution-threads instead.\033[0m')
|
||||
modules.globals.execution_threads = args.cpu_cores_deprecated
|
||||
@@ -141,7 +98,7 @@ def handle_deprecated_args(args) -> None:
|
||||
print('\033[33mArgument --gpu-vendor nvidia is deprecated. Use --execution-provider cuda instead.\033[0m')
|
||||
modules.globals.execution_providers = decode_execution_providers(['cuda'])
|
||||
if args.gpu_vendor_deprecated == 'amd':
|
||||
print('\033[33mArgument --gpu-vendor amd is deprecated. Use --execution-provider rocm instead.\033[0m')
|
||||
print('\033[33mArgument --gpu-vendor amd is deprecated. Use --execution-provider cuda instead.\033[0m')
|
||||
modules.globals.execution_providers = decode_execution_providers(['rocm'])
|
||||
if args.gpu_threads_deprecated:
|
||||
print('\033[33mArgument --gpu-threads is deprecated. Use --execution-threads instead.\033[0m')
|
||||
@@ -149,22 +106,18 @@ def handle_deprecated_args(args) -> None:
|
||||
|
||||
|
||||
def encode_execution_providers(execution_providers: List[str]) -> List[str]:
|
||||
return [provider.replace('ExecutionProvider', '').lower() for provider in execution_providers]
|
||||
return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers]
|
||||
|
||||
|
||||
def decode_execution_providers(execution_providers: List[str]) -> List[str]:
|
||||
available_providers = onnxruntime.get_available_providers()
|
||||
encoded_providers = encode_execution_providers(available_providers)
|
||||
|
||||
selected_providers = [available_providers[encoded_providers.index(req)] for req in execution_providers
|
||||
if req in encoded_providers]
|
||||
|
||||
# Default to CPU if no suitable providers are found
|
||||
return selected_providers if selected_providers else ['CPUExecutionProvider']
|
||||
return [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
|
||||
if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)]
|
||||
|
||||
|
||||
def suggest_max_memory() -> int:
|
||||
return 4 if platform.system().lower() == 'darwin' else 16
|
||||
if platform.system().lower() == 'darwin':
|
||||
return 4
|
||||
return 16
|
||||
|
||||
|
||||
def suggest_execution_providers() -> List[str]:
|
||||
@@ -172,43 +125,34 @@ def suggest_execution_providers() -> List[str]:
|
||||
|
||||
|
||||
def suggest_execution_threads() -> int:
|
||||
if 'dml' in modules.globals.execution_providers:
|
||||
if 'DmlExecutionProvider' in modules.globals.execution_providers:
|
||||
return 1
|
||||
if 'rocm' in modules.globals.execution_providers:
|
||||
if 'ROCMExecutionProvider' in modules.globals.execution_providers:
|
||||
return 1
|
||||
return 8
|
||||
|
||||
|
||||
def limit_resources() -> None:
|
||||
# Prevent TensorFlow memory leak
|
||||
# prevent tensorflow memory leak
|
||||
gpus = tensorflow.config.experimental.list_physical_devices('GPU')
|
||||
for gpu in gpus:
|
||||
tensorflow.config.experimental.set_memory_growth(gpu, True)
|
||||
|
||||
# Limit memory usage
|
||||
# limit memory usage
|
||||
if modules.globals.max_memory:
|
||||
memory = modules.globals.max_memory * 1024 ** 3
|
||||
if platform.system().lower() == 'darwin':
|
||||
memory = modules.globals.max_memory * 1024 ** 3
|
||||
elif platform.system().lower() == 'windows':
|
||||
memory = modules.globals.max_memory * 1024 ** 6
|
||||
if platform.system().lower() == 'windows':
|
||||
import ctypes
|
||||
kernel32 = ctypes.windll.kernel32
|
||||
kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
|
||||
else:
|
||||
import resource
|
||||
try:
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_DATA)
|
||||
if memory > hard:
|
||||
print(
|
||||
f"Warning: Requested memory limit {memory / (1024 ** 3)} GB exceeds system's hard limit. Setting to maximum allowed {hard / (1024 ** 3)} GB.")
|
||||
memory = hard
|
||||
resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
|
||||
except ValueError as e:
|
||||
print(f"Warning: Could not set memory limit: {e}. Continuing with default limits.")
|
||||
resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
|
||||
|
||||
|
||||
def release_resources() -> None:
|
||||
if 'cuda' in modules.globals.execution_providers:
|
||||
if 'CUDAExecutionProvider' in modules.globals.execution_providers:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
@@ -219,86 +163,52 @@ def pre_check() -> bool:
|
||||
if not shutil.which('ffmpeg'):
|
||||
update_status('ffmpeg is not installed.')
|
||||
return False
|
||||
if 'cuda' in modules.globals.execution_providers and not torch.cuda.is_available():
|
||||
update_status('CUDA is not available. Please check your GPU or CUDA installation.')
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def update_status(message: str, scope: str = 'DLC.CORE') -> None:
|
||||
print(f'[{scope}] {message}')
|
||||
if not modules.globals.headless and ui.status_label:
|
||||
if not modules.globals.headless:
|
||||
ui.update_status(message)
|
||||
|
||||
|
||||
def start() -> None:
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
if not frame_processor.pre_start():
|
||||
return
|
||||
|
||||
# Process image to image
|
||||
update_status('Processing...')
|
||||
# process image to image
|
||||
if has_image_extension(modules.globals.target_path):
|
||||
process_image_to_image()
|
||||
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
|
||||
return
|
||||
try:
|
||||
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
|
||||
except Exception as e:
|
||||
print("Error copying file:", str(e))
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
update_status('Progressing...', frame_processor.NAME)
|
||||
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path)
|
||||
release_resources()
|
||||
if is_image(modules.globals.target_path):
|
||||
update_status('Processing to image succeed!')
|
||||
else:
|
||||
update_status('Processing to image failed!')
|
||||
return
|
||||
# process image to videos
|
||||
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
|
||||
return
|
||||
|
||||
# Process image to video
|
||||
process_image_to_video()
|
||||
if not modules.globals.map_faces:
|
||||
update_status('Creating temp resources...')
|
||||
create_temp(modules.globals.target_path)
|
||||
update_status('Extracting frames...')
|
||||
extract_frames(modules.globals.target_path)
|
||||
|
||||
|
||||
def process_image_to_image() -> None:
|
||||
if modules.globals.nsfw:
|
||||
from modules.predicter import predict_image
|
||||
if predict_image(modules.globals.target_path):
|
||||
destroy(to_quit=False)
|
||||
update_status('Processing to image ignored!')
|
||||
return
|
||||
|
||||
try:
|
||||
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
|
||||
except Exception as e:
|
||||
print("Error copying file:", str(e))
|
||||
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
update_status('Processing...', frame_processor.NAME)
|
||||
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path,
|
||||
modules.globals.output_path)
|
||||
release_resources()
|
||||
|
||||
if is_image(modules.globals.target_path):
|
||||
update_status('Processing to image succeeded!')
|
||||
else:
|
||||
update_status('Processing to image failed!')
|
||||
|
||||
|
||||
def process_image_to_video() -> None:
|
||||
if modules.globals.nsfw:
|
||||
from modules.predicter import predict_video
|
||||
if predict_video(modules.globals.target_path):
|
||||
destroy(to_quit=False)
|
||||
update_status('Processing to video ignored!')
|
||||
return
|
||||
|
||||
update_status('Creating temporary resources...')
|
||||
create_temp(modules.globals.target_path)
|
||||
update_status('Extracting frames...')
|
||||
extract_frames(modules.globals.target_path)
|
||||
temp_frame_paths = get_temp_frame_paths(modules.globals.target_path)
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
update_status('Processing...', frame_processor.NAME)
|
||||
update_status('Progressing...', frame_processor.NAME)
|
||||
frame_processor.process_video(modules.globals.source_path, temp_frame_paths)
|
||||
release_resources()
|
||||
|
||||
handle_video_fps()
|
||||
handle_video_audio()
|
||||
clean_temp(modules.globals.target_path)
|
||||
|
||||
if is_video(modules.globals.target_path):
|
||||
update_status('Processing to video succeeded!')
|
||||
else:
|
||||
update_status('Processing to video failed!')
|
||||
|
||||
|
||||
def handle_video_fps() -> None:
|
||||
# handles fps
|
||||
if modules.globals.keep_fps:
|
||||
update_status('Detecting fps...')
|
||||
fps = detect_fps(modules.globals.target_path)
|
||||
@@ -307,9 +217,7 @@ def handle_video_fps() -> None:
|
||||
else:
|
||||
update_status('Creating video with 30.0 fps...')
|
||||
create_video(modules.globals.target_path)
|
||||
|
||||
|
||||
def handle_video_audio() -> None:
|
||||
# handle audio
|
||||
if modules.globals.keep_audio:
|
||||
if modules.globals.keep_fps:
|
||||
update_status('Restoring audio...')
|
||||
@@ -318,6 +226,12 @@ def handle_video_audio() -> None:
|
||||
restore_audio(modules.globals.target_path, modules.globals.output_path)
|
||||
else:
|
||||
move_temp(modules.globals.target_path, modules.globals.output_path)
|
||||
# clean and validate
|
||||
clean_temp(modules.globals.target_path)
|
||||
if is_video(modules.globals.target_path):
|
||||
update_status('Processing to video succeed!')
|
||||
else:
|
||||
update_status('Processing to video failed!')
|
||||
|
||||
|
||||
def destroy(to_quit=True) -> None:
|
||||
@@ -327,20 +241,15 @@ def destroy(to_quit=True) -> None:
|
||||
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
parse_args()
|
||||
if not pre_check():
|
||||
parse_args()
|
||||
if not pre_check():
|
||||
return
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
if not frame_processor.pre_check():
|
||||
return
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
if not frame_processor.pre_check():
|
||||
return
|
||||
limit_resources()
|
||||
if modules.globals.headless:
|
||||
start()
|
||||
else:
|
||||
window = ui.init(start, destroy)
|
||||
window.mainloop()
|
||||
except Exception as e:
|
||||
print(f"UI initialization failed: {str(e)}")
|
||||
update_status(f"UI initialization failed: {str(e)}")
|
||||
destroy() # Ensure any resources are cleaned up on failure
|
||||
limit_resources()
|
||||
if modules.globals.headless:
|
||||
start()
|
||||
else:
|
||||
window = ui.init(start, destroy)
|
||||
window.mainloop()
|
||||
|
||||
+176
-14
@@ -1,27 +1,189 @@
|
||||
from typing import Any, Optional
|
||||
import os
|
||||
import shutil
|
||||
from typing import Any
|
||||
import insightface
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import modules.globals
|
||||
from tqdm import tqdm
|
||||
from modules.typing import Frame
|
||||
from modules.cluster_analysis import find_cluster_centroids, find_closest_centroid
|
||||
from modules.utilities import get_temp_directory_path, create_temp, extract_frames, clean_temp, get_temp_frame_paths
|
||||
from pathlib import Path
|
||||
|
||||
FACE_ANALYSER: Optional[insightface.app.FaceAnalysis] = None
|
||||
FACE_ANALYSER = None
|
||||
|
||||
def get_face_analyser() -> insightface.app.FaceAnalysis:
|
||||
|
||||
def get_face_analyser() -> Any:
|
||||
global FACE_ANALYSER
|
||||
|
||||
if FACE_ANALYSER is None:
|
||||
FACE_ANALYSER = insightface.app.FaceAnalysis(
|
||||
name='buffalo_l',
|
||||
providers=modules.globals.execution_providers
|
||||
)
|
||||
FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=modules.globals.execution_providers)
|
||||
FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640))
|
||||
|
||||
return FACE_ANALYSER
|
||||
|
||||
def get_one_face(frame: Frame) -> Optional[Any]:
|
||||
faces = get_face_analyser().get(frame)
|
||||
return min(faces, key=lambda x: x.bbox[0], default=None)
|
||||
|
||||
def get_many_faces(frame: Frame) -> Optional[Any]:
|
||||
faces = get_face_analyser().get(frame)
|
||||
return faces if faces else None
|
||||
def get_one_face(frame: Frame) -> Any:
|
||||
face = get_face_analyser().get(frame)
|
||||
try:
|
||||
return min(face, key=lambda x: x.bbox[0])
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def get_many_faces(frame: Frame) -> Any:
|
||||
try:
|
||||
return get_face_analyser().get(frame)
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def has_valid_map() -> bool:
|
||||
for map in modules.globals.souce_target_map:
|
||||
if "source" in map and "target" in map:
|
||||
return True
|
||||
return False
|
||||
|
||||
def default_source_face() -> Any:
|
||||
for map in modules.globals.souce_target_map:
|
||||
if "source" in map:
|
||||
return map['source']['face']
|
||||
return None
|
||||
|
||||
def simplify_maps() -> Any:
|
||||
centroids = []
|
||||
faces = []
|
||||
for map in modules.globals.souce_target_map:
|
||||
if "source" in map and "target" in map:
|
||||
centroids.append(map['target']['face'].normed_embedding)
|
||||
faces.append(map['source']['face'])
|
||||
|
||||
modules.globals.simple_map = {'source_faces': faces, 'target_embeddings': centroids}
|
||||
return None
|
||||
|
||||
def add_blank_map() -> Any:
|
||||
try:
|
||||
max_id = -1
|
||||
if len(modules.globals.souce_target_map) > 0:
|
||||
max_id = max(modules.globals.souce_target_map, key=lambda x: x['id'])['id']
|
||||
|
||||
modules.globals.souce_target_map.append({
|
||||
'id' : max_id + 1
|
||||
})
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
def get_unique_faces_from_target_image() -> Any:
|
||||
try:
|
||||
modules.globals.souce_target_map = []
|
||||
target_frame = cv2.imread(modules.globals.target_path)
|
||||
many_faces = get_many_faces(target_frame)
|
||||
i = 0
|
||||
|
||||
for face in many_faces:
|
||||
x_min, y_min, x_max, y_max = face['bbox']
|
||||
modules.globals.souce_target_map.append({
|
||||
'id' : i,
|
||||
'target' : {
|
||||
'cv2' : target_frame[int(y_min):int(y_max), int(x_min):int(x_max)],
|
||||
'face' : face
|
||||
}
|
||||
})
|
||||
i = i + 1
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def get_unique_faces_from_target_video() -> Any:
|
||||
try:
|
||||
modules.globals.souce_target_map = []
|
||||
frame_face_embeddings = []
|
||||
face_embeddings = []
|
||||
|
||||
print('Creating temp resources...')
|
||||
clean_temp(modules.globals.target_path)
|
||||
create_temp(modules.globals.target_path)
|
||||
print('Extracting frames...')
|
||||
extract_frames(modules.globals.target_path)
|
||||
|
||||
temp_frame_paths = get_temp_frame_paths(modules.globals.target_path)
|
||||
|
||||
i = 0
|
||||
for temp_frame_path in tqdm(temp_frame_paths, desc="Extracting face embeddings from frames"):
|
||||
temp_frame = cv2.imread(temp_frame_path)
|
||||
many_faces = get_many_faces(temp_frame)
|
||||
|
||||
for face in many_faces:
|
||||
face_embeddings.append(face.normed_embedding)
|
||||
|
||||
frame_face_embeddings.append({'frame': i, 'faces': many_faces, 'location': temp_frame_path})
|
||||
i += 1
|
||||
|
||||
centroids = find_cluster_centroids(face_embeddings)
|
||||
|
||||
for frame in frame_face_embeddings:
|
||||
for face in frame['faces']:
|
||||
closest_centroid_index, _ = find_closest_centroid(centroids, face.normed_embedding)
|
||||
face['target_centroid'] = closest_centroid_index
|
||||
|
||||
for i in range(len(centroids)):
|
||||
modules.globals.souce_target_map.append({
|
||||
'id' : i
|
||||
})
|
||||
|
||||
temp = []
|
||||
for frame in tqdm(frame_face_embeddings, desc=f"Mapping frame embeddings to centroids-{i}"):
|
||||
temp.append({'frame': frame['frame'], 'faces': [face for face in frame['faces'] if face['target_centroid'] == i], 'location': frame['location']})
|
||||
|
||||
modules.globals.souce_target_map[i]['target_faces_in_frame'] = temp
|
||||
|
||||
# dump_faces(centroids, frame_face_embeddings)
|
||||
default_target_face()
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def default_target_face():
|
||||
for map in modules.globals.souce_target_map:
|
||||
best_face = None
|
||||
best_frame = None
|
||||
for frame in map['target_faces_in_frame']:
|
||||
if len(frame['faces']) > 0:
|
||||
best_face = frame['faces'][0]
|
||||
best_frame = frame
|
||||
break
|
||||
|
||||
for frame in map['target_faces_in_frame']:
|
||||
for face in frame['faces']:
|
||||
if face['det_score'] > best_face['det_score']:
|
||||
best_face = face
|
||||
best_frame = frame
|
||||
|
||||
x_min, y_min, x_max, y_max = best_face['bbox']
|
||||
|
||||
target_frame = cv2.imread(best_frame['location'])
|
||||
map['target'] = {
|
||||
'cv2' : target_frame[int(y_min):int(y_max), int(x_min):int(x_max)],
|
||||
'face' : best_face
|
||||
}
|
||||
|
||||
|
||||
def dump_faces(centroids: Any, frame_face_embeddings: list):
|
||||
temp_directory_path = get_temp_directory_path(modules.globals.target_path)
|
||||
|
||||
for i in range(len(centroids)):
|
||||
if os.path.exists(temp_directory_path + f"/{i}") and os.path.isdir(temp_directory_path + f"/{i}"):
|
||||
shutil.rmtree(temp_directory_path + f"/{i}")
|
||||
Path(temp_directory_path + f"/{i}").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for frame in tqdm(frame_face_embeddings, desc=f"Copying faces to temp/./{i}"):
|
||||
temp_frame = cv2.imread(frame['location'])
|
||||
|
||||
j = 0
|
||||
for face in frame['faces']:
|
||||
if face['target_centroid'] == i:
|
||||
x_min, y_min, x_max, y_max = face['bbox']
|
||||
|
||||
if temp_frame[int(y_min):int(y_max), int(x_min):int(x_max)].size > 0:
|
||||
cv2.imwrite(temp_directory_path + f"/{i}/{frame['frame']}_{j}.png", temp_frame[int(y_min):int(y_max), int(x_min):int(x_max)])
|
||||
j += 1
|
||||
+8
-6
@@ -1,5 +1,5 @@
|
||||
import os
|
||||
from typing import List, Dict
|
||||
from typing import List, Dict, Any
|
||||
|
||||
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
WORKFLOW_DIR = os.path.join(ROOT_DIR, 'workflow')
|
||||
@@ -9,6 +9,9 @@ file_types = [
|
||||
('Video', ('*.mp4','*.mkv'))
|
||||
]
|
||||
|
||||
souce_target_map = []
|
||||
simple_map = {}
|
||||
|
||||
source_path = None
|
||||
target_path = None
|
||||
output_path = None
|
||||
@@ -17,6 +20,9 @@ keep_fps = None
|
||||
keep_audio = None
|
||||
keep_frames = None
|
||||
many_faces = None
|
||||
map_faces = None
|
||||
color_correction = None # New global variable for color correction toggle
|
||||
nsfw_filter = None
|
||||
video_encoder = None
|
||||
video_quality = None
|
||||
live_mirror = None
|
||||
@@ -27,9 +33,5 @@ execution_threads = None
|
||||
headless = None
|
||||
log_level = 'error'
|
||||
fp_ui: Dict[str, bool] = {}
|
||||
nsfw = None
|
||||
camera_input_combobox = None
|
||||
webcam_preview_running = False
|
||||
enhancer_upscale_factor = 1
|
||||
source_image_scaling_factor = 2
|
||||
sr_scale_factor = 4
|
||||
webcam_preview_running = False
|
||||
@@ -2,68 +2,58 @@ from typing import Any, List
|
||||
import cv2
|
||||
import insightface
|
||||
import threading
|
||||
import os
|
||||
|
||||
import modules.globals
|
||||
import modules.processors.frame.core
|
||||
from modules.core import update_status
|
||||
from modules.face_analyser import get_one_face, get_many_faces
|
||||
from modules.face_analyser import get_one_face, get_many_faces, default_source_face
|
||||
from modules.typing import Face, Frame
|
||||
from modules.utilities import conditional_download, resolve_relative_path, is_image, is_video
|
||||
import numpy as np
|
||||
from modules.cluster_analysis import find_closest_centroid
|
||||
|
||||
FACE_SWAPPER = None
|
||||
THREAD_LOCK = threading.Lock()
|
||||
NAME = 'DLC.FACE-SWAPPER'
|
||||
|
||||
|
||||
def pre_check() -> bool:
|
||||
download_directory_path = resolve_relative_path('../models')
|
||||
conditional_download(download_directory_path, [
|
||||
'https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128.onnx'
|
||||
])
|
||||
conditional_download(download_directory_path, ['https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128_fp16.onnx'])
|
||||
return True
|
||||
|
||||
|
||||
def pre_start() -> bool:
|
||||
if not is_image(modules.globals.source_path):
|
||||
if not modules.globals.map_faces and not is_image(modules.globals.source_path):
|
||||
update_status('Select an image for source path.', NAME)
|
||||
return False
|
||||
elif not get_one_face(cv2.imread(modules.globals.source_path)):
|
||||
update_status('No face detected in the source path.', NAME)
|
||||
elif not modules.globals.map_faces and not get_one_face(cv2.imread(modules.globals.source_path)):
|
||||
update_status('No face in source path detected.', NAME)
|
||||
return False
|
||||
if not is_image(modules.globals.target_path) and not is_video(modules.globals.target_path):
|
||||
update_status('Select an image or video for target path.', NAME)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_face_swapper() -> Any:
|
||||
global FACE_SWAPPER
|
||||
|
||||
with THREAD_LOCK:
|
||||
if FACE_SWAPPER is None:
|
||||
model_path = resolve_relative_path('../models/inswapper_128.onnx')
|
||||
model_path = resolve_relative_path('../models/inswapper_128_fp16.onnx')
|
||||
FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=modules.globals.execution_providers)
|
||||
return FACE_SWAPPER
|
||||
|
||||
def upscale_image(image: np.ndarray, scaling_factor: int = modules.globals.source_image_scaling_factor) -> np.ndarray:
|
||||
"""
|
||||
Upscales the given image by the specified scaling factor.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The input image to upscale.
|
||||
scaling_factor (int): The factor by which to upscale the image.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The upscaled image.
|
||||
"""
|
||||
height, width = image.shape[:2]
|
||||
new_size = (width * scaling_factor, height * scaling_factor)
|
||||
upscaled_image = cv2.resize(image, new_size, interpolation=cv2.INTER_CUBIC)
|
||||
return upscaled_image
|
||||
|
||||
def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
|
||||
return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)
|
||||
|
||||
|
||||
def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
|
||||
# Ensure the frame is in RGB format if color correction is enabled
|
||||
if modules.globals.color_correction:
|
||||
temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
if modules.globals.many_faces:
|
||||
many_faces = get_many_faces(temp_frame)
|
||||
if many_faces:
|
||||
@@ -75,30 +65,99 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
return temp_frame
|
||||
|
||||
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
|
||||
source_image = cv2.imread(source_path)
|
||||
if source_image is None:
|
||||
print(f"Failed to load source image from {source_path}")
|
||||
return
|
||||
# Upscale the source image for better quality
|
||||
source_image_upscaled = upscale_image(source_image, scaling_factor=2)
|
||||
source_face = get_one_face(source_image_upscaled)
|
||||
|
||||
for temp_frame_path in temp_frame_paths:
|
||||
temp_frame = cv2.imread(temp_frame_path)
|
||||
try:
|
||||
result = process_frame(source_face, temp_frame)
|
||||
cv2.imwrite(temp_frame_path, result)
|
||||
except Exception as exception:
|
||||
print(f"Error processing frame {temp_frame_path}: {exception}")
|
||||
if progress:
|
||||
progress.update(1)
|
||||
def process_frame_v2(temp_frame: Frame, temp_frame_path: str = "") -> Frame:
|
||||
if is_image(modules.globals.target_path):
|
||||
if modules.globals.many_faces:
|
||||
source_face = default_source_face()
|
||||
for map in modules.globals.souce_target_map:
|
||||
target_face = map['target']['face']
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
|
||||
elif not modules.globals.many_faces:
|
||||
for map in modules.globals.souce_target_map:
|
||||
if "source" in map:
|
||||
source_face = map['source']['face']
|
||||
target_face = map['target']['face']
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
|
||||
elif is_video(modules.globals.target_path):
|
||||
if modules.globals.many_faces:
|
||||
source_face = default_source_face()
|
||||
for map in modules.globals.souce_target_map:
|
||||
target_frame = [f for f in map['target_faces_in_frame'] if f['location'] == temp_frame_path]
|
||||
|
||||
for frame in target_frame:
|
||||
for target_face in frame['faces']:
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
|
||||
elif not modules.globals.many_faces:
|
||||
for map in modules.globals.souce_target_map:
|
||||
if "source" in map:
|
||||
target_frame = [f for f in map['target_faces_in_frame'] if f['location'] == temp_frame_path]
|
||||
source_face = map['source']['face']
|
||||
|
||||
for frame in target_frame:
|
||||
for target_face in frame['faces']:
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
else:
|
||||
many_faces = get_many_faces(temp_frame)
|
||||
if modules.globals.many_faces:
|
||||
source_face = default_source_face()
|
||||
if many_faces:
|
||||
for target_face in many_faces:
|
||||
temp_frame = swap_face(source_face, target_face, temp_frame)
|
||||
|
||||
elif not modules.globals.many_faces:
|
||||
if many_faces:
|
||||
for target_face in many_faces:
|
||||
closest_centroid_index, _ = find_closest_centroid(modules.globals.simple_map['target_embeddings'], target_face.normed_embedding)
|
||||
|
||||
temp_frame = swap_face(modules.globals.simple_map['source_faces'][closest_centroid_index], target_face, temp_frame)
|
||||
return temp_frame
|
||||
|
||||
|
||||
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
|
||||
if not modules.globals.map_faces:
|
||||
source_face = get_one_face(cv2.imread(source_path))
|
||||
for temp_frame_path in temp_frame_paths:
|
||||
temp_frame = cv2.imread(temp_frame_path)
|
||||
try:
|
||||
result = process_frame(source_face, temp_frame)
|
||||
cv2.imwrite(temp_frame_path, result)
|
||||
except Exception as exception:
|
||||
print(exception)
|
||||
pass
|
||||
if progress:
|
||||
progress.update(1)
|
||||
else:
|
||||
for temp_frame_path in temp_frame_paths:
|
||||
temp_frame = cv2.imread(temp_frame_path)
|
||||
try:
|
||||
result = process_frame_v2(temp_frame, temp_frame_path)
|
||||
cv2.imwrite(temp_frame_path, result)
|
||||
except Exception as exception:
|
||||
print(exception)
|
||||
pass
|
||||
if progress:
|
||||
progress.update(1)
|
||||
|
||||
|
||||
def process_image(source_path: str, target_path: str, output_path: str) -> None:
|
||||
source_face = get_one_face(cv2.imread(source_path))
|
||||
target_frame = cv2.imread(target_path)
|
||||
result = process_frame(source_face, target_frame)
|
||||
cv2.imwrite(output_path, result)
|
||||
if not modules.globals.map_faces:
|
||||
source_face = get_one_face(cv2.imread(source_path))
|
||||
target_frame = cv2.imread(target_path)
|
||||
result = process_frame(source_face, target_frame)
|
||||
cv2.imwrite(output_path, result)
|
||||
else:
|
||||
if modules.globals.many_faces:
|
||||
update_status('Many faces enabled. Using first source image. Progressing...', NAME)
|
||||
target_frame = cv2.imread(output_path)
|
||||
result = process_frame_v2(target_frame)
|
||||
cv2.imwrite(output_path, result)
|
||||
|
||||
|
||||
def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
|
||||
if modules.globals.map_faces and modules.globals.many_faces:
|
||||
update_status('Many faces enabled. Using first source image. Progressing...', NAME)
|
||||
modules.processors.frame.core.process_video(source_path, temp_frame_paths, process_frames)
|
||||
|
||||
+372
-29
@@ -19,7 +19,7 @@ if platform.system() == 'Windows' or platform.system() == 'Linux': # Windows or
|
||||
|
||||
import modules.globals
|
||||
import modules.metadata
|
||||
from modules.face_analyser import get_one_face
|
||||
from modules.face_analyser import get_one_face, get_unique_faces_from_target_image, get_unique_faces_from_target_video, add_blank_map, has_valid_map, simplify_maps
|
||||
from modules.capturer import get_video_frame, get_video_frame_total
|
||||
from modules.processors.frame.core import get_frame_processors_modules
|
||||
from modules.utilities import is_image, is_video, resolve_relative_path
|
||||
@@ -34,6 +34,22 @@ PREVIEW_MAX_WIDTH = 1200
|
||||
PREVIEW_DEFAULT_WIDTH = 960
|
||||
PREVIEW_DEFAULT_HEIGHT = 540
|
||||
|
||||
POPUP_WIDTH = 750
|
||||
POPUP_HEIGHT = 810
|
||||
POPUP_SCROLL_WIDTH = 740,
|
||||
POPUP_SCROLL_HEIGHT = 700
|
||||
|
||||
POPUP_LIVE_WIDTH = 900
|
||||
POPUP_LIVE_HEIGHT = 820
|
||||
POPUP_LIVE_SCROLL_WIDTH = 890,
|
||||
POPUP_LIVE_SCROLL_HEIGHT = 700
|
||||
|
||||
MAPPER_PREVIEW_MAX_HEIGHT = 100
|
||||
MAPPER_PREVIEW_MAX_WIDTH = 100
|
||||
|
||||
DEFAULT_BUTTON_WIDTH = 200
|
||||
DEFAULT_BUTTON_HEIGHT = 40
|
||||
|
||||
RECENT_DIRECTORY_SOURCE = None
|
||||
RECENT_DIRECTORY_TARGET = None
|
||||
RECENT_DIRECTORY_OUTPUT = None
|
||||
@@ -43,6 +59,11 @@ preview_slider = None
|
||||
source_label = None
|
||||
target_label = None
|
||||
status_label = None
|
||||
popup_status_label = None
|
||||
popup_status_label_live = None
|
||||
source_label_dict = {}
|
||||
source_label_dict_live = {}
|
||||
target_label_dict_live = {}
|
||||
|
||||
img_ft, vid_ft = modules.globals.file_types
|
||||
|
||||
@@ -148,8 +169,12 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C
|
||||
nsfw_switch = ctk.CTkSwitch(root, text='NSFW', variable=nsfw_value, cursor='hand2', command=lambda: setattr(modules.globals, 'nsfw', nsfw_value.get()))
|
||||
nsfw_switch.place(relx=0.6, rely=0.6125)
|
||||
|
||||
start_button = ctk.CTkButton(root, text='Start', cursor='hand2', command=lambda: select_output_path(start))
|
||||
start_button.place(relx=0.15, rely=0.7, relwidth=0.2, relheight=0.05)
|
||||
map_faces = ctk.BooleanVar(value=modules.globals.map_faces)
|
||||
map_faces_switch = ctk.CTkSwitch(root, text='Map faces', variable=map_faces, cursor='hand2', command=lambda: setattr(modules.globals, 'map_faces', map_faces.get()))
|
||||
map_faces_switch.place(relx=0.1, rely=0.75)
|
||||
|
||||
start_button = ctk.CTkButton(root, text='Start', cursor='hand2', command=lambda: analyze_target(start, root))
|
||||
start_button.place(relx=0.15, rely=0.80, relwidth=0.2, relheight=0.05)
|
||||
|
||||
stop_button = ctk.CTkButton(root, text='Destroy', cursor='hand2', command=destroy)
|
||||
stop_button.place(relx=0.4, rely=0.7, relwidth=0.2, relheight=0.05)
|
||||
@@ -157,22 +182,8 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C
|
||||
preview_button = ctk.CTkButton(root, text='Preview', cursor='hand2', command=toggle_preview)
|
||||
preview_button.place(relx=0.65, rely=0.7, relwidth=0.2, relheight=0.05)
|
||||
|
||||
camera_label = ctk.CTkLabel(root, text="Select Camera:")
|
||||
camera_label.place(relx=0.4, rely=0.7525, relwidth=0.2, relheight=0.05)
|
||||
|
||||
available_cameras = get_available_cameras()
|
||||
available_camera_strings = [str(cam) for cam in available_cameras]
|
||||
|
||||
camera_variable = ctk.StringVar(value=available_camera_strings[0] if available_camera_strings else "No cameras found")
|
||||
camera_optionmenu = ctk.CTkOptionMenu(root, variable=camera_variable, values=available_camera_strings)
|
||||
camera_optionmenu.place(relx=0.65, rely=0.7525, relwidth=0.2, relheight=0.05)
|
||||
|
||||
virtual_cam_out_value = ctk.BooleanVar(value=False)
|
||||
virtual_cam_out_switch = ctk.CTkSwitch(root, text='Virtual Cam Output (OBS)', variable=virtual_cam_out_value, cursor='hand2')
|
||||
virtual_cam_out_switch.place(relx=0.4, rely=0.805)
|
||||
|
||||
live_button = ctk.CTkButton(root, text='Live', cursor='hand2', command=lambda: webcam_preview(camera_variable.get(), virtual_cam_out_value.get()))
|
||||
live_button.place(relx=0.15, rely=0.7525, relwidth=0.2, relheight=0.05)
|
||||
live_button = ctk.CTkButton(root, text='Live', cursor='hand2', command=lambda: webcam_preview(root))
|
||||
live_button.place(relx=0.40, rely=0.86, relwidth=0.2, relheight=0.05)
|
||||
|
||||
status_label = ctk.CTkLabel(root, text=None, justify='center')
|
||||
status_label.place(relx=0.1, relwidth=0.8, rely=0.875)
|
||||
@@ -184,6 +195,109 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C
|
||||
|
||||
return root
|
||||
|
||||
def analyze_target(start: Callable[[], None], root: ctk.CTk):
|
||||
if POPUP != None and POPUP.winfo_exists():
|
||||
update_status("Please complete pop-up or close it.")
|
||||
return
|
||||
|
||||
if modules.globals.map_faces:
|
||||
modules.globals.souce_target_map = []
|
||||
|
||||
if is_image(modules.globals.target_path):
|
||||
update_status('Getting unique faces')
|
||||
get_unique_faces_from_target_image()
|
||||
elif is_video(modules.globals.target_path):
|
||||
update_status('Getting unique faces')
|
||||
get_unique_faces_from_target_video()
|
||||
|
||||
if len(modules.globals.souce_target_map) > 0:
|
||||
create_source_target_popup(start, root, modules.globals.souce_target_map)
|
||||
else:
|
||||
update_status("No faces found in target")
|
||||
else:
|
||||
select_output_path(start)
|
||||
|
||||
def create_source_target_popup(start: Callable[[], None], root: ctk.CTk, map: list) -> None:
|
||||
global POPUP, popup_status_label
|
||||
|
||||
POPUP = ctk.CTkToplevel(root)
|
||||
POPUP.title("Source x Target Mapper")
|
||||
POPUP.geometry(f"{POPUP_WIDTH}x{POPUP_HEIGHT}")
|
||||
POPUP.focus()
|
||||
|
||||
def on_submit_click(start):
|
||||
if has_valid_map():
|
||||
POPUP.destroy()
|
||||
select_output_path(start)
|
||||
else:
|
||||
update_pop_status("Atleast 1 source with target is required!")
|
||||
|
||||
scrollable_frame = ctk.CTkScrollableFrame(POPUP, width=POPUP_SCROLL_WIDTH, height=POPUP_SCROLL_HEIGHT)
|
||||
scrollable_frame.grid(row=0, column=0, padx=0, pady=0, sticky='nsew')
|
||||
|
||||
def on_button_click(map, button_num):
|
||||
map = update_popup_source(scrollable_frame, map, button_num)
|
||||
|
||||
for item in map:
|
||||
id = item['id']
|
||||
|
||||
button = ctk.CTkButton(scrollable_frame, text="Select source image", command=lambda id=id: on_button_click(map, id), width=DEFAULT_BUTTON_WIDTH, height=DEFAULT_BUTTON_HEIGHT)
|
||||
button.grid(row=id, column=0, padx=50, pady=10)
|
||||
|
||||
x_label = ctk.CTkLabel(scrollable_frame, text=f"X", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
x_label.grid(row=id, column=2, padx=10, pady=10)
|
||||
|
||||
image = Image.fromarray(cv2.cvtColor(item['target']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
target_image = ctk.CTkLabel(scrollable_frame, text=f"T-{id}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
target_image.grid(row=id, column=3, padx=10, pady=10)
|
||||
target_image.configure(image=tk_image)
|
||||
|
||||
popup_status_label = ctk.CTkLabel(POPUP, text=None, justify='center')
|
||||
popup_status_label.grid(row=1, column=0, pady=15)
|
||||
|
||||
close_button = ctk.CTkButton(POPUP, text="Submit", command=lambda: on_submit_click(start))
|
||||
close_button.grid(row=2, column=0, pady=10)
|
||||
|
||||
|
||||
def update_popup_source(scrollable_frame: ctk.CTkScrollableFrame, map: list, button_num: int) -> list:
|
||||
global source_label_dict
|
||||
|
||||
source_path = ctk.filedialog.askopenfilename(title='select an source image', initialdir=RECENT_DIRECTORY_SOURCE, filetypes=[img_ft])
|
||||
|
||||
if "source" in map[button_num]:
|
||||
map[button_num].pop("source")
|
||||
source_label_dict[button_num].destroy()
|
||||
del source_label_dict[button_num]
|
||||
|
||||
if source_path == "":
|
||||
return map
|
||||
else:
|
||||
cv2_img = cv2.imread(source_path)
|
||||
face = get_one_face(cv2_img)
|
||||
|
||||
if face:
|
||||
x_min, y_min, x_max, y_max = face['bbox']
|
||||
|
||||
map[button_num]['source'] = {
|
||||
'cv2' : cv2_img[int(y_min):int(y_max), int(x_min):int(x_max)],
|
||||
'face' : face
|
||||
}
|
||||
|
||||
image = Image.fromarray(cv2.cvtColor(map[button_num]['source']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
source_image = ctk.CTkLabel(scrollable_frame, text=f"S-{button_num}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
source_image.grid(row=button_num, column=1, padx=10, pady=10)
|
||||
source_image.configure(image=tk_image)
|
||||
source_label_dict[button_num] = source_image
|
||||
else:
|
||||
update_pop_status("Face could not be detected in last upload!")
|
||||
return map
|
||||
|
||||
|
||||
def create_preview(parent: ctk.CTk) -> ctk.CTkToplevel:
|
||||
global preview_label, preview_slider
|
||||
@@ -206,6 +320,11 @@ def update_status(text: str) -> None:
|
||||
status_label.configure(text=text)
|
||||
ROOT.update()
|
||||
|
||||
def update_pop_status(text: str) -> None:
|
||||
popup_status_label.configure(text=text)
|
||||
|
||||
def update_pop_live_status(text: str) -> None:
|
||||
popup_status_label_live.configure(text=text)
|
||||
|
||||
def update_tumbler(var: str, value: bool) -> None:
|
||||
modules.globals.fp_ui[var] = value
|
||||
@@ -394,10 +513,75 @@ def fit_image_to_size(image, width: int, height: int):
|
||||
new_size = (int(ratio * w), int(ratio * h))
|
||||
return cv2.resize(image, dsize=new_size)
|
||||
|
||||
def webcam_preview(camera_name: str, virtual_cam_output: bool):
|
||||
if modules.globals.source_path is None:
|
||||
return
|
||||
|
||||
def render_image_preview(image_path: str, size: Tuple[int, int]) -> ctk.CTkImage:
|
||||
image = Image.open(image_path)
|
||||
if size:
|
||||
image = ImageOps.fit(image, size, Image.LANCZOS)
|
||||
return ctk.CTkImage(image, size=image.size)
|
||||
|
||||
|
||||
def render_video_preview(video_path: str, size: Tuple[int, int], frame_number: int = 0) -> ctk.CTkImage:
|
||||
capture = cv2.VideoCapture(video_path)
|
||||
if frame_number:
|
||||
capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
||||
has_frame, frame = capture.read()
|
||||
if has_frame:
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
if size:
|
||||
image = ImageOps.fit(image, size, Image.LANCZOS)
|
||||
return ctk.CTkImage(image, size=image.size)
|
||||
capture.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def toggle_preview() -> None:
|
||||
if PREVIEW.state() == 'normal':
|
||||
PREVIEW.withdraw()
|
||||
elif modules.globals.source_path and modules.globals.target_path:
|
||||
init_preview()
|
||||
update_preview()
|
||||
|
||||
|
||||
def init_preview() -> None:
|
||||
if is_image(modules.globals.target_path):
|
||||
preview_slider.pack_forget()
|
||||
if is_video(modules.globals.target_path):
|
||||
video_frame_total = get_video_frame_total(modules.globals.target_path)
|
||||
preview_slider.configure(to=video_frame_total)
|
||||
preview_slider.pack(fill='x')
|
||||
preview_slider.set(0)
|
||||
|
||||
|
||||
def update_preview(frame_number: int = 0) -> None:
|
||||
if modules.globals.source_path and modules.globals.target_path:
|
||||
update_status('Processing...')
|
||||
temp_frame = get_video_frame(modules.globals.target_path, frame_number)
|
||||
if modules.globals.nsfw_filter and check_and_ignore_nsfw(temp_frame):
|
||||
return
|
||||
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
|
||||
temp_frame = frame_processor.process_frame(
|
||||
get_one_face(cv2.imread(modules.globals.source_path)),
|
||||
temp_frame
|
||||
)
|
||||
image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))
|
||||
image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
image = ctk.CTkImage(image, size=image.size)
|
||||
preview_label.configure(image=image)
|
||||
update_status('Processing succeed!')
|
||||
PREVIEW.deiconify()
|
||||
|
||||
def webcam_preview(root: ctk.CTk):
|
||||
if not modules.globals.map_faces:
|
||||
if modules.globals.source_path is None:
|
||||
# No image selected
|
||||
return
|
||||
create_webcam_preview()
|
||||
else:
|
||||
modules.globals.souce_target_map = []
|
||||
create_source_target_popup_for_webcam(root, modules.globals.souce_target_map)
|
||||
|
||||
def create_webcam_preview():
|
||||
global preview_label, PREVIEW
|
||||
|
||||
WIDTH = 960
|
||||
@@ -440,11 +624,7 @@ def webcam_preview(camera_name: str, virtual_cam_output: bool):
|
||||
while preview_running:
|
||||
preview_running = webcam_preview_loop(camera, source_image, frame_processors, virtual_cam)
|
||||
|
||||
while preview_running:
|
||||
preview_running = webcam_preview_loop(camera, source_image, frame_processors)
|
||||
|
||||
if camera: camera.release()
|
||||
PREVIEW.withdraw()
|
||||
temp_frame = frame.copy() #Create a copy of the frame
|
||||
|
||||
|
||||
def get_camera_index_by_name(camera_name: str) -> int:
|
||||
@@ -459,6 +639,18 @@ def get_camera_index_by_name(camera_name: str) -> int:
|
||||
return get_available_cameras().index(camera_name)
|
||||
return -1
|
||||
|
||||
if not modules.globals.map_faces:
|
||||
# Select and save face image only once
|
||||
if source_image is None and modules.globals.source_path:
|
||||
source_image = get_one_face(cv2.imread(modules.globals.source_path))
|
||||
|
||||
for frame_processor in frame_processors:
|
||||
temp_frame = frame_processor.process_frame(source_image, temp_frame)
|
||||
else:
|
||||
modules.globals.target_path = None
|
||||
|
||||
for frame_processor in frame_processors:
|
||||
temp_frame = frame_processor.process_frame_v2(temp_frame)
|
||||
|
||||
def get_available_cameras():
|
||||
"""Get available camera names (cross-platform)."""
|
||||
@@ -490,5 +682,156 @@ def get_available_cameras():
|
||||
cap.release()
|
||||
index += 1
|
||||
|
||||
available_cameras = devices
|
||||
return available_cameras
|
||||
if PREVIEW.state() == 'withdrawn':
|
||||
break
|
||||
|
||||
camera.release()
|
||||
PREVIEW.withdraw() # Close preview window when loop is finished
|
||||
|
||||
|
||||
def create_source_target_popup_for_webcam(root: ctk.CTk, map: list) -> None:
|
||||
global POPUP_LIVE, popup_status_label_live
|
||||
|
||||
POPUP_LIVE = ctk.CTkToplevel(root)
|
||||
POPUP_LIVE.title("Source x Target Mapper")
|
||||
POPUP_LIVE.geometry(f"{POPUP_LIVE_WIDTH}x{POPUP_LIVE_HEIGHT}")
|
||||
POPUP_LIVE.focus()
|
||||
|
||||
def on_submit_click():
|
||||
if has_valid_map():
|
||||
POPUP_LIVE.destroy()
|
||||
simplify_maps()
|
||||
create_webcam_preview()
|
||||
else:
|
||||
update_pop_live_status("Atleast 1 source with target is required!")
|
||||
|
||||
def on_add_click():
|
||||
add_blank_map()
|
||||
refresh_data(map)
|
||||
update_pop_live_status("Please provide mapping!")
|
||||
|
||||
popup_status_label_live = ctk.CTkLabel(POPUP_LIVE, text=None, justify='center')
|
||||
popup_status_label_live.grid(row=1, column=0, pady=15)
|
||||
|
||||
add_button = ctk.CTkButton(POPUP_LIVE, text="Add", command=lambda: on_add_click())
|
||||
add_button.place(relx=0.2, rely=0.92, relwidth=0.2, relheight=0.05)
|
||||
|
||||
close_button = ctk.CTkButton(POPUP_LIVE, text="Submit", command=lambda: on_submit_click())
|
||||
close_button.place(relx=0.6, rely=0.92, relwidth=0.2, relheight=0.05)
|
||||
|
||||
|
||||
def refresh_data(map: list):
|
||||
global POPUP_LIVE
|
||||
|
||||
scrollable_frame = ctk.CTkScrollableFrame(POPUP_LIVE, width=POPUP_LIVE_SCROLL_WIDTH, height=POPUP_LIVE_SCROLL_HEIGHT)
|
||||
scrollable_frame.grid(row=0, column=0, padx=0, pady=0, sticky='nsew')
|
||||
|
||||
def on_sbutton_click(map, button_num):
|
||||
map = update_webcam_source(scrollable_frame, map, button_num)
|
||||
|
||||
def on_tbutton_click(map, button_num):
|
||||
map = update_webcam_target(scrollable_frame, map, button_num)
|
||||
|
||||
for item in map:
|
||||
id = item['id']
|
||||
|
||||
button = ctk.CTkButton(scrollable_frame, text="Select source image", command=lambda id=id: on_sbutton_click(map, id), width=DEFAULT_BUTTON_WIDTH, height=DEFAULT_BUTTON_HEIGHT)
|
||||
button.grid(row=id, column=0, padx=30, pady=10)
|
||||
|
||||
x_label = ctk.CTkLabel(scrollable_frame, text=f"X", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
x_label.grid(row=id, column=2, padx=10, pady=10)
|
||||
|
||||
button = ctk.CTkButton(scrollable_frame, text="Select target image", command=lambda id=id: on_tbutton_click(map, id), width=DEFAULT_BUTTON_WIDTH, height=DEFAULT_BUTTON_HEIGHT)
|
||||
button.grid(row=id, column=3, padx=20, pady=10)
|
||||
|
||||
if "source" in item:
|
||||
image = Image.fromarray(cv2.cvtColor(item['source']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
source_image = ctk.CTkLabel(scrollable_frame, text=f"S-{id}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
source_image.grid(row=id, column=1, padx=10, pady=10)
|
||||
source_image.configure(image=tk_image)
|
||||
|
||||
if "target" in item:
|
||||
image = Image.fromarray(cv2.cvtColor(item['target']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
target_image = ctk.CTkLabel(scrollable_frame, text=f"T-{id}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
target_image.grid(row=id, column=4, padx=20, pady=10)
|
||||
target_image.configure(image=tk_image)
|
||||
|
||||
|
||||
def update_webcam_source(scrollable_frame: ctk.CTkScrollableFrame, map: list, button_num: int) -> list:
|
||||
global source_label_dict_live
|
||||
|
||||
source_path = ctk.filedialog.askopenfilename(title='select an source image', initialdir=RECENT_DIRECTORY_SOURCE, filetypes=[img_ft])
|
||||
|
||||
if "source" in map[button_num]:
|
||||
map[button_num].pop("source")
|
||||
source_label_dict_live[button_num].destroy()
|
||||
del source_label_dict_live[button_num]
|
||||
|
||||
if source_path == "":
|
||||
return map
|
||||
else:
|
||||
cv2_img = cv2.imread(source_path)
|
||||
face = get_one_face(cv2_img)
|
||||
|
||||
if face:
|
||||
x_min, y_min, x_max, y_max = face['bbox']
|
||||
|
||||
map[button_num]['source'] = {
|
||||
'cv2' : cv2_img[int(y_min):int(y_max), int(x_min):int(x_max)],
|
||||
'face' : face
|
||||
}
|
||||
|
||||
image = Image.fromarray(cv2.cvtColor(map[button_num]['source']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
source_image = ctk.CTkLabel(scrollable_frame, text=f"S-{button_num}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
source_image.grid(row=button_num, column=1, padx=10, pady=10)
|
||||
source_image.configure(image=tk_image)
|
||||
source_label_dict_live[button_num] = source_image
|
||||
else:
|
||||
update_pop_live_status("Face could not be detected in last upload!")
|
||||
return map
|
||||
|
||||
def update_webcam_target(scrollable_frame: ctk.CTkScrollableFrame, map: list, button_num: int) -> list:
|
||||
global target_label_dict_live
|
||||
|
||||
target_path = ctk.filedialog.askopenfilename(title='select an target image', initialdir=RECENT_DIRECTORY_SOURCE, filetypes=[img_ft])
|
||||
|
||||
if "target" in map[button_num]:
|
||||
map[button_num].pop("target")
|
||||
target_label_dict_live[button_num].destroy()
|
||||
del target_label_dict_live[button_num]
|
||||
|
||||
if target_path == "":
|
||||
return map
|
||||
else:
|
||||
cv2_img = cv2.imread(target_path)
|
||||
face = get_one_face(cv2_img)
|
||||
|
||||
if face:
|
||||
x_min, y_min, x_max, y_max = face['bbox']
|
||||
|
||||
map[button_num]['target'] = {
|
||||
'cv2' : cv2_img[int(y_min):int(y_max), int(x_min):int(x_max)],
|
||||
'face' : face
|
||||
}
|
||||
|
||||
image = Image.fromarray(cv2.cvtColor(map[button_num]['target']['cv2'], cv2.COLOR_BGR2RGB))
|
||||
image = image.resize((MAPPER_PREVIEW_MAX_WIDTH, MAPPER_PREVIEW_MAX_HEIGHT), Image.LANCZOS)
|
||||
tk_image = ctk.CTkImage(image, size=image.size)
|
||||
|
||||
target_image = ctk.CTkLabel(scrollable_frame, text=f"T-{button_num}", width=MAPPER_PREVIEW_MAX_WIDTH, height=MAPPER_PREVIEW_MAX_HEIGHT)
|
||||
target_image.grid(row=button_num, column=4, padx=20, pady=10)
|
||||
target_image.configure(image=tk_image)
|
||||
target_label_dict_live[button_num] = target_image
|
||||
else:
|
||||
update_pop_live_status("Face could not be detected in last upload!")
|
||||
return map
|
||||
|
||||
|
||||
Reference in New Issue
Block a user