diff --git a/.github/workflows/code_check.yaml b/.github/workflows/code_check.yaml index 5955aab..472e956 100644 --- a/.github/workflows/code_check.yaml +++ b/.github/workflows/code_check.yaml @@ -1,5 +1,6 @@ --- name: code-check + on: push: branches: diff --git a/README.md b/README.md index fc48987..35caec8 100644 --- a/README.md +++ b/README.md @@ -119,12 +119,9 @@ Run `dot --help` to get a full list of available options. ```bash dot \ - --swap_type simswap \ + -c ./configs/simswap.yaml \ --target 0 \ --source "./data" \ - --parsing_model_path ./saved_models/simswap/parsing_model/checkpoint/79999_iter.pth \ - --arcface_model_path ./saved_models/simswap/arcface_model/arcface_checkpoint.tar \ - --checkpoints_dir ./saved_models/simswap/checkpoints \ --show_fps \ --use_gpu ``` @@ -133,27 +130,20 @@ Run `dot --help` to get a full list of available options. ```bash dot \ - --swap_type simswap \ + -c ./configs/simswap_hq.yaml \ --target 0 \ --source "./data" \ - --parsing_model_path ./saved_models/simswap/parsing_model/checkpoint/79999_iter.pth \ - --arcface_model_path ./saved_models/simswap/arcface_model/arcface_checkpoint.tar \ - --checkpoints_dir ./saved_models/simswap/checkpoints \ - --crop_size 512 \ --show_fps \ --use_gpu ``` - Additionally, to enable face superresolution, use the flag `--gpen_type gpen_256` or `--gpen_type gpen_512`. - 3. FOMM ```bash dot \ - --swap_type fomm \ + -c ./configs/fomm.yaml \ --target 0 \ --source "./data" \ - --model_path ./saved_models/fomm/vox-adv-cpk.pth.tar \ --show_fps \ --use_gpu ``` @@ -162,21 +152,20 @@ Run `dot --help` to get a full list of available options. ```bash dot \ - --swap_type faceswap_cv2 \ + -c ./configs/faceswap.yaml \ --target 0 \ --source "./data" \ - --model_path ./saved_models/faceswap_cv/shape_predictor_68_face_landmarks.dat \ --show_fps \ --use_gpu ``` -**Note**: To use *dot* on CPU (not recommended), do not pass the `--use_gpu` flag. +**Note**: To enable face superresolution, use the flag `--gpen_type gpen_256` or `--gpen_type gpen_512`. To use *dot* on CPU (not recommended), do not pass the `--use_gpu` flag. ### Controlling dot > **Disclaimer**: We use the `SimSwap` technique for the following demonstration -Running *dot* via any of the above methods generates real-time Deepfake on the input video feed using source images from the `./data` folder. +Running *dot* via any of the above methods generates real-time Deepfake on the input video feed using source images from the `data/` folder.

diff --git a/configs/faceswap_cv2.yaml b/configs/faceswap_cv2.yaml new file mode 100644 index 0000000..f350097 --- /dev/null +++ b/configs/faceswap_cv2.yaml @@ -0,0 +1,3 @@ +--- +swap_type: faceswap_cv2 +model_path: ./saved_models/faceswap_cv/shape_predictor_68_face_landmarks.dat diff --git a/configs/fomm.yaml b/configs/fomm.yaml new file mode 100644 index 0000000..d50b557 --- /dev/null +++ b/configs/fomm.yaml @@ -0,0 +1,3 @@ +--- +swap_type: fomm +model_path: ./saved_models/fomm/vox-adv-cpk.pth.tar diff --git a/configs/simswap.yaml b/configs/simswap.yaml new file mode 100644 index 0000000..db1bf89 --- /dev/null +++ b/configs/simswap.yaml @@ -0,0 +1,5 @@ +--- +swap_type: simswap +parsing_model_path: ./saved_models/simswap/parsing_model/checkpoint/79999_iter.pth +arcface_model_path: ./saved_models/simswap/arcface_model/arcface_checkpoint.tar +checkpoints_dir: ./saved_models/simswap/checkpoints diff --git a/configs/simswaphq.yaml b/configs/simswaphq.yaml new file mode 100644 index 0000000..8c523e2 --- /dev/null +++ b/configs/simswaphq.yaml @@ -0,0 +1,6 @@ +--- +swap_type: simswap +parsing_model_path: ./saved_models/simswap/parsing_model/checkpoint/79999_iter.pth +arcface_model_path: ./saved_models/simswap/arcface_model/arcface_checkpoint.tar +checkpoints_dir: ./saved_models/simswap/checkpoints +crop_size: 512 diff --git a/docs/run_without_camera.md b/docs/run_without_camera.md index 21c1c7d..90d5ab3 100644 --- a/docs/run_without_camera.md +++ b/docs/run_without_camera.md @@ -5,26 +5,22 @@ ```bash dot \ ---swap_type simswap \ ---target ./data \ ---use_gpu \ ---source "./data" \ ---parsing_model_path ./saved_models/simswap/parsing_model/checkpoint/79999_iter.pth \ ---arcface_model_path ./saved_models/simswap/arcface_model/arcface_checkpoint.tar \ ---checkpoints_dir ./saved_models/simswap/checkpoints \ +-c ./configs/simswap.yaml \ +--target data/ \ +--source "data/" \ --save_folder test_local/ ---use_image +--use_image \ +--use_gpu ``` ```bash dot \ ---swap_type faceswap_cv2 \ ---target ./data \ ---use_gpu \ ---source "./data" \ ---model_path ./saved_models/faceswap_cv/shape_predictor_68_face_landmarks.dat \ +-c ./configs/faceswap.yaml \ +--target data/ \ +--source "data/" \ --save_folder test_local/ ---use_image +--use_image \ +--use_gpu ``` ## Faceswap images from directory(Simswap) diff --git a/dot/__main__.py b/dot/__main__.py index b32ac0b..e91fd0f 100644 --- a/dot/__main__.py +++ b/dot/__main__.py @@ -7,16 +7,81 @@ from typing import Union import click +import yaml from .dot import DOT +def run( + swap_type: str, + source: str, + target: Union[int, str], + model_path: str = None, + parsing_model_path: str = None, + arcface_model_path: str = None, + checkpoints_dir: str = None, + gpen_type: str = None, + gpen_path: str = "./saved_models/gpen", + crop_size: int = 224, + save_folder: str = None, + show_fps: bool = False, + use_gpu: bool = False, + use_video: bool = False, + use_image: bool = False, + limit: int = None, +): + """Builds a DOT object and runs it. + + Args: + swap_type (str): The type of swap to run. + source (str): The source image or video. + target (Union[int, str]): The target image or video. + model_path (str, optional): The path to the model's weights. Defaults to None. + parsing_model_path (str, optional): The path to the parsing model. Defaults to None. + arcface_model_path (str, optional): The path to the arcface model. Defaults to None. + checkpoints_dir (str, optional): The path to the checkpoints directory. Defaults to None. + gpen_type (str, optional): The type of gpen model to use. Defaults to None. + gpen_path (str, optional): The path to the gpen models. Defaults to "./saved_models/gpen". + crop_size (int, optional): The size to crop the images to. Defaults to 224. + save_folder (str, optional): The path to the save folder. Defaults to None. + show_fps (bool, optional): Pass flag to show fps value. Defaults to False. + use_gpu (bool, optional): Pass flag to use GPU else use CPU. Defaults to False. + use_video (bool, optional): Pass flag to use video-swap pipeline. Defaults to False. + use_image (bool, optional): Pass flag to use image-swap pipeline. Defaults to False. + limit (int, optional): The number of frames to process. Defaults to None. + """ + # initialize dot + _dot = DOT(use_video=use_video, use_image=use_image, save_folder=save_folder) + + # build dot + option = _dot.build_option( + swap_type=swap_type, + use_gpu=use_gpu, + gpen_type=gpen_type, + gpen_path=gpen_path, + crop_size=crop_size, + ) + + # run dot + _dot.generate( + option=option, + source=source, + target=target, + show_fps=show_fps, + model_path=model_path, + limit=limit, + parsing_model_path=parsing_model_path, + arcface_model_path=arcface_model_path, + checkpoints_dir=checkpoints_dir, + opt_crop_size=crop_size, + ) + + @click.command() @click.option( "--swap_type", "swap_type", type=click.Choice(["fomm", "faceswap_cv2", "simswap"], case_sensitive=False), - required=True, ) @click.option( "--source", @@ -101,6 +166,14 @@ help="Pass flag to use image-swap pipeline.", ) @click.option("--limit", "limit", type=int, default=None) +@click.option( + "-c", + "--config", + "config_file", + help="Configuration file. Overrides duplicate options passed.", + required=False, + default=None, +) def main( swap_type: str, source: str, @@ -118,34 +191,33 @@ def main( use_video: bool = False, use_image: bool = False, limit: int = None, + config_file: str = None, ): """CLI entrypoint for dot.""" - # initialize dot - _dot = DOT( - use_video=use_video, use_image=use_image, save_folder=save_folder, target=target - ) - - # build dot - option = _dot.build_option( - swap_type=swap_type, - use_gpu=use_gpu, - gpen_type=gpen_type, - gpen_path=gpen_path, - crop_size=crop_size, - ) + # load config, if provided + config = {} + if config_file is not None: + with open(config_file) as f: + config = yaml.safe_load(f) # run dot - _dot.generate( - option=option, + run( + swap_type=config.get("swap_type", swap_type), source=source, target=target, + model_path=config.get("model_path", model_path), + parsing_model_path=config.get("parsing_model_path", parsing_model_path), + arcface_model_path=config.get("arcface_model_path", arcface_model_path), + checkpoints_dir=config.get("checkpoints_dir", checkpoints_dir), + gpen_type=config.get("gpen_type", gpen_type), + gpen_path=config.get("gpen_path", gpen_path), + crop_size=config.get("crop_size", crop_size), + save_folder=save_folder, show_fps=show_fps, - model_path=model_path, + use_gpu=use_gpu, + use_video=use_video, + use_image=use_image, limit=limit, - parsing_model_path=parsing_model_path, - arcface_model_path=arcface_model_path, - checkpoints_dir=checkpoints_dir, - opt_crop_size=crop_size, )