- allow processing all faces within a frame

This commit is contained in:
chris 2023-05-31 11:28:13 -04:00
parent 90e2ba0fc7
commit 6fbc2d32b3
4 changed files with 32 additions and 6 deletions

View File

@ -11,3 +11,11 @@ def get_face(img_data):
return sorted(analysed, key=lambda x: x.bbox[0])[0] return sorted(analysed, key=lambda x: x.bbox[0])[0]
except IndexError: except IndexError:
return None return None
def get_all_faces(img_data):
analysed = face_analyser.get(img_data)
try:
return analysed
except IndexError:
return None

View File

@ -2,3 +2,4 @@ import onnxruntime
use_gpu = False use_gpu = False
providers = onnxruntime.get_available_providers() providers = onnxruntime.get_available_providers()
all_faces = False

View File

@ -16,13 +16,26 @@ def process_video(source_img, frame_paths):
for frame_path in frame_paths: for frame_path in frame_paths:
frame = cv2.imread(frame_path) frame = cv2.imread(frame_path)
try: try:
face = get_face(frame) if core.globals.all_faces:
if face: all_faces = get_all_faces(frame)
result = face_swapper.get(frame, face, source_face, paste_back=True) result = frame
cv2.imwrite(frame_path, result) for singleFace in all_faces:
print('.', end='', flush=True) if singleFace:
result = face_swapper.get(result, singleFace, source_face, paste_back=True)
print('.', end='', flush=True)
else:
print('S', end='', flush=True)
if result is not None:
cv2.imwrite(frame_path, result)
else: else:
print('S', end='', flush=True) face = get_face(frame)
if face:
result = face_swapper.get(frame, face, source_face, paste_back=True)
cv2.imwrite(frame_path, result)
print('.', end='', flush=True)
else:
print('S', end='', flush=True)
except Exception as e: except Exception as e:
print('E', end='', flush=True) print('E', end='', flush=True)
pass pass

4
run.py
View File

@ -13,6 +13,9 @@ elif 'ROCMExecutionProvider' not in core.globals.providers:
import torch import torch
if not torch.cuda.is_available(): if not torch.cuda.is_available():
quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.") quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
if '--all-faces' in sys.argv or '-a' in sys.argv:
core.globals.all_faces = True
import glob import glob
import argparse import argparse
@ -41,6 +44,7 @@ parser.add_argument('-o', '--output', help='save output to this file', dest='out
parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False)
parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
parser.add_argument('-a', '--all-faces', help='swap all faces in frame', dest='all_faces', default=False)
for name, value in vars(parser.parse_args()).items(): for name, value in vars(parser.parse_args()).items():
args[name] = value args[name] = value