diff --git a/roop/uis/__components__/preview.py b/roop/uis/__components__/preview.py index f2728bf..806ca08 100644 --- a/roop/uis/__components__/preview.py +++ b/roop/uis/__components__/preview.py @@ -29,12 +29,12 @@ def render() -> None: if is_image(roop.globals.target_path): temp_frame = cv2.imread(roop.globals.target_path) preview_frame = get_preview_frame(temp_frame) - preview_image_args['value'] = normalize_preview_frame(preview_frame) + preview_image_args['value'] = ui.normalize_frame(preview_frame) preview_image_args['visible'] = True if is_video(roop.globals.target_path): temp_frame = get_video_frame(roop.globals.target_path, roop.globals.reference_frame_number) preview_frame = get_preview_frame(temp_frame) - preview_image_args['value'] = normalize_preview_frame(preview_frame) + preview_image_args['value'] = ui.normalize_frame(preview_frame) preview_image_args['visible'] = True preview_slider_args['value'] = roop.globals.reference_frame_number preview_slider_args['maximum'] = get_video_frame_total(roop.globals.target_path) @@ -54,12 +54,12 @@ def update(frame_number: int = 0) -> Tuple[Dict[Any, Any], Dict[Any, Any]]: if is_image(roop.globals.target_path): temp_frame = cv2.imread(roop.globals.target_path) preview_frame = get_preview_frame(temp_frame) - return gradio.update(value=normalize_preview_frame(preview_frame), visible=True), gradio.update(value=0, maximum=1, visible=False) + return gradio.update(value=ui.normalize_frame(preview_frame), visible=True), gradio.update(value=0, maximum=1, visible=False) if is_video(roop.globals.target_path): video_frame_total = get_video_frame_total(roop.globals.target_path) temp_frame = get_video_frame(roop.globals.target_path, frame_number) preview_frame = get_preview_frame(temp_frame) - return gradio.update(value=normalize_preview_frame(preview_frame), visible=True), gradio.update(value=frame_number, maximum=video_frame_total, visible=True) + return gradio.update(value=ui.normalize_frame(preview_frame), visible=True), gradio.update(value=frame_number, maximum=video_frame_total, visible=True) return gradio.update(value=None, visible=False), gradio.update(value=0, maximum=1, visible=False) @@ -83,5 +83,3 @@ def get_preview_frame(temp_frame: Frame) -> Frame: return temp_frame -def normalize_preview_frame(preview_frame: Frame) -> Frame: - return cv2.cvtColor(preview_frame, cv2.COLOR_BGR2RGB) diff --git a/roop/uis/__components__/reference.py b/roop/uis/__components__/reference.py new file mode 100644 index 0000000..118ae5d --- /dev/null +++ b/roop/uis/__components__/reference.py @@ -0,0 +1,42 @@ +from time import sleep +from typing import Dict, Any + +import gradio + +import roop.globals +from roop.capturer import get_video_frame +from roop.face_analyser import get_many_faces +from roop.uis import core as ui +from roop.utilities import is_video + + +def render() -> None: + with gradio.Box(): + reference_face_gallery_args: Dict[str, Any] = { + 'label': 'reference_faces', + 'columns': 6, + 'height': 200, + 'allow_preview': False, + 'visible': True + } + if is_video(roop.globals.target_path): + target_frame = get_video_frame(roop.globals.target_path, roop.globals.reference_frame_number) + faces = get_many_faces(target_frame) + if faces: + value = [] + for face in faces: + start_x, start_y, end_x, end_y = map(int, face['bbox']) + crop_frame = target_frame[start_y:end_y, start_x:end_x] + value.append(ui.normalize_frame(crop_frame)) + reference_face_gallery_args['value'] = value + else: + reference_face_gallery_args['value'] = [] + reference_face_gallery = gradio.Gallery(**reference_face_gallery_args) + similar_face_distance_slider = gradio.Slider( + label='similar_face_distance', + value=roop.globals.similar_face_distance, + minimum=0, + maximum=2, + step=0.05 + ) + ui.register_component('reference_face_gallery', reference_face_gallery) diff --git a/roop/uis/__components__/settings.py b/roop/uis/__components__/settings.py index bf83539..4bbad90 100644 --- a/roop/uis/__components__/settings.py +++ b/roop/uis/__components__/settings.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Literal +from typing import Any, Dict, List import gradio import onnxruntime diff --git a/roop/uis/__layouts__/default.py b/roop/uis/__layouts__/default.py index 3f6556d..28bffdd 100644 --- a/roop/uis/__layouts__/default.py +++ b/roop/uis/__layouts__/default.py @@ -1,6 +1,6 @@ import gradio -from roop.uis.__components__ import settings, source, target, preview, output +from roop.uis.__components__ import settings, source, target, preview, reference, output def render() -> gradio.Blocks: @@ -13,6 +13,7 @@ def render() -> gradio.Blocks: target.render() with gradio.Column(scale=2): preview.render() + reference.render() with gradio.Row(): output.render() return layout diff --git a/roop/uis/core.py b/roop/uis/core.py index db9519d..a6d0e7c 100644 --- a/roop/uis/core.py +++ b/roop/uis/core.py @@ -1,8 +1,11 @@ import importlib from typing import Dict, Optional, List + +import cv2 import gradio import roop.globals +from roop.typing import Frame from roop.uis.typing import Component, ComponentName from roop.utilities import list_module_names @@ -34,3 +37,7 @@ def register_component(name: ComponentName, component: Component) -> None: def list_ui_layouts_names() -> Optional[List[str]]: return list_module_names('roop/uis/__layouts__') + + +def normalize_frame(frame: Frame) -> Frame: + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) diff --git a/roop/uis/typing.py b/roop/uis/typing.py index 6b3c6fd..536fad1 100644 --- a/roop/uis/typing.py +++ b/roop/uis/typing.py @@ -5,6 +5,7 @@ Component = gradio.File or gradio.Image or gradio.Video ComponentName = Literal[ 'source_file', 'target_file', + 'reference_face_gallery', 'frame_processors_checkbox_group', 'many_faces_checkbox' ]