Add Feature to Use Config Files (#17)

* ⚡ update dot main to accept config file

* 🔧 add config files for all available options

* 📝 update usage instructions in README

* 📝 update doc to run dot without camera

.github/workflows/code_check.yaml

@@ -1,5 +1,6 @@
 ---
 name: code-check
+
 on:
     push:
         branches:

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.
 
 <p align="center">
 <img src="./assets/dot_run.gif" width="500"/>

configs/faceswap_cv2.yaml

@@ -0,0 +1,3 @@
+---
+swap_type: faceswap_cv2
+model_path: ./saved_models/faceswap_cv/shape_predictor_68_face_landmarks.dat

configs/fomm.yaml

@@ -0,0 +1,3 @@
+---
+swap_type: fomm
+model_path: ./saved_models/fomm/vox-adv-cpk.pth.tar

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

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

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)

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,
     )