From 981027fe52f794f419faeb7fe344acff354bde50 Mon Sep 17 00:00:00 2001 From: Ghassen Date: Mon, 27 May 2024 16:25:48 +0200 Subject: [PATCH] Add video and image swap to the GUI (#116) * Add video and image swap to the GUI * Update CHANGELOG.md * Restore CHANGELOG.md * Restore ui.py * Update ui * Update ui * Update CHANGELOG.md * Update executable --- CHANGELOG.md | 4 + README.md | 2 +- src/dot/commons/video/video_utils.py | 13 +- src/dot/ui/ui.py | 418 +++++++++++++++++---------- 4 files changed, 268 insertions(+), 169 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8449f10..73f98c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +* Add video and image swap to the GUI by @Ghassen-Chaabouni in https://github.com/sensity-ai/dot/pull/116 + ## [1.3.0] - 2024-02-19 ## What's Changed diff --git a/README.md b/README.md index d4308ff..10bb434 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ Supported methods: Download and run the dot executable for your OS: - Windows (Tested on Windows 10 and 11): - - Download `dot.zip` from [here](https://drive.google.com/file/d/10LXgtE721YPXdHfzcDUp6ba-9klueocR/view), unzip it and then run `dot.exe` + - Download `dot.zip` from [here](https://drive.google.com/file/d/1IgMaaKzFw4lBKa8MWnsH7nKwBQtrtLGJ/view), unzip it and then run `dot.exe` - Ubuntu: - ToDo - Mac (Tested on Apple M2 Sonoma 14.0): diff --git a/src/dot/commons/video/video_utils.py b/src/dot/commons/video/video_utils.py index 3afc1b5..73a28f9 100644 --- a/src/dot/commons/video/video_utils.py +++ b/src/dot/commons/video/video_utils.py @@ -12,7 +12,6 @@ import mediapipe as mp import numpy as np from mediapipe.python.solutions.drawing_utils import _normalized_to_pixel_coordinates -from tqdm import tqdm from ..pose.head_pose import pose_estimation from ..utils import expand_bbox, find_files_from_path @@ -120,7 +119,7 @@ def video_pipeline( print("Total number of face-swaps: ", len(swaps_combination)) # iterate on each source-target pair - for (source, target) in tqdm(swaps_combination): + for (source, target) in swaps_combination: img_a_whole = cv2.imread(source) img_a_whole = _crop_and_pose(img_a_whole, estimate_pose=head_pose) if isinstance(img_a_whole, int): @@ -138,11 +137,9 @@ def video_pipeline( cap = cv2.VideoCapture(target) fps = int(cap.get(cv2.CAP_PROP_FPS)) - if crop_size == 256: # fomm - frame_width = frame_height = crop_size - else: - frame_width = int(cap.get(3)) - frame_height = int(cap.get(4)) + + frame_width = int(cap.get(3)) + frame_height = int(cap.get(4)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # trim original video length @@ -164,7 +161,7 @@ def video_pipeline( ) # process each frame individually - for _ in tqdm(range(total_frames)): + for _ in range(total_frames): ret, frame = cap.read() if ret is True: frame = cv2.flip(frame, 1) diff --git a/src/dot/ui/ui.py b/src/dot/ui/ui.py index e239c15..b394bdb 100644 --- a/src/dot/ui/ui.py +++ b/src/dot/ui/ui.py @@ -128,42 +128,33 @@ def __init__(self, *args, **kwargs): self.textbox.configure(state=tkinter.DISABLED) -class App(customtkinter.CTk): +class TabView: """ - The main class of the ui interface + A class to handle the layout and functionality for each tab. """ - def __init__(self): - super().__init__() - - # configure window - self.title("Deepfake Offensive Toolkit") - self.geometry(f"{835}x{600}") - self.resizable(False, False) - - self.grid_columnconfigure((0, 1), weight=1) - self.grid_rowconfigure((0, 1, 2, 3), weight=1) - - # create menubar - menubar = tkinter.Menu(self) - - filemenu = tkinter.Menu(menubar, tearoff=0) - filemenu.add_command(label="Exit", command=self.quit) - menubar.add_cascade(label="File", menu=filemenu) + def __init__(self, tab_view, target_tip_text, use_image=False, use_video=False): + self.tab_view = tab_view + self.target_tip_text = target_tip_text + self.use_image = use_image + self.use_video = use_video + self.save_folder = None - helpmenu = tkinter.Menu(menubar, tearoff=0) - helpmenu.add_command(label="Usage", command=self.usage_window) - helpmenu.add_separator() - helpmenu.add_command(label="About DOT", command=self.about_window) - menubar.add_cascade(label="Help", menu=helpmenu) + self.resources_path = "" - self.config(menu=menubar) + # MacOS bundle has different resource directory structure + if sys.platform == "darwin": + if getattr(sys, "frozen", False): + self.resources_path = os.path.join( + str(Path(sys.executable).resolve().parents[0]).replace("MacOS", ""), + "Resources", + ) - self.toplevel_usage_window = None - self.toplevel_about_window = None + self.setup_ui() + def setup_ui(self): # create entry text for source, target and config - self.source_target_config_frame = customtkinter.CTkFrame(self) + self.source_target_config_frame = customtkinter.CTkFrame(self.tab_view) self.source_target_config_frame.grid( row=0, column=0, padx=(20, 20), pady=(20, 0), sticky="nsew" ) @@ -182,7 +173,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.source), + command=lambda: self.upload_file_action(self.source), width=10, ) @@ -192,6 +183,32 @@ def __init__(self): self.target_label = customtkinter.CTkLabel( master=self.source_target_config_frame, text="target" ) + if (self.use_image) or (self.use_video): + self.target_button = customtkinter.CTkButton( + master=self.source_target_config_frame, + fg_color="gray", + text_color="white", + text="Open", + command=lambda: self.upload_file_action(self.target), + width=10, + ) + + self.save_folder = customtkinter.CTkEntry( + master=self.source_target_config_frame, + placeholder_text="save_folder", + width=85, + ) + self.save_folder_label = customtkinter.CTkLabel( + master=self.source_target_config_frame, text="save_folder" + ) + self.save_folder_button = customtkinter.CTkButton( + master=self.source_target_config_frame, + fg_color="gray", + text_color="white", + text="Open", + command=lambda: self.upload_folder_action(self.save_folder), + width=10, + ) self.config_file_var = customtkinter.StringVar( value="Select" @@ -241,19 +258,41 @@ def __init__(self): ) self.target.grid(row=2, column=0, pady=10, padx=(80, 20), sticky="w") + if (self.use_image) or (self.use_video): + self.target_button.grid( + row=2, + column=0, + pady=(10, 10), + padx=(175, 20), + sticky="w", + ) self.target_label.grid(row=2, column=0, pady=10, padx=(35, 20), sticky="w") - self.target.insert(0, 0) - self.CreateToolTip( - self.target, text="The camera id. Usually 0 is the correct id" - ) + if (not self.use_image) and (not self.use_video): + self.target.insert(0, 0) + + self.CreateToolTip(self.target, text=self.target_tip_text) + + if (self.use_image) or (self.use_video): + self.save_folder.grid(row=3, column=0, pady=10, padx=(80, 20), sticky="w") + + self.save_folder_button.grid( + row=3, + column=0, + pady=(10, 10), + padx=(175, 20), + sticky="w", + ) + self.save_folder_label.grid(row=3, column=0, pady=10, padx=5, sticky="w") + + self.CreateToolTip(self.save_folder, text="The path to the save folder") self.config_file_combobox.grid( - row=3, column=0, pady=10, padx=(80, 20), sticky="w" + row=4, column=0, pady=10, padx=(80, 20), sticky="w" ) - self.config_file_label.grid(row=3, column=0, pady=10, padx=10, sticky="w") + self.config_file_label.grid(row=4, column=0, pady=10, padx=10, sticky="w") self.config_file_button.grid( - row=3, + row=4, column=0, pady=10, padx=(175, 20), @@ -264,7 +303,7 @@ def __init__(self): ) # create entry text for dot options - self.option_entry_frame = customtkinter.CTkFrame(self) + self.option_entry_frame = customtkinter.CTkFrame(self.tab_view) self.option_entry_frame.grid( row=1, column=0, columnspan=4, padx=(20, 20), pady=(20, 0), sticky="nsew" ) @@ -324,7 +363,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.model_path), + command=lambda: self.upload_file_action(self.model_path), width=10, ) self.parsing_model_path_button = customtkinter.CTkButton( @@ -332,7 +371,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.parsing_model_path), + command=lambda: self.upload_file_action(self.parsing_model_path), width=10, ) self.arcface_model_path_button = customtkinter.CTkButton( @@ -340,7 +379,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.arcface_model_path), + command=lambda: self.upload_file_action(self.arcface_model_path), width=10, ) self.checkpoints_dir_button = customtkinter.CTkButton( @@ -348,7 +387,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.checkpoints_dir), + command=lambda: self.upload_file_action(self.checkpoints_dir), width=10, ) self.gpen_path_button = customtkinter.CTkButton( @@ -356,7 +395,7 @@ def __init__(self): fg_color="gray", text_color="white", text="Open", - command=lambda: self.UploadAction(self.gpen_path), + command=lambda: self.upload_file_action(self.gpen_path), width=10, ) @@ -444,7 +483,7 @@ def __init__(self): ) # create radiobutton frame for swap_type - self.swap_type_frame = customtkinter.CTkFrame(self) + self.swap_type_frame = customtkinter.CTkFrame(self.tab_view) self.swap_type_frame.grid( row=0, column=1, padx=(20, 20), pady=(20, 0), sticky="nsew" ) @@ -489,7 +528,7 @@ def __init__(self): ) # create radiobutton frame for gpen_type - self.gpen_type_frame = customtkinter.CTkFrame(self) + self.gpen_type_frame = customtkinter.CTkFrame(self.tab_view) self.gpen_type_frame.grid( row=0, column=2, padx=(20, 20), pady=(20, 0), sticky="nsew" ) @@ -528,7 +567,7 @@ def __init__(self): ) # create checkbox and switch frame - self.checkbox_slider_frame = customtkinter.CTkFrame(self) + self.checkbox_slider_frame = customtkinter.CTkFrame(self.tab_view) self.checkbox_slider_frame.grid( row=0, column=3, padx=(20, 20), pady=(20, 0), sticky="nsew" ) @@ -569,18 +608,19 @@ def __init__(self): # create run button self.error_label = customtkinter.CTkLabel( - master=self, text_color="red", text="" + master=self.tab_view, text_color="red", text="" ) self.error_label.grid( row=4, column=0, columnspan=4, padx=(20, 20), pady=(0, 20), sticky="nsew" ) self.run_button = customtkinter.CTkButton( - master=self, + master=self.tab_view, fg_color="white", border_width=2, text_color="black", text="RUN", + height=40, command=lambda: self.start_button_event(self.error_label), ) self.run_button.grid( @@ -589,7 +629,7 @@ def __init__(self): self.CreateToolTip(self.run_button, text="Start running the deepfake") self.run_label = customtkinter.CTkLabel( - master=self, + master=self.tab_view, text="The initial execution of dot may require a few minutes to complete.", text_color="gray", ) @@ -597,67 +637,6 @@ def __init__(self): row=3, column=0, columnspan=3, padx=(180, 0), pady=(0, 20), sticky="nsew" ) - self.resources_path = "" - - # MacOS bundle has different resource directory structure - if sys.platform == "darwin": - if getattr(sys, "frozen", False): - self.resources_path = os.path.join( - str(Path(sys.executable).resolve().parents[0]).replace("MacOS", ""), - "Resources", - ) - - def CreateToolTip(self, widget, text): - toolTip = ToolTip(widget) - - def enter(event): - toolTip.showtip(text) - - def leave(event): - toolTip.hidetip() - - widget.bind("", enter) - widget.bind("", leave) - - def usage_window(self): - """ - Open the usage window - """ - - if ( - self.toplevel_usage_window is None - or not self.toplevel_usage_window.winfo_exists() - ): - self.toplevel_usage_window = ToplevelUsageWindow( - self - ) # create window if its None or destroyed - self.toplevel_usage_window.focus() - - def about_window(self): - """ - Open the about window - """ - - if ( - self.toplevel_about_window is None - or not self.toplevel_about_window.winfo_exists() - ): - self.toplevel_about_window = ToplevelAboutWindow( - self - ) # create window if its None or destroyed - self.toplevel_about_window.focus() - - def UploadAction(self, entry_element: customtkinter.CTkOptionMenu): - """ - Action for the upload buttons to update the value of a CTkEntry - - Args: - entry_element (customtkinter.CTkOptionMenu): The CTkEntry element. - """ - - filename = tkinter.filedialog.askopenfilename() - self.modify_entry(entry_element, filename) - def modify_entry(self, entry_element: customtkinter.CTkEntry, text: str): """ Modify the value of the CTkEntry @@ -712,62 +691,29 @@ def upload_action_config_file( element.set(config_file_var) - for key in config.keys(): + for key, value in config.items(): if key in entry_list: - self.modify_entry(eval(f"self.{key}"), config[key]) + self.modify_entry(getattr(self, key), value) elif key in radio_list: - self.swap_type_radio_var = tkinter.StringVar(value=config[key]) - eval(f"self.{config[key]}_radio_button").invoke() + self.swap_type_radio_var = tkinter.StringVar(value=value) + radio_button = getattr(self, f"{value}_radio_button") + radio_button.invoke() for entry in entry_list: - if entry not in ["source", "target"]: - if entry not in config.keys(): - self.modify_entry(eval(f"self.{entry}"), "") + if (entry not in ["source", "target"]) and (entry not in config): + self.modify_entry(getattr(self, entry), "") - def optionmenu_callback(self, choice: str): - """ - Set the configurations for the swap_type using the optionmenu - - Args: - choice (str): The type of swap to run. - """ - - entry_list = ["source", "target", "crop_size"] - radio_list = ["swap_type", "gpen_type"] - model_list = [ - "model_path", - "parsing_model_path", - "arcface_model_path", - "checkpoints_dir", - "gpen_path", - ] - - config_file = os.path.join(self.resources_path, f"configs/{choice}.yaml") + def CreateToolTip(self, widget, text): + toolTip = ToolTip(widget) - if os.path.isfile(config_file): - config = {} - with open(config_file) as f: - config = yaml.safe_load(f) + def enter(event): + toolTip.showtip(text) - for key in config.keys(): - if key in entry_list: - self.modify_entry(eval(f"self.{key}"), config[key]) - elif key in radio_list: - if key == "swap_type": - self.swap_type_radio_var = tkinter.StringVar(value=config[key]) - elif key == "gpen_type": - self.gpen_type_radio_var = tkinter.StringVar(value=config[key]) - eval(f"self.{config[key]}_radio_button").invoke() - elif key in model_list: - self.modify_entry( - eval(f"self.{key}"), - os.path.join(self.resources_path, config[key]), - ) + def leave(event): + toolTip.hidetip() - for entry in entry_list: - if entry not in ["source", "target"]: - if entry not in config.keys(): - self.modify_entry(eval(f"self.{entry}"), "") + widget.bind("", enter) + widget.bind("", leave) def start_button_event(self, error_label): """ @@ -813,11 +759,13 @@ def start_button_event(self, error_label): or 224, ), head_pose=config.get("head_pose", int(self.head_pose_checkbox.get())), - save_folder=None, + save_folder=self.save_folder.get() + if self.save_folder is not None + else None, show_fps=config.get("show_fps", int(self.show_fps_checkbox.get())), use_gpu=config.get("use_gpu", int(self.use_gpu_checkbox.get())), - use_video=False, - use_image=False, + use_video=self.use_video, + use_image=self.use_image, limit=None, ) except Exception as e: @@ -825,6 +773,156 @@ def start_button_event(self, error_label): print(traceback.format_exc()) error_label.configure(text=e) + def upload_folder_action(self, entry_element: customtkinter.CTkOptionMenu): + """ + Action for the upload folder buttons to update the value of a CTkEntry + + Args: + entry_element (customtkinter.CTkOptionMenu): The CTkEntry element. + """ + + foldername = tkinter.filedialog.askdirectory() + self.modify_entry(entry_element, foldername) + + def upload_file_action(self, entry_element: customtkinter.CTkOptionMenu): + """ + Action for the upload file buttons to update the value of a CTkEntry + + Args: + entry_element (customtkinter.CTkOptionMenu): The CTkEntry element. + """ + + filename = tkinter.filedialog.askopenfilename() + self.modify_entry(entry_element, filename) + + def optionmenu_callback(self, choice: str): + """ + Set the configurations for the swap_type using the optionmenu + + Args: + choice (str): The type of swap to run. + """ + + entry_list = ["source", "target", "crop_size"] + radio_list = ["swap_type", "gpen_type"] + model_list = [ + "model_path", + "parsing_model_path", + "arcface_model_path", + "checkpoints_dir", + "gpen_path", + ] + + config_file = os.path.join(self.resources_path, f"configs/{choice}.yaml") + + if os.path.isfile(config_file): + config = {} + with open(config_file) as f: + config = yaml.safe_load(f) + + for key in config.keys(): + if key in entry_list: + self.modify_entry(eval(f"self.{key}"), config[key]) + elif key in radio_list: + if key == "swap_type": + self.swap_type_radio_var = tkinter.StringVar(value=config[key]) + elif key == "gpen_type": + self.gpen_type_radio_var = tkinter.StringVar(value=config[key]) + eval(f"self.{config[key]}_radio_button").invoke() + elif key in model_list: + self.modify_entry( + eval(f"self.{key}"), + os.path.join(self.resources_path, config[key]), + ) + + for entry in entry_list: + if entry not in ["source", "target"]: + if entry not in config.keys(): + self.modify_entry(eval(f"self.{entry}"), "") + + +class App(customtkinter.CTk): + """ + The main class of the ui interface + """ + + def __init__(self): + super().__init__() + + # configure window + self.title("Deepfake Offensive Toolkit") + self.geometry(f"{835}x{600}") + self.resizable(False, False) + + self.grid_columnconfigure((0, 1), weight=1) + self.grid_rowconfigure((0, 1, 2, 3), weight=1) + + # create menubar + menubar = tkinter.Menu(self) + + filemenu = tkinter.Menu(menubar, tearoff=0) + filemenu.add_command(label="Exit", command=self.quit) + menubar.add_cascade(label="File", menu=filemenu) + + helpmenu = tkinter.Menu(menubar, tearoff=0) + helpmenu.add_command(label="Usage", command=self.usage_window) + helpmenu.add_separator() + helpmenu.add_command(label="About DOT", command=self.about_window) + menubar.add_cascade(label="Help", menu=helpmenu) + + self.config(menu=menubar) + + self.toplevel_usage_window = None + self.toplevel_about_window = None + + tabview = customtkinter.CTkTabview(self) + tabview.pack(padx=0, pady=0) + live_tab = tabview.add("Live") + image_tab = tabview.add("Image") + video_tab = tabview.add("Video") + + self.live_tab_view = TabView( + live_tab, target_tip_text="The camera id. Usually 0 is the correct id" + ) + self.image_tab_view = TabView( + image_tab, + target_tip_text="target images folder or certain image file", + use_image=True, + ) + self.video_tab_view = TabView( + video_tab, + target_tip_text="target videos folder or certain video file", + use_video=True, + ) + + def usage_window(self): + """ + Open the usage window + """ + + if ( + self.toplevel_usage_window is None + or not self.toplevel_usage_window.winfo_exists() + ): + self.toplevel_usage_window = ToplevelUsageWindow( + self + ) # create window if its None or destroyed + self.toplevel_usage_window.focus() + + def about_window(self): + """ + Open the about window + """ + + if ( + self.toplevel_about_window is None + or not self.toplevel_about_window.winfo_exists() + ): + self.toplevel_about_window = ToplevelAboutWindow( + self + ) # create window if its None or destroyed + self.toplevel_about_window.focus() + @click.command() def main():