add support for two faces

pull/372/head
ivideogameboss 2024-08-16 10:43:03 -05:00
parent d602d2e3c6
commit fc52401379
3 changed files with 63 additions and 15 deletions

View File

@ -22,6 +22,20 @@ def get_one_face(frame: Frame) -> Any:
return min(face, key=lambda x: x.bbox[0])
except ValueError:
return None
def get_one_face_left(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_one_face_right(frame: Frame) -> Any:
face = get_face_analyser().get(frame)
try:
return max(face, key=lambda x: x.bbox[0])
except ValueError:
return None
def get_many_faces(frame: Frame) -> Any:

View File

@ -6,7 +6,7 @@ import threading
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,get_one_face_left,get_one_face_right
from modules.typing import Face, Frame
from modules.utilities import conditional_download, resolve_relative_path, is_image, is_video
@ -47,26 +47,50 @@ def get_face_swapper() -> Any:
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 get_two_faces(frame: Frame) -> List[Face]:
faces = get_many_faces(frame)
if faces:
# Sort faces from left to right based on the x-coordinate of the bounding box
sorted_faces = sorted(faces, key=lambda x: x.bbox[0])
return sorted_faces[:2] # Return up to two faces, leftmost and rightmost
return []
def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
def process_frame(source_face: List[Face], temp_frame: Frame) -> Frame:
if modules.globals.many_faces:
many_faces = get_many_faces(temp_frame)
if many_faces:
for target_face in many_faces:
temp_frame = swap_face(source_face, target_face, temp_frame)
temp_frame = swap_face(source_face[0], target_face, temp_frame)
else:
target_face = get_one_face(temp_frame)
if target_face:
temp_frame = swap_face(source_face, target_face, temp_frame)
target_faces = get_two_faces(temp_frame)
if len(target_faces) >= 2:
# Swap the first face
temp_frame = swap_face(source_face[0], target_faces[0], temp_frame)
# Swap the second face
temp_frame = swap_face(source_face[1], target_faces[1], temp_frame)
elif len(target_faces) == 1:
# If only one face is found, swap with the first source face
temp_frame = swap_face(source_face[0], target_faces[0], temp_frame)
return temp_frame
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
source_image_left = None # Initialize variable for the selected face image
source_image_right = None # Initialize variable for the selected face image
if source_image_left is None and source_path:
source_image_left = get_one_face_left(cv2.imread(source_path))
if source_image_right is None and source_path:
source_image_right = get_one_face_right(cv2.imread(source_path))
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)
result = process_frame([source_image_left,source_image_right], temp_frame)
cv2.imwrite(temp_frame_path, result)
except Exception as exception:
print(exception)

View File

@ -7,7 +7,7 @@ from PIL import Image, ImageOps
import modules.globals
import modules.metadata
from modules.face_analyser import get_one_face
from modules.face_analyser import get_one_face,get_one_face_left,get_one_face_right
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
@ -61,7 +61,7 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C
target_label = ctk.CTkLabel(root, text=None)
target_label.place(relx=0.6, rely=0.1, relwidth=0.3, relheight=0.25)
select_face_button = ctk.CTkButton(root, text='Select a face', cursor='hand2', command=lambda: select_source_path())
select_face_button = ctk.CTkButton(root, text='Select a face/s \n(left face)(right face)', cursor='hand2', command=lambda: select_source_path())
select_face_button.place(relx=0.1, rely=0.4, relwidth=0.3, relheight=0.1)
select_target_button = ctk.CTkButton(root, text='Select a target', cursor='hand2', command=lambda: select_target_path())
@ -239,9 +239,16 @@ def update_preview(frame_number: int = 0) -> None:
from modules.predicter import predict_frame
if predict_frame(temp_frame):
quit()
source_image_left = None # Initialize variable for the selected face image
source_image_right = None # Initialize variable for the selected face image
if source_image_left is None and modules.globals.source_path:
source_image_left = get_one_face_left(cv2.imread(modules.globals.source_path))
if source_image_right is None and modules.globals.source_path:
source_image_right = get_one_face_right(cv2.imread(modules.globals.source_path))
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 = frame_processor.process_frame([source_image_left,source_image_right],
temp_frame
)
image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))
@ -269,7 +276,8 @@ def webcam_preview():
frame_processors = get_frame_processors_modules(modules.globals.frame_processors)
source_image = None # Initialize variable for the selected face image
source_image_left = None # Initialize variable for the selected face image
source_image_right = None # Initialize variable for the selected face image
while True:
ret, frame = cap.read()
@ -277,13 +285,15 @@ def webcam_preview():
break
# 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))
if source_image_left is None and modules.globals.source_path:
source_image_left = get_one_face_left(cv2.imread(modules.globals.source_path))
if source_image_right is None and modules.globals.source_path:
source_image_right = get_one_face_right(cv2.imread(modules.globals.source_path))
temp_frame = frame.copy() #Create a copy of the frame
for frame_processor in frame_processors:
temp_frame = frame_processor.process_frame(source_image, temp_frame)
temp_frame = frame_processor.process_frame([source_image_left,source_image_right], temp_frame)
image = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) # Convert the image to RGB format to display it with Tkinter
image = Image.fromarray(image)