diff --git a/plugins/extract/align/.cache/.keep b/.fs_cache/.keep similarity index 100% rename from plugins/extract/align/.cache/.keep rename to .fs_cache/.keep diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml new file mode 100644 index 0000000000..1172b91594 --- /dev/null +++ b/.github/workflows/pytest.yml @@ -0,0 +1,161 @@ +name: ci/build + +on: + push: + pull_request: + paths-ignore: + - docs/** + - "**/README.md" + +jobs: + build_conda: + name: conda (${{ matrix.os }}, ${{ matrix.backend }}) + runs-on: ${{ matrix.os }} + defaults: + run: + shell: bash -el {0} + strategy: + fail-fast: false + matrix: + os: ["ubuntu-latest", "macos-latest", "windows-latest"] + backend: ["nvidia", "cpu"] + include: + - os: "ubuntu-latest" + backend: "rocm" + - os: "windows-latest" + backend: "directml" + steps: + - uses: actions/checkout@v3 + - name: Set cache date + run: echo "DATE=$(date +'%Y%m%d')" >> $GITHUB_ENV + - name: Cache conda + uses: actions/cache@v3 + env: + # Increase this value to manually reset cache + CACHE_NUMBER: 1 + REQ_FILE: ./requirements/requirements_${{ matrix.backend }}.txt + with: + path: ~/conda_pkgs_dir + key: ${{ runner.os }}-${{ matrix.backend }}-conda-${{ env.CACHE_NUMBER }}-${{ env.DATE }}-${{ hashFiles('./requirements/requirements.txt', env.REQ_FILE) }} + - name: Set up Conda + uses: conda-incubator/setup-miniconda@v2 + with: + python-version: "3.10" + auto-update-conda: true + activate-environment: faceswap + - name: Conda info + run: conda info && conda list + - name: Install + run: | + python setup.py --installer --${{ matrix.backend }} + pip install flake8 pylint mypy pytest pytest-mock wheel pytest-xvfb + pip install types-attrs types-cryptography types-pyOpenSSL types-PyYAML types-setuptools + - name: Lint with flake8 + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --select=E9,F63,F7,F82 --show-source + flake8 . --exit-zero + - name: MyPy Typing + continue-on-error: true + run: | + mypy . + - name: SysInfo + run: python -c "from lib.sysinfo import sysinfo ; print(sysinfo)" + - name: Simple Tests + # These backends will fail as GPU drivers not available + if: matrix.backend != 'rocm' && matrix.backend != 'nvidia' && matrix.backend != 'directml' + run: | + FACESWAP_BACKEND="${{ matrix.backend }}" py.test -v tests/; + - name: End to End Tests + # These backends will fail as GPU drivers not available + # macOS fails on first extract test with 'died with ' + if: matrix.backend != 'rocm' && matrix.backend != 'nvidia' && matrix.backend != 'directml' && matrix.os != 'macos-latest' + run: | + FACESWAP_BACKEND="${{ matrix.backend }}" python tests/simple_tests.py; + + build_linux: + name: "pip (ubuntu-latest, ${{ matrix.backend }})" + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10"] + backend: ["cpu"] + include: + - backend: "cpu" + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v3 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + cache-dependency-path: './requirements/requirements_${{ matrix.backend }}.txt' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install flake8 pylint mypy pytest pytest-mock pytest-xvfb wheel + pip install types-attrs types-cryptography types-pyOpenSSL types-PyYAML types-setuptools + pip install -r ./requirements/requirements_${{ matrix.backend }}.txt + - name: List installed packages + run: pip freeze + - name: Lint with flake8 + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --select=E9,F63,F7,F82 --show-source + # exit-zero treats all errors as warnings. + flake8 . --exit-zero + - name: MyPy Typing + continue-on-error: true + run: | + mypy . + - name: Simple Tests + run: | + FACESWAP_BACKEND="${{ matrix.backend }}" py.test -v tests/; + - name: End to End Tests + run: | + FACESWAP_BACKEND="${{ matrix.backend }}" python tests/simple_tests.py; + + build_windows: + name: "pip (windows-latest, ${{ matrix.backend }})" + runs-on: windows-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10"] + backend: ["cpu", "directml"] + include: + - backend: "cpu" + - backend: "directml" + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v3 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + cache-dependency-path: './requirements/requirements_${{ matrix.backend }}.txt' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install flake8 pylint mypy pytest pytest-mock wheel + pip install types-attrs types-cryptography types-pyOpenSSL types-PyYAML types-setuptools + pip install -r ./requirements/requirements_${{ matrix.backend }}.txt + - name: List installed packages + run: pip freeze + - name: Set Backend EnvVar + run: echo "FACESWAP_BACKEND=${{ matrix.backend }}" | Out-File -FilePath $env:GITHUB_ENV -Append + - name: Lint with flake8 + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --select=E9,F63,F7,F82 --show-source + # exit-zero treats all errors as warnings. + flake8 . --exit-zero + - name: MyPy Typing + continue-on-error: true + run: | + mypy . + - name: Simple Tests + run: py.test -v tests + - name: End to End Tests + run: python tests/simple_tests.py diff --git a/.gitignore b/.gitignore index 580d05d5fc..63f21e29a7 100644 --- a/.gitignore +++ b/.gitignore @@ -19,8 +19,9 @@ !/update_deps.py # Support files -!_travis/ -!_travis/*.py +!/.github/ +!/.github/workflows/ +!/.github/workflows/*.yml !.install/ !.install/** !config/ @@ -35,11 +36,13 @@ !tests/**/*.py # Core files +!.fs_cache !lib/ !lib/**/ !lib/**/*.py !lib/gui/**/icons/*.png !lib/gui/**/themes/default.json +!lib/gui/**/presets/**/*.json !plugins/ !plugins/**/ !plugins/**/*.py diff --git a/.install/linux/faceswap_setup_x64.sh b/.install/linux/faceswap_setup_x64.sh index 980540fb98..26d3997fc0 100644 --- a/.install/linux/faceswap_setup_x64.sh +++ b/.install/linux/faceswap_setup_x64.sh @@ -12,7 +12,7 @@ DIR_CONDA="$HOME/miniconda3" CONDA_EXECUTABLE="${DIR_CONDA}/bin/conda" CONDA_TO_PATH=false ENV_NAME="faceswap" -PYENV_VERSION="3.9" +PYENV_VERSION="3.10" DIR_FACESWAP="$HOME/faceswap" VERSION="nvidia" @@ -35,6 +35,13 @@ info () { done <<< "$(echo "$1" | fmt -cu -w 70)" } +warn () { + # output warning message + while read -r line ; do + echo -e "\e[33mWARNING\e[97m $line" + done <<< "$(echo "$1" | fmt -cu -w 70)" +} + error () { # output error message. while read -r line ; do @@ -127,11 +134,11 @@ ask_version() { # Ask which version of faceswap to install while true; do default=1 - read -rp $'\e[36m'"Select: 1 (NVIDIA), 2 (AMD), 3 (CPU) [default: $default]: "$'\e[97m' vers + read -rp $'\e[36mSelect:\t1: NVIDIA\n\t2: AMD (ROCm)\n\t3: CPU\n'"[default: $default]: "$'\e[97m' vers vers="${vers:-${default}}" case $vers in 1) VERSION="nvidia" ; break ;; - 2) VERSION="amd" ; PYENV_VERSION="3.8" ; break ;; + 2) VERSION="rocm" ; break ;; 3) VERSION="cpu" ; break ;; * ) echo "Invalid selection." ;; esac @@ -259,7 +266,7 @@ conda_opts () { echo "" info "Faceswap will be installed inside a Conda Environment. If an environment already\ exists with the name specified then it will be deleted." - ask "Please specify a name for the Faceswap Conda Environmnet" "ENV_NAME" + ask "Please specify a name for the Faceswap Conda Environment" "ENV_NAME" } faceswap_opts () { @@ -273,6 +280,12 @@ faceswap_opts () { latest graphics card drivers installed from the relevant vendor. Please select the version\ of Faceswap you wish to install." ask_version + if [ $VERSION == "rocm" ] ; then + warn "ROCm support is experimental. Please make sure that your GPU is supported by ROCm and that \ + ROCm has been installed on your system before proceeding. Installation instructions: \ + https://docs.amd.com/bundle/ROCm_Installation_Guidev5.0/page/Overview_of_ROCm_Installation_Methods.html" + sleep 2 + fi } post_install_opts() { @@ -309,6 +322,10 @@ review() { fi echo " - Faceswap will be installed in '$DIR_FACESWAP'" echo " - Installing for '$VERSION'" + if [ $VERSION == "rocm" ] ; then + echo -e " \e[33m- Note: Please ensure that ROCm is supported by your GPU\e[97m" + echo -e " \e[33m and is installed prior to proceeding.\e[97m" + fi if $DESKTOP ; then echo " - A Desktop shortcut will be created" ; fi if ! ask_yesno "Do you wish to continue?" "No" ; then exit ; fi } @@ -346,7 +363,7 @@ delete_env() { } create_env() { - # Create Python 3.8 env for faceswap + # Create Python 3.10 env for faceswap delete_env info "Creating Conda Virtual Environment..." yellow ; "$CONDA_EXECUTABLE" create -n "$ENV_NAME" -q python="$PYENV_VERSION" -y @@ -363,7 +380,9 @@ activate_env() { install_git() { # Install git inside conda environment info "Installing Git..." - yellow ; conda install git -q -y + # TODO On linux version 2.45.2 makes the font fixed TK pull in Python from + # graalpy, which breaks pretty much everything + yellow ; conda install "git<2.45" -q -y } delete_faceswap() { @@ -384,8 +403,7 @@ clone_faceswap() { setup_faceswap() { # Run faceswap setup script info "Setting up Faceswap..." - if [ $VERSION != "cpu" ] ; then args="--$VERSION" ; else args="" ; fi - python "$DIR_FACESWAP/setup.py" --installer $args + python -u "$DIR_FACESWAP/setup.py" --installer --$VERSION } create_gui_launcher () { diff --git a/.install/macos/app.zip b/.install/macos/app.zip new file mode 100644 index 0000000000..9ce64629d6 Binary files /dev/null and b/.install/macos/app.zip differ diff --git a/.install/macos/faceswap_setup_macos.sh b/.install/macos/faceswap_setup_macos.sh new file mode 100644 index 0000000000..804b828c3e --- /dev/null +++ b/.install/macos/faceswap_setup_macos.sh @@ -0,0 +1,491 @@ +#!/bin/bash + +TMP_DIR="/tmp/faceswap_install" + +URL_CONDA="https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-" +DL_CONDA="${URL_CONDA}x86_64.sh" +DL_FACESWAP="https://github.com/deepfakes/faceswap.git" +DL_XQUARTZ="https://github.com/XQuartz/XQuartz/releases/latest/download/XQuartz-2.8.5.pkg" + +CONDA_PATHS=("/opt" "$HOME") +CONDA_NAMES=("anaconda" "miniconda" "miniforge") +CONDA_VERSIONS=("3" "2") +CONDA_BINS=("/bin/conda" "/condabin/conda") +DIR_CONDA="$HOME/miniconda3" +CONDA_EXECUTABLE="${DIR_CONDA}/bin/conda" +CONDA_TO_PATH=false +ENV_NAME="faceswap" +PYENV_VERSION="3.10" + +DIR_FACESWAP="$HOME/faceswap" +VERSION="nvidia" + +DESKTOP=false +XQUARTZ=false + +header() { + # Format header text + length=${#1} + padding=$(( (72 - length) / 2)) + sep=$(printf '=%.0s' $(seq 1 $padding)) + echo "" + echo $'\e[32m'$sep $1 $sep +} + +info () { + # output info message + while read -r line ; do + echo $'\e[32mINFO\e[39m '$line + done <<< "$(echo "$1" | fmt -s -w 70)" +} + +warn () { + # output warning message + while read -r line ; do + echo $'\e[33mWARNING\e[39m '$line + done <<< "$(echo "$1" | fmt -s -w 70)" +} + +error () { + # output error message. + while read -r line ; do + echo $'\e[31mERROR\e[39m '$line + done <<< "$(echo "$1" | fmt -s -w 70)" +} + +yellow () { + # Change text color to yellow + echo $'\e[33m' +} + +check_file_exists () { + # Check whether a file exists and return true or false + test -f "$1" +} + +check_folder_exists () { + # Check whether a folder exists and return true or false + test -d "$1" +} + +download_file () { + # Download a file to the temp folder + fname=$(basename -- "$1") + curl -L "$1" --output "$TMP_DIR/$fname" --progress-bar +} + +check_for_sudo() { + # Ensure user isn't running as sudo/root. We don't want to screw up any system install + if [ "$EUID" == 0 ] ; then + error "This install script should not be run with root privileges. Please run as a normal user." + exit 1 + fi +} + +check_for_curl() { + # Ensure that curl is available on the system + if ! command -V curl &> /dev/null ; then + error "'curl' is required for running the Faceswap installer, but could not be found. \ + Please install 'curl' before proceeding." + exit 1 + fi +} + +check_for_xcode() { + # Ensure that xcode command line tools are available on the system + if xcode-select -p 2>&1 | grep -q "xcode-select: error" ; then + error "Xcode is required to install faceswap. Please install Xcode Command Line Tools \ + before proceeding. If the Xcode installer does not automatically open, then \ + you can run the command:" + error "xcode-select --install" + echo "" + xcode-select --install + exit 1 + fi +} + +create_tmp_dir() { + TMP_DIR="$(mktemp -d)" + if [ -z "$TMP_DIR" -o ! -d "$TMP_DIR" ]; then + # This shouldn't happen, but just in case to prevent the tmp cleanup function to mess things up. + error "Failed creating the temporary install directory." + exit 2 + fi + trap cleanup_tmp_dir EXIT +} + +cleanup_tmp_dir() { + rm -rf "$TMP_DIR" +} + +ask () { + # Ask for input. First parameter: Display text, 2nd parameter variable name + default="${!2}" + read -rp $'\e[35m'"$1 [default: '$default']: "$'\e[39m' inp + inp="${inp:-${default}}" + if [ "$inp" == "\n" ] ; then inp=${!2} ; fi + printf -v $2 "$inp" +} + +ask_yesno () { + # Ask yes or no. First Param: Question, 2nd param: Default + # Returns True for yes, False for No + case $2 in + [Yy]* ) opts="[YES/no]" ;; + [Nn]* ) opts="[yes/NO]" ;; + esac + while true; do + read -rp $'\e[35m'"$1 $opts: "$'\e[39m' yn + yn="${yn:-${2}}" + case $yn in + [Yy]* ) retval=true ; break ;; + [Nn]* ) retval=false ; break ;; + * ) echo "Please answer yes or no." ;; + esac + done + $retval +} + + +ask_version() { + # Ask which version of faceswap to install + while true; do + default=1 + read -rp $'\e[35mSelect:\t1: Apple Silicon\n\t2: NVIDIA\n\t3: CPU\n'"[default: $default]: "$'\e[39m' vers + vers="${vers:-${default}}" + case $vers in + 1) VERSION="apple_silicon" ; break ;; + 2) VERSION="nvidia" ; break ;; + 3) VERSION="cpu" ; break ;; + * ) echo "Invalid selection." ;; + esac + done +} + +banner () { + echo $' \e[32m 001' + echo $' \e[32m 11 10 010' + echo $' \e[39m @@@@\e[32m 10' + echo $' \e[39m @@@@@@@@\e[32m 00 1' + echo $' \e[39m @@@@@@@@@@\e[32m 1 1 0' + echo $' \e[39m @@@@@@@@\e[32m 0000 01111' + echo $' \e[39m @@@@@@@@@@\e[32m 01 110 01 1' + echo $' \e[39m@@@@@@@@@@@@\e[32m 111 010 0' + echo $' \e[39m@@@@@@@@@@@@@@@@\e[32m 10 0' + echo $' \e[39m@@@@@@@@@@@@@\e[32m 0010 1' + echo $' \e[39m@@@@@@@@@ @@@\e[32m 100 1' + echo $' \e[39m@@@@@@@ .@@@@\e[32m 10 1' + echo $' \e[39m #@@@@@@@@@@@\e[32m 001 0' + echo $' \e[39m @@@@@@@@@@@ ,' + echo ' @@@@@@@@ @@@@@' + echo ' @@@@@@@@ @@@@@@@@ _' + echo ' @@@@@@@@@,@@@@@@@@ / _|' + echo ' %@@@@@@@@@@@@@@@@@ | |_ ___ ' + echo ' @@@@@@@@@@@@@@ | _|/ __|' + echo ' @@@@@@@@@@@@ | | \__ \' + echo ' @@@@@@@@@@( |_| |___/' + echo ' @@@@@@' + echo ' @@@@' + sleep 2 +} + +find_conda_install() { + if check_conda_path; + then true + elif check_conda_locations ; then true + else false + fi +} + +set_conda_dir_from_bin() { + # Set the DIR_CONDA variable from the bin file + pth="$(dirname "$1")/.." + DIR_CONDA=$(python -c "import os, sys; print(os.path.realpath('$pth'))") + info "Found existing conda install at: $DIR_CONDA" +} + +check_conda_path() { + # Check if conda is in PATH + conda_bin="$(which conda 2>/dev/null)" + if [[ "$?" == "0" ]]; then + set_conda_dir_from_bin "$conda_bin" + CONDA_EXECUTABLE="$conda_bin" + true + else + false + fi +} + +check_conda_locations() { + # Check common conda install locations + retval=false + for path in "${CONDA_PATHS[@]}"; do + for name in "${CONDA_NAMES[@]}" ; do + foldername="$path/$name" + for vers in "${CONDA_VERSIONS[@]}" ; do + for bin in "${CONDA_BINS[@]}" ; do + condabin="$foldername$vers$bin" + if check_file_exists "$condabin" ; then + set_conda_dir_from_bin "$condabin" + CONDA_EXECUTABLE="$condabin"; + retval=true + break 4 + fi + done + done + done + done + $retval +} + +user_input() { + # Get user options for install + header "Welcome to the macOS Faceswap Installer" + info "To get setup we need to gather some information about where you would like Faceswap\ + and Conda to be installed." + info "To accept the default values just hit the 'ENTER' key for each option. You will have\ + an opportunity to review your responses prior to commencing the install." + echo "" + info "IMPORTANT: Make sure that the user '$USER' has full permissions for all of the\ + destinations that you select." + read -rp $'\e[35m'"Press 'ENTER' to continue with the setup..."$'\e[39m' + apps_opts + conda_opts + faceswap_opts + post_install_opts +} + +apps_opts () { + # Options pertaining to additional apps that are required + if ! command -V xquartz &> /dev/null ; then + header "APPS" + info "XQuartz is required to use the Faceswap GUI but was not detected. " + if ask_yesno "Install XQuartz for GUI support?" "Yes" ; then + XQUARTZ=true + fi + fi +} + +conda_opts () { + # Options pertaining to the installation of conda + header "CONDA" + info "Faceswap uses Conda as it handles the installation of all prerequisites." + if find_conda_install && ask_yesno "Use the pre installed conda?" "Yes"; then + info "Using Conda install at $DIR_CONDA" + else + echo "" + info "If you have an existing Conda install then enter the location here,\ + otherwise Miniconda3 will be installed in the given location." + err_msg="The location for Conda must not contain spaces (this is a specific\ + limitation of Conda)." + tmp_dir_conda="$DIR_CONDA" + while true ; do + ask "Please specify a location for Conda." "DIR_CONDA" + case ${DIR_CONDA} in + *\ * ) error "$err_msg" ; DIR_CONDA=$tmp_dir_conda ;; + * ) break ;; + esac + CONDA_EXECUTABLE="${DIR_CONDA}/bin/conda" + done + fi + if ! check_file_exists "$CONDA_EXECUTABLE" ; then + echo "" + info "The Conda executable can be added to your PATH. This makes it easier to run Conda\ + commands directly. If you already have a pre-existing Conda install then you should\ + probably not enable this, otherwise this should be fine." + if ask_yesno "Add Conda executable to path?" "Yes" ; then CONDA_TO_PATH=true ; fi + fi + echo "" + info "Faceswap will be installed inside a Conda Environment. If an environment already\ + exists with the name specified then it will be deleted." + ask "Please specify a name for the Faceswap Conda Environment" "ENV_NAME" +} + +faceswap_opts () { + # Options pertaining to the installation of faceswap + header "FACESWAP" + info "Faceswap will be installed in the given location. If a folder exists at the\ + location you specify, then it will be deleted." + ask "Please specify a location for Faceswap" "DIR_FACESWAP" + echo "" + info "Faceswap can be run on Apple Silicon (M1, M2 etc.), compatible NVIDIA gpus, or on CPU. You should make sure that any \ + drivers are up to date. Please select the version of Faceswap you wish to install." + ask_version + if [ $VERSION == "apple_silicon" ] ; then + DL_CONDA="${URL_CONDA}arm64.sh" + fi +} + +post_install_opts() { + # Post installation options + header "POST INSTALLATION ACTIONS" + info "Launching Faceswap requires activating your Conda Environment and then running\ + Faceswap. The installer can simplify this by creating an Application Launcher file and placing it \ + on your desktop to launch straight into the Faceswap GUI" + if ask_yesno "Create FaceswapGUI Launcher?" "Yes" ; then + DESKTOP=true + fi +} + +review() { + # Review user options and ask continue + header "Review install options" + info "Please review the selected installation options before proceeding:" + echo "" + if $XQUARTZ ; then echo " - The XQuartz installer will be downloaded and launched" ; fi + if ! check_folder_exists "$DIR_CONDA" + then + echo " - MiniConda3 will be installed in '$DIR_CONDA'" + else + echo " - Existing Conda install at '$DIR_CONDA' will be used" + fi + if $CONDA_TO_PATH ; then echo " - MiniConda3 will be added to your PATH" ; fi + if check_env_exists ; then + echo $' \e[33m- Existing Conda Environment '$ENV_NAME $' will be removed\e[39m' + fi + echo " - Conda Environment '$ENV_NAME' will be created." + if check_folder_exists "$DIR_FACESWAP" ; then + echo $' \e[33m- Existing Faceswap folder '$DIR_FACESWAP $' will be removed\e[39m' + fi + echo " - Faceswap will be installed in '$DIR_FACESWAP'" + echo " - Installing for '$VERSION'" + if [ $VERSION == "nvidia" ] ; then + echo $' \e[33m- Note: Please ensure that Nvidia drivers are installed prior to proceeding\e[39m' + fi + if $DESKTOP ; then echo " - An Application Launcher will be created" ; fi + if ! ask_yesno "Do you wish to continue?" "No" ; then exit ; fi +} + +xquartz_install() { + # Download and install XQuartz + if $XQUARTZ ; then + info "Downloading XQuartz..." + yellow ; download_file $DL_XQUARTZ + echo "" + + info "Installing XQuartz..." + info "Admin password required to install XQuartz:" + fname="$(basename -- $DL_XQUARTZ)" + yellow ; sudo installer -pkg "$TMP_DIR/$fname" -target / + echo "" + fi +} + +conda_install() { + # Download and install Mini Conda3 + if ! check_folder_exists "$DIR_CONDA" ; then + info "Downloading Miniconda3..." + yellow ; download_file $DL_CONDA + info "Installing Miniconda3..." + yellow ; fname="$(basename -- $DL_CONDA)" + bash "$TMP_DIR/$fname" -b -p "$DIR_CONDA" + if $CONDA_TO_PATH ; then + info "Adding Miniconda3 to PATH..." + yellow ; "$CONDA_EXECUTABLE" init zsh bash + "$CONDA_EXECUTABLE" config --set auto_activate_base false + fi + fi +} + +check_env_exists() { + # Check if an environment with the given name exists + if check_file_exists "$CONDA_EXECUTABLE" ; then + "$CONDA_EXECUTABLE" env list | grep -qE "^${ENV_NAME}\W" + else false + fi +} + +delete_env() { + # Delete the env if it previously exists + if check_env_exists ; then + info "Removing pre-existing Virtual Environment" + yellow ; "$CONDA_EXECUTABLE" env remove -n "$ENV_NAME" + fi +} + +create_env() { + # Create Python 3.10 env for faceswap + delete_env + info "Creating Conda Virtual Environment..." + yellow ; "$CONDA_EXECUTABLE" create -n "$ENV_NAME" -q python="$PYENV_VERSION" -y +} + + +activate_env() { + # Activate the conda environment + # shellcheck source=/dev/null + source "$DIR_CONDA/etc/profile.d/conda.sh" activate + conda activate "$ENV_NAME" +} + +delete_faceswap() { + # Delete existing faceswap folder + if check_folder_exists "$DIR_FACESWAP" ; then + info "Removing Faceswap folder: '$DIR_FACESWAP'" + rm -rf "$DIR_FACESWAP" + fi +} + +clone_faceswap() { + # Clone the faceswap repo + delete_faceswap + info "Downloading Faceswap..." + yellow ; git clone --depth 1 --no-single-branch "$DL_FACESWAP" "$DIR_FACESWAP" +} + +setup_faceswap() { + # Run faceswap setup script + info "Setting up Faceswap..." + python -u "$DIR_FACESWAP/setup.py" --installer --$VERSION +} + +create_gui_launcher () { + # Create a shortcut to launch into the GUI + launcher="$DIR_FACESWAP/faceswap_gui_launcher.command" + launch_script="#!/bin/bash\n" + launch_script+="source \"$DIR_CONDA/etc/profile.d/conda.sh\" activate && \n" + launch_script+="conda activate '$ENV_NAME' && \n" + launch_script+="python \"$DIR_FACESWAP/faceswap.py\" gui" + printf "$launch_script" > "$launcher" + chmod +x "$launcher" +} + +create_app_on_desktop () { + # Create a simple .app wrapper to launch GUI + if $DESKTOP ; then + app_name="FaceswapGUI" + app_dir="$TMP_DIR/$app_name.app" + + unzip -qq "$DIR_FACESWAP/.install/macos/app.zip" -d "$TMP_DIR" + + script="#!/bin/bash\n" + script+="bash \"$DIR_FACESWAP/faceswap_gui_launcher.command\"" + printf "$script" > "$app_dir/Contents/Resources/script" + chmod +x "$app_dir/Contents/Resources/script" + + rm -rf "$HOME/Desktop/$app_name.app" + mv "$app_dir" "$HOME/Desktop" + fi ; +} + +check_for_sudo +check_for_curl +check_for_xcode +banner +user_input +review +create_tmp_dir +xquartz_install +conda_install +create_env +activate_env +clone_faceswap +setup_faceswap +create_gui_launcher +create_app_on_desktop +info "Faceswap installation is complete!" +if $CONDA_TO_PATH ; then + info "You should close the terminal before proceeding" ; fi +if $DESKTOP ; then info "You can launch Faceswap from the icon on your desktop" ; fi +if $XQUARTZ ; then + warn "XQuartz has been installed. You must log out and log in again to be able to use the GUI" ; fi diff --git a/.install/windows/install.nsi b/.install/windows/install.nsi index 277a0228ec..d407807872 100644 --- a/.install/windows/install.nsi +++ b/.install/windows/install.nsi @@ -22,7 +22,7 @@ InstallDir $PROFILE\faceswap # Install cli flags !define flagsConda "/S /RegisterPython=0 /AddToPath=0 /D=$PROFILE\MiniConda3" !define flagsRepo "--depth 1 --no-single-branch ${wwwRepo}" -!define flagsEnv "-y python=3." +!define flagsEnv "-y python=3.10" # Folders Var ProgramData @@ -138,9 +138,9 @@ Function pgPrereqCreate ${NSD_AddStyle} $ctlRadio ${WS_GROUP} nsDialogs::SetUserData $ctlRadio "nvidia" ${NSD_OnClick} $ctlRadio RadioClick - ${NSD_CreateRadioButton} 40% $lblPos% 25% 11u "Setup for AMD GPU" + ${NSD_CreateRadioButton} 40% $lblPos% 25% 11u "Setup for DirectML" Pop $ctlRadio - nsDialogs::SetUserData $ctlRadio "amd" + nsDialogs::SetUserData $ctlRadio "directml" ${NSD_OnClick} $ctlRadio RadioClick ${NSD_CreateRadioButton} 70% $lblPos% 20% 11u "Setup for CPU" Pop $ctlRadio @@ -200,7 +200,7 @@ FunctionEnd Function CheckSetupType ${If} $setupType == "" - MessageBox MB_OK "Please specify whether to setup for Nvidia, AMD or CPU." + MessageBox MB_OK "Please specify whether to setup for Nvidia, DirectML or CPU." Abort ${EndIf} StrCpy $Log "$log(check) Setting up for: $setupType$\n" @@ -397,11 +397,7 @@ Function SetEnvironment CreateEnv: SetDetailsPrint listonly - ${If} $setupType == "amd" - StrCpy $0 "${flagsEnv}8" - ${else} - StrCpy $0 "${flagsEnv}9" - ${EndIf} + StrCpy $0 "${flagsEnv}" ExecDos::exec /NOUNLOAD /ASYNC /DETAILED "$\"$dirConda\scripts\activate.bat$\" && conda create $0 -n $\"$envName$\" && conda deactivate" pop $0 ExecDos::wait $0 @@ -444,11 +440,9 @@ FunctionEnd Function SetupFaceSwap DetailPrint "Setting up FaceSwap Environment... This may take a while" StrCpy $0 "${flagsSetup}" - ${If} $setupType != "cpu" - StrCpy $0 "$0 --$setupType" - ${EndIf} + StrCpy $0 "$0 --$setupType" SetDetailsPrint listonly - ExecDos::exec /NOUNLOAD /ASYNC /DETAILED "$\"$dirConda\scripts\activate.bat$\" && conda activate $\"$envName$\" && python $\"$INSTDIR\setup.py$\" $0 && conda deactivate" + ExecDos::exec /NOUNLOAD /ASYNC /DETAILED "$\"$dirConda\scripts\activate.bat$\" && conda activate $\"$envName$\" && python -u $\"$INSTDIR\setup.py$\" $0 && conda deactivate" pop $0 ExecDos::wait $0 pop $0 diff --git a/.pylintrc b/.pylintrc index 5c52c8bf7d..5f6e63247f 100644 --- a/.pylintrc +++ b/.pylintrc @@ -60,85 +60,14 @@ confidence= # --enable=similarities". If you want to run only the classes checker, but have # no Warning level messages displayed, use "--disable=all --enable=classes # --disable=W". -disable=print-statement, - parameter-unpacking, - unpacking-in-except, - old-raise-syntax, - backtick, - long-suffix, - old-ne-operator, - old-octal-literal, - import-star-module-level, - non-ascii-bytes-literal, - raw-checker-failed, +disable=raw-checker-failed, bad-inline-option, locally-disabled, file-ignored, suppressed-message, useless-suppression, deprecated-pragma, - use-symbolic-message-instead, - apply-builtin, - basestring-builtin, - buffer-builtin, - cmp-builtin, - coerce-builtin, - execfile-builtin, - file-builtin, - long-builtin, - raw_input-builtin, - reduce-builtin, - standarderror-builtin, - unicode-builtin, - xrange-builtin, - coerce-method, - delslice-method, - getslice-method, - setslice-method, - no-absolute-import, - old-division, - dict-iter-method, - dict-view-method, - next-method-called, - metaclass-assignment, - indexing-exception, - raising-string, - reload-builtin, - oct-method, - hex-method, - nonzero-method, - cmp-method, - input-builtin, - round-builtin, - intern-builtin, - unichr-builtin, - map-builtin-not-iterating, - zip-builtin-not-iterating, - range-builtin-not-iterating, - filter-builtin-not-iterating, - using-cmp-argument, - eq-without-hash, - div-method, - idiv-method, - rdiv-method, - exception-message-attribute, - invalid-str-codec, - sys-max-int, - bad-python3-import, - deprecated-string-function, - deprecated-str-translate-call, - deprecated-itertools-function, - deprecated-types-field, - next-method-defined, - dict-items-not-iterating, - dict-keys-not-iterating, - dict-values-not-iterating, - deprecated-operator-function, - deprecated-urllib-function, - xreadlines-attribute, - deprecated-sys-function, - exception-escape, - comprehension-escape + use-symbolic-message-instead # Enable the message, report, category or checker with the given id(s). You can # either give multiple identifier separated by comma (,) or put this option @@ -360,13 +289,6 @@ max-line-length=100 # Maximum number of lines in a module. max-module-lines=1000 -# List of optional constructs for which whitespace checking is disabled. `dict- -# separator` is used to allow tabulation in dicts, etc.: {1 : 1,\n222: 2}. -# `trailing-comma` allows a space between comma and closing bracket: (a, ). -# `empty-line` allows space-only lines. -no-space-check=trailing-comma, - dict-separator - # Allow the body of a class to be on the same line as the declaration if body # contains single statement. single-line-class-stmt=no @@ -424,7 +346,7 @@ contextmanager-decorators=contextlib.contextmanager # List of members which are set dynamically and missed by pylint inference # system, and so shouldn't trigger E1101 when accessed. Python regular # expressions are accepted. -generated-members= +generated-members=cv2.* # Tells whether missing members accessed in mixin class should be ignored. A # mixin class is detected if its name ends with "mixin" (case insensitive). @@ -480,7 +402,7 @@ notes=FIXME, max-args=10 # Maximum number of attributes for a class (see R0902). -max-attributes=10 +max-attributes=12 # Maximum number of boolean expressions in an if statement. max-bool-expr=5 @@ -504,7 +426,7 @@ max-returns=6 max-statements=50 # Minimum number of public methods for a class (see R0903). -min-public-methods=2 +min-public-methods=1 [CLASSES] @@ -566,5 +488,5 @@ known-third-party=enchant # Exceptions that will emit a warning when being caught. Defaults to # "BaseException, Exception". -overgeneral-exceptions=BaseException, - Exception +overgeneral-exceptions=builtins.BaseException, + builtins.Exception diff --git a/.readthedocs.yml b/.readthedocs.yml index 2aa3c9934b..8d199514eb 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -7,9 +7,9 @@ version: 2 # Set the version of Python and other tools you might need build: - os: ubuntu-20.04 + os: ubuntu-22.04 tools: - python: "3.8" + python: "3.10" # Build documentation in the docs/ directory with Sphinx sphinx: diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index 9075e4af8e..0000000000 --- a/.travis.yml +++ /dev/null @@ -1,102 +0,0 @@ -# Adapted from https://github.com/kangwonlee/travis-yml-conda-posix-nt/blob/master/.travis.yml - -language: shell - -env: - global: - - CONDA_PYTHON=3.8 - - CONDA_BLD_PATH=${HOME}/conda-bld - -os: - - linux - # - windows - # - osx - - -cache: - # More time is needed for caching due to the sheer size of the conda env. - timeout: 1000 - directories: - - ${HOME}/cache - -before_cache: - # adapted from https://github.com/theochem/cgrid/blob/master/.travis.yml - - echo "Cleaning stuff in miniconda path. (${MINICONDA_PATH})" - - ls -la ${MINICONDA_PATH} - - rm -rf ${MINICONDA_PATH}/conda-bld - - rm -rf ${MINICONDA_PATH}/locks - - rm -rf ${MINICONDA_PATH}/pkgs - - rm -rf ${MINICONDA_PATH}/var - - rm -rf ${MINICONDA_PATH}/envs/*/conda-bld - - rm -rf ${MINICONDA_PATH}/envs/*/locks - - rm -rf ${MINICONDA_PATH}/envs/*/pkgs - - rm -rf ${MINICONDA_PATH}/envs/*/var - - ls -la ${MINICONDA_PATH} - - echo "Cleaning test results" - - rm -rf ${HOME}/cache/tests/*/faces - - rm -rf ${HOME}/cache/tests/*/conv - - rm -rf ${HOME}/cache/tests/*/*.json - - rm -rf ${HOME}/cache/tests/vid/faces_sorted - - rm -rf ${HOME}/cache/tests/vid/model - -before_install: - # set conda path info - - | - if [[ "$TRAVIS_OS_NAME" != "windows" ]]; then - MINICONDA_PATH=${HOME}/cache/miniconda; - MINICONDA_SUB_PATH=$MINICONDA_PATH/bin; - elif [[ "$TRAVIS_OS_NAME" == "windows" ]]; then - MINICONDA_PATH=${HOME}/cache/miniconda3/; - MINICONDA_PATH_WIN=`cygpath --windows $MINICONDA_PATH`; - MINICONDA_SUB_PATH=$MINICONDA_PATH/Scripts; - fi; - # obtain miniconda installer - - | - if [[ "$TRAVIS_OS_NAME" == "linux" ]]; then - wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh; - # Not used at the moment, but in case we want to also run test on osx, we can. - elif [[ "$TRAVIS_OS_NAME" == "osx" ]]; then - wget https://repo.continuum.io/miniconda/Miniconda3-latest-MacOSX-x86_64.sh -O miniconda.sh; - fi; - -install: - # install miniconda - # pip and conda will also need OpenSSL for Windows - - | - if test -e "$MINICONDA_PATH"; then - echo "Conda already installed"; - else - echo "Installing conda"; - if [[ "$TRAVIS_OS_NAME" != "windows" ]]; then - bash miniconda.sh -b -p $MINICONDA_PATH; - elif [[ "$TRAVIS_OS_NAME" == "windows" ]]; then - choco install openssl.light; - choco install miniconda3 --params="'/AddToPath:1 /D:$MINICONDA_PATH_WIN'"; - fi; - fi; - - export PATH="$MINICONDA_PATH:$MINICONDA_SUB_PATH:$PATH"; - # for conda version 4.4 or later - - source $MINICONDA_PATH/etc/profile.d/conda.sh; - - hash -r; - - conda config --set always_yes yes --set changeps1 no; - - conda update -q conda; - # Useful for debugging any issues with conda - - conda info -a - - echo "Python $CONDA_PYTHON running on $TRAVIS_OS_NAME"; - # Only create the environment if we don't have it already - - conda env list | grep faceswap || conda create -q --name faceswap python=$CONDA_PYTHON; - - conda activate faceswap; - - conda --version ; python --version ; pip --version; - # We set up for plaidML as we can then use both the plaidML and Tensorflow backends for testing - - python setup.py --installer --amd; - - conda install pytest; - # For debugging purposes - - df -h - -script: - - rm -f ~/.plaidml; - - echo "{\"PLAIDML_DEVICE_IDS\":[\"llvm_cpu.0\"],\"PLAIDML_EXPERIMENTAL\":true}" > ~/.plaidml; - - FACESWAP_BACKEND="amd" KERAS_BACKEND="plaidml.keras.backend" PYTHONPATH=$PWD:$PYTHONPATH py.test -v tests/; - - rm -f ~/.plaidml; - - FACESWAP_BACKEND="cpu" KERAS_BACKEND="tensorflow" PYTHONPATH=$PWD:$PYTHONPATH py.test -v tests/; - - FACESWAP_BACKEND="cpu" KERAS_BACKEND="tensorflow" python _travis/simple_tests.py; diff --git a/Dockerfile.cpu b/Dockerfile.cpu index 8b9d297737..0c27ec9b69 100755 --- a/Dockerfile.cpu +++ b/Dockerfile.cpu @@ -1,19 +1,19 @@ -FROM tensorflow/tensorflow:2.8.2 +FROM ubuntu:22.04 # To disable tzdata and others from asking for input ENV DEBIAN_FRONTEND noninteractive +ENV FACESWAP_BACKEND cpu -RUN apt-get update -qq -y \ - && apt-get install -y software-properties-common \ - && add-apt-repository -y ppa:jonathonf/ffmpeg-4 \ - && apt-get update -qq -y \ - && apt-get install -y libsm6 libxrender1 libxext-dev python3-tk ffmpeg git \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* +RUN apt-get update -qq -y +RUN apt-get upgrade -y +RUN apt-get install -y libgl1 libglib2.0-0 python3 python3-pip python3-tk git -COPY ./requirements/_requirements_base.txt /opt/ -RUN pip3 install --upgrade pip -RUN pip3 --no-cache-dir install -r /opt/_requirements_base.txt && rm /opt/_requirements_base.txt +RUN ln -s $(which python3) /usr/local/bin/python + +RUN git clone --depth 1 --no-single-branch https://github.com/deepfakes/faceswap.git +WORKDIR "/faceswap" + +RUN python -m pip install --upgrade pip +RUN python -m pip --no-cache-dir install -r ./requirements/requirements_cpu.txt -WORKDIR "/srv" CMD ["/bin/bash"] diff --git a/Dockerfile.gpu b/Dockerfile.gpu index 078875f5ed..5b9c0abd0a 100755 --- a/Dockerfile.gpu +++ b/Dockerfile.gpu @@ -1,29 +1,18 @@ -FROM nvidia/cuda:11.7.0-runtime-ubuntu18.04 -ARG DEBIAN_FRONTEND=noninteractive +FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 -#install python3.8 -RUN apt-get update -RUN apt-get install software-properties-common -y -RUN add-apt-repository ppa:deadsnakes/ppa -y -RUN apt-get update -RUN apt-get install python3.8 -y -RUN apt-get install python3.8-distutils -y -RUN apt-get install python3.8-tk -y -RUN apt-get install curl -y -RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py -RUN python3.8 get-pip.py -RUN rm get-pip.py +ENV DEBIAN_FRONTEND=noninteractive +ENV FACESWAP_BACKEND nvidia -# install requirements -RUN apt-get install ffmpeg git -y -COPY ./requirements/_requirements_base.txt /opt/ -COPY ./requirements/requirements_nvidia.txt /opt/ -RUN python3.8 -m pip --no-cache-dir install -r /opt/requirements_nvidia.txt && rm /opt/_requirements_base.txt && rm /opt/requirements_nvidia.txt +RUN apt-get update -qq -y +RUN apt-get upgrade -y +RUN apt-get install -y libgl1 libglib2.0-0 python3 python3-pip python3-tk git -RUN python3.8 -m pip install jupyter matplotlib tqdm -RUN python3.8 -m pip install jupyter_http_over_ws -RUN jupyter serverextension enable --py jupyter_http_over_ws -RUN alias python=python3.8 -RUN echo "alias python=python3.8" >> /root/.bashrc -WORKDIR "/notebooks" -CMD ["jupyter-notebook", "--allow-root" ,"--port=8888" ,"--no-browser" ,"--ip=0.0.0.0"] +RUN ln -s $(which python3) /usr/local/bin/python + +RUN git clone --depth 1 --no-single-branch https://github.com/deepfakes/faceswap.git +WORKDIR "/faceswap" + +RUN python -m pip install --upgrade pip +RUN python -m pip --no-cache-dir install -r ./requirements/requirements_nvidia.txt + +CMD ["/bin/bash"] diff --git a/INSTALL.md b/INSTALL.md index 2adb73ec15..752c272929 100755 --- a/INSTALL.md +++ b/INSTALL.md @@ -4,7 +4,7 @@ - [Hardware Requirements](#hardware-requirements) - [Supported operating systems](#supported-operating-systems) - [Important before you proceed](#important-before-you-proceed) -- [Linux and Windows Install Guide](#linux-and-windows-install-guide) +- [Linux, Windows and macOS Install Guide](#linux-windows-and-macos-install-guide) - [Installer](#installer) - [Manual Install](#manual-install) - [Prerequisites](#prerequisites-1) @@ -39,12 +39,9 @@ - [Setup](#setup-2) - [About some of the options](#about-some-of-the-options) - [Docker Install Guide](#docker-install-guide) - - [Docker General](#docker-general) - - [CUDA with Docker in 20 minutes.](#cuda-with-docker-in-20-minutes) - - [CUDA with Docker on Arch Linux](#cuda-with-docker-on-arch-linux) - - [Install docker](#install-docker) - - [A successful setup log, without docker.](#a-successful-setup-log-without-docker) - - [Run the project](#run-the-project) + - [Docker CPU](#docker-cpu) + - [Docker Nvidia](#docker-nvidia) +- [Run the project](#run-the-project) - [Notes](#notes) # Prerequisites @@ -58,15 +55,20 @@ The type of computations that the process does are well suited for graphics card - **A powerful CPU** - Laptop CPUs can often run the software, but will not be fast enough to train at reasonable speeds - **A powerful GPU** - - Currently, Nvidia GPUs are fully supported. and AMD graphics cards are partially supported through plaidML. + - Currently, Nvidia GPUs are fully supported + - DirectX 12 AMD GPUs are supported on Windows through DirectML. + - More modern AMD GPUs are supported on Linux through ROCm. + - M-series Macs are supported through Tensorflow-Metal - If using an Nvidia GPU, then it needs to support at least CUDA Compute Capability 3.5. (Release 1.0 will work on Compute Capability 3.0) To see which version your GPU supports, consult this list: https://developer.nvidia.com/cuda-gpus Desktop cards later than the 7xx series are most likely supported. - **A lot of patience** ## Supported operating systems -- **Windows 10** - Windows 7 and 8 might work. Your mileage may vary. Windows has an installer which will set up everything you need. See: https://github.com/deepfakes/faceswap/releases +- **Windows 10/11** + Windows 7 and 8 might work for Nvidia. Your mileage may vary. + DirectML support is only available in Windows 10 onwards. + Windows has an installer which will set up everything you need. See: https://github.com/deepfakes/faceswap/releases - **Linux** Most Ubuntu/Debian or CentOS based Linux distributions will work. There is a Linux install script that will install and set up everything you need. See: https://github.com/deepfakes/faceswap/releases - **macOS** @@ -81,10 +83,10 @@ Alternatively, there is a docker image that is based on Debian. The developers are also not responsible for any damage you might cause to your own computer. -# Linux and Windows Install Guide +# Linux, Windows and macOS Install Guide ## Installer -Windows and Linux now both have an installer which installs everything for you and creates a desktop shortcut to launch straight into the GUI. You can download the installer from https://github.com/deepfakes/faceswap/releases. +Windows, Linux and macOS all have installers which set up everything for you. You can download the installer from https://github.com/deepfakes/faceswap/releases. If you have issues with the installer then read on for the more manual way to install faceswap on Windows. @@ -110,7 +112,7 @@ Reboot your PC, so that everything you have just installed gets registered. - Select "Create" at the bottom - In the pop up: - Give it the name: faceswap - - **IMPORTANT**: Select python version 3.8 + - **IMPORTANT**: Select python version 3.10 - Hit "Create" (NB: This may take a while as it will need to download Python) ![Anaconda virtual env setup](https://i.imgur.com/CLIDDfa.png) @@ -136,7 +138,6 @@ If you are using an Nvidia card make sure you have the correct versions of Cuda/ - Install tkinter (required for the GUI) by typing: `conda install tk` - Install requirements: - For Nvidia GPU users: `pip install -r ./requirements/requirements_nvidia.txt` - - For AMD GPU users: `pip install -r ./requirements/requirements_amd.txt` - For CPU users: `pip install -r ./requirements/requirements_cpu.txt` ## Running faceswap @@ -164,6 +165,8 @@ It's good to keep faceswap up to date as new features are added and bugs are fix # macOS (Apple Silicon) Install Guide +macOS now has [an installer](#linux-windows-and-macos-install-guide) which sets everything up for you, but if you run into difficulties and need to set things up manually, the steps are as follows: + ## Prerequisites ### OS @@ -191,7 +194,7 @@ $ source ~/miniforge3/bin/activate ## Setup ### Create and Activate the Environment ```sh -$ conda create --name faceswap python=3.9 +$ conda create --name faceswap python=3.10 $ conda activate faceswap ``` @@ -221,7 +224,7 @@ Obtain git for your distribution from the [git website](https://git-scm.com/down The recommended install method is to use a Conda3 Environment as this will handle the installation of Nvidia's CUDA and cuDNN straight into your Conda Environment. This is by far the easiest and most reliable way to setup the project. - MiniConda3 is recommended: [MiniConda3](https://docs.conda.io/en/latest/miniconda.html) -Alternatively you can install Python (>= 3.7-3.9 64-bit) for your distribution (links below.) If you go down this route and are using an Nvidia GPU you should install CUDA (https://developer.nvidia.com/cuda-zone) and cuDNN (https://developer.nvidia.com/cudnn). for your system. If you do not plan to build Tensorflow yourself, make sure you install the correct Cuda and cuDNN package for the currently installed version of Tensorflow (Current release: Tensorflow 2.8. Release v1.0: Tensorflow 1.15). You can check for the compatible versions here: (https://www.tensorflow.org/install/source#gpu). +Alternatively you can install Python (3.10 64-bit) for your distribution (links below.) If you go down this route and are using an Nvidia GPU you should install CUDA (https://developer.nvidia.com/cuda-zone) and cuDNN (https://developer.nvidia.com/cudnn). for your system. If you do not plan to build Tensorflow yourself, make sure you install the correct Cuda and cuDNN package for the currently installed version of Tensorflow (Current release: Tensorflow 2.9. Release v1.0: Tensorflow 1.15). You can check for the compatible versions here: (https://www.tensorflow.org/install/source#gpu). - Python distributions: - apt/yum install python3 (Linux) - [Installer](https://www.python.org/downloads/release/python-368/) (Windows) @@ -256,153 +259,83 @@ If setup fails for any reason you can still manually install the packages listed # Docker Install Guide -## Docker General -
- Click to expand! - - ### CUDA with Docker in 20 minutes. - - 1. Install Docker - https://www.docker.com/community-edition - - 2. Install Nvidia-Docker & Restart Docker Service - https://github.com/NVIDIA/nvidia-docker +This Faceswap repo contains Docker build scripts for CPU and Nvidia backends. The scripts will set up a Docker container for you and install the latest version of the Faceswap software. - 3. Build Docker Image For faceswap - - ```bash - docker build -t deepfakes-gpu -f Dockerfile.gpu . - ``` +You must first ensure that Docker is installed and running on your system. Follow the guide for downloading and installing Docker from their website: - 4. Mount faceswap volume and Run it - a). without `gui.tools.py` gui not working. + - https://www.docker.com/get-started - ```bash - nvidia-docker run --rm -it -p 8888:8888 \ - --hostname faceswap-gpu --name faceswap-gpu \ - -v /opt/faceswap:/srv \ - deepfakes-gpu - ``` +Once Docker is installed and running, follow the relevant steps for your chosen backend +## Docker CPU +To run the CPU version of Faceswap follow these steps: - b). with gui. tools.py gui working. - -Enable local access to X11 server - -```bash -xhost +local: +1. Build the Docker image For faceswap: ``` - -Enable nvidia device if working under bumblebee - -```bash -echo ON > /proc/acpi/bbswitch -``` - -Create container -```bash -nvidia-docker run -p 8888:8888 \ - --hostname faceswap-gpu --name faceswap-gpu \ - -v /opt/faceswap:/srv \ - -v /tmp/.X11-unix:/tmp/.X11-unix \ - -e DISPLAY=unix$DISPLAY \ - -e AUDIO_GID=`getent group audio | cut -d: -f3` \ - -e VIDEO_GID=`getent group video | cut -d: -f3` \ - -e GID=`id -g` \ - -e UID=`id -u` \ - deepfakes-gpu - +docker build \ +-t faceswap-cpu \ +https://raw.githubusercontent.com/deepfakes/faceswap/master/Dockerfile.cpu ``` - -Open a new terminal to interact with the project - -```bash -docker exec -it deepfakes-gpu /bin/bash +2. Launch and enter the Faceswap container: + + a. For the **headless/command line** version of Faceswap run: + ``` + docker run --rm -it faceswap-cpu + ``` + You can then execute faceswap the standard way: + ``` + python faceswap.py --help + ``` + b. For the **GUI** version of Faceswap run: + ``` + xhost +local: && \ + docker run --rm -it \ + -v /tmp/.X11-unix:/tmp/.X11-unix \ + -e DISPLAY=${DISPLAY} \ + faceswap-cpu + ``` + You can then launch the GUI with + ``` + python faceswap.py gui + ``` + ## Docker Nvidia +To build the NVIDIA GPU version of Faceswap, follow these steps: + +1. Nvidia Docker builds need extra resources to provide the Docker container with access to your GPU. + + a. Follow the instructions to install and apply the `Nvidia Container Toolkit` for your distribution from: + - https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html + + b. If Docker is already running, restart it to pick up the changes made by the Nvidia Container Toolkit. + +2. Build the Docker image For faceswap ``` - -Launch deepfakes gui (Answer 3 for NVIDIA at the prompt) - -```bash -python3.8 /srv/faceswap.py gui +docker build \ +-t faceswap-gpu \ +https://raw.githubusercontent.com/deepfakes/faceswap/master/Dockerfile.gpu ``` -
- -## CUDA with Docker on Arch Linux - -
- Click to expand! - -### Install docker - -```bash -sudo pacman -S docker -``` - -The steps are same but Arch linux doesn't use nvidia-docker - -create container - -```bash -docker run -p 8888:8888 --gpus all --privileged -v /dev:/dev \ - --hostname faceswap-gpu --name faceswap-gpu \ - -v /mnt/hdd2/faceswap:/srv \ - -v /tmp/.X11-unix:/tmp/.X11-unix \ - -e DISPLAY=unix$DISPLAY \ - -e AUDIO_GID=`getent group audio | cut -d: -f3` \ - -e VIDEO_GID=`getent group video | cut -d: -f3` \ - -e GID=`id -g` \ - -e UID=`id -u` \ - deepfakes-gpu -``` - -Open a new terminal to interact with the project - -```bash -docker exec -it deepfakes-gpu /bin/bash -``` - -Launch deepfakes gui (Answer 3 for NVIDIA at the prompt) - -**With `gui.tools.py` gui working.** - Enable local access to X11 server - - ```bash -xhost +local: -``` - - ```bash - python3.8 /srv/faceswap.py gui - ``` - -
- ---- -## A successful setup log, without docker. -``` -INFO The tool provides tips for installation - and installs required python packages -INFO Setup in Linux 4.14.39-1-MANJARO -INFO Installed Python: 3.7.5 64bit -INFO Installed PIP: 10.0.1 -Enable Docker? [Y/n] n -INFO Docker Disabled -Enable CUDA? [Y/n] -INFO CUDA Enabled -INFO CUDA version: 9.1 -INFO cuDNN version: 7 -WARNING Tensorflow has no official prebuild for CUDA 9.1 currently. - To continue, You have to build your own tensorflow-gpu. - Help: https://www.tensorflow.org/install/install_sources -Are System Dependencies met? [y/N] y -INFO Installing Missing Python Packages... -INFO Installing tensorflow-gpu -...... -INFO Installing tqdm -INFO Installing matplotlib -INFO All python3 dependencies are met. - You are good to go. -``` - -## Run the project +1. Launch and enter the Faceswap container: + + a. For the **headless/command line** version of Faceswap run: + ``` + docker run --runtime=nvidia --rm -it faceswap-gpu + ``` + You can then execute faceswap the standard way: + ``` + python faceswap.py --help + ``` + b. For the **GUI** version of Faceswap run: + ``` + xhost +local: && \ + docker run --runtime=nvidia --rm -it \ + -v /tmp/.X11-unix:/tmp/.X11-unix \ + -e DISPLAY=${DISPLAY} \ + faceswap-gpu + ``` + You can then launch the GUI with + ``` + python faceswap.py gui + ``` +# Run the project Once all these requirements are installed, you can attempt to run the faceswap tools. Use the `-h` or `--help` options for a list of options. ```bash diff --git a/README.md b/README.md index c4cf43bbda..2b1e829214 100755 --- a/README.md +++ b/README.md @@ -10,12 +10,18 @@

    

+ +

+ +
Emma Stone/Scarlett Johansson FaceSwap using the Phaze-A model +

+


Jennifer Lawrence/Steve Buscemi FaceSwap using the Villain model

-[![Build Status](https://travis-ci.org/deepfakes/faceswap.svg?branch=master)](https://travis-ci.org/deepfakes/faceswap) [![Documentation Status](https://readthedocs.org/projects/faceswap/badge/?version=latest)](https://faceswap.readthedocs.io/en/latest/?badge=latest) +![Build Status](https://github.com/deepfakes/faceswap/actions/workflows/pytest.yml/badge.svg) [![Documentation Status](https://readthedocs.org/projects/faceswap/badge/?version=latest)](https://faceswap.readthedocs.io/en/latest/?badge=latest) Make sure you check out [INSTALL.md](INSTALL.md) before getting started. @@ -73,7 +79,7 @@ We are very troubled by the fact that FaceSwap can be used for unethical and dis # How To setup and run the project FaceSwap is a Python program that will run on multiple Operating Systems including Windows, Linux, and MacOS. -See [INSTALL.md](INSTALL.md) for full installation instructions. You will need a modern GPU with CUDA support for best performance. AMD GPUs are partially supported. +See [INSTALL.md](INSTALL.md) for full installation instructions. You will need a modern GPU with CUDA support for best performance. Many AMD GPUs are supported through DirectML (Windows) and ROCm (Linux). # Overview The project has multiple entry points. You will have to: diff --git a/USAGE.md b/USAGE.md index 49476007ad..de78f5abf8 100755 --- a/USAGE.md +++ b/USAGE.md @@ -2,24 +2,24 @@ **Before attempting any of this, please make sure you have read, understood and completed the [installation instructions](../master/INSTALL.md). If you are experiencing issues, please raise them in the [faceswap Forum](https://faceswap.dev/forum) or the [FaceSwap Discord server](https://discord.gg/FdEwxXd) instead of the main repo.** -- [Workflow](#Workflow) -- [Introduction](#Introduction) - - [Disclaimer](#Disclaimer) - - [Getting Started](#Getting-Started) -- [Extract](#Extract) - - [Gathering raw data](#Gathering-raw-data) - - [Extracting Faces](#Extracting-Faces) - - [General Tips](#General-Tips) -- [Training a model](#Training-a-model) - - [General Tips](#General-Tips-1) -- [Converting a video](#Converting-a-video) - - [General Tips](#General-Tips-2) -- [GUI](#GUI) -- [Video's](#Videos) -- [EFFMPEG](#EFFMPEG) -- [Extracting video frames with FFMPEG](#Extracting-video-frames-with-FFMPEG) -- [Generating a video](#Generating-a-video) -- [Notes](#Notes) +- [Workflow](#workflow) +- [Introduction](#introduction) + - [Disclaimer](#disclaimer) + - [Getting Started](#getting-started) +- [Extract](#extract) + - [Gathering raw data](#gathering-raw-data) + - [Extracting Faces](#extracting-faces) + - [General Tips](#general-tips) +- [Training a model](#training-a-model) + - [General Tips](#general-tips-1) +- [Converting a video](#converting-a-video) + - [General Tips](#general-tips-2) +- [GUI](#gui) +- [Video's](#videos) +- [EFFMPEG](#effmpeg) +- [Extracting video frames with FFMPEG](#extracting-video-frames-with-ffmpeg) +- [Generating a video](#generating-a-video) +- [Notes](#notes) # Introduction diff --git a/docs/_static/theme_overrides.css b/docs/_static/theme_overrides.css deleted file mode 100644 index abc9c0fcee..0000000000 --- a/docs/_static/theme_overrides.css +++ /dev/null @@ -1,14 +0,0 @@ -/* override table width restrictions */ -@media screen and (min-width: 767px) { - - .wy-table-responsive table td { - /* !important prevents the common CSS stylesheets from overriding - this as on RTD they are loaded after this stylesheet */ - white-space: normal !important; - } - - .wy-table-responsive { - overflow: visible !important; - } - } - \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index bbd311d5ac..f547655839 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -12,13 +12,21 @@ # import os import sys +from unittest import mock + +os.environ["FACESWAP_BACKEND"] = "nvidia" sys.path.insert(0, os.path.abspath('../')) sys.setrecursionlimit(1500) + +MOCK_MODULES = ["pynvx", "ctypes.windll", "comtypes"] +for mod_name in MOCK_MODULES: + sys.modules[mod_name] = mock.Mock() + # -- Project information ----------------------------------------------------- project = 'faceswap' -copyright = '2019, faceswap.dev' +copyright = '2022, faceswap.dev' author = 'faceswap.dev' # The full version, including alpha/beta/rc tags @@ -31,6 +39,7 @@ # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = ['sphinx.ext.napoleon', "sphinx.ext.autosummary", ] +napoleon_custom_sections = ['License'] # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] @@ -61,12 +70,6 @@ # so a file named "default.css" will overwrite the builtin "default.css". html_static_path = ['_static'] -html_context = { - 'css_files': [ - '_static/theme_overrides.css', # override wide tables in RTD theme - ], - } - master_doc = 'index' autosummary_generate = True diff --git a/docs/full/lib/align.rst b/docs/full/lib/align.rst index bb449f805d..feebc7af4a 100644 --- a/docs/full/lib/align.rst +++ b/docs/full/lib/align.rst @@ -7,6 +7,7 @@ The align Package handles detected faces, their alignments and masks. .. contents:: Contents :local: + aligned\_face module ==================== @@ -16,10 +17,9 @@ Handles aligned faces and corresponding pose estimates .. autosummary:: :nosignatures: - + ~lib.align.aligned_face.AlignedFace ~lib.align.aligned_face.get_matrix_scaling - ~lib.align.aligned_face.PoseEstimate ~lib.align.aligned_face.transform_image .. rubric:: Module @@ -29,6 +29,29 @@ Handles aligned faces and corresponding pose estimates :undoc-members: :show-inheritance: + +aligned\_mask module +==================== + +Handles aligned storage and retrieval of Faceswap generated masks + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~lib.align.aligned_mask.BlurMask + ~lib.align.aligned_mask.LandmarksMask + ~lib.align.aligned_mask.Mask + +.. rubric:: Module + +.. automodule:: lib.align.aligned_mask + :members: + :undoc-members: + :show-inheritance: + + alignments module ================= @@ -38,7 +61,7 @@ Handles alignments stored in a serialized alignments.fsa file .. autosummary:: :nosignatures: - + ~lib.align.alignments.Alignments ~lib.align.alignments.Thumbnails @@ -49,6 +72,17 @@ Handles alignments stored in a serialized alignments.fsa file :undoc-members: :show-inheritance: + +constants module +================ +Holds various constants for use in generating and manipulating aligned face images + +.. automodule:: lib.align.constants + :members: + :undoc-members: + :show-inheritance: + + detected\_face module ===================== @@ -58,10 +92,8 @@ Handles detected face objects and their associated masks. .. autosummary:: :nosignatures: - - ~lib.align.detected_face.BlurMask + ~lib.align.detected_face.DetectedFace - ~lib.align.detected_face.Mask ~lib.align.detected_face.update_legacy_png_header .. rubric:: Module @@ -70,3 +102,33 @@ Handles detected face objects and their associated masks. :members: :undoc-members: :show-inheritance: + + +pose module +=========== +Handles pose estimates based on aligned face data + +.. automodule:: lib.align.pose + :members: + :undoc-members: + :show-inheritance: + + +thumbnails module +================= +Handles creation of jpg thumbnails for storage in alignment files/png headers + +.. automodule:: lib.align.thumbnails + :members: + :undoc-members: + :show-inheritance: + + +updater module +============== +Handles the update of alignments files to the latest version + +.. automodule:: lib.align.updater + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/config.rst b/docs/full/lib/config.rst new file mode 100755 index 0000000000..dcd5dcb8b8 --- /dev/null +++ b/docs/full/lib/config.rst @@ -0,0 +1,7 @@ +config module +============= + +.. automodule:: lib.config + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/git.rst b/docs/full/lib/git.rst new file mode 100644 index 0000000000..3f8d8de585 --- /dev/null +++ b/docs/full/lib/git.rst @@ -0,0 +1,10 @@ +********** +git module +********** + +Handles interfacing with the git executable + +.. automodule:: lib.git + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/gpu_stats.rst b/docs/full/lib/gpu_stats.rst index 6f6aaa309a..9ba5ce3c97 100755 --- a/docs/full/lib/gpu_stats.rst +++ b/docs/full/lib/gpu_stats.rst @@ -38,6 +38,14 @@ gpu_stats.cpu module :undoc-members: :show-inheritance: +gpu_stats.directml module +------------------------- + +.. automodule:: lib.gpu_stats.directml + :members: + :undoc-members: + :show-inheritance: + gpu_stats.nvidia_apple module ----------------------------- @@ -53,3 +61,11 @@ gpu_stats.nvidia module :members: :undoc-members: :show-inheritance: + +gpu_stats.rocm module +---------------------- + +.. automodule:: lib.gpu_stats.rocm + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/gui.rst b/docs/full/lib/gui.rst index aa8f54d18c..07981633a4 100755 --- a/docs/full/lib/gui.rst +++ b/docs/full/lib/gui.rst @@ -19,7 +19,7 @@ stats module .. autosummary:: :nosignatures: - + ~lib.gui.analysis.stats.Calculations ~lib.gui.analysis.stats.GlobalSession ~lib.gui.analysis.stats.SessionsSummary @@ -47,7 +47,7 @@ custom\_widgets module .. autosummary:: :nosignatures: - + ~lib.gui.custom_widgets.ConsoleOut ~lib.gui.custom_widgets.ContextMenu ~lib.gui.custom_widgets.MultiOption @@ -77,7 +77,7 @@ display\_analysis module .. autosummary:: :nosignatures: - + ~lib.gui.display_analysis.Analysis ~lib.gui.display_analysis.StatsData @@ -88,6 +88,14 @@ display\_analysis module :undoc-members: :show-inheritance: +display\_command module +======================= + +.. automodule:: lib.gui.display_command + :members: + :undoc-members: + :show-inheritance: + display\_graph module ===================== @@ -96,6 +104,20 @@ display\_graph module :undoc-members: :show-inheritance: +menu module +=========== +.. automodule:: lib.gui.menu + :members: + :undoc-members: + :show-inheritance: + +options module +============== +.. automodule:: lib.gui.options + :members: + :undoc-members: + :show-inheritance: + popup_configure module ====================== .. automodule:: lib.gui.popup_configure @@ -117,7 +139,7 @@ project module .. autosummary:: :nosignatures: - + ~lib.gui.project.LastSession ~lib.gui.project.Project ~lib.gui.project.Tasks @@ -139,28 +161,61 @@ theme module :undoc-members: :show-inheritance: -utils module -============ +utils package +============= -.. rubric:: Module Summary +.. rubric:: Package Summary .. autosummary:: :nosignatures: - - ~lib.gui.utils.Config - ~lib.gui.utils.FileHandler - ~lib.gui.utils.Images - ~lib.gui.utils.LongRunningTask - ~lib.gui.utils.get_config - ~lib.gui.utils.get_images - ~lib.gui.utils.initialize_config - ~lib.gui.utils.initialize_images -.. rubric:: Module + ~lib.gui.utils.config.Config + ~lib.gui.utils.config.initialize_config + ~lib.gui.utils.config.get_config + ~lib.gui.utils.file_handler.FileHandler + ~lib.gui.utils.image.Images + ~lib.gui.utils.image.get_images + ~lib.gui.utils.image.initialize_images + ~lib.gui.utils.misc.LongRunningTask + + +.. rubric:: config Module -.. automodule:: lib.gui.utils +.. automodule:: lib.gui.utils.config :members: :undoc-members: :show-inheritance: +.. rubric:: file_handler Module + +.. automodule:: lib.gui.utils.file_handler + :members: + :undoc-members: + :show-inheritance: + + +.. rubric:: image Module + +.. automodule:: lib.gui.utils.image + :members: + :undoc-members: + :show-inheritance: + + +.. rubric:: misc Module + +.. automodule:: lib.gui.utils.misc + :members: + :undoc-members: + :show-inheritance: + +wrapper module +============== + +.. rubric:: Module + +.. automodule:: lib.gui.wrapper + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/keras_utils.rst b/docs/full/lib/keras_utils.rst new file mode 100644 index 0000000000..1dda86a3ec --- /dev/null +++ b/docs/full/lib/keras_utils.rst @@ -0,0 +1,8 @@ +****************** +keras_utils module +****************** + +.. automodule:: lib.keras_utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/model.rst b/docs/full/lib/model.rst index c253d9a8b3..e01f9fd02f 100755 --- a/docs/full/lib/model.rst +++ b/docs/full/lib/model.rst @@ -1,5 +1,6 @@ +************* model package -============= +************* The Model Package handles interfacing with the neural network backend and holds custom objects. @@ -7,7 +8,7 @@ The Model Package handles interfacing with the neural network backend and holds :local: model.backup_restore module ---------------------------- +=========================== .. automodule:: lib.model.backup_restore :members: @@ -15,13 +16,13 @@ model.backup_restore module :show-inheritance: model.initializers module -------------------------- +========================= .. rubric:: Module Summary .. autosummary:: :nosignatures: - + ~lib.model.initializers.ConvolutionAware ~lib.model.initializers.ICNR ~lib.model.initializers.compute_fans @@ -32,51 +33,82 @@ model.initializers module :show-inheritance: model.layers module -------------------- +=================== .. rubric:: Module Summary .. autosummary:: :nosignatures: - + ~lib.model.layers.GlobalMinPooling2D ~lib.model.layers.GlobalStdDevPooling2D + ~lib.model.layers.KResizeImages ~lib.model.layers.L2_normalize ~lib.model.layers.PixelShuffler + ~lib.model.layers.QuickGELU ~lib.model.layers.ReflectionPadding2D ~lib.model.layers.SubPixelUpscaling - + ~lib.model.layers.Swish + .. automodule:: lib.model.layers :members: :undoc-members: :show-inheritance: model.losses module -------------------- +=================== -The losses listed here are generated from the docstrings in :mod:`lib.model.losses_tf`, however -the functions are excactly the same for :mod:`lib.model.losses_plaid`. The correct loss module will -be imported as :mod:`lib.model.losses` depending on the backend in use. +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~lib.model.loss.loss_tf.FocalFrequencyLoss + ~lib.model.loss.loss_tf.GeneralizedLoss + ~lib.model.loss.loss_tf.GradientLoss + ~lib.model.loss.loss_tf.LaplacianPyramidLoss + ~lib.model.loss.loss_tf.LInfNorm + ~lib.model.loss.loss_tf.LossWrapper + ~lib.model.loss.feature_loss_tf.LPIPSLoss + ~lib.model.loss.perceptual_loss_tf.DSSIMObjective + ~lib.model.loss.perceptual_loss_tf.GMSDLoss + ~lib.model.loss.perceptual_loss_tf.LDRFLIPLoss + ~lib.model.loss.perceptual_loss_tf.MSSIMLoss + +.. automodule:: lib.model.loss.loss_tf + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: lib.model.loss.feature_loss_tf + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: lib.model.loss.perceptual_loss_tf + :members: + :undoc-members: + :show-inheritance: + + +model.nets module +================= .. rubric:: Module Summary .. autosummary:: :nosignatures: - ~lib.model.losses_tf.DSSIMObjective - ~lib.model.losses_tf.GeneralizedLoss - ~lib.model.losses_tf.GMSDLoss - ~lib.model.losses_tf.GradientLoss - ~lib.model.losses_tf.LInfNorm - ~lib.model.losses_tf.LossWrapper + ~lib.model.nets.AlexNet + ~lib.model.nets.SqueezeNet -.. automodule:: lib.model.losses_tf +.. automodule:: lib.model.nets :members: :undoc-members: :show-inheritance: model.nn_blocks module ----------------------- +====================== .. rubric:: Module Summary @@ -98,26 +130,22 @@ model.nn_blocks module :show-inheritance: model.normalization module --------------------------- +========================== .. rubric:: Module Summary .. autosummary:: :nosignatures: - + ~lib.model.normalization.InstanceNormalization - + .. automodule:: lib.model.normalization :members: :undoc-members: :show-inheritance: model.optimizers module ------------------------ - -The optimizers listed here are generated from the docstrings in :mod:`lib.model.optimizers_tf`, however -the functions are excactly the same for :mod:`lib.model.optimizers_plaid`. The correct optimizers module will -be imported as :mod:`lib.model.optimizers` depending on the backend in use. +======================= .. rubric:: Module Summary @@ -132,10 +160,9 @@ be imported as :mod:`lib.model.optimizers` depending on the backend in use. :show-inheritance: model.session module ---------------------- +===================== .. automodule:: lib.model.session :members: :undoc-members: :show-inheritance: - diff --git a/docs/full/lib/multithreading.rst b/docs/full/lib/multithreading.rst new file mode 100644 index 0000000000..e786abe83a --- /dev/null +++ b/docs/full/lib/multithreading.rst @@ -0,0 +1,7 @@ +multithreading module +===================== + +.. automodule:: lib.multithreading + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/plaidml_utils.rst b/docs/full/lib/plaidml_utils.rst new file mode 100644 index 0000000000..256e96ed7a --- /dev/null +++ b/docs/full/lib/plaidml_utils.rst @@ -0,0 +1,8 @@ +******************** +plaidml_utils module +******************** + +.. automodule:: lib.plaidml_utils + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/lib/training.rst b/docs/full/lib/training.rst index ff8c3cde04..f3cb797f40 100644 --- a/docs/full/lib/training.rst +++ b/docs/full/lib/training.rst @@ -7,18 +7,51 @@ The training Package handles the processing of faces for feeding into a Faceswap .. contents:: Contents :local: -augmentation module -=================== +training.augmentation module +============================ .. automodule:: lib.training.augmentation :members: :undoc-members: :show-inheritance: -generator module -================ +training.cache module +===================== + +.. automodule:: lib.training.cache + :members: + :undoc-members: + :show-inheritance: + + +training.generator module +========================= .. automodule:: lib.training.generator :members: :undoc-members: :show-inheritance: + +training.lr_finder module +========================= + +.. automodule:: lib.training.lr_finder + :members: + :undoc-members: + :show-inheritance: + +training.preview_cv module +========================== + +.. automodule:: lib.training.preview_cv + :members: + :undoc-members: + :show-inheritance: + +training.preview_tk module +========================== + +.. automodule:: lib.training.preview_tk + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/modules.rst b/docs/full/modules.rst index 8dd0581022..877e28ef66 100644 --- a/docs/full/modules.rst +++ b/docs/full/modules.rst @@ -7,4 +7,7 @@ faceswap lib/lib plugins/plugins scripts + tests/tests tools/tools + setup + update_deps diff --git a/docs/full/plugins/convert.rst b/docs/full/plugins/convert.rst index 05f3a110ef..103d650ba7 100755 --- a/docs/full/plugins/convert.rst +++ b/docs/full/plugins/convert.rst @@ -44,3 +44,27 @@ writer.gif module :members: :undoc-members: :show-inheritance: + +writer.opencv module +-------------------- + +.. automodule:: plugins.convert.writer.opencv + :members: + :undoc-members: + :show-inheritance: + +writer.patch module +-------------------- + +.. automodule:: plugins.convert.writer.patch + :members: + :undoc-members: + :show-inheritance: + +writer.pillow module +-------------------- + +.. automodule:: plugins.convert.writer.pillow + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/plugins/extract.rst b/docs/full/plugins/extract.rst index 59da672e2d..f7e441a106 100755 --- a/docs/full/plugins/extract.rst +++ b/docs/full/plugins/extract.rst @@ -7,66 +7,121 @@ The Extract Package handles the various plugins available for extracting face se .. contents:: Contents :local: -pipeline module -=============== -.. rubric:: Module Summary - -.. autosummary:: - :nosignatures: - - ~plugins.extract.pipeline.ExtractMedia - ~plugins.extract.pipeline.Extractor +extract\_media module +===================== +.. automodule:: plugins.extract.extract_media + :members: + :undoc-members: + :show-inheritance: -.. rubric:: Module +pipeline module +=============== .. automodule:: plugins.extract.pipeline :members: :undoc-members: :show-inheritance: -extract plugins package -======================= +_base module +============ +.. automodule:: plugins.extract._base + :members: + :undoc-members: + :show-inheritance: + + +align plugins package +===================== .. contents:: Contents :local: -_base module ------------- +align._base.aligner module +-------------------------- +.. automodule:: plugins.extract.align._base.aligner + :members: + :undoc-members: + :show-inheritance: -.. automodule:: plugins.extract._base +align._base.processing module +----------------------------- +.. automodule:: plugins.extract.align._base.processing :members: :undoc-members: :show-inheritance: +align.cv2_dnn module +-------------------- +.. automodule:: plugins.extract.align.cv2_dnn + :members: + :undoc-members: + :show-inheritance: + +align.fan module +---------------- +.. automodule:: plugins.extract.align.fan + :members: + :undoc-members: + :show-inheritance: + + +detect plugins package +====================== +.. contents:: Contents + :local: + detect._base module ------------------- - .. automodule:: plugins.extract.detect._base :members: :undoc-members: :show-inheritance: -align._base module ------------------- - -.. automodule:: plugins.extract.align._base +detect.mtcnn module +------------------- +.. automodule:: plugins.extract.detect.mtcnn :members: :undoc-members: :show-inheritance: + +mask plugins package +==================== +.. contents:: Contents + :local: + mask._base module ----------------- - .. automodule:: plugins.extract.mask._base :members: :undoc-members: :show-inheritance: -vgg\_face2\_keras module +mask.bisenet_fp module +---------------------- +.. automodule:: plugins.extract.mask.bisenet_fp + :members: + :undoc-members: + :show-inheritance: + + +recognition plugins package +=========================== +.. contents:: Contents + :local: + +recognition._base module ------------------------ +.. automodule:: plugins.extract.recognition._base + :members: + :undoc-members: + :show-inheritance: + -.. automodule:: plugins.extract.recognition.vgg_face2_keras +recognition.vgg_face2 module +---------------------------- +.. automodule:: plugins.extract.recognition.vgg_face2 :members: :undoc-members: - :show-inheritance: \ No newline at end of file + :show-inheritance: diff --git a/docs/full/plugins/train.rst b/docs/full/plugins/train.rst index 4f0675d23b..cfaf89bd45 100755 --- a/docs/full/plugins/train.rst +++ b/docs/full/plugins/train.rst @@ -8,37 +8,59 @@ The Train Package handles the Model and Trainer plugins for training models in F .. contents:: Contents :local: -model._base module -================== +model package +============= + +This package contains various helper functions that plugins can inherit from .. rubric:: Module Summary .. autosummary:: :nosignatures: - ~plugins.train.model._base.KerasModel - ~plugins.train.model._base.ModelBase - ~plugins.train.model._base.State + ~plugins.train.model._base.model + ~plugins.train.model._base.settings + ~plugins.train.model._base.io -.. rubric:: Module +model._base.model module +------------------------ -.. automodule:: plugins.train.model._base +.. automodule:: plugins.train.model._base.model + :members: + :undoc-members: + :show-inheritance: + +model._base.settings module +--------------------------- + +.. automodule:: plugins.train.model._base.settings + :members: + :undoc-members: + :show-inheritance: + +model._base.io module +--------------------- + +.. automodule:: plugins.train.model._base.io :members: :undoc-members: :show-inheritance: model.original module -===================== +---------------------- .. automodule:: plugins.train.model.original :members: :undoc-members: :show-inheritance: +trainer package +=============== + trainer._base module -==================== +---------------------- .. automodule:: plugins.train.trainer._base :members: :undoc-members: - :show-inheritance: \ No newline at end of file + :show-inheritance: diff --git a/docs/full/scripts.rst b/docs/full/scripts.rst index eb8871b4ca..1736d5a308 100644 --- a/docs/full/scripts.rst +++ b/docs/full/scripts.rst @@ -51,7 +51,6 @@ fsmedia module ~scripts.fsmedia.Alignments ~scripts.fsmedia.DebugLandmarks - ~scripts.fsmedia.FaceFilter ~scripts.fsmedia.Images ~scripts.fsmedia.PostProcess ~scripts.fsmedia.finalize diff --git a/docs/full/setup.rst b/docs/full/setup.rst new file mode 100644 index 0000000000..baa29b9f19 --- /dev/null +++ b/docs/full/setup.rst @@ -0,0 +1,8 @@ +************ +setup module +************ + +.. automodule:: setup + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tests/lib.gpu_stats.rst b/docs/full/tests/lib.gpu_stats.rst new file mode 100644 index 0000000000..dbca67ef7f --- /dev/null +++ b/docs/full/tests/lib.gpu_stats.rst @@ -0,0 +1,15 @@ +***************** +gpu_stats package +***************** + +.. contents:: Contents + :local: + +_base_test module +***************** +Unittests for the :class:`~lib.gpu_stats._base` module + +.. automodule:: tests.lib.gpu_stats._base_test + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tests/lib.gui.rst b/docs/full/tests/lib.gui.rst new file mode 100644 index 0000000000..4ec4258e7b --- /dev/null +++ b/docs/full/tests/lib.gui.rst @@ -0,0 +1,15 @@ +*********** +gui package +*********** + +.. contents:: Contents + :local: + +gui.analysis.event_reader module +******************************** +Unittests for the :class:`~lib.gui.analysis.event_reader` module + +.. automodule:: tests.lib.gui.analysis.event_reader_test + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tests/lib.rst b/docs/full/tests/lib.rst new file mode 100644 index 0000000000..a020342f8b --- /dev/null +++ b/docs/full/tests/lib.rst @@ -0,0 +1,39 @@ +*********** +lib package +*********** + +.. contents:: Contents + :local: + +Subpackages +=========== + +.. toctree:: + :maxdepth: 1 + + lib.gpu_stats + lib.gui + + +sysinfo module +************** +Unit tests for :class:`~lib.sysinfo` module + +.. rubric:: Module + +.. automodule:: tests.lib.sysinfo_test + :members: + :undoc-members: + :show-inheritance: + + +utils_test module +***************** +Unit tests for :class:`~lib.utils` module + +.. rubric:: Module + +.. automodule:: tests.lib.utils_test + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/tests/tests.rst b/docs/full/tests/tests.rst new file mode 100644 index 0000000000..a24f36aa6f --- /dev/null +++ b/docs/full/tests/tests.rst @@ -0,0 +1,17 @@ +************* +tests package +************* + +The Tests Package provides Faceswap's Unit Tests. + +.. contents:: Contents + :local: + +Subpackages +=========== + +.. toctree:: + :maxdepth: 1 + + lib + tools diff --git a/docs/full/tests/tools.alignments.rst b/docs/full/tests/tools.alignments.rst new file mode 100644 index 0000000000..cb5e3d1018 --- /dev/null +++ b/docs/full/tests/tools.alignments.rst @@ -0,0 +1,15 @@ +****************** +alignments package +****************** + +.. contents:: Contents + :local: + +media_test module +***************** +Unittests for the :class:`~tools.alignments.media` module + +.. automodule:: tests.tools.alignments.media_test + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/tests/tools.preview.rst b/docs/full/tests/tools.preview.rst new file mode 100644 index 0000000000..7c744b12f0 --- /dev/null +++ b/docs/full/tests/tools.preview.rst @@ -0,0 +1,15 @@ +*************** +preview package +*************** + +.. contents:: Contents + :local: + +viewer_test module +****************** +Unittests for the :class:`~tools.preview.viewer` module + +.. automodule:: tests.tools.preview.viewer_test + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tests/tools.rst b/docs/full/tests/tools.rst new file mode 100644 index 0000000000..d19ad80daf --- /dev/null +++ b/docs/full/tests/tools.rst @@ -0,0 +1,15 @@ +************* +tools package +************* + +.. contents:: Contents + :local: + +Subpackages +=========== + +.. toctree:: + :maxdepth: 1 + + tools.alignments + tools.preview diff --git a/docs/full/tools/alignments.rst b/docs/full/tools/alignments.rst new file mode 100644 index 0000000000..33119bfa64 --- /dev/null +++ b/docs/full/tools/alignments.rst @@ -0,0 +1,50 @@ +****************** +alignments package +****************** + +.. contents:: Contents + :local: + + +alignments module +***************** +The Alignments Module is the main entry point into the Alignments Tool. + +.. automodule:: tools.alignments.alignments + :members: + :undoc-members: + :show-inheritance: + + +jobs_faces module +================= + +.. automodule:: tools.alignments.jobs_faces + :members: + :undoc-members: + :show-inheritance: + + +jobs_frames module +================== + +.. automodule:: tools.alignments.jobs_frames + :members: + :undoc-members: + :show-inheritance: + +jobs module +=========== + +.. automodule:: tools.alignments.jobs + :members: + :undoc-members: + :show-inheritance: + +media module +============ + +.. automodule:: tools.alignments.media + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/tools/manual.faceviewer.rst b/docs/full/tools/manual.faceviewer.rst index ed3209f4f1..5589e144a1 100644 --- a/docs/full/tools/manual.faceviewer.rst +++ b/docs/full/tools/manual.faceviewer.rst @@ -7,6 +7,7 @@ Handles the display of faces in the Face Viewer section of Faceswap's Manual Too .. contents:: Contents :local: + frame module ============ @@ -28,6 +29,27 @@ frame module :undoc-members: :show-inheritance: + +interact module +=============== + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.faceviewer.interact.ActiveFrame + ~tools.manual.faceviewer.interact.Asset + ~tools.manual.faceviewer.interact.HoverBox + +.. rubric:: Module + +.. automodule:: tools.manual.faceviewer.interact + :members: + :undoc-members: + :show-inheritance: + + viewport module =============== @@ -36,8 +58,6 @@ viewport module .. autosummary:: :nosignatures: - ~tools.manual.faceviewer.viewport.ActiveFrame - ~tools.manual.faceviewer.viewport.HoverBox ~tools.manual.faceviewer.viewport.TKFace ~tools.manual.faceviewer.viewport.Viewport ~tools.manual.faceviewer.viewport.VisibleObjects diff --git a/docs/full/tools/manual.rst b/docs/full/tools/manual.rst index f2428395c3..4b35542559 100644 --- a/docs/full/tools/manual.rst +++ b/docs/full/tools/manual.rst @@ -23,11 +23,10 @@ The Manual Module is the main entry point into the Manual Editor Tool. .. autosummary:: :nosignatures: - + ~tools.manual.manual.Aligner ~tools.manual.manual.FrameLoader ~tools.manual.manual.Manual - ~tools.manual.manual.TkGlobals .. rubric:: Module @@ -43,11 +42,10 @@ detected_faces module .. autosummary:: :nosignatures: - + ~tools.manual.detected_faces.DetectedFaces ~tools.manual.detected_faces.FaceUpdate ~tools.manual.detected_faces.Filter - ~tools.manual.detected_faces.ThumbsCreator .. rubric:: Module @@ -55,3 +53,32 @@ detected_faces module :members: :undoc-members: :show-inheritance: + +globals module +============== + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.globals.CurrentFrame + ~tools.manual.globals.TkGlobals + ~tools.manual.globals.TKVars + +.. rubric:: Module + +.. automodule:: tools.manual.globals + :members: + :undoc-members: + :show-inheritance: + + +thumbnails module +================== + +.. automodule:: tools.manual.thumbnails + :members: + :undoc-members: + :show-inheritance: + diff --git a/docs/full/tools/preview.rst b/docs/full/tools/preview.rst new file mode 100644 index 0000000000..5c0c76d790 --- /dev/null +++ b/docs/full/tools/preview.rst @@ -0,0 +1,44 @@ +*************** +preview package +*************** + +.. contents:: Contents + :local: + + +preview module +============== +The Preview Module is the main entry point into the Preview Tool. + +.. automodule:: tools.preview.preview + :members: + :undoc-members: + :show-inheritance: + + +cli module +========== + +.. automodule:: tools.preview.cli + :members: + :undoc-members: + :show-inheritance: + + +control_panels module +===================== + +.. automodule:: tools.preview.control_panels + :members: + :undoc-members: + :show-inheritance: + + +viewer module +============= + +.. automodule:: tools.preview.viewer + :members: + :undoc-members: + :show-inheritance: + diff --git a/docs/full/tools/sort.rst b/docs/full/tools/sort.rst new file mode 100644 index 0000000000..6335339799 --- /dev/null +++ b/docs/full/tools/sort.rst @@ -0,0 +1,34 @@ +************ +sort package +************ + +.. contents:: Contents + :local: + + +sort module +=========== +The Sort Module is the main entry point into the Sort Tool. + +.. automodule:: tools.sort.sort + :members: + :undoc-members: + :show-inheritance: + + +sort_methods module +=================== + +.. automodule:: tools.sort.sort_methods + :members: + :undoc-members: + :show-inheritance: + + +sort_methods_aligned module +=========================== + +.. automodule:: tools.sort.sort_methods_aligned + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/tools/tools.rst b/docs/full/tools/tools.rst index 07468d3cea..0381c4b9c3 100644 --- a/docs/full/tools/tools.rst +++ b/docs/full/tools/tools.rst @@ -13,14 +13,10 @@ Subpackages .. toctree:: :maxdepth: 1 + alignments manual - -alignments module -================= -.. automodule:: tools.alignments.alignments - :members: - :undoc-members: - :show-inheritance: + preview + sort mask module =========== @@ -30,27 +26,10 @@ mask module :undoc-members: :show-inheritance: -preview module -=============== - -.. rubric:: Module Summary - -.. autosummary:: - :nosignatures: - - ~tools.preview.preview.ActionFrame - ~tools.preview.preview.ConfigFrame - ~tools.preview.preview.ConfigTools - ~tools.preview.preview.FacesDisplay - ~tools.preview.preview.ImagesCanvas - ~tools.preview.preview.OptionsBook - ~tools.preview.preview.Patch - ~tools.preview.preview.Preview - ~tools.preview.preview.Samples - -.. rubric:: Module +model module +============ -.. automodule:: tools.preview.preview +.. automodule:: tools.model.model :members: :undoc-members: :show-inheritance: diff --git a/docs/full/update_deps.rst b/docs/full/update_deps.rst new file mode 100644 index 0000000000..aea4753eaf --- /dev/null +++ b/docs/full/update_deps.rst @@ -0,0 +1,8 @@ +****************** +update_deps module +****************** + +.. automodule:: update_deps + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/sphinx_requirements.txt b/docs/sphinx_requirements.txt index 83f2db1e56..59a4b9b8f6 100755 --- a/docs/sphinx_requirements.txt +++ b/docs/sphinx_requirements.txt @@ -1,19 +1,21 @@ # NB Do not install from this requirements file # It is for documentation purposes only -tqdm==4.62 -psutil==5.8.0 -numpy==1.18.0 -opencv-python>4.5.3.0,<4.5.4.0 -pillow==8.3.1 -scikit-learn==0.24.2 -fastcluster==1.1.26 -matplotlib==3.5.1 -imageio==2.9.0 -imageio-ffmpeg==0.4.5 -ffmpy==0.2.3 -nvidia-ml-py3 -pywin32==228 ; sys_platform == "win32" -pynvx==1.0.0 ; sys_platform == "darwin" -plaidml-keras==0.7.0 -tensorflow==2.2.0 +sphinx>=6.0.0,<7.0.0 +sphinx_rtd_theme==1.2.2 +tqdm==4.65 +psutil==5.9.0 +numexpr>=2.8.7 +numpy>=1.26.0 +opencv-python>=4.9.0.0 +pillow==9.4.0 +scikit-learn>=1.3.0 +fastcluster>=1.2.6 +matplotlib==3.8.0 +imageio==2.33.1 +imageio-ffmpeg==0.4.9 +ffmpy==0.3.0 +nvidia-ml-py>=12.535,<12.536 +pytest==7.2.0 +pytest-mock==3.10.0 +tensorflow>=2.10.0,<2.11.0 diff --git a/faceswap.py b/faceswap.py index b8857eb697..5f27ba1792 100755 --- a/faceswap.py +++ b/faceswap.py @@ -1,23 +1,25 @@ #!/usr/bin/env python3 """ The master faceswap.py script """ import gettext +import locale +import os import sys -from lib.cli import args as cli_args -from lib.config import generate_configs -from lib.utils import get_backend +# Translations don't work by default in Windows, so hack in environment variable +if sys.platform.startswith("win"): + os.environ["LANG"], _ = locale.getdefaultlocale() +from lib.cli import args as cli_args # pylint:disable=wrong-import-position +from lib.cli.args_train import TrainArgs # pylint:disable=wrong-import-position +from lib.cli.args_extract_convert import ConvertArgs, ExtractArgs # noqa:E501 pylint:disable=wrong-import-position +from lib.config import generate_configs # pylint:disable=wrong-import-position # LOCALES _LANG = gettext.translation("faceswap", localedir="locales", fallback=True) _ = _LANG.gettext - -if sys.version_info < (3, 7): - raise Exception("This program requires at least python3.7") -if get_backend() == "amd" and sys.version_info >= (3, 9): - raise Exception("The AMD version of Faceswap cannot run on versions of Python higher than 3.8") - +if sys.version_info < (3, 10): + raise ValueError("This program requires at least python 3.10") _PARSER = cli_args.FullHelpArgumentParser() @@ -41,11 +43,11 @@ def _main() -> None: generate_configs() subparser = _PARSER.add_subparsers() - cli_args.ExtractArgs(subparser, "extract", _("Extract the faces from pictures or a video")) - cli_args.TrainArgs(subparser, "train", _("Train a model for the two faces A and B")) - cli_args.ConvertArgs(subparser, - "convert", - _("Convert source pictures or video to a new one with the face swapped")) + ExtractArgs(subparser, "extract", _("Extract the faces from pictures or a video")) + TrainArgs(subparser, "train", _("Train a model for the two faces A and B")) + ConvertArgs(subparser, + "convert", + _("Convert source pictures or video to a new one with the face swapped")) cli_args.GuiArgs(subparser, "gui", _("Launch the Faceswap Graphical User Interface")) _PARSER.set_defaults(func=_bad_args) arguments = _PARSER.parse_args() diff --git a/lib/__init__.py b/lib/__init__.py index e69de29bb2..c87f4c4316 100644 --- a/lib/__init__.py +++ b/lib/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python3 +""" Initialization for faceswap's lib section """ +# Import logger here so our custom loglevels are set for when executing code outside of FS +from . import logger diff --git a/lib/align/__init__.py b/lib/align/__init__.py index 12b82f283a..3f5887bcd6 100644 --- a/lib/align/__init__.py +++ b/lib/align/__init__.py @@ -1,6 +1,9 @@ #!/usr/bin/env python3 """ Package for handling alignments files, detected faces and aligned faces along with their associated objects. """ -from .aligned_face import AlignedFace, _EXTRACT_RATIOS, get_matrix_scaling, get_centered_size, PoseEstimate, transform_image # noqa -from .alignments import Alignments # noqa -from .detected_face import BlurMask, DetectedFace, Mask, update_legacy_png_header # noqa +from .aligned_face import (AlignedFace, get_adjusted_center, get_matrix_scaling, + get_centered_size, transform_image) +from .aligned_mask import BlurMask, LandmarksMask, Mask +from .alignments import Alignments +from .constants import CenteringType, EXTRACT_RATIOS, LANDMARK_PARTS, LandmarkType +from .detected_face import DetectedFace, update_legacy_png_header diff --git a/lib/align/aligned_face.py b/lib/align/aligned_face.py index b5dd48687d..41f2eed8c3 100644 --- a/lib/align/aligned_face.py +++ b/lib/align/aligned_face.py @@ -1,64 +1,25 @@ #!/usr/bin/env python3 """ Aligner for faceswap.py """ +from __future__ import annotations +from dataclasses import dataclass, field import logging +import typing as T + from threading import Lock import cv2 import numpy as np -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - - -_MEAN_FACE = np.array([[0.010086, 0.106454], [0.085135, 0.038915], [0.191003, 0.018748], - [0.300643, 0.034489], [0.403270, 0.077391], [0.596729, 0.077391], - [0.699356, 0.034489], [0.808997, 0.018748], [0.914864, 0.038915], - [0.989913, 0.106454], [0.500000, 0.203352], [0.500000, 0.307009], - [0.500000, 0.409805], [0.500000, 0.515625], [0.376753, 0.587326], - [0.435909, 0.609345], [0.500000, 0.628106], [0.564090, 0.609345], - [0.623246, 0.587326], [0.131610, 0.216423], [0.196995, 0.178758], - [0.275698, 0.179852], [0.344479, 0.231733], [0.270791, 0.245099], - [0.192616, 0.244077], [0.655520, 0.231733], [0.724301, 0.179852], - [0.803005, 0.178758], [0.868389, 0.216423], [0.807383, 0.244077], - [0.729208, 0.245099], [0.264022, 0.780233], [0.350858, 0.745405], - [0.438731, 0.727388], [0.500000, 0.742578], [0.561268, 0.727388], - [0.649141, 0.745405], [0.735977, 0.780233], [0.652032, 0.864805], - [0.566594, 0.902192], [0.500000, 0.909281], [0.433405, 0.902192], - [0.347967, 0.864805], [0.300252, 0.784792], [0.437969, 0.778746], - [0.500000, 0.785343], [0.562030, 0.778746], [0.699747, 0.784792], - [0.563237, 0.824182], [0.500000, 0.831803], [0.436763, 0.824182]]) - -_MEAN_FACE_3D = np.array([[4.056931, -11.432347, 1.636229], # 8 chin LL - [1.833492, -12.542305, 4.061275], # 7 chin L - [0.0, -12.901019, 4.070434], # 6 chin C - [-1.833492, -12.542305, 4.061275], # 5 chin R - [-4.056931, -11.432347, 1.636229], # 4 chin RR - [6.825897, 1.275284, 4.402142], # 33 L eyebrow L - [1.330353, 1.636816, 6.903745], # 29 L eyebrow R - [-1.330353, 1.636816, 6.903745], # 34 R eyebrow L - [-6.825897, 1.275284, 4.402142], # 38 R eyebrow R - [1.930245, -5.060977, 5.914376], # 54 nose LL - [0.746313, -5.136947, 6.263227], # 53 nose L - [0.0, -5.485328, 6.76343], # 52 nose C - [-0.746313, -5.136947, 6.263227], # 51 nose R - [-1.930245, -5.060977, 5.914376], # 50 nose RR - [5.311432, 0.0, 3.987654], # 13 L eye L - [1.78993, -0.091703, 4.413414], # 17 L eye R - [-1.78993, -0.091703, 4.413414], # 25 R eye L - [-5.311432, 0.0, 3.987654], # 21 R eye R - [2.774015, -7.566103, 5.048531], # 43 mouth L - [0.509714, -7.056507, 6.566167], # 42 mouth top L - [0.0, -7.131772, 6.704956], # 41 mouth top C - [-0.509714, -7.056507, 6.566167], # 40 mouth top R - [-2.774015, -7.566103, 5.048531], # 39 mouth R - [-0.589441, -8.443925, 6.109526], # 46 mouth bottom R - [0.0, -8.601736, 6.097667], # 45 mouth bottom C - [0.589441, -8.443925, 6.109526]]) # 44 mouth bottom L - -_EXTRACT_RATIOS = dict(legacy=0.375, face=0.5, head=0.625) - - -def get_matrix_scaling(matrix): +from lib.logger import parse_class_init + +from .constants import CenteringType, EXTRACT_RATIOS, LandmarkType, _MEAN_FACE +from .pose import PoseEstimate + +logger = logging.getLogger(__name__) + + +def get_matrix_scaling(matrix: np.ndarray) -> tuple[int, int]: """ Given a matrix, return the cv2 Interpolation method and inverse interpolation method for applying the matrix on an image. @@ -80,11 +41,15 @@ def get_matrix_scaling(matrix): interpolators = cv2.INTER_CUBIC, cv2.INTER_AREA else: interpolators = cv2.INTER_AREA, cv2.INTER_CUBIC - logger.trace("interpolator: %s, inverse interpolator: %s", interpolators[0], interpolators[1]) + logger.trace("interpolator: %s, inverse interpolator: %s", # type:ignore[attr-defined] + interpolators[0], interpolators[1]) return interpolators -def transform_image(image, matrix, size, padding=0): +def transform_image(image: np.ndarray, + matrix: np.ndarray, + size: int, + padding: int = 0) -> np.ndarray: """ Perform transformation on an image, applying the given size and padding to the matrix. Parameters @@ -103,7 +68,7 @@ def transform_image(image, matrix, size, padding=0): :class:`numpy.ndarray` The transformed image """ - logger.trace("image shape: %s, matrix: %s, size: %s. padding: %s", + logger.trace("image shape: %s, matrix: %s, size: %s. padding: %s", # type:ignore[attr-defined] image.shape, matrix, size, padding) # transform the matrix for size and padding mat = matrix * (size - 2 * padding) @@ -112,10 +77,167 @@ def transform_image(image, matrix, size, padding=0): # transform image interpolators = get_matrix_scaling(mat) retval = cv2.warpAffine(image, mat, (size, size), flags=interpolators[0]) - logger.trace("transformed matrix: %s, final image shape: %s", mat, image.shape) + logger.trace("transformed matrix: %s, final image shape: %s", # type:ignore[attr-defined] + mat, image.shape) + return retval + + +def get_adjusted_center(image_size: int, + source_offset: np.ndarray, + target_offset: np.ndarray, + source_centering: CenteringType) -> np.ndarray: + """ Obtain the correct center of a face extracted image to translate between two different + extract centerings. + + Parameters + ---------- + image_size: int + The size of the image at the given :attr:`source_centering` + source_offset: :class:`numpy.ndarray` + The pose offset to translate a base extracted face to source centering + target_offset: :class:`numpy.ndarray` + The pose offset to translate a base extracted face to target centering + source_centering: ["face", "head", "legacy"] + The centering of the source image + + Returns + ------- + :class:`numpy.ndarray` + The center point of the image at the given size for the target centering + """ + source_size = image_size - (image_size * EXTRACT_RATIOS[source_centering]) + offset = target_offset - source_offset + offset *= source_size + center = np.rint(offset + image_size / 2).astype("int32") + logger.trace( # type:ignore[attr-defined] + "image_size: %s, source_offset: %s, target_offset: %s, source_centering: '%s', " + "adjusted_offset: %s, center: %s", + image_size, source_offset, target_offset, source_centering, offset, center) + return center + + +def get_centered_size(source_centering: CenteringType, + target_centering: CenteringType, + size: int, + coverage_ratio: float = 1.0) -> int: + """ Obtain the size of a cropped face from an aligned image. + + Given an image of a certain dimensions, returns the dimensions of the sub-crop within that + image for the requested centering at the requested coverage ratio + + Notes + ----- + `"legacy"` places the nose in the center of the image (the original method for aligning). + `"face"` aligns for the nose to be in the center of the face (top to bottom) but the center + of the skull for left to right. `"head"` places the center in the middle of the skull in 3D + space. + + The ROI in relation to the source image is calculated by rounding the padding of one side + to the nearest integer then applying this padding to the center of the crop, to ensure that + any dimensions always have an even number of pixels. + + Parameters + ---------- + source_centering: ["head", "face", "legacy"] + The centering that the original image is aligned at + target_centering: ["head", "face", "legacy"] + The centering that the sub-crop size should be obtained for + size: int + The size of the source image to obtain the cropped size for + coverage_ratio: float, optional + The coverage ratio to be applied to the target image. Default: `1.0` + + Returns + ------- + int + The pixel size of a sub-crop image from a full head aligned image with the given coverage + ratio + """ + if source_centering == target_centering and coverage_ratio == 1.0: + retval = size + else: + src_size = size - (size * EXTRACT_RATIOS[source_centering]) + retval = 2 * int(np.rint((src_size / (1 - EXTRACT_RATIOS[target_centering]) + * coverage_ratio) / 2)) + logger.trace( # type:ignore[attr-defined] + "source_centering: %s, target_centering: %s, size: %s, coverage_ratio: %s, " + "source_size: %s, crop_size: %s", + source_centering, target_centering, size, coverage_ratio, src_size, retval) return retval +@dataclass +class _FaceCache: # pylint:disable=too-many-instance-attributes + """ Cache for storing items related to a single aligned face. + + Items are cached so that they are only created the first time they are called. + Each item includes a threading lock to make cache creation thread safe. + + Parameters + ---------- + pose: :class:`lib.align.PoseEstimate`, optional + The estimated pose in 3D space. Default: ``None`` + original_roi: :class:`numpy.ndarray`, optional + The location of the extracted face box within the original frame. Default: ``None`` + landmarks: :class:`numpy.ndarray`, optional + The 68 point facial landmarks aligned to the extracted face box. Default: ``None`` + landmarks_normalized: :class:`numpy.ndarray`: + The 68 point facial landmarks normalized to 0.0 - 1.0 as aligned by Umeyama. + Default: ``None`` + average_distance: float, optional + The average distance of the core landmarks (18-67) from the mean face that was used for + aligning the image. Default: `0.0` + relative_eye_mouth_position: float, optional + A float value representing the relative position of the lowest eye/eye-brow point to the + highest mouth point. Positive values indicate that eyes/eyebrows are aligned above the + mouth, negative values indicate that eyes/eyebrows are misaligned below the mouth. + Default: `0.0` + adjusted_matrix: :class:`numpy.ndarray`, optional + The 3x2 transformation matrix for extracting and aligning the core face area out of the + original frame with padding and sizing applied. Default: ``None`` + interpolators: tuple, optional + (`interpolator` and `reverse interpolator`) for the :attr:`adjusted matrix`. + Default: `(0, 0)` + cropped_roi, dict, optional + The (`left`, `top`, `right`, `bottom` location of the region of interest within an + aligned face centered for each centering. Default: `{}` + cropped_slices: dict, optional + The slices for an input full head image and output cropped image. Default: `{}` + """ + pose: PoseEstimate | None = None + original_roi: np.ndarray | None = None + landmarks: np.ndarray | None = None + landmarks_normalized: np.ndarray | None = None + average_distance: float = 0.0 + relative_eye_mouth_position: float = 0.0 + adjusted_matrix: np.ndarray | None = None + interpolators: tuple[int, int] = (0, 0) + cropped_roi: dict[CenteringType, np.ndarray] = field(default_factory=dict) + cropped_slices: dict[CenteringType, dict[T.Literal["in", "out"], + tuple[slice, slice]]] = field(default_factory=dict) + + _locks: dict[str, Lock] = field(default_factory=dict) + + def __post_init__(self): + """ Initialize the locks for the class parameters """ + self._locks = {name: Lock() for name in self.__dict__} + + def lock(self, name: str) -> Lock: + """ Obtain the lock for the given property + + Parameters + ---------- + name: str + The name of a parameter within the cache + + Returns + ------- + :class:`threading.Lock` + The lock associated with the requested parameter + """ + return self._locks[name] + + class AlignedFace(): """ Class to align a face. @@ -147,179 +269,225 @@ class AlignedFace(): is_aligned_face: bool, optional Indicates that the :attr:`image` is an aligned face rather than a frame. Default: ``False`` + is_legacy: bool, optional + Only used if `is_aligned` is ``True``. ``True`` indicates that the aligned image being + loaded is a legacy extracted face rather than a current head extracted face """ - def __init__(self, landmarks, image=None, centering="face", size=64, coverage_ratio=1.0, - dtype=None, is_aligned=False): - logger.trace("Initializing: %s (image shape: %s, centering: '%s', size: %s, " - "coverage_ratio: %s, dtype: %s, is_aligned: %s)", self.__class__.__name__, - image if image is None else image.shape, centering, size, coverage_ratio, - dtype, is_aligned) + def __init__(self, + landmarks: np.ndarray, + image: np.ndarray | None = None, + centering: CenteringType = "face", + size: int = 64, + coverage_ratio: float = 1.0, + dtype: str | None = None, + is_aligned: bool = False, + is_legacy: bool = False) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] self._frame_landmarks = landmarks + self._landmark_type = LandmarkType.from_shape(landmarks.shape) self._centering = centering self._size = size + self._coverage_ratio = coverage_ratio self._dtype = dtype self._is_aligned = is_aligned - self._matrices = dict(legacy=_umeyama(landmarks[17:], _MEAN_FACE, True)[0:2], - face=None, - head=None) + self._source_centering: CenteringType = "legacy" if is_legacy and is_aligned else "head" self._padding = self._padding_from_coverage(size, coverage_ratio) - self._cache = self._set_cache() + lookup = self._landmark_type + self._mean_lookup = LandmarkType.LM_2D_51 if lookup == LandmarkType.LM_2D_68 else lookup + + self._cache = _FaceCache() + self._matrices: dict[CenteringType, np.ndarray] = {"legacy": self._get_default_matrix()} self._face = self.extract_face(image) - logger.trace("Initialized: %s (matrix: %s, padding: %s, face shape: %s)", - self.__class__.__name__, self._matrices["legacy"], self._padding, + logger.trace("Initialized: %s (padding: %s, face shape: %s)", # type:ignore[attr-defined] + self.__class__.__name__, self._padding, self._face if self._face is None else self._face.shape) @property - def size(self): + def centering(self) -> T.Literal["legacy", "head", "face"]: + """ str: The centering of the Aligned Face. One of `"legacy"`, `"head"`, `"face"`. """ + return self._centering + + @property + def size(self) -> int: """ int: The size (in pixels) of one side of the square extracted face image. """ return self._size @property - def padding(self): + def padding(self) -> int: """ int: The amount of padding (in pixels) that is applied to each side of the extracted face image for the selected extract type. """ return self._padding[self._centering] @property - def matrix(self): + def matrix(self) -> np.ndarray: """ :class:`numpy.ndarray`: The 3x2 transformation matrix for extracting and aligning the core face area out of the original frame, with no padding or sizing applied. The returned matrix is offset for the given :attr:`centering`. """ - if self._matrices[self._centering] is None: + if self._centering not in self._matrices: matrix = self._matrices["legacy"].copy() matrix[:, 2] -= self.pose.offset[self._centering] self._matrices[self._centering] = matrix - logger.trace("original matrix: %s, new matrix: %s", self._matrices["legacy"], matrix) + logger.trace("original matrix: %s, new matrix: %s", # type:ignore[attr-defined] + self._matrices["legacy"], matrix) return self._matrices[self._centering] @property - def _head_size(self): - """ int: The size of the full head extract image calculated from the required - centering. """ - with self._cache["head_size"][1]: - if self._centering not in self._cache["head_size"][0]: - self._cache["head_size"][0][self._centering] = get_centered_size(self._centering, - "head", - self.size) - return self._cache["head_size"][0][self._centering] - - @property - def pose(self): + def pose(self) -> PoseEstimate: """ :class:`lib.align.PoseEstimate`: The estimated pose in 3D space. """ - with self._cache["pose"][1]: - if self._cache["pose"][0] is None: - lms = cv2.transform(np.expand_dims(self._frame_landmarks, axis=1), - self._matrices["legacy"]).squeeze() - self._cache["pose"][0] = PoseEstimate(lms) - return self._cache["pose"][0] + with self._cache.lock("pose"): + if self._cache.pose is None: + lms = np.nan_to_num(cv2.transform(np.expand_dims(self._frame_landmarks, axis=1), + self._matrices["legacy"]).squeeze()) + self._cache.pose = PoseEstimate(lms, self._landmark_type) + return self._cache.pose @property - def adjusted_matrix(self): + def adjusted_matrix(self) -> np.ndarray: """ :class:`numpy.ndarray`: The 3x2 transformation matrix for extracting and aligning the core face area out of the original frame with padding and sizing applied. """ - with self._cache["adjusted_matrix"][1]: - if self._cache["adjusted_matrix"][0] is None: + with self._cache.lock("adjusted_matrix"): + if self._cache.adjusted_matrix is None: matrix = self.matrix.copy() mat = matrix * (self._size - 2 * self.padding) mat[:, 2] += self.padding - logger.trace("adjusted_matrix: %s", mat) - self._cache["adjusted_matrix"][0] = mat - return self._cache["adjusted_matrix"][0] + logger.trace("adjusted_matrix: %s", mat) # type:ignore[attr-defined] + self._cache.adjusted_matrix = mat + return self._cache.adjusted_matrix @property - def face(self): + def face(self) -> np.ndarray | None: """ :class:`numpy.ndarray`: The aligned face at the given :attr:`size` at the specified :attr:`coverage` in the given :attr:`dtype`. If an :attr:`image` has not been provided then an the attribute will return ``None``. """ return self._face @property - def original_roi(self): + def original_roi(self) -> np.ndarray: """ :class:`numpy.ndarray`: The location of the extracted face box within the original frame. """ - with self._cache["original_roi"][1]: - if self._cache["original_roi"][0] is None: + with self._cache.lock("original_roi"): + if self._cache.original_roi is None: roi = np.array([[0, 0], [0, self._size - 1], [self._size - 1, self._size - 1], [self._size - 1, 0]]) roi = np.rint(self.transform_points(roi, invert=True)).astype("int32") - logger.trace("original roi: %s", roi) - self._cache["original_roi"][0] = roi - return self._cache["original_roi"][0] + logger.trace("original roi: %s", roi) # type:ignore[attr-defined] + self._cache.original_roi = roi + return self._cache.original_roi @property - def landmarks(self): + def landmarks(self) -> np.ndarray: """ :class:`numpy.ndarray`: The 68 point facial landmarks aligned to the extracted face box. """ - with self._cache["landmarks"][1]: - if self._cache["landmarks"][0] is None: + with self._cache.lock("landmarks"): + if self._cache.landmarks is None: lms = self.transform_points(self._frame_landmarks) - logger.trace("aligned landmarks: %s", lms) - self._cache["landmarks"][0] = lms - return self._cache["landmarks"][0] + logger.trace("aligned landmarks: %s", lms) # type:ignore[attr-defined] + self._cache.landmarks = lms + return self._cache.landmarks @property - def normalized_landmarks(self): + def landmark_type(self) -> LandmarkType: + """:class:`~LandmarkType`: The type of landmarks that generated this aligned face """ + return self._landmark_type + + @property + def normalized_landmarks(self) -> np.ndarray: """ :class:`numpy.ndarray`: The 68 point facial landmarks normalized to 0.0 - 1.0 as aligned by Umeyama. """ - with self._cache["landmarks_normalized"][1]: - if self._cache["landmarks_normalized"][0] is None: + with self._cache.lock("landmarks_normalized"): + if self._cache.landmarks_normalized is None: lms = np.expand_dims(self._frame_landmarks, axis=1) - lms = cv2.transform(lms, self._matrices["legacy"], lms.shape).squeeze() - logger.trace("normalized landmarks: %s", lms) - self._cache["landmarks_normalized"][0] = lms - return self._cache["landmarks_normalized"][0] + lms = cv2.transform(lms, self._matrices["legacy"]).squeeze() + logger.trace("normalized landmarks: %s", lms) # type:ignore[attr-defined] + self._cache.landmarks_normalized = lms + return self._cache.landmarks_normalized @property - def interpolators(self): + def interpolators(self) -> tuple[int, int]: """ tuple: (`interpolator` and `reverse interpolator`) for the :attr:`adjusted matrix`. """ - with self._cache["interpolators"][1]: - if self._cache["interpolators"][0] is None: + with self._cache.lock("interpolators"): + if not any(self._cache.interpolators): interpolators = get_matrix_scaling(self.adjusted_matrix) - logger.trace("interpolators: %s", interpolators) - self._cache["interpolators"][0] = interpolators - return self._cache["interpolators"][0] + logger.trace("interpolators: %s", interpolators) # type:ignore[attr-defined] + self._cache.interpolators = interpolators + return self._cache.interpolators @property - def average_distance(self): + def average_distance(self) -> float: """ float: The average distance of the core landmarks (18-67) from the mean face that was used for aligning the image. """ - with self._cache["average_distance"][1]: - if self._cache["average_distance"][0] is None: - # pylint:disable=unsubscriptable-object - average_distance = np.mean(np.abs(self.normalized_landmarks[17:] - _MEAN_FACE)) - logger.trace("average_distance: %s", average_distance) - self._cache["average_distance"][0] = average_distance - return self._cache["average_distance"][0] + with self._cache.lock("average_distance"): + if not self._cache.average_distance: + mean_face = _MEAN_FACE[self._mean_lookup] + lms = self.normalized_landmarks + if self._landmark_type == LandmarkType.LM_2D_68: + lms = lms[17:] # 68 point landmarks only use core face items + average_distance = np.mean(np.abs(lms - mean_face)) + logger.trace("average_distance: %s", average_distance) # type:ignore[attr-defined] + self._cache.average_distance = average_distance + return self._cache.average_distance + + @property + def relative_eye_mouth_position(self) -> float: + """ float: Value representing the relative position of the lowest eye/eye-brow point to the + highest mouth point. Positive values indicate that eyes/eyebrows are aligned above the + mouth, negative values indicate that eyes/eyebrows are misaligned below the mouth. """ + with self._cache.lock("relative_eye_mouth_position"): + if not self._cache.relative_eye_mouth_position: + if self._landmark_type != LandmarkType.LM_2D_68: + position = 1.0 # arbitrary positive value + else: + lowest_eyes = np.max(self.normalized_landmarks[np.r_[17:27, 36:48], 1]) + highest_mouth = np.min(self.normalized_landmarks[48:68, 1]) + position = highest_mouth - lowest_eyes + logger.trace("lowest_eyes: %s, highest_mouth: %s, " # type:ignore[attr-defined] + "relative_eye_mouth_position: %s", lowest_eyes, highest_mouth, + position) + self._cache.relative_eye_mouth_position = position + return self._cache.relative_eye_mouth_position @classmethod - def _set_cache(cls): - """ Set the cache items. + def _padding_from_coverage(cls, size: int, coverage_ratio: float) -> dict[CenteringType, int]: + """ Return the image padding for a face from coverage_ratio set against a + pre-padded training image. - Items are cached so that they are only created the first time they are called. - Each item includes a threading lock to make cache creation thread safe. + Parameters + ---------- + size: int + The final size of the aligned image in pixels + coverage_ratio: float + The ratio of the final image to pad to Returns ------- dict - The Aligned Face cache + The padding required, in pixels for 'head', 'face' and 'legacy' face types """ - return dict(pose=[None, Lock()], - original_roi=[None, Lock()], - landmarks=[None, Lock()], - landmarks_normalized=[None, Lock()], - average_distance=[None, Lock()], - adjusted_matrix=[None, Lock()], - interpolators=[None, Lock()], - head_size=[dict(), Lock()], - cropped_roi=[dict(), Lock()], - cropped_size=[dict(), Lock()], - cropped_slices=[dict(), Lock()]) - - def transform_points(self, points, invert=False): + retval = {_type: round((size * (coverage_ratio - (1 - EXTRACT_RATIOS[_type]))) / 2) + for _type in T.get_args(T.Literal["legacy", "face", "head"])} + logger.trace(retval) # type:ignore[attr-defined] + return retval + + def _get_default_matrix(self) -> np.ndarray: + """ Get the default (legacy) matrix. All subsequent matrices are calculated from this + + Returns + ------- + :class:`numpy.ndarray` + The default 'legacy' matrix + """ + lms = self._frame_landmarks + if self._landmark_type == LandmarkType.LM_2D_68: + lms = lms[17:] # 68 point landmarks only use core face items + retval = _umeyama(lms, _MEAN_FACE[self._mean_lookup], True)[0:2] + logger.trace("Default matrix: %s", retval) # type:ignore[attr-defined] + return retval + + def transform_points(self, points: np.ndarray, invert: bool = False) -> np.ndarray: """ Perform transformation on a series of (x, y) co-ordinates in world space into aligned face space. @@ -338,12 +506,12 @@ def transform_points(self, points, invert=False): """ retval = np.expand_dims(points, axis=1) mat = cv2.invertAffineTransform(self.adjusted_matrix) if invert else self.adjusted_matrix - retval = cv2.transform(retval, mat, retval.shape).squeeze() - logger.trace("invert: %s, Original points: %s, transformed points: %s", - invert, points, retval) + retval = cv2.transform(retval, mat).squeeze() + logger.trace( # type:ignore[attr-defined] + "invert: %s, Original points: %s, transformed points: %s", invert, points, retval) return retval - def extract_face(self, image): + def extract_face(self, image: np.ndarray | None) -> np.ndarray | None: """ Extract the face from a source image and populate :attr:`face`. If an image is not provided then ``None`` is returned. @@ -360,10 +528,13 @@ def extract_face(self, image): ``None`` if no image has been provided. """ if image is None: - logger.trace("_extract_face called without a loaded image. Returning empty face.") + logger.trace("_extract_face called without a loaded " # type:ignore[attr-defined] + "image. Returning empty face.") return None - if self._is_aligned and self._centering != "head": # Crop out the sub face from full head + if self._is_aligned and (self._centering != self._source_centering or + self._coverage_ratio != 1.0): + # Crop out the sub face from full head image = self._convert_centering(image) if self._is_aligned and image.shape[0] != self._size: # Resize the given aligned face @@ -376,13 +547,13 @@ def extract_face(self, image): retval = retval if self._dtype is None else retval.astype(self._dtype) return retval - def _convert_centering(self, image): + def _convert_centering(self, image: np.ndarray) -> np.ndarray: """ When the face being loaded is pre-aligned, the loaded image will have 'head' centering so it needs to be cropped out to the appropriate centering. This function temporarily converts this object to a full head aligned face, extracts the sub-cropped face to the correct centering, reverse the sub crop and returns the cropped - face. + face at the selected coverage ratio. Parameters ---------- @@ -392,50 +563,74 @@ def _convert_centering(self, image): Returns ------- :class:`numpy.ndarray` - The aligned image with the correct centering + The aligned image with the correct centering, scaled to image input size """ - # Input image is sized up because of integer rounding - logger.trace("head_size: %s, image_size: %s, target_size: %s", - self._head_size, image.shape[0], self.size) - if self._head_size != image.shape[0]: - interp = cv2.INTER_CUBIC if image.shape[0] < self._head_size else cv2.INTER_AREA - image = cv2.resize(image, (self._head_size, self._head_size), interpolation=interp) - - out = np.zeros((self.size, self.size, image.shape[-1]), dtype=image.dtype) - slices = self._get_cropped_slices() + logger.trace( # type:ignore[attr-defined] + "image_size: %s, target_size: %s, coverage_ratio: %s", + image.shape[0], self.size, self._coverage_ratio) + + img_size = image.shape[0] + target_size = get_centered_size(self._source_centering, + self._centering, + img_size, + self._coverage_ratio) + out = np.zeros((target_size, target_size, image.shape[-1]), dtype=image.dtype) + + slices = self._get_cropped_slices(img_size, target_size) out[slices["out"][0], slices["out"][1], :] = image[slices["in"][0], slices["in"][1], :] - logger.trace("Cropped from aligned extract: (centering: %s, in shape: %s, out shape: %s)", - self._centering, image.shape, out.shape) + logger.trace( # type:ignore[attr-defined] + "Cropped from aligned extract: (centering: %s, in shape: %s, out shape: %s)", + self._centering, image.shape, out.shape) return out - @classmethod - def _padding_from_coverage(cls, size, coverage_ratio): - """ Return the image padding for a face from coverage_ratio set against a - pre-padded training image. + def _get_cropped_slices(self, + image_size: int, + target_size: int, + ) -> dict[T.Literal["in", "out"], tuple[slice, slice]]: + """ Obtain the slices to turn a full head extract into an alternatively centered extract. Parameters ---------- - size: int - The final size of the aligned image in pixels - coverage_ratio: float - The ratio of the final image to pad to + image_size: int + The size of the full head extracted image loaded from disk + target_size: int + The size of the target centered face with coverage ratio applied in relation to the + original image size Returns ------- dict - The padding required, in pixels for 'head', 'face' and 'legacy' face types + The slices for an input full head image and output cropped image """ - retval = {_type: round((size * (coverage_ratio - (1 - _EXTRACT_RATIOS[_type]))) / 2) - for _type in ("legacy", "face", "head")} - logger.trace(retval) - return retval - - def get_cropped_roi(self, centering): + with self._cache.lock("cropped_slices"): + if not self._cache.cropped_slices.get(self._centering): + roi = self.get_cropped_roi(image_size, target_size, self._centering) + slice_in = (slice(max(roi[1], 0), max(roi[3], 0)), + slice(max(roi[0], 0), max(roi[2], 0))) + slice_out = (slice(max(roi[1] * -1, 0), + target_size - min(target_size, max(0, roi[3] - image_size))), + slice(max(roi[0] * -1, 0), + target_size - min(target_size, max(0, roi[2] - image_size)))) + self._cache.cropped_slices[self._centering] = {"in": slice_in, "out": slice_out} + logger.trace("centering: %s, cropped_slices: %s", # type:ignore[attr-defined] + self._centering, self._cache.cropped_slices[self._centering]) + return self._cache.cropped_slices[self._centering] + + def get_cropped_roi(self, + image_size: int, + target_size: int, + centering: CenteringType) -> np.ndarray: """ Obtain the region of interest within an aligned face set to centered coverage for an alternative centering Parameters ---------- + image_size: int + The size of the full head extracted image loaded from disk + target_size: int + The size of the target centered face with coverage ratio applied in relation to the + original image size + centering: ["legacy", "face"] The type of centering to obtain the region of interest for. "legacy" places the nose in the center of the image (the original method for aligning). "face" aligns for the @@ -448,218 +643,41 @@ def get_cropped_roi(self, centering): The (`left`, `top`, `right`, `bottom` location of the region of interest within an aligned face centered on the head for the given centering """ - with self._cache["cropped_roi"][1]: - if centering not in self._cache["cropped_roi"][0]: - offset = self.pose.offset.get(centering, np.float32((0, 0))) # legacy = 0.0 - adjusted = offset - self.pose.offset["head"] - adjusted *= (self._head_size - (self._head_size * _EXTRACT_RATIOS["head"])) - - center = np.rint(adjusted + self._head_size / 2).astype("int32") - padding = self.size // 2 + with self._cache.lock("cropped_roi"): + if centering not in self._cache.cropped_roi: + center = get_adjusted_center(image_size, + self.pose.offset[self._source_centering], + self.pose.offset[centering], + self._source_centering) + padding = target_size // 2 roi = np.array([center - padding, center + padding]).ravel() - logger.trace("centering: '%s', center: %s, padding: %s, sub roi: %s", - centering, center, padding, roi) - self._cache["cropped_roi"][0][centering] = roi - return self._cache["cropped_roi"][0][centering] - - def _get_cropped_slices(self): - """ Obtain the slices to turn a full head extract into an alternatively centered extract. - - Returns - ------- - dict - The slices for an input full head image and output cropped image - """ - with self._cache["cropped_slices"][1]: - if not self._cache["cropped_slices"][0].get(self._centering): - roi = self.get_cropped_roi(self._centering) - slice_in = [slice(max(roi[1], 0), max(roi[3], 0)), - slice(max(roi[0], 0), max(roi[2], 0))] - slice_out = [slice(max(roi[1] * -1, 0), - self._size - min(self._size, max(0, roi[3] - self._head_size))), - slice(max(roi[0] * -1, 0), - self._size - min(self._size, max(0, roi[2] - self._head_size)))] - self._cache["cropped_slices"][0][self._centering] = {"in": slice_in, - "out": slice_out} - logger.trace("centering: %s, cropped_slices: %s", - self._centering, self._cache["cropped_slices"][0][self._centering]) - return self._cache["cropped_slices"][0][self._centering] - - -class PoseEstimate(): - """ Estimates pose from a generic 3D head model for the given 2D face landmarks. - - Parameters - ---------- - landmarks: :class:`numpy.ndarry` - The original 68 point landmarks aligned to 0.0 - 1.0 range - - References - ---------- - Head Pose Estimation using OpenCV and Dlib - https://www.learnopencv.com/tag/solvepnp/ - 3D Model points - http://aifi.isr.uc.pt/Downloads/OpenGL/glAnthropometric3DModel.cpp - """ - def __init__(self, landmarks): - self._distortion_coefficients = np.zeros((4, 1)) # Assuming no lens distortion - self._xyz_2d = None - - self._camera_matrix = self._get_camera_matrix() - self._rotation, self._translation = self._solve_pnp(landmarks) - self._offset = self._get_offset() - self._pitch_yaw = None + logger.trace( # type:ignore[attr-defined] + "centering: '%s', center: %s, padding: %s, sub roi: %s", + centering, center, padding, roi) + self._cache.cropped_roi[centering] = roi + return self._cache.cropped_roi[centering] - @property - def xyz_2d(self): - """ :class:`numpy.ndarray` projected (x, y) coordinates for each x, y, z point at a - constant distance from adjusted center of the skull (0.5, 0.5) in the 2D space. """ - if self._xyz_2d is None: - xyz = cv2.projectPoints(np.float32([[6, 0, -2.3], [0, 6, -2.3], [0, 0, 3.7]]), - self._rotation, - self._translation, - self._camera_matrix, - self._distortion_coefficients)[0].squeeze() - self._xyz_2d = xyz - self._offset["head"] - return self._xyz_2d - - @property - def offset(self): - """ dict: The amount to offset a standard 0.0 - 1.0 umeyama transformation matrix for a - from the center of the face (between the eyes) or center of the head (middle of skull) - rather than the nose area. """ - return self._offset - - @property - def pitch(self): - """ float: The pitch of the aligned face in eular angles """ - if not self._pitch_yaw: - self._get_pitch_yaw() - return self._pitch_yaw[0] - - @property - def yaw(self): - """ float: The yaw of the aligned face in eular angles """ - if not self._pitch_yaw: - self._get_pitch_yaw() - return self._pitch_yaw[1] - - def _get_pitch_yaw(self): - """ Obtain the yaw and pitch from the :attr:`_rotation` in eular angles. """ - proj_matrix = np.zeros((3, 4), dtype="float32") - proj_matrix[:3, :3] = cv2.Rodrigues(self._rotation)[0] - euler = cv2.decomposeProjectionMatrix(proj_matrix)[-1] - self._pitch_yaw = (euler[0][0], euler[1][0]) - logger.trace("yaw_pitch: %s", self._pitch_yaw) - - @classmethod - def _get_camera_matrix(cls): - """ Obtain an estimate of the camera matrix based off the original frame dimensions. + def split_mask(self) -> np.ndarray: + """ Remove the mask from the alpha channel of :attr:`face` and return the mask Returns ------- :class:`numpy.ndarray` - An estimated camera matrix - """ - focal_length = 4 - camera_matrix = np.array([[focal_length, 0, 0.5], - [0, focal_length, 0.5], - [0, 0, 1]], dtype="double") - logger.trace("camera_matrix: %s", camera_matrix) - return camera_matrix - - def _solve_pnp(self, landmarks): - """ Solve the Perspective-n-Point for the given landmarks. + The mask that was stored in the :attr:`face`'s alpha channel - Takes 2D landmarks in world space and estimates the rotation and translation vectors - in 3D space. - - Parameters - ---------- - landmarks: :class:`numpy.ndarry` - The original 68 point landmark co-ordinates relating to the original frame - - Returns - ------- - rotation: :class:`numpy.ndarray` - The solved rotation vector - translation: :class:`numpy.ndarray` - The solved translation vector - """ - points = landmarks[[6, 7, 8, 9, 10, 17, 21, 22, 26, 31, 32, 33, 34, - 35, 36, 39, 42, 45, 48, 50, 51, 52, 54, 56, 57, 58]] - _, rotation, translation = cv2.solvePnP(_MEAN_FACE_3D, - points, - self._camera_matrix, - self._distortion_coefficients, - flags=cv2.SOLVEPNP_ITERATIVE) - logger.trace("points: %s, rotation: %s, translation: %s", points, rotation, translation) - return rotation, translation - - def _get_offset(self): - """ Obtain the offset between the original center of the extracted face to the new center - of the head in 2D space. - - Returns - ------- - :class:`numpy.ndarray` - The x, y offset of the new center from the old center. + Raises + ------ + AssertionError + If :attr:`face` does not contain a mask in the alpha channel """ - offset = dict(legacy=np.array([0.0, 0.0])) - points = dict(head=(0, 0, -2.3), face=(0, -1.5, 4.2)) - - for key, pnts in points.items(): - center = cv2.projectPoints(np.float32([pnts]), - self._rotation, - self._translation, - self._camera_matrix, - self._distortion_coefficients)[0].squeeze() - logger.trace("center %s: %s", key, center) - offset[key] = center - (0.5, 0.5) - logger.trace("offset: %s", offset) - return offset - - -def get_centered_size(source_centering, target_centering, size): - """ Obtain the size of a cropped face from an aligned image. + assert self._face is not None + assert self._face.shape[-1] == 4, "No mask stored in the alpha channel" + mask = self._face[..., 3] + self._face = self._face[..., :3] + return mask - Given an image of a certain dimensions, returns the dimensions of the sub-crop within that - image for the requested centering. - Notes - ----- - `"legacy"` places the nose in the center of the image (the original method for aligning). - `"face"` aligns for the nose to be in the center of the face (top to bottom) but the center - of the skull for left to right. `"head"` places the center in the middle of the skull in 3D - space. - - The ROI in relation to the source image is calculated by rounding the padding of one side - to the nearest integer then applying this padding to the center of the crop, to ensure that - any dimensions always have an even number of pixels. - - Parameters - ---------- - source_centering: ["head", "face", "legacy"] - The centering that the original image is aligned at - target_centering: ["head", "face", "legacy"] - The centering that the sub-crop size should be obtained for - size: int - The size of the source image to obtain the cropped size for - - Returns - ------- - int - The pixel size of a sub-crop image from a full head aligned image - """ - if source_centering == target_centering: - retval = size - else: - src_size = size - (size * _EXTRACT_RATIOS[source_centering]) - retval = 2 * int(np.rint(src_size / (1 - _EXTRACT_RATIOS[target_centering]) / 2)) - logger.trace("source_centering: %s, target_centering: %s, size: %s, crop_size: %s", - source_centering, target_centering, size, retval) - return retval - - -def _umeyama(source, destination, estimate_scale): +def _umeyama(source: np.ndarray, destination: np.ndarray, estimate_scale: bool) -> np.ndarray: """Estimate N-D similarity transformation with or without scaling. Imported, and slightly adapted, directly from: @@ -706,24 +724,24 @@ def _umeyama(source, destination, estimate_scale): if np.linalg.det(A) < 0: d[dim - 1] = -1 - T = np.eye(dim + 1, dtype=np.double) + retval = np.eye(dim + 1, dtype=np.double) U, S, V = np.linalg.svd(A) # Eq. (40) and (43). rank = np.linalg.matrix_rank(A) if rank == 0: - return np.nan * T + return np.nan * retval if rank == dim - 1: if np.linalg.det(U) * np.linalg.det(V) > 0: - T[:dim, :dim] = U @ V + retval[:dim, :dim] = U @ V else: s = d[dim - 1] d[dim - 1] = -1 - T[:dim, :dim] = U @ np.diag(d) @ V + retval[:dim, :dim] = U @ np.diag(d) @ V d[dim - 1] = s else: - T[:dim, :dim] = U @ np.diag(d) @ V + retval[:dim, :dim] = U @ np.diag(d) @ V if estimate_scale: # Eq. (41) and (42). @@ -731,7 +749,7 @@ def _umeyama(source, destination, estimate_scale): else: scale = 1.0 - T[:dim, dim] = dst_mean - scale * (T[:dim, :dim] @ src_mean.T) - T[:dim, :dim] *= scale + retval[:dim, dim] = dst_mean - scale * (retval[:dim, :dim] @ src_mean.T) + retval[:dim, :dim] *= scale - return T + return retval diff --git a/lib/align/aligned_mask.py b/lib/align/aligned_mask.py new file mode 100644 index 0000000000..6a34060653 --- /dev/null +++ b/lib/align/aligned_mask.py @@ -0,0 +1,599 @@ +#!/usr/bin python3 +""" Handles retrieval and storage of Faceswap aligned masks """ + +from __future__ import annotations +import logging +import typing as T + +from zlib import compress, decompress + +import cv2 +import numpy as np + +from lib.logger import parse_class_init + +from .alignments import MaskAlignmentsFileDict +from . import get_adjusted_center, get_centered_size + +if T.TYPE_CHECKING: + from collections.abc import Callable + from .aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +class Mask(): + """ Face Mask information and convenience methods + + Holds a Faceswap mask as generated from :mod:`plugins.extract.mask` and the information + required to transform it to its original frame. + + Holds convenience methods to handle the warping, storing and retrieval of the mask. + + Parameters + ---------- + storage_size: int, optional + The size (in pixels) that the mask should be stored at. Default: 128. + storage_centering, str (optional): + The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. + Default: `"face"` + + Attributes + ---------- + stored_size: int + The size, in pixels, of the stored mask across its height and width. + stored_centering: str + The centering that the mask is stored at. One of `"legacy"`, `"face"`, `"head"` + """ + def __init__(self, + storage_size: int = 128, + storage_centering: CenteringType = "face") -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self.stored_size = storage_size + self.stored_centering = storage_centering + + self._mask: bytes | None = None + self._affine_matrix: np.ndarray | None = None + self._interpolator: int | None = None + + self._blur_type: T.Literal["gaussian", "normalized"] | None = None + self._blur_passes: int = 0 + self._blur_kernel: float | int = 0 + self._threshold = 0.0 + self._dilation: tuple[T.Literal["erode", "dilate"], np.ndarray | None] = ("erode", None) + self._sub_crop_size = 0 + self._sub_crop_slices: dict[T.Literal["in", "out"], list[slice]] = {} + + self.set_blur_and_threshold() + logger.trace("Initialized: %s", self.__class__.__name__) # type:ignore[attr-defined] + + @property + def mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` with any requested + blurring, threshold amount and centering applied.""" + mask = self.stored_mask + if self._dilation[-1] is not None or self._threshold != 0.0 or self._blur_kernel != 0: + mask = mask.copy() + self._dilate_mask(mask) + if self._threshold != 0.0: + mask[mask < self._threshold] = 0.0 + mask[mask > 255.0 - self._threshold] = 255.0 + if self._blur_kernel != 0 and self._blur_type is not None: + mask = BlurMask(self._blur_type, + mask, + self._blur_kernel, + passes=self._blur_passes).blurred + if self._sub_crop_size: # Crop the mask to the given centering + out = np.zeros((self._sub_crop_size, self._sub_crop_size, 1), dtype=mask.dtype) + slice_in, slice_out = self._sub_crop_slices["in"], self._sub_crop_slices["out"] + out[slice_out[0], slice_out[1], :] = mask[slice_in[0], slice_in[1], :] + mask = out + logger.trace("mask shape: %s", mask.shape) # type:ignore[attr-defined] + return mask + + @property + def stored_mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` as it is stored + (i.e. with no blurring/centering applied). """ + assert self._mask is not None + dims = (self.stored_size, self.stored_size, 1) + mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims) + logger.trace("stored mask shape: %s", mask.shape) # type:ignore[attr-defined] + return mask + + @property + def original_roi(self) -> np.ndarray: + """ :class: `numpy.ndarray`: The original region of interest of the mask in the + source frame. """ + points = np.array([[0, 0], + [0, self.stored_size - 1], + [self.stored_size - 1, self.stored_size - 1], + [self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2)) + matrix = cv2.invertAffineTransform(self.affine_matrix) + roi = cv2.transform(points, matrix).reshape((4, 2)) + logger.trace("Returning: %s", roi) # type:ignore[attr-defined] + return roi + + @property + def affine_matrix(self) -> np.ndarray: + """ :class: `numpy.ndarray`: The affine matrix to transpose the mask to a full frame. """ + assert self._affine_matrix is not None + return self._affine_matrix + + @property + def interpolator(self) -> int: + """ int: The cv2 interpolator required to transpose the mask to a full frame. """ + assert self._interpolator is not None + return self._interpolator + + def _dilate_mask(self, mask: np.ndarray) -> None: + """ Erode/Dilate the mask. The action is performed in-place on the given mask. + + No action is performed if a dilation amount has not been set + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask to be eroded/dilated + """ + if self._dilation[-1] is None: + return + + func = cv2.erode if self._dilation[0] == "erode" else cv2.dilate + func(mask, self._dilation[-1], dst=mask, iterations=1) + + def get_full_frame_mask(self, width: int, height: int) -> np.ndarray: + """ Return the stored mask in a full size frame of the given dimensions + + Parameters + ---------- + width: int + The width of the original frame that the mask was extracted from + height: int + The height of the original frame that the mask was extracted from + + Returns + ------- + :class:`numpy.ndarray`: The mask affined to the original full frame of the given dimensions + """ + frame = np.zeros((width, height, 1), dtype="uint8") + mask = cv2.warpAffine(self.mask, + self.affine_matrix, + (width, height), + frame, + flags=cv2.WARP_INVERSE_MAP | self.interpolator, + borderMode=cv2.BORDER_CONSTANT) + logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] + "mask max: %s", mask.shape, mask.dtype, mask.min(), mask.max()) + return mask + + def add(self, mask: np.ndarray, affine_matrix: np.ndarray, interpolator: int) -> None: + """ Add a Faceswap mask to this :class:`Mask`. + + The mask should be the original output from :mod:`plugins.extract.mask` + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask that is to be added as output from :mod:`plugins.extract.mask` + It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` + affine_matrix: :class:`numpy.ndarray` + The transformation matrix required to transform the mask to the original frame. + interpolator, int: + The CV2 interpolator required to transform this mask to it's original frame + """ + logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] + "mask max: %s, affine_matrix: %s, interpolator: %s)", + mask.shape, mask.dtype, mask.min(), affine_matrix, mask.max(), interpolator) + self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix) + self._interpolator = interpolator + self.replace_mask(mask) + + def replace_mask(self, mask: np.ndarray) -> None: + """ Replace the existing :attr:`_mask` with the given mask. + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask that is to be added as output from :mod:`plugins.extract.mask`. + It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` + """ + mask = (cv2.resize(mask * 255.0, + (self.stored_size, self.stored_size), + interpolation=cv2.INTER_AREA)).astype("uint8") + self._mask = compress(mask.tobytes()) + + def set_dilation(self, amount: float) -> None: + """ Set the internal dilation object for returned masks + + Parameters + ---------- + amount: float + The amount of erosion/dilation to apply as a percentage of the total mask size. + Negative values erode the mask. Positive values dilate the mask + """ + if amount == 0: + self._dilation = ("erode", None) + return + + action: T.Literal["erode", "dilate"] = "erode" if amount < 0 else "dilate" + kernel = int(round(self.stored_size * abs(amount / 100.), 0)) + self._dilation = (action, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel, kernel))) + + logger.trace("action: '%s', amount: %s, kernel: %s, ", # type:ignore[attr-defined] + action, amount, kernel) + + def set_blur_and_threshold(self, + blur_kernel: int = 0, + blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian", + blur_passes: int = 1, + threshold: int = 0) -> None: + """ Set the internal blur kernel and threshold amount for returned masks + + Parameters + ---------- + blur_kernel: int, optional + The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no + blurring. Should be odd, if an even number is passed in (outside of 0) then it is + rounded up to the next odd number. Default: 0 + blur_type: ["gaussian", "normalized"], optional + The blur type to use. ``gaussian`` or ``normalized`` box filter. Default: ``gaussian`` + blur_passes: int, optional + The number of passed to perform when blurring. Default: 1 + threshold: int, optional + The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value. + Default: 0 + """ + logger.trace("blur_kernel: %s, blur_type: %s, " # type:ignore[attr-defined] + "blur_passes: %s, threshold: %s", + blur_kernel, blur_type, blur_passes, threshold) + if blur_type is not None: + blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1 + self._blur_kernel = blur_kernel + self._blur_type = blur_type + self._blur_passes = blur_passes + self._threshold = (threshold / 100.0) * 255.0 + + def set_sub_crop(self, + source_offset: np.ndarray, + target_offset: np.ndarray, + centering: CenteringType, + coverage_ratio: float = 1.0) -> None: + """ Set the internal crop area of the mask to be returned. + + This impacts the returned mask from :attr:`mask` if the requested mask is required for + different face centering than what has been stored. + + Parameters + ---------- + source_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at its stored centering + target_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at the requested target centering + centering: str + The centering to set the sub crop area for. One of `"legacy"`, `"face"`. `"head"` + coverage_ratio: float, optional + The coverage ratio to be applied to the target image. ``None`` for default (1.0). + Default: ``None`` + """ + if centering == self.stored_centering and coverage_ratio == 1.0: + return + + center = get_adjusted_center(self.stored_size, + source_offset, + target_offset, + self.stored_centering) + crop_size = get_centered_size(self.stored_centering, + centering, + self.stored_size, + coverage_ratio=coverage_ratio) + roi = np.array([center - crop_size // 2, center + crop_size // 2]).ravel() + + self._sub_crop_size = crop_size + self._sub_crop_slices["in"] = [slice(max(roi[1], 0), max(roi[3], 0)), + slice(max(roi[0], 0), max(roi[2], 0))] + self._sub_crop_slices["out"] = [ + slice(max(roi[1] * -1, 0), + crop_size - min(crop_size, max(0, roi[3] - self.stored_size))), + slice(max(roi[0] * -1, 0), + crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))] + + logger.trace("src_size: %s, coverage_ratio: %s, " # type:ignore[attr-defined] + "sub_crop_size: %s, sub_crop_slices: %s", + roi, coverage_ratio, self._sub_crop_size, self._sub_crop_slices) + + def _adjust_affine_matrix(self, mask_size: int, affine_matrix: np.ndarray) -> np.ndarray: + """ Adjust the affine matrix for the mask's storage size + + Parameters + ---------- + mask_size: int + The original size of the mask. + affine_matrix: :class:`numpy.ndarray` + The affine matrix to transform the mask at original size to the parent frame. + + Returns + ------- + affine_matrix: :class:`numpy,ndarray` + The affine matrix adjusted for the mask at its stored dimensions. + """ + zoom = self.stored_size / mask_size + zoom_mat = np.array([[zoom, 0, 0.], [0, zoom, 0.]]) + adjust_mat = np.dot(zoom_mat, np.concatenate((affine_matrix, np.array([[0., 0., 1.]])))) + logger.trace("storage_size: %s, mask_size: %s, zoom: %s, " # type:ignore[attr-defined] + "original matrix: %s, adjusted_matrix: %s", self.stored_size, mask_size, zoom, + affine_matrix.shape, adjust_mat.shape) + return adjust_mat + + def to_dict(self, is_png=False) -> MaskAlignmentsFileDict: + """ Convert the mask to a dictionary for saving to an alignments file + + Parameters + ---------- + is_png: bool + ``True`` if the dictionary is being created for storage in a png header otherwise + ``False``. Default: ``False`` + + Returns + ------- + dict: + The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + assert self._mask is not None + affine_matrix = self.affine_matrix.tolist() if is_png else self.affine_matrix + retval = MaskAlignmentsFileDict(mask=self._mask, + affine_matrix=affine_matrix, + interpolator=self.interpolator, + stored_size=self.stored_size, + stored_centering=self.stored_centering) + logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] + for k, v in retval.items()}) + return retval + + def to_png_meta(self) -> MaskAlignmentsFileDict: + """ Convert the mask to a dictionary supported by png itxt headers. + + Returns + ------- + dict: + The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + return self.to_dict(is_png=True) + + def from_dict(self, mask_dict: MaskAlignmentsFileDict) -> None: + """ Populates the :class:`Mask` from a dictionary loaded from an alignments file. + + Parameters + ---------- + mask_dict: dict + A dictionary stored in an alignments file containing the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + self._mask = mask_dict["mask"] + affine_matrix = mask_dict["affine_matrix"] + self._affine_matrix = (affine_matrix if isinstance(affine_matrix, np.ndarray) + else np.array(affine_matrix, dtype="float64")) + self._interpolator = mask_dict["interpolator"] + self.stored_size = mask_dict["stored_size"] + centering = mask_dict.get("stored_centering") + self.stored_centering = "face" if centering is None else centering + logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] + for k, v in mask_dict.items()}) + + +class LandmarksMask(Mask): + """ Create a single channel mask from aligned landmark points. + + Landmarks masks are created on the fly, so the stored centering and size should be the same as + the aligned face that the mask will be applied to. As the masks are created on the fly, blur + + dilation is applied to the mask at creation (prior to compression) rather than after + decompression when requested. + + Note + ---- + Threshold is not used for Landmarks mask as the mask is binary + + Parameters + ---------- + points: list + A list of landmark points that correspond to the given storage_size to create + the mask. Each item in the list should be a :class:`numpy.ndarray` that a filled + convex polygon will be created from + storage_size: int, optional + The size (in pixels) that the compressed mask should be stored at. Default: 128. + storage_centering, str (optional): + The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. + Default: `"face"` + dilation: float, optional + The amount of dilation to apply to the mask. as a percentage of the mask size. Default: 0.0 + """ + def __init__(self, + points: list[np.ndarray], + storage_size: int = 128, + storage_centering: CenteringType = "face", + dilation: float = 0.0) -> None: + super().__init__(storage_size=storage_size, storage_centering=storage_centering) + self._points = points + self.set_dilation(dilation) + + @property + def mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: Overrides the default mask property, creating the processed + mask at first call and compressing it. The decompressed mask is returned from this + property. """ + return self.stored_mask + + def generate_mask(self, affine_matrix: np.ndarray, interpolator: int) -> None: + """ Generate the mask. + + Creates the mask applying any requested dilation and blurring and assigns compressed mask + to :attr:`_mask` + + Parameters + ---------- + affine_matrix: :class:`numpy.ndarray` + The transformation matrix required to transform the mask to the original frame. + interpolator, int: + The CV2 interpolator required to transform this mask to it's original frame + """ + mask = np.zeros((self.stored_size, self.stored_size, 1), dtype="float32") + for landmarks in self._points: + lms = np.rint(landmarks).astype("int") + cv2.fillConvexPoly(mask, cv2.convexHull(lms), [1.0], lineType=cv2.LINE_AA) + if self._dilation[-1] is not None: + self._dilate_mask(mask) + if self._blur_kernel != 0 and self._blur_type is not None: + mask = BlurMask(self._blur_type, + mask, + self._blur_kernel, + passes=self._blur_passes).blurred + logger.trace("mask: (shape: %s, dtype: %s)", # type:ignore[attr-defined] + mask.shape, mask.dtype) + self.add(mask, affine_matrix, interpolator) + + +class BlurMask(): + """ Factory class to return the correct blur object for requested blur type. + + Works for square images only. Currently supports Gaussian and Normalized Box Filters. + + Parameters + ---------- + blur_type: ["gaussian", "normalized"] + The type of blur to use + mask: :class:`numpy.ndarray` + The mask to apply the blur to + kernel: int or float + Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size + is_ratio: bool, optional + Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the + actual kernel size will be calculated from the given ratio and the mask size. If + ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. + Default: ``False`` + passes: int, optional + The number of passes to perform when blurring. Default: ``1`` + + Example + ------- + >>> print(mask.shape) + (128, 128, 1) + >>> new_mask = BlurMask("gaussian", mask, 3, is_ratio=False, passes=1).blurred + >>> print(new_mask.shape) + (128, 128, 1) + """ + def __init__(self, + blur_type: T.Literal["gaussian", "normalized"], + mask: np.ndarray, + kernel: int | float, + is_ratio: bool = False, + passes: int = 1) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self._blur_type = blur_type + self._mask = mask + self._passes = passes + kernel_size = self._get_kernel_size(kernel, is_ratio) + self._kernel_size = self._get_kernel_tuple(kernel_size) + logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] + + @property + def blurred(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The final mask with blurring applied. """ + func = self._func_mapping[self._blur_type] + kwargs = self._get_kwargs() + blurred = self._mask + for i in range(self._passes): + assert isinstance(kwargs["ksize"], tuple) + ksize = int(kwargs["ksize"][0]) + logger.trace("Pass: %s, kernel_size: %s", # type:ignore[attr-defined] + i + 1, (ksize, ksize)) + blurred = func(blurred, **kwargs) + ksize = int(round(ksize * self._multipass_factor)) + kwargs["ksize"] = self._get_kernel_tuple(ksize) + blurred = blurred[..., None] + logger.trace("Returning blurred mask. Shape: %s", # type:ignore[attr-defined] + blurred.shape) + return blurred + + @property + def _multipass_factor(self) -> float: + """ For multiple passes the kernel must be scaled down. This value is + different for box filter and gaussian """ + factor = {"gaussian": 0.8, "normalized": 0.5} + return factor[self._blur_type] + + @property + def _sigma(self) -> T.Literal[0]: + """ int: The Sigma for Gaussian Blur. Returns 0 to force calculation from kernel size. """ + return 0 + + @property + def _func_mapping(self) -> dict[T.Literal["gaussian", "normalized"], Callable]: + """ dict: :attr:`_blur_type` mapped to cv2 Function name. """ + return {"gaussian": cv2.GaussianBlur, "normalized": cv2.blur} + + @property + def _kwarg_requirements(self) -> dict[T.Literal["gaussian", "normalized"], list[str]]: + """ dict: :attr:`_blur_type` mapped to cv2 Function required keyword arguments. """ + return {"gaussian": ['ksize', 'sigmaX'], "normalized": ['ksize']} + + @property + def _kwarg_mapping(self) -> dict[str, int | tuple[int, int]]: + """ dict: cv2 function keyword arguments mapped to their parameters. """ + return {"ksize": self._kernel_size, "sigmaX": self._sigma} + + def _get_kernel_size(self, kernel: int | float, is_ratio: bool) -> int: + """ Set the kernel size to absolute value. + + If :attr:`is_ratio` is ``True`` then the kernel size is calculated from the given ratio and + the :attr:`_mask` size, otherwise the given kernel size is just returned. + + Parameters + ---------- + kernel: int or float + Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size + is_ratio: bool, optional + Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the + actual kernel size will be calculated from the given ratio and the mask size. If + ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. + + Returns + ------- + int + The size (in pixels) of the blur kernel + """ + if not is_ratio: + return int(kernel) + + mask_diameter = np.sqrt(np.sum(self._mask)) + radius = round(max(1., mask_diameter * kernel / 100.)) + kernel_size = int(radius * 2 + 1) + logger.trace("kernel_size: %s", kernel_size) # type:ignore[attr-defined] + return kernel_size + + @staticmethod + def _get_kernel_tuple(kernel_size: int) -> tuple[int, int]: + """ Make sure kernel_size is odd and return it as a tuple. + + Parameters + ---------- + kernel_size: int + The size in pixels of the blur kernel + + Returns + ------- + tuple + The kernel size as a tuple of ('int', 'int') + """ + kernel_size += 1 if kernel_size % 2 == 0 else 0 + retval = (kernel_size, kernel_size) + logger.trace(retval) # type:ignore[attr-defined] + return retval + + def _get_kwargs(self) -> dict[str, int | tuple[int, int]]: + """ dict: the valid keyword arguments for the requested :attr:`_blur_type` """ + retval = {kword: self._kwarg_mapping[kword] + for kword in self._kwarg_requirements[self._blur_type]} + logger.trace("BlurMask kwargs: %s", retval) # type:ignore[attr-defined] + return retval diff --git a/lib/align/alignments.py b/lib/align/alignments.py index 9e0f1aee9a..ff0f2207d3 100644 --- a/lib/align/alignments.py +++ b/lib/align/alignments.py @@ -1,9 +1,10 @@ #!/usr/bin/env python3 """ Alignments file functions for reading, writing and manipulating the data stored in a serialized alignments file. """ - +from __future__ import annotations import logging import os +import typing as T from datetime import datetime import numpy as np @@ -11,10 +12,16 @@ from lib.serializer import get_serializer, get_serializer_from_filename from lib.utils import FaceswapError -logger = logging.getLogger(__name__) # pylint: disable=invalid-name -_VERSION = 2.2 +from .thumbnails import Thumbnails +from .updater import (FileStructure, IdentityAndVideoMeta, LandmarkRename, Legacy, ListToNumpy, + MaskCentering, VideoExtension) +if T.TYPE_CHECKING: + from collections.abc import Generator + from .aligned_face import CenteringType +logger = logging.getLogger(__name__) +_VERSION = 2.4 # VERSION TRACKING # 1.0 - Never really existed. Basically any alignments file prior to version 2.0 # 2.0 - Implementation of full head extract. Any alignments version below this will have used @@ -22,6 +29,59 @@ # 2.1 - Alignments data to extracted face PNG header. SHA1 hashes of faces no longer calculated # or stored in alignments file # 2.2 - Add support for differently centered masks (i.e. not all masks stored as face centering) +# 2.3 - Add 'identity' key to alignments file. May or may not be populated, to contain vggface2 +# embeddings. Make 'video_meta' key a standard key. Can be unpopulated +# 2.4 - Update video file alignment keys to end in the video extension rather than '.png' + + +# TODO Convert these to Dataclasses +class MaskAlignmentsFileDict(T.TypedDict): + """ Typed Dictionary for storing Masks. """ + mask: bytes + affine_matrix: list[float] | np.ndarray + interpolator: int + stored_size: int + stored_centering: CenteringType + + +class PNGHeaderAlignmentsDict(T.TypedDict): + """ Base Dictionary for storing a single faces' Alignment Information in Alignments files and + PNG Headers. """ + x: int + y: int + w: int + h: int + landmarks_xy: list[float] | np.ndarray + mask: dict[str, MaskAlignmentsFileDict] + identity: dict[str, list[float]] + + +class AlignmentFileDict(PNGHeaderAlignmentsDict): + """ Typed Dictionary for storing a single faces' Alignment Information in alignments files. """ + thumb: np.ndarray | None + + +class PNGHeaderSourceDict(T.TypedDict): + """ Dictionary for storing additional meta information in PNG headers """ + alignments_version: float + original_filename: str + face_index: int + source_filename: str + source_is_video: bool + source_frame_dims: tuple[int, int] | None + + +class AlignmentDict(T.TypedDict): + """ Dictionary for holding all of the alignment information within a single alignment file """ + faces: list[AlignmentFileDict] + video_meta: dict[str, float | int] + + +class PNGHeaderDict(T.TypedDict): + """ Dictionary for storing all alignment and meta information in PNG Headers """ + alignments: PNGHeaderAlignmentsDict + source: PNGHeaderSourceDict + class Alignments(): """ The alignments file is a custom serialized ``.fsa`` file that holds information for each @@ -42,78 +102,63 @@ class Alignments(): The filename of the ``.fsa`` alignments file. If not provided then the given folder will be checked for a default alignments file filename. Default: "alignments" """ - def __init__(self, folder, filename="alignments"): + def __init__(self, folder: str, filename: str = "alignments") -> None: logger.debug("Initializing %s: (folder: '%s', filename: '%s')", self.__class__.__name__, folder, filename) - self._version = _VERSION - self._serializer = get_serializer("compressed") - self._file = self._get_location(folder, filename) - self._meta = None + self._io = _IO(self, folder, filename) self._data = self._load() - self._update_legacy() - self._hashes_to_frame = dict() - self._hashes_to_alignment = dict() + self._io.update_legacy() + + self._legacy = Legacy(self) self._thumbnails = Thumbnails(self) logger.debug("Initialized %s", self.__class__.__name__) # << PROPERTIES >> # @property - def frames_count(self): + def frames_count(self) -> int: """ int: The number of frames that appear in the alignments :attr:`data`. """ retval = len(self._data) - logger.trace(retval) + logger.trace(retval) # type:ignore[attr-defined] return retval @property - def faces_count(self): + def faces_count(self) -> int: """ int: The total number of faces that appear in the alignments :attr:`data`. """ retval = sum(len(val["faces"]) for val in self._data.values()) - logger.trace(retval) + logger.trace(retval) # type:ignore[attr-defined] return retval @property - def file(self): + def file(self) -> str: """ str: The full path to the currently loaded alignments file. """ - return self._file + return self._io.file @property - def data(self): + def data(self) -> dict[str, AlignmentDict]: """ dict: The loaded alignments :attr:`file` in dictionary form. """ return self._data @property - def have_alignments_file(self): + def have_alignments_file(self) -> bool: """ bool: ``True`` if an alignments file exists at location :attr:`file` otherwise ``False``. """ - retval = os.path.exists(self._file) - logger.trace(retval) - return retval + return self._io.have_alignments_file @property - def hashes_to_frame(self): + def hashes_to_frame(self) -> dict[str, dict[str, int]]: """ dict: The SHA1 hash of the face mapped to the frame(s) and face index within the frame - that the hash corresponds to. The structure of the dictionary is: - - {**SHA1_hash** (`str`): {**filename** (`str`): **face_index** (`int`)}}. + that the hash corresponds to. Notes ----- This method is depractated and exists purely for updating legacy hash based alignments to new png header storage in :class:`lib.align.update_legacy_png_header`. - - The first time this property is referenced, the dictionary will be created and cached. - Subsequent references will be made to this cached dictionary. """ - if not self._hashes_to_frame: - logger.debug("Generating hashes to frame") - for frame_name, val in self._data.items(): - for idx, face in enumerate(val["faces"]): - self._hashes_to_frame.setdefault(face["hash"], dict())[frame_name] = idx - return self._hashes_to_frame + return self._legacy.hashes_to_frame @property - def hashes_to_alignment(self): + def hashes_to_alignment(self) -> dict[str, AlignmentFileDict]: """ dict: The SHA1 hash of the face mapped to the alignment for the face that the hash corresponds to. The structure of the dictionary is: @@ -121,105 +166,54 @@ def hashes_to_alignment(self): ----- This method is depractated and exists purely for updating legacy hash based alignments to new png header storage in :class:`lib.align.update_legacy_png_header`. - - The first time this property is referenced, the dictionary will be created and cached. - Subsequent references will be made to this cached dictionary. """ - if not self._hashes_to_alignment: - logger.debug("Generating hashes to alignment") - self._hashes_to_alignment = {face["hash"]: face - for val in self._data.values() - for face in val["faces"]} - return self._hashes_to_alignment + return self._legacy.hashes_to_alignment @property - def mask_summary(self): + def mask_summary(self) -> dict[str, int]: """ dict: The mask type names stored in the alignments :attr:`data` as key with the number of faces which possess the mask type as value. """ - masks = dict() + masks: dict[str, int] = {} for val in self._data.values(): for face in val["faces"]: if face.get("mask", None) is None: masks["none"] = masks.get("none", 0) + 1 - for key in face.get("mask", dict()): + for key in face.get("mask", {}): masks[key] = masks.get(key, 0) + 1 return masks @property - def video_meta_data(self): + def video_meta_data(self) -> dict[str, list[int] | list[float] | None]: """ dict: The frame meta data stored in the alignments file. If data does not exist in the alignments file then ``None`` is returned for each Key """ - retval = dict(pts_time=None, keyframes=None) - pts_time = [] - keyframes = [] + retval: dict[str, list[int] | list[float] | None] = {"pts_time": None, "keyframes": None} + pts_time: list[float] = [] + keyframes: list[int] = [] for idx, key in enumerate(sorted(self.data)): - if "video_meta" not in self.data[key]: + if not self.data[key].get("video_meta", {}): return retval meta = self.data[key]["video_meta"] - pts_time.append(meta["pts_time"]) + pts_time.append(T.cast(float, meta["pts_time"])) if meta["keyframe"]: keyframes.append(idx) - retval = dict(pts_time=pts_time, keyframes=keyframes) + retval = {"pts_time": pts_time, "keyframes": keyframes} return retval @property - def thumbnails(self): - """ :class:`~lib.align.Thumbnails`: The low resolution thumbnail images that exist - within the alignments file """ + def thumbnails(self) -> Thumbnails: + """ :class:`~lib.align.thumbnails.Thumbnails`: The low resolution thumbnail images that + exist within the alignments file """ return self._thumbnails @property - def version(self): + def version(self) -> float: """ float: The alignments file version number. """ - return self._version - - # << INIT FUNCTIONS >> # - - def _get_location(self, folder, filename): - """ Obtains the location of an alignments file. - - If a legacy alignments file is provided/discovered, then the alignments file will be - updated to the custom ``.fsa`` format and saved. - - Parameters - ---------- - folder: str - The folder that the alignments file is located in - filename: str - The filename of the alignments file - - Returns - ------- - str - The full path to the alignments file - """ - logger.debug("Getting location: (folder: '%s', filename: '%s')", folder, filename) - noext_name, extension = os.path.splitext(filename) - if extension in (".json", ".p", ".pickle", ".yaml", ".yml"): - # Reformat legacy alignments file - filename = self._update_file_format(folder, filename) - logger.debug("Updated legacy alignments. New filename: '%s'", filename) - if extension[1:] == self._serializer.file_extension: - logger.debug("Valid Alignments filename provided: '%s'", filename) - else: - filename = "{}.{}".format(noext_name, self._serializer.file_extension) - logger.debug("File extension set from serializer: '%s'", - self._serializer.file_extension) - location = os.path.join(str(folder), filename) - if not os.path.exists(location): - # Test for old format alignments files and reformat if they exist. This will be - # executed if an alignments file has not been explicitly provided therefore it will not - # have been picked up in the extension test - self._test_for_legacy(location) - logger.verbose("Alignments filepath: '%s'", location) - return location + return self._io.version - # << I/O >> # - - def _load(self): + def _load(self) -> dict[str, AlignmentDict]: """ Load the alignments data from the serialized alignments :attr:`file`. - Populates :attr:`_meta` with the alignment file's meta information as well as returning + Populates :attr:`_version` with the alignment file's loaded version as well as returning the serialized data. Returns @@ -227,49 +221,23 @@ def _load(self): dict: The loaded alignments data """ - logger.debug("Loading alignments") - if not self.have_alignments_file: - raise FaceswapError("Error: Alignments file not found at " - "{}".format(self._file)) - - logger.info("Reading alignments from: '%s'", self._file) - data = self._serializer.load(self._file) - self._meta = data.get("__meta__", dict(version=1.0)) - self._version = self._meta["version"] - data = data.get("__data__", data) - logger.debug("Loaded alignments") - return data + return self._io.load() - def save(self): + def save(self) -> None: """ Write the contents of :attr:`data` and :attr:`_meta` to a serialized ``.fsa`` file at the location :attr:`file`. """ - logger.debug("Saving alignments") - logger.info("Writing alignments to: '%s'", self._file) - data = dict(__meta__=dict(version=self._version), - __data__=self._data) - self._serializer.save(self._file, data) - logger.debug("Saved alignments") + return self._io.save() - def backup(self): + def backup(self) -> None: """ Create a backup copy of the alignments :attr:`file`. Creates a copy of the serialized alignments :attr:`file` appending a timestamp onto the end of the file name and storing in the same folder as the original :attr:`file`. """ - logger.debug("Backing up alignments") - if not os.path.isfile(self._file): - logger.debug("No alignments to back up") - return - now = datetime.now().strftime("%Y%m%d_%H%M%S") - src = self._file - split = os.path.splitext(src) - dst = split[0] + "_" + now + split[1] - logger.info("Backing up original alignments to '%s'", dst) - os.rename(src, dst) - logger.debug("Backed up alignments") + return self._io.backup() - def save_video_meta_data(self, pts_time, keyframes): + def save_video_meta_data(self, pts_time: list[float], keyframes: list[int]) -> None: """ Save video meta data to the alignments file. If the alignments file does not have an entry for every frame (e.g. if Extract Every N @@ -289,34 +257,37 @@ def save_video_meta_data(self, pts_time, keyframes): sample_filename = next(fname for fname in self.data) basename = sample_filename[:sample_filename.rfind("_")] - logger.debug("sample filename: %s, base filename: %s", sample_filename, basename) + ext = os.path.splitext(sample_filename)[-1] + logger.debug("sample filename: '%s', base filename: '%s' extension: '%s'", + sample_filename, basename, ext) logger.info("Saving video meta information to Alignments file") for idx, pts in enumerate(pts_time): - meta = dict(pts_time=pts, keyframe=idx in keyframes) - key = "{}_{:06d}.png".format(basename, idx + 1) + meta: dict[str, float | int] = {"pts_time": pts, "keyframe": idx in keyframes} + key = f"{basename}_{idx + 1:06d}{ext}" if key not in self.data: - self.data[key] = dict(video_meta=meta, faces=[]) + self.data[key] = {"video_meta": meta, "faces": []} else: self.data[key]["video_meta"] = meta logger.debug("Alignments count: %s, timestamp count: %s", len(self.data), len(pts_time)) if len(self.data) != len(pts_time): raise FaceswapError( - "There is a mismatch between the number of frames found in the video file ({}) " - "and the number of frames found in the alignments file ({})." - "\nThis can be caused by a number of issues:" + "There is a mismatch between the number of frames found in the video file " + f"({len(pts_time)}) and the number of frames found in the alignments file " + f"({len(self.data)}).\nThis can be caused by a number of issues:" "\n - The video has a Variable Frame Rate and FFMPEG is having a hard time " "calculating the correct number of frames." "\n - You are working with a Merged Alignments file. This is not supported for " "your current use case." "\nYou should either extract the video to individual frames, re-encode the " "video at a constant frame rate and re-run extraction or work with a dedicated " - "alignments file for your requested video.".format(len(pts_time), len(self.data))) - self.save() + "alignments file for your requested video.") + self._io.save() @classmethod - def _pad_leading_frames(cls, pts_time, keyframes): + def _pad_leading_frames(cls, pts_time: list[float], keyframes: list[int]) -> tuple[list[float], + list[int]]: """ Calculate the number of frames to pad the video by when the first frame is not a key frame. @@ -326,9 +297,11 @@ def _pad_leading_frames(cls, pts_time, keyframes): Parameters ---------- - pts_time: list + pts_time: list A list of presentation timestamps (`float`) in frame index order for every frame in the input video + keyframes: list + A list of keyframes (`int`) for the input video Returns ------- @@ -338,7 +311,7 @@ def _pad_leading_frames(cls, pts_time, keyframes): """ start_pts = pts_time[0] logger.debug("Video not cut on keyframe. Start pts: %s", start_pts) - gaps = [] + gaps: list[float] = [] prev_time = None for item in pts_time: if prev_time is not None: @@ -355,8 +328,7 @@ def _pad_leading_frames(cls, pts_time, keyframes): return pts_time, keyframes # << VALIDATION >> # - - def frame_exists(self, frame_name): + def frame_exists(self, frame_name: str) -> bool: """ Check whether a given frame_name exists within the alignments :attr:`data`. Parameters @@ -371,10 +343,10 @@ def frame_exists(self, frame_name): otherwise ``False`` """ retval = frame_name in self._data.keys() - logger.trace("'%s': %s", frame_name, retval) + logger.trace("'%s': %s", frame_name, retval) # type:ignore[attr-defined] return retval - def frame_has_faces(self, frame_name): + def frame_has_faces(self, frame_name: str) -> bool: """ Check whether a given frame_name exists within the alignments :attr:`data` and contains at least 1 face. @@ -389,11 +361,12 @@ def frame_has_faces(self, frame_name): ``True`` if the given frame_name exists within the alignments :attr:`data` and has at least 1 face associated with it, otherwise ``False`` """ - retval = bool(self._data.get(frame_name, dict()).get("faces", [])) - logger.trace("'%s': %s", frame_name, retval) + frame_data = self._data.get(frame_name, T.cast(AlignmentDict, {})) + retval = bool(frame_data.get("faces", [])) + logger.trace("'%s': %s", frame_name, retval) # type:ignore[attr-defined] return retval - def frame_has_multiple_faces(self, frame_name): + def frame_has_multiple_faces(self, frame_name: str) -> bool: """ Check whether a given frame_name exists within the alignments :attr:`data` and contains more than 1 face. @@ -412,11 +385,12 @@ def frame_has_multiple_faces(self, frame_name): if not frame_name: retval = False else: - retval = bool(len(self._data.get(frame_name, dict()).get("faces", [])) > 1) - logger.trace("'%s': %s", frame_name, retval) + frame_data = self._data.get(frame_name, T.cast(AlignmentDict, {})) + retval = bool(len(frame_data.get("faces", [])) > 1) + logger.trace("'%s': %s", frame_name, retval) # type:ignore[attr-defined] return retval - def mask_is_valid(self, mask_type): + def mask_is_valid(self, mask_type: str) -> bool: """ Ensure the given ``mask_type`` is valid for the alignments :attr:`data`. Every face in the alignments :attr:`data` must have the given mask type to successfully @@ -433,16 +407,15 @@ def mask_is_valid(self, mask_type): ``True`` if all faces in the current alignments possess the given ``mask_type`` otherwise ``False`` """ - retval = any([(face.get("mask", None) is not None and - face["mask"].get(mask_type, None) is not None) - for val in self._data.values() - for face in val["faces"]]) + retval = all((face.get("mask") is not None and + face["mask"].get(mask_type) is not None) + for val in self._data.values() + for face in val["faces"]) logger.debug(retval) return retval # << DATA >> # - - def get_faces_in_frame(self, frame_name): + def get_faces_in_frame(self, frame_name: str) -> list[AlignmentFileDict]: """ Obtain the faces from :attr:`data` associated with a given frame_name. Parameters @@ -456,10 +429,11 @@ def get_faces_in_frame(self, frame_name): list The list of face dictionaries that appear within the requested frame_name """ - logger.trace("Getting faces for frame_name: '%s'", frame_name) - return self._data.get(frame_name, dict()).get("faces", []) + logger.trace("Getting faces for frame_name: '%s'", frame_name) # type:ignore[attr-defined] + frame_data = self._data.get(frame_name, T.cast(AlignmentDict, {})) + return frame_data.get("faces", T.cast(list[AlignmentFileDict], [])) - def _count_faces_in_frame(self, frame_name): + def count_faces_in_frame(self, frame_name: str) -> int: """ Return number of faces that appear within :attr:`data` for the given frame_name. Parameters @@ -473,13 +447,13 @@ def _count_faces_in_frame(self, frame_name): int The number of faces that appear in the given frame_name """ - retval = len(self._data.get(frame_name, dict()).get("faces", [])) - logger.trace(retval) + frame_data = self._data.get(frame_name, T.cast(AlignmentDict, {})) + retval = len(frame_data.get("faces", [])) + logger.trace(retval) # type:ignore[attr-defined] return retval # << MANIPULATION >> # - - def delete_face_at_index(self, frame_name, face_index): + def delete_face_at_index(self, frame_name: str, face_index: int) -> bool: """ Delete the face for the given frame_name at the given face index from :attr:`data`. Parameters @@ -497,7 +471,7 @@ def delete_face_at_index(self, frame_name, face_index): """ logger.debug("Deleting face %s for frame_name '%s'", face_index, frame_name) face_index = int(face_index) - if face_index + 1 > self._count_faces_in_frame(frame_name): + if face_index + 1 > self.count_faces_in_frame(frame_name): logger.debug("No face to delete: (frame_name: '%s', face_index %s)", frame_name, face_index) return False @@ -505,7 +479,7 @@ def delete_face_at_index(self, frame_name, face_index): logger.debug("Deleted face: (frame_name: '%s', face_index %s)", frame_name, face_index) return True - def add_face(self, frame_name, face): + def add_face(self, frame_name: str, face: AlignmentFileDict) -> int: """ Add a new face for the given frame_name in :attr:`data` and return it's index. Parameters @@ -524,13 +498,13 @@ def add_face(self, frame_name, face): """ logger.debug("Adding face to frame_name: '%s'", frame_name) if frame_name not in self._data: - self._data[frame_name] = dict(faces=[]) + self._data[frame_name] = {"faces": [], "video_meta": {}} self._data[frame_name]["faces"].append(face) - retval = self._count_faces_in_frame(frame_name) - 1 + retval = self.count_faces_in_frame(frame_name) - 1 logger.debug("Returning new face index: %s", retval) return retval - def update_face(self, frame_name, face_index, face): + def update_face(self, frame_name: str, face_index: int, face: AlignmentFileDict) -> None: """ Update the face for the given frame_name at the given face index in :attr:`data`. Parameters @@ -547,7 +521,7 @@ def update_face(self, frame_name, face_index, face): logger.debug("Updating face %s for frame_name '%s'", face_index, frame_name) self._data[frame_name]["faces"][face_index] = face - def filter_faces(self, filter_dict, filter_out=False): + def filter_faces(self, filter_dict: dict[str, list[int]], filter_out: bool = False) -> None: """ Remove faces from :attr:`data` based on a given filter list. Parameters @@ -568,15 +542,28 @@ def filter_faces(self, filter_dict, filter_out=False): else: filter_list = [idx for idx in range(len(frame_data["faces"])) if idx not in face_indices] - logger.trace("frame: '%s', filter_list: %s", source_frame, filter_list) + logger.trace("frame: '%s', filter_list: %s", # type:ignore[attr-defined] + source_frame, filter_list) for face_idx in reversed(sorted(filter_list)): - logger.verbose("Filtering out face: (filename: %s, index: %s)", - source_frame, face_idx) + logger.verbose( # type:ignore[attr-defined] + "Filtering out face: (filename: %s, index: %s)", source_frame, face_idx) del frame_data["faces"][face_idx] + def update_from_dict(self, data: dict[str, AlignmentDict]) -> None: + """ Replace all alignments with the contents of the given dictionary + + Parameters + ---------- + data: dict[str, AlignmentDict] + The alignments, in correctly formatted dictionary form, to be populated into this + :class:`Alignments` + """ + logger.debug("Populating alignments with %s entries", len(data)) + self._data = data + # << GENERATORS >> # - def yield_faces(self): + def yield_faces(self) -> Generator[tuple[str, list[AlignmentFileDict], int, str], None, None]: """ Generator to obtain all faces with meta information from :attr:`data`. The results are yielded by frame. @@ -599,58 +586,65 @@ def yield_faces(self): for frame_fullname, val in self._data.items(): frame_name = os.path.splitext(frame_fullname)[0] face_count = len(val["faces"]) - logger.trace("Yielding: (frame: '%s', faces: %s, frame_fullname: '%s')", - frame_name, face_count, frame_fullname) + logger.trace( # type:ignore[attr-defined] + "Yielding: (frame: '%s', faces: %s, frame_fullname: '%s')", + frame_name, face_count, frame_fullname) yield frame_name, val["faces"], face_count, frame_fullname - # << LEGACY FUNCTIONS >> # + def update_legacy_has_source(self, filename: str) -> None: + """ Update legacy alignments files when we have the source filename available. - def _update_legacy(self): - """ Check whether the alignments are legacy, and if so update them to current alignments - format. """ - updated = False - if self._has_legacy_structure(): - self._update_legacy_structure() - - if self._has_legacy_landmarksxy(): - logger.info("Updating legacy landmarksXY to landmarks_xy") - self._update_legacy_landmarksxy() - updated = True - if self._has_legacy_landmarks_list(): - logger.info("Updating legacy landmarks from list to numpy array") - self._update_legacy_landmarks_list() - updated = True - if self._version < 2.2: - logger.info("Updating legacy mask centering") - self._update_mask_centering() - updated = True - if updated: - self._version = _VERSION - self.save() - - # # - # Serializer is now a compressed pickle custom format. This used to be any number - # of serializers - def _test_for_legacy(self, location): - """ For alignments filenames passed in without an extension, test for legacy - serialization formats and update to current ``.fsa`` format if any are found. + Updates here can only be performed when we have the source filename Parameters ---------- - location: str - The folder location to check for legacy alignments + filename: str: + The filename/folder of the original source images/video for the current alignments """ - logger.debug("Checking for legacy alignments file formats: '%s'", location) - filename = os.path.splitext(location)[0] - for ext in (".json", ".p", ".pickle", ".yaml"): - legacy_filename = "{}{}".format(filename, ext) - if os.path.exists(legacy_filename): - logger.debug("Legacy alignments file exists: '%s'", legacy_filename) - _ = self._update_file_format(*os.path.split(legacy_filename)) - break - logger.debug("Legacy alignments file does not exist: '%s'", legacy_filename) + updates = [updater.is_updated for updater in (VideoExtension(self, filename), )] + if any(updates): + self._io.update_version() + self.save() + + +class _IO(): + """ Class to handle the saving/loading of an alignments file. + + Parameters + ---------- + alignments: :class:'~Alignments` + The parent alignments class that these IO operations belong to + folder: str + The folder that contains the alignments ``.fsa`` file + filename: str + The filename of the ``.fsa`` alignments file. + """ + def __init__(self, alignments: Alignments, folder: str, filename: str) -> None: + logger.debug("Initializing %s: (alignments: %s)", self.__class__.__name__, alignments) + self._alignments = alignments + self._serializer = get_serializer("compressed") + self._file = self._get_location(folder, filename) + self._version: float = _VERSION + + @property + def file(self) -> str: + """ str: The full path to the currently loaded alignments file. """ + return self._file - def _update_file_format(self, folder, filename): + @property + def version(self) -> float: + """ float: The alignments file version number. """ + return self._version + + @property + def have_alignments_file(self) -> bool: + """ bool: ``True`` if an alignments file exists at location :attr:`file` otherwise + ``False``. """ + retval = os.path.exists(self._file) + logger.trace(retval) # type:ignore[attr-defined] + return retval + + def _update_file_format(self, folder: str, filename: str) -> str: """ Convert old style serialized alignments to new ``.fsa`` format. Parameters @@ -667,8 +661,7 @@ def _update_file_format(self, folder, filename): """ logger.info("Reformatting legacy alignments file...") old_location = os.path.join(str(folder), filename) - new_location = "{}.{}".format(os.path.splitext(old_location)[0], - self._serializer.file_extension) + new_location = f"{os.path.splitext(old_location)[0]}.{self._serializer.file_extension}" if os.path.exists(old_location): if os.path.exists(new_location): logger.info("Using existing updated alignments file found at '%s'. If you do not " @@ -681,155 +674,137 @@ def _update_file_format(self, folder, filename): self._serializer.save(new_location, data) return os.path.basename(new_location) - # # - # Alignments were structured: {frame_name: }. We need to be able to store - # information at the frame level, so new structure is: {frame_name: {faces: }} - def _has_legacy_structure(self): - """ Test whether the alignments file is laid out in the old structure of - `{frame_name: [faces]}` + def _test_for_legacy(self, location: str) -> None: + """ For alignments filenames passed in without an extension, test for legacy + serialization formats and update to current ``.fsa`` format if any are found. - Returns - ------- - bool - ``True`` if the file has legacy structure otherwise ``False`` + Parameters + ---------- + location: str + The folder location to check for legacy alignments """ - retval = any(isinstance(val, list) for val in self._data.values()) - logger.debug("legacy structure: %s", retval) - return retval - - def _update_legacy_structure(self): - """ Update legacy alignments files from the format `{frame_name: [faces}` to the - format `{frame_name: {faces: [faces]}`.""" - for key, val in self._data.items(): - self._data[key] = dict(faces=val) - logger.debug("Updated alignments file structure") - - # # - # Landmarks renamed from landmarksXY to landmarks_xy for PEP compliance - def _has_legacy_landmarksxy(self): - """ check for legacy landmarksXY keys. + logger.debug("Checking for legacy alignments file formats: '%s'", location) + filename = os.path.splitext(location)[0] + for ext in (".json", ".p", ".pickle", ".yaml"): + legacy_filename = f"{filename}{ext}" + if os.path.exists(legacy_filename): + logger.debug("Legacy alignments file exists: '%s'", legacy_filename) + _ = self._update_file_format(*os.path.split(legacy_filename)) + break + logger.debug("Legacy alignments file does not exist: '%s'", legacy_filename) - Returns - ------- - bool - ``True`` if the alignments file contains legacy `landmarksXY` keys otherwise ``False`` - """ - logger.debug("checking legacy landmarksXY") - retval = (any(key == "landmarksXY" - for val in self._data.values() - for alignment in val["faces"] - for key in alignment)) - logger.debug("legacy landmarksXY: %s", retval) - return retval + def _get_location(self, folder: str, filename: str) -> str: + """ Obtains the location of an alignments file. - def _update_legacy_landmarksxy(self): - """ Update legacy `landmarksXY` keys to PEP compliant `landmarks_xy` keys. """ - update_count = 0 - for val in self._data.values(): - for alignment in val["faces"]: - alignment["landmarks_xy"] = alignment.pop("landmarksXY") - update_count += 1 - logger.debug("Updated landmarks_xy: %s", update_count) + If a legacy alignments file is provided/discovered, then the alignments file will be + updated to the custom ``.fsa`` format and saved. - # Landmarks stored as list instead of numpy array - def _has_legacy_landmarks_list(self): - """ check for legacy landmarks stored as `list` rather than :class:`numpy.ndarray`. + Parameters + ---------- + folder: str + The folder that the alignments file is located in + filename: str + The filename of the alignments file Returns ------- - bool - ``True`` if not all landmarks are :class:`numpy.ndarray` otherwise ``False`` + str + The full path to the alignments file """ - logger.debug("checking legacy landmarks as list") - retval = not all(isinstance(face["landmarks_xy"], np.ndarray) - for val in self._data.values() - for face in val["faces"]) - return retval - - def _update_legacy_landmarks_list(self): - """ Update landmarks stored as `list` to :class:`numpy.ndarray`. """ - update_count = 0 - for val in self._data.values(): - for alignment in val["faces"]: - test = alignment["landmarks_xy"] - if not isinstance(test, np.ndarray): - alignment["landmarks_xy"] = np.array(test, dtype="float32") - update_count += 1 - logger.debug("Updated landmarks_xy: %s", update_count) - - # Masks not containing the stored_centering parameters. Prior to this implementation all masks - # were stored with face centering - def _update_mask_centering(self): - update_count = 0 - for val in self._data.values(): - for alignment in val["faces"]: - if "mask" not in alignment: - alignment["mask"] = {} - for mask in alignment["mask"].values(): - mask["stored_centering"] = "face" - update_count += 1 - logger.debug("Updated legacy mask centering: %s", update_count) - - -class Thumbnails(): - """ Thumbnail images stored in the alignments file. - - The thumbnails are stored as low resolution (64px), low quality jpg in the alignments file - and are used for the Manual Alignments tool. + logger.debug("Getting location: (folder: '%s', filename: '%s')", folder, filename) + noext_name, extension = os.path.splitext(filename) + if extension in (".json", ".p", ".pickle", ".yaml", ".yml"): + # Reformat legacy alignments file + filename = self._update_file_format(folder, filename) + logger.debug("Updated legacy alignments. New filename: '%s'", filename) + if extension[1:] == self._serializer.file_extension: + logger.debug("Valid Alignments filename provided: '%s'", filename) + else: + filename = f"{noext_name}.{self._serializer.file_extension}" + logger.debug("File extension set from serializer: '%s'", + self._serializer.file_extension) + location = os.path.join(str(folder), filename) + if not os.path.exists(location): + # Test for old format alignments files and reformat if they exist. This will be + # executed if an alignments file has not been explicitly provided therefore it will not + # have been picked up in the extension test + self._test_for_legacy(location) + logger.verbose("Alignments filepath: '%s'", location) # type:ignore[attr-defined] + return location - Parameters - ---------- - alignments: :class:'~lib.align.Alignments` - The parent alignments class that these thumbs belong to - """ - def __init__(self, alignments): - logger.debug("Initializing %s: (alignments: %s)", self.__class__.__name__, alignments) - self._alignments_dict = alignments.data - self._frame_list = list(sorted(self._alignments_dict)) - logger.debug("Initialized %s", self.__class__.__name__) + def update_legacy(self) -> None: + """ Check whether the alignments are legacy, and if so update them to current alignments + format. """ + updates = [updater.is_updated for updater in (FileStructure(self._alignments), + LandmarkRename(self._alignments), + ListToNumpy(self._alignments), + MaskCentering(self._alignments), + IdentityAndVideoMeta(self._alignments))] + if any(updates): + self.update_version() + self.save() - @property - def has_thumbnails(self): - """ bool: ``True`` if all faces in the alignments file contain thumbnail images - otherwise ``False``. """ - retval = all("thumb" in face - for frame in self._alignments_dict.values() - for face in frame["faces"]) - logger.trace(retval) - return retval + def update_version(self) -> None: + """ Update the version of the alignments file to the latest version """ + self._version = _VERSION + logger.info("Updating alignments file to version %s", self._version) - def get_thumbnail_by_index(self, frame_index, face_index): - """ Obtain a jpg thumbnail from the given frame index for the given face index + def load(self) -> dict[str, AlignmentDict]: + """ Load the alignments data from the serialized alignments :attr:`file`. - Parameters - ---------- - frame_index: int - The frame index that contains the thumbnail - face_index: int - The face index within the frame to retrieve the thumbnail for + Populates :attr:`_version` with the alignment file's loaded version as well as returning + the serialized data. Returns ------- - :class:`numpy.ndarray` - The encoded jpg thumbnail + dict: + The loaded alignments data """ - retval = self._alignments_dict[self._frame_list[frame_index]]["faces"][face_index]["thumb"] - logger.trace("frame index: %s, face_index: %s, thumb shape: %s", - frame_index, face_index, retval.shape) - return retval + logger.debug("Loading alignments") + if not self.have_alignments_file: + raise FaceswapError(f"Error: Alignments file not found at {self._file}") - def add_thumbnail(self, frame, face_index, thumb): - """ Add a thumbnail for the given face index for the given frame. + logger.info("Reading alignments from: '%s'", self._file) + data = self._serializer.load(self._file) + meta = data.get("__meta__", {"version": 1.0}) + self._version = meta["version"] + data = data.get("__data__", data) + logger.debug("Loaded alignments") + return data - Parameters - ---------- - frame: str - The name of the frame to add the thumbnail for - face_index: int - The face index within the given frame to add the thumbnail for - thumb: :class:`numpy.ndarray` - The encoded jpg thumbnail at 64px to add to the alignments file + def save(self) -> None: + """ Write the contents of :attr:`data` and :attr:`_meta` to a serialized ``.fsa`` file at + the location :attr:`file`. """ + logger.debug("Saving alignments") + logger.info("Writing alignments to: '%s'", self._file) + data = {"__meta__": {"version": self._version}, + "__data__": self._alignments.data} + self._serializer.save(self._file, data) + logger.debug("Saved alignments") + + def backup(self) -> None: + """ Create a backup copy of the alignments :attr:`file`. + + Creates a copy of the serialized alignments :attr:`file` appending a + timestamp onto the end of the file name and storing in the same folder as + the original :attr:`file`. """ - logger.debug("frame: %s, face_index: %s, thumb shape: %s thumb dtype: %s", - frame, face_index, thumb.shape, thumb.dtype) - self._alignments_dict[frame]["faces"][face_index]["thumb"] = thumb + logger.debug("Backing up alignments") + if not os.path.isfile(self._file): + logger.debug("No alignments to back up") + return + now = datetime.now().strftime("%Y%m%d_%H%M%S") + src = self._file + split = os.path.splitext(src) + dst = f"{split[0]}_{now}{split[1]}" + idx = 1 + while True: + if not os.path.exists(dst): + break + logger.debug("Backup file %s exists. Incrementing", dst) + dst = f"{split[0]}_{now}({idx}){split[1]}" + idx += 1 + + logger.info("Backing up original alignments to '%s'", dst) + os.rename(src, dst) + logger.debug("Backed up alignments") diff --git a/lib/align/constants.py b/lib/align/constants.py new file mode 100644 index 0000000000..6d5b484e59 --- /dev/null +++ b/lib/align/constants.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +""" Constants that are required across faceswap's lib.align package """ +from __future__ import annotations + +import typing as T +from enum import Enum + +import numpy as np + +CenteringType = T.Literal["face", "head", "legacy"] + +EXTRACT_RATIOS: dict[CenteringType, float] = {"legacy": 0.375, "face": 0.5, "head": 0.625} +"""dict[Literal["legacy", "face", head"] float]: The amount of padding applied to each +centering type when generating aligned faces """ + + +class LandmarkType(Enum): + """ Enumeration for the landmark types that Faceswap supports """ + LM_2D_4 = 1 + LM_2D_51 = 2 + LM_2D_68 = 3 + LM_3D_26 = 4 + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> LandmarkType: + """ The landmark type for a given shape + + Parameters + ---------- + shape: tuple[int, ...] + The shape to get the landmark type for + + Returns + ------- + Type[LandmarkType] + The enum for the given shape + + Raises + ------ + ValueError + If the requested shape is not valid + """ + shapes: dict[tuple[int, ...], LandmarkType] = {(4, 2): cls.LM_2D_4, + (51, 2): cls.LM_2D_51, + (68, 2): cls.LM_2D_68, + (26, 3): cls.LM_3D_26} + if shape not in shapes: + raise ValueError(f"The given shape {shape} is not valid. Valid shapes: {list(shapes)}") + return shapes[shape] + + +_MEAN_FACE: dict[LandmarkType, np.ndarray] = { + LandmarkType.LM_2D_4: np.array( + [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]), # Clockwise from TL + LandmarkType.LM_2D_51: np.array([ + [0.010086, 0.106454], [0.085135, 0.038915], [0.191003, 0.018748], [0.300643, 0.034489], + [0.403270, 0.077391], [0.596729, 0.077391], [0.699356, 0.034489], [0.808997, 0.018748], + [0.914864, 0.038915], [0.989913, 0.106454], [0.500000, 0.203352], [0.500000, 0.307009], + [0.500000, 0.409805], [0.500000, 0.515625], [0.376753, 0.587326], [0.435909, 0.609345], + [0.500000, 0.628106], [0.564090, 0.609345], [0.623246, 0.587326], [0.131610, 0.216423], + [0.196995, 0.178758], [0.275698, 0.179852], [0.344479, 0.231733], [0.270791, 0.245099], + [0.192616, 0.244077], [0.655520, 0.231733], [0.724301, 0.179852], [0.803005, 0.178758], + [0.868389, 0.216423], [0.807383, 0.244077], [0.729208, 0.245099], [0.264022, 0.780233], + [0.350858, 0.745405], [0.438731, 0.727388], [0.500000, 0.742578], [0.561268, 0.727388], + [0.649141, 0.745405], [0.735977, 0.780233], [0.652032, 0.864805], [0.566594, 0.902192], + [0.500000, 0.909281], [0.433405, 0.902192], [0.347967, 0.864805], [0.300252, 0.784792], + [0.437969, 0.778746], [0.500000, 0.785343], [0.562030, 0.778746], [0.699747, 0.784792], + [0.563237, 0.824182], [0.500000, 0.831803], [0.436763, 0.824182]]), + LandmarkType.LM_3D_26: np.array([ + [4.056931, -11.432347, 1.636229], # 8 chin LL + [1.833492, -12.542305, 4.061275], # 7 chin L + [0.0, -12.901019, 4.070434], # 6 chin C + [-1.833492, -12.542305, 4.061275], # 5 chin R + [-4.056931, -11.432347, 1.636229], # 4 chin RR + [6.825897, 1.275284, 4.402142], # 33 L eyebrow L + [1.330353, 1.636816, 6.903745], # 29 L eyebrow R + [-1.330353, 1.636816, 6.903745], # 34 R eyebrow L + [-6.825897, 1.275284, 4.402142], # 38 R eyebrow R + [1.930245, -5.060977, 5.914376], # 54 nose LL + [0.746313, -5.136947, 6.263227], # 53 nose L + [0.0, -5.485328, 6.76343], # 52 nose C + [-0.746313, -5.136947, 6.263227], # 51 nose R + [-1.930245, -5.060977, 5.914376], # 50 nose RR + [5.311432, 0.0, 3.987654], # 13 L eye L + [1.78993, -0.091703, 4.413414], # 17 L eye R + [-1.78993, -0.091703, 4.413414], # 25 R eye L + [-5.311432, 0.0, 3.987654], # 21 R eye R + [2.774015, -7.566103, 5.048531], # 43 mouth L + [0.509714, -7.056507, 6.566167], # 42 mouth top L + [0.0, -7.131772, 6.704956], # 41 mouth top C + [-0.509714, -7.056507, 6.566167], # 40 mouth top R + [-2.774015, -7.566103, 5.048531], # 39 mouth R + [-0.589441, -8.443925, 6.109526], # 46 mouth bottom R + [0.0, -8.601736, 6.097667], # 45 mouth bottom C + [0.589441, -8.443925, 6.109526]])} # 44 mouth bottom L +"""dict[:class:`~LandmarkType, np.ndarray]: 'Mean' landmark points for various landmark types. Used +for aligning faces """ + +LANDMARK_PARTS: dict[LandmarkType, dict[str, tuple[int, int, bool]]] = { + LandmarkType.LM_2D_68: {"mouth_outer": (48, 60, True), + "mouth_inner": (60, 68, True), + "right_eyebrow": (17, 22, False), + "left_eyebrow": (22, 27, False), + "right_eye": (36, 42, True), + "left_eye": (42, 48, True), + "nose": (27, 36, False), + "jaw": (0, 17, False), + "chin": (8, 11, False)}, + LandmarkType.LM_2D_4: {"face": (0, 4, True)}} +"""dict[:class:`LandmarkType`, dict[str, tuple[int, int, bool]]: For each landmark type, stores +the (start index, end index, is polygon) information about each part of the face. """ diff --git a/lib/align/detected_face.py b/lib/align/detected_face.py index d82155f161..efd4475ab2 100644 --- a/lib/align/detected_face.py +++ b/lib/align/detected_face.py @@ -1,19 +1,28 @@ #!/usr/bin python3 """ Face and landmarks detection for faceswap.py """ +from __future__ import annotations import logging import os +import typing as T from hashlib import sha1 from zlib import compress, decompress -import cv2 import numpy as np from lib.image import encode_image, read_image +from lib.logger import parse_class_init from lib.utils import FaceswapError -from . import AlignedFace, _EXTRACT_RATIOS, get_centered_size +from .alignments import (Alignments, AlignmentFileDict, PNGHeaderAlignmentsDict, + PNGHeaderDict, PNGHeaderSourceDict) +from .aligned_face import AlignedFace +from .aligned_mask import LandmarksMask, Mask +from .constants import LANDMARK_PARTS -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from .aligned_face import CenteringType + +logger = logging.getLogger(__name__) class DetectedFace(): @@ -28,16 +37,16 @@ class DetectedFace(): ---------- image: numpy.ndarray, optional Original frame that holds this face. Optional (not required if just storing coordinates) - x: int + left: int The left most point (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - w: int + width: int The width (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - y: int + top: int The top most point (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - h: int + height: int The height (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` landmarks_xy: list @@ -45,7 +54,7 @@ class DetectedFace(): of 68 `(x, y)` ``tuples`` with each of the landmark co-ordinates. mask: dict The generated mask(s) for the face as generated in :mod:`plugins.extract.mask`. Must be a - dict of {**name** (`str`): :class:`Mask`}. + dict of {**name** (`str`): :class:`~lib.align.aligned_mask.Mask`}. Attributes ---------- @@ -53,68 +62,84 @@ class DetectedFace(): This is a generic image placeholder that should not be relied on to be holding a particular image. It may hold the source frame that holds the face, a cropped face or a scaled image depending on the method using this object. - x: int + left: int The left most point (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - w: int + width: int The width (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - y: int + top: int The top most point (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` - h: int + height: int The height (in pixels) of the face's bounding box as discovered in :mod:`plugins.extract.detect` landmarks_xy: list The 68 point landmarks as discovered in :mod:`plugins.extract.align`. mask: dict The generated mask(s) for the face as generated in :mod:`plugins.extract.mask`. Is a - dict of {**name** (`str`): :class:`Mask`}. + dict of {**name** (`str`): :class:`~lib.align.aligned_mask.Mask`}. """ - def __init__(self, image=None, x=None, w=None, y=None, h=None, landmarks_xy=None, mask=None, - filename=None): - logger.trace("Initializing %s: (image: %s, x: %s, w: %s, y: %s, h:%s, landmarks_xy: %s, " - "mask: %s, filename: %s)", - self.__class__.__name__, - image.shape if image is not None and image.any() else image, - x, w, y, h, landmarks_xy, - {k: v.shape for k, v in mask} if mask is not None else mask, - filename) + def __init__(self, + image: np.ndarray | None = None, + left: int | None = None, + width: int | None = None, + top: int | None = None, + height: int | None = None, + landmarks_xy: np.ndarray | None = None, + mask: dict[str, Mask] | None = None) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] self.image = image - self.x = x # pylint:disable=invalid-name - self.w = w # pylint:disable=invalid-name - self.y = y # pylint:disable=invalid-name - self.h = h # pylint:disable=invalid-name - self.landmarks_xy = landmarks_xy - self.thumbnail = None - self.mask = dict() if mask is None else mask - - self.aligned = None - logger.trace("Initialized %s", self.__class__.__name__) + self.left = left + self.width = width + self.top = top + self.height = height + self._landmarks_xy = landmarks_xy + self._identity: dict[str, np.ndarray] = {} + self.thumbnail: np.ndarray | None = None + self.mask = {} if mask is None else mask + self._training_masks: tuple[bytes, tuple[int, int, int]] | None = None + + self._aligned: AlignedFace | None = None + logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] @property - def left(self): - """int: Left point (in pixels) of face detection bounding box within the parent image """ - return self.x + def aligned(self) -> AlignedFace: + """ The aligned face connected to this detected face. """ + assert self._aligned is not None + return self._aligned @property - def top(self): - """int: Top point (in pixels) of face detection bounding box within the parent image """ - return self.y + def landmarks_xy(self) -> np.ndarray: + """ The aligned face connected to this detected face. """ + assert self._landmarks_xy is not None + return self._landmarks_xy @property - def right(self): + def right(self) -> int: """int: Right point (in pixels) of face detection bounding box within the parent image """ - return self.x + self.w + assert self.left is not None and self.width is not None + return self.left + self.width @property - def bottom(self): + def bottom(self) -> int: """int: Bottom point (in pixels) of face detection bounding box within the parent image """ - return self.y + self.h + assert self.top is not None and self.height is not None + return self.top + self.height - def add_mask(self, name, mask, affine_matrix, interpolator, - storage_size=128, storage_centering="face"): - """ Add a :class:`Mask` to this detected face + @property + def identity(self) -> dict[str, np.ndarray]: + """ dict: Identity mechanism as key, identity embedding as value. """ + return self._identity + + def add_mask(self, + name: str, + mask: np.ndarray, + affine_matrix: np.ndarray, + interpolator: int, + storage_size: int = 128, + storage_centering: CenteringType = "face") -> None: + """ Add a :class:`~lib.align.aligned_mask.Mask` to this detected face The mask should be the original output from :mod:`plugins.extract.mask` If a mask with this name already exists it will be overwritten by the given @@ -138,68 +163,140 @@ def add_mask(self, name, mask, affine_matrix, interpolator, The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. Default: `"face"` """ - logger.trace("name: '%s', mask shape: %s, affine_matrix: %s, interpolator: %s, " - "storage_size: %s, storage_centering: %s)", name, mask.shape, affine_matrix, - interpolator, storage_size, storage_centering) + logger.trace("name: '%s', mask shape: %s, affine_matrix: %s, " # type:ignore[attr-defined] + "interpolator: %s, storage_size: %s, storage_centering: %s)", name, + mask.shape, affine_matrix, interpolator, storage_size, storage_centering) fsmask = Mask(storage_size=storage_size, storage_centering=storage_centering) fsmask.add(mask, affine_matrix, interpolator) self.mask[name] = fsmask - def get_landmark_mask(self, size, area, - aligned=True, centering="face", dilation=0, blur_kernel=0, as_zip=False): - """ Obtain a single channel mask based on the face's landmark points. + def add_landmarks_xy(self, landmarks: np.ndarray) -> None: + """ Add landmarks to the detected face object. If landmarks alread exist, they will be + overwritten. + + Parameters + ---------- + landmarks: :class:`numpy.ndarray` + The 68 point face landmarks to add for the face + """ + logger.trace("landmarks shape: '%s'", landmarks.shape) # type:ignore[attr-defined] + self._landmarks_xy = landmarks + + def add_identity(self, name: str, embedding: np.ndarray, ) -> None: + """ Add an identity embedding to this detected face. If an identity already exists for the + given :attr:`name` it will be overwritten Parameters ---------- - size: int or tuple - The size of the aligned mask to retrieve. Should be an `int` if an aligned face is - being requested, or a ('height', 'width') shape tuple if a full frame is being - requested - area: ["mouth", "eyes"] + name: str + The name of the mechanism that calculated the identity + embedding: numpy.ndarray + The identity embedding + """ + logger.trace("name: '%s', embedding shape: %s", # type:ignore[attr-defined] + name, embedding.shape) + assert name == "vggface2" + assert embedding.shape[0] == 512 + self._identity[name] = embedding + + def clear_all_identities(self) -> None: + """ Remove all stored identity embeddings """ + self._identity = {} + + def get_landmark_mask(self, + area: T.Literal["eye", "face", "mouth"], + blur_kernel: int, + dilation: float) -> np.ndarray: + """ Add a :class:`L~lib.align.aligned_mask.LandmarksMask` to this detected face + + Landmark based masks are generated from face Aligned Face landmark points. An aligned + face must be loaded. As the data is coming from the already aligned face, no further mask + cropping is required. + + Parameters + ---------- + area: ["face", "mouth", "eye"] The type of mask to obtain. `face` is a full face mask the others are masks for those specific areas - aligned: bool, optional - ``True`` if the returned mask should be for an aligned face. ``False`` if a full frame - mask should be returned. Default ``True`` - centering: ["legacy", "face", "head"], optional - Only used if `aligned`=``True``. The centering for the landmarks based mask. Should be - the same as the centering used for the extracted face that this mask will be applied - to. "legacy" places the nose in the center of the image (the original method for - aligning). "face" aligns for the nose to be in the center of the face (top to bottom) - but the center of the skull for left to right. "head" aligns for the center of the - skull (in 3D space) being the center of the extracted image, with the crop holding the - full head. Default: `"face"` - dilation: int, optional - The amount of dilation to apply to the mask. `0` for none. Default: `0` - blur_kernel: int, optional - The kernel size for applying gaussian blur to apply to the mask. `0` for none. - Default: `0` - as_zip: bool, optional - ``True`` if the mask should be returned zipped otherwise ``False`` + blur_kernel: int + The size of the kernel for blurring the mask edges + dilation: float + The amount of dilation to apply to the mask. as a percentage of the mask size Returns ------- - :class:`numpy.ndarray` or zipped array - The mask as a single channel image of the given :attr:`size` dimension. If - :attr:`as_zip` is ``True`` then the :class:`numpy.ndarray` will be contained within a - zipped container + :class:`numpy.ndarray` + The generated landmarks mask for the selected area + + Raises + ------ + FaceSwapError + If the aligned face does not contain the correct landmarks to generate a landmark mask """ # TODO Face mask generation from landmarks - logger.trace("size: %s, area: %s, aligned: %s, dilation: %s, blur_kernel: %s, as_zip: %s", - size, area, aligned, dilation, blur_kernel, as_zip) - areas = dict(mouth=[slice(48, 60)], eyes=[slice(36, 42), slice(42, 48)]) - if aligned: - face = AlignedFace(self.landmarks_xy, centering=centering, size=size) - landmarks = face.landmarks - size = (size, size) - else: - landmarks = self.landmarks_xy - points = [landmarks[zone] for zone in areas[area]] # pylint:disable=unsubscriptable-object - mask = _LandmarksMask(size, points, dilation=dilation, blur_kernel=blur_kernel) - retval = mask.get(as_zip=as_zip) - return retval + logger.trace("area: %s, dilation: %s", area, dilation) # type:ignore[attr-defined] + + lm_type = self.aligned.landmark_type + if lm_type not in LANDMARK_PARTS: + raise FaceswapError(f"Landmark based masks cannot be created for {lm_type.name}") + + lm_parts = LANDMARK_PARTS[self.aligned.landmark_type] + mapped = {"mouth": ["mouth_outer"], "eye": ["right_eye", "left_eye"]} + if not all(part in lm_parts for parts in mapped.values() for part in parts): + raise FaceswapError(f"Landmark based masks cannot be created for {lm_type.name}") + + areas = {key: [slice(*lm_parts[v][:2]) for v in val]for key, val in mapped.items()} + points = [self.aligned.landmarks[zone] for zone in areas[area]] + + lmmask = LandmarksMask(points, + storage_size=self.aligned.size, + storage_centering=self.aligned.centering, + dilation=dilation) + lmmask.set_blur_and_threshold(blur_kernel=blur_kernel) + lmmask.generate_mask( + self.aligned.adjusted_matrix, + self.aligned.interpolators[1]) + return lmmask.mask + + def store_training_masks(self, + masks: list[np.ndarray | None], + delete_masks: bool = False) -> None: + """ Concatenate and compress the given training masks and store for retrieval. + + Parameters + ---------- + masks: list + A list of training mask. Must be all be uint-8 3D arrays of the same size in + 0-255 range + delete_masks: bool, optional + ``True`` to delete any of the :class:`~lib.align.aligned_mask.Mask` objects owned by + this detected face. Use to free up unrequired memory usage. Default: ``False`` + """ + if delete_masks: + del self.mask + self.mask = {} + + valid = [msk for msk in masks if msk is not None] + if not valid: + return + combined = np.concatenate(valid, axis=-1) + self._training_masks = (compress(combined), combined.shape) - def to_alignment(self): + def get_training_masks(self) -> np.ndarray | None: + """ Obtain the decompressed combined training masks. + + Returns + ------- + :class:`numpy.ndarray` + A 3D array containing the decompressed training masks as uint8 in 0-255 range if + training masks are present otherwise ``None`` + """ + if not self._training_masks: + return None + return np.frombuffer(decompress(self._training_masks[0]), + dtype="uint8").reshape(self._training_masks[1]) + + def to_alignment(self) -> AlignmentFileDict: """ Return the detected face formatted for an alignments file returns @@ -209,18 +306,22 @@ def to_alignment(self): ``landmarks_xy``, ``mask``. The additional key ``thumb`` will be provided if the detected face object contains a thumbnail. """ - alignment = dict(x=self.x, - w=self.w, - y=self.y, - h=self.h, - landmarks_xy=self.landmarks_xy, - mask={name: mask.to_dict() for name, mask in self.mask.items()}) - if self.thumbnail is not None: - alignment["thumb"] = self.thumbnail - logger.trace("Returning: %s", alignment) + if (self.left is None or self.width is None or self.top is None or self.height is None): + raise AssertionError("Some detected face variables have not been initialized") + alignment = AlignmentFileDict(x=self.left, + w=self.width, + y=self.top, + h=self.height, + landmarks_xy=self.landmarks_xy, + mask={name: mask.to_dict() + for name, mask in self.mask.items()}, + identity={k: v.tolist() for k, v in self._identity.items()}, + thumb=self.thumbnail) + logger.trace("Returning: %s", alignment) # type:ignore[attr-defined] return alignment - def from_alignment(self, alignment, image=None, with_thumb=False): + def from_alignment(self, alignment: AlignmentFileDict, + image: np.ndarray | None = None, with_thumb: bool = False) -> None: """ Set the attributes of this class from an alignments file and optionally load the face into the ``image`` attribute. @@ -241,50 +342,56 @@ def from_alignment(self, alignment, image=None, with_thumb=False): Default: ``False`` """ - logger.trace("Creating from alignment: (alignment: %s, has_image: %s)", - alignment, bool(image is not None)) - self.x = alignment["x"] - self.w = alignment["w"] - self.y = alignment["y"] - self.h = alignment["h"] + logger.trace("Creating from alignment: (alignment: %s," # type:ignore[attr-defined] + " has_image: %s)", alignment, bool(image is not None)) + self.left = alignment["x"] + self.width = alignment["w"] + self.top = alignment["y"] + self.height = alignment["h"] landmarks = alignment["landmarks_xy"] if not isinstance(landmarks, np.ndarray): landmarks = np.array(landmarks, dtype="float32") - self.landmarks_xy = landmarks.copy() + self._identity = {T.cast(T.Literal["vggface2"], k): np.array(v, dtype="float32") + for k, v in alignment.get("identity", {}).items()} + self._landmarks_xy = landmarks.copy() if with_thumb: # Thumbnails currently only used for manual tool. Default to None - self.thumbnail = alignment.get("thumb", None) + self.thumbnail = alignment.get("thumb") # Manual tool and legacy alignments will not have a mask - self.aligned = None + self._aligned = None if alignment.get("mask", None) is not None: - self.mask = dict() + self.mask = {} for name, mask_dict in alignment["mask"].items(): self.mask[name] = Mask() self.mask[name].from_dict(mask_dict) if image is not None and image.any(): self._image_to_face(image) - logger.trace("Created from alignment: (x: %s, w: %s, y: %s. h: %s, " - "landmarks: %s, mask: %s)", - self.x, self.w, self.y, self.h, self.landmarks_xy, self.mask) + logger.trace("Created from alignment: (left: %s, width: %s, " # type:ignore[attr-defined] + "top: %s, height: %s, landmarks: %s, mask: %s)", + self.left, self.width, self.top, self.height, self.landmarks_xy, self.mask) - def to_png_meta(self): + def to_png_meta(self) -> PNGHeaderAlignmentsDict: """ Return the detected face formatted for insertion into a png itxt header. returns: dict The alignments dict will be returned with the keys ``x``, ``w``, ``y``, ``h``, ``landmarks_xy`` and ``mask`` """ - alignment = dict(x=self.x, - w=self.w, - y=self.y, - h=self.h, - landmarks_xy=self.landmarks_xy.tolist(), - mask={name: mask.to_png_meta() for name, mask in self.mask.items()}) + if (self.left is None or self.width is None or self.top is None or self.height is None): + raise AssertionError("Some detected face variables have not been initialized") + alignment = PNGHeaderAlignmentsDict( + x=self.left, + w=self.width, + y=self.top, + h=self.height, + landmarks_xy=self.landmarks_xy.tolist(), + mask={name: mask.to_png_meta() for name, mask in self.mask.items()}, + identity={k: v.tolist() for k, v in self._identity.items()}) return alignment - def from_png_meta(self, alignment): + def from_png_meta(self, alignment: PNGHeaderAlignmentsDict) -> None: """ Set the attributes of this class from alignments stored in a png exif header. Parameters @@ -293,26 +400,40 @@ def from_png_meta(self, alignment): A dictionary entry for a face from alignments stored in a png exif header containing the keys ``x``, ``w``, ``y``, ``h``, ``landmarks_xy`` and ``mask`` """ - self.x = alignment["x"] - self.w = alignment["w"] - self.y = alignment["y"] - self.h = alignment["h"] - self.landmarks_xy = np.array(alignment["landmarks_xy"], dtype="float32") - self.mask = dict() + self.left = alignment["x"] + self.width = alignment["w"] + self.top = alignment["y"] + self.height = alignment["h"] + self._landmarks_xy = np.array(alignment["landmarks_xy"], dtype="float32") + self.mask = {} for name, mask_dict in alignment["mask"].items(): self.mask[name] = Mask() self.mask[name].from_dict(mask_dict) - logger.trace("Created from png exif header: (x: %s, w: %s, y: %s. h: %s, landmarks: %s, " - "mask: %s)", self.x, self.w, self.y, self.h, self.landmarks_xy, self.mask) - - def _image_to_face(self, image): + self._identity = {} + for key, val in alignment.get("identity", {}).items(): + assert key in ["vggface2"] + self._identity[T.cast(T.Literal["vggface2"], key)] = np.array(val, dtype="float32") + logger.trace("Created from png exif header: (left: %s, " # type:ignore[attr-defined] + "width: %s, top: %s height: %s, landmarks: %s, mask: %s, identity: %s)", + self.left, self.width, self.top, self.height, self.landmarks_xy, self.mask, + {k: v.shape for k, v in self._identity.items()}) + + def _image_to_face(self, image: np.ndarray) -> None: """ set self.image to be the cropped face from detected bounding box """ - logger.trace("Cropping face from image") + logger.trace("Cropping face from image") # type:ignore[attr-defined] self.image = image[self.top: self.bottom, self.left: self.right] # <<< Aligned Face methods and properties >>> # - def load_aligned(self, image, size=256, dtype=None, centering="head", force=False): + def load_aligned(self, + image: np.ndarray | None, + size: int = 256, + dtype: str | None = None, + centering: CenteringType = "head", + coverage_ratio: float = 1.0, + force: bool = False, + is_aligned: bool = False, + is_legacy: bool = False) -> None: """ Align a face from a given image. Aligning a face is a relatively expensive task and is not required for all uses of @@ -338,553 +459,44 @@ def load_aligned(self, image, size=256, dtype=None, centering="head", force=Fals right. "head" aligns for the center of the skull (in 3D space) being the center of the extracted image, with the crop holding the full head. Default: `"head"` + coverage_ratio: float, optional + The amount of the aligned image to return. A ratio of 1.0 will return the full contents + of the aligned image. A ratio of 0.5 will return an image of the given size, but will + crop to the central 50%% of the image. Default: `1.0` force: bool, optional Force an update of the aligned face, even if it is already loaded. Default: ``False`` - + is_aligned: bool, optional + Indicates that the :attr:`image` is an aligned face rather than a frame. + Default: ``False`` + is_legacy: bool, optional + Only used if `is_aligned` is ``True``. ``True`` indicates that the aligned image being + loaded is a legacy extracted face rather than a current head extracted face Notes ----- This method must be executed to get access to the following an :class:`AlignedFace` object """ - if self.aligned and not force: + if self._aligned and not force: # Don't reload an already aligned face - logger.trace("Skipping alignment calculation for already aligned face") + logger.trace("Skipping alignment calculation for already " # type:ignore[attr-defined] + "aligned face") else: - logger.trace("Loading aligned face: (size: %s, dtype: %s)", size, dtype) - self.aligned = AlignedFace(self.landmarks_xy, - image=image, - centering=centering, - size=size, - coverage_ratio=1.0, - dtype=dtype, - is_aligned=False) - - -class _LandmarksMask(): # pylint:disable=too-few-public-methods - """ Create a single channel mask from aligned landmark points. - - size: tuple - The (height, width) shape tuple that the mask should be returned as - points: list - A list of landmark points that correspond to the given shape tuple to create - the mask. Each item in the list should be a :class:`numpy.ndarray` that a filled - convex polygon will be created from - dilation: int, optional - The amount of dilation to apply to the mask. `0` for none. Default: `0` - blur_kernel: int, optional - The kernel size for applying gaussian blur to apply to the mask. `0` for none. Default: `0` - """ - def __init__(self, size, points, dilation=0, blur_kernel=0): - logger.trace("Initializing: %s: (size: %s, points: %s, dilation: %s, blur_kernel: %s)", - self.__class__.__name__, size, points, dilation, blur_kernel) - self._size = size - self._points = points - self._dilation = dilation - self._blur_kernel = blur_kernel - self._mask = None - logger.trace("Initialized: %s", self.__class__.__name__) - - def get(self, as_zip=False): - """ Obtain the mask. - - Parameters - ---------- - as_zip: bool, optional - ``True`` if the mask should be returned zipped otherwise ``False`` - - Returns - ------- - :class:`numpy.ndarray` or zipped array - The mask as a single channel image of the given :attr:`size` dimension. If - :attr:`as_zip` is ``True`` then the :class:`numpy.ndarray` will be contained within a - zipped container - """ - if not np.any(self._mask): - self._generate_mask() - retval = compress(self._mask) if as_zip else self._mask - logger.trace("as_zip: %s, retval type: %s", as_zip, type(retval)) - return retval - - def _generate_mask(self): - """ Generate the mask. - - Creates the mask applying any requested dilation and blurring and assigns to - :attr:`_mask` - - Returns - ------- - :class:`numpy.ndarray` - The mask as a single channel image of the given :attr:`size` dimension. - """ - mask = np.zeros((self._size) + (1, ), dtype="float32") - for landmarks in self._points: - lms = np.rint(landmarks).astype("int") - cv2.fillConvexPoly(mask, cv2.convexHull(lms), 1.0, lineType=cv2.LINE_AA) - if self._dilation != 0: - mask = cv2.dilate(mask, - cv2.getStructuringElement(cv2.MORPH_ELLIPSE, - (self._dilation, self._dilation)), - iterations=1) - if self._blur_kernel != 0: - mask = BlurMask("gaussian", mask, self._blur_kernel).blurred - logger.trace("mask: (shape: %s, dtype: %s)", mask.shape, mask.dtype) - self._mask = (mask * 255.0).astype("uint8") - - -class Mask(): - """ Face Mask information and convenience methods - - Holds a Faceswap mask as generated from :mod:`plugins.extract.mask` and the information - required to transform it to its original frame. - - Holds convenience methods to handle the warping, storing and retrieval of the mask. - - Parameters - ---------- - storage_size: int, optional - The size (in pixels) that the mask should be stored at. Default: 128. - storage_centering, str (optional): - The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. - Default: `"face"` - - Attributes - ---------- - stored_size: int - The size, in pixels, of the stored mask across its height and width. - stored_centering: str - The centering that the mask is stored at. One of `"legacy"`, `"face"`, `"head"` - """ - def __init__(self, storage_size=128, storage_centering="face"): - logger.trace("Initializing: %s (storage_size: %s, storage_centering: %s)", - self.__class__.__name__, storage_size, storage_centering) - self.stored_size = storage_size - self.stored_centering = storage_centering - - self._mask = None - self._affine_matrix = None - self._interpolator = None - - self._blur = dict() - self._blur_kernel = 0 - self._threshold = 0.0 - self._sub_crop = dict(size=None, slice_in=[], slice_out=[]) - self.set_blur_and_threshold() - logger.trace("Initialized: %s", self.__class__.__name__) - - @property - def mask(self): - """ numpy.ndarray: The mask at the size of :attr:`stored_size` with any requested blurring, - threshold amount and centering applied.""" - mask = self.stored_mask - if self._threshold != 0.0 or self._blur["kernel"] != 0: - mask = mask.copy() - if self._threshold != 0.0: - mask[mask < self._threshold] = 0.0 - mask[mask > 255.0 - self._threshold] = 255.0 - if self._blur["kernel"] != 0: - mask = BlurMask(self._blur["type"], - mask, - self._blur["kernel"], - passes=self._blur["passes"]).blurred - if self._sub_crop["size"]: # Crop the mask to the given centering - out = np.zeros((self._sub_crop["size"], self._sub_crop["size"], 1), dtype=mask.dtype) - slice_in, slice_out = self._sub_crop["slice_in"], self._sub_crop["slice_out"] - out[slice_out[0], slice_out[1], :] = mask[slice_in[0], slice_in[1], :] - mask = out - logger.trace("mask shape: %s", mask.shape) - return mask - - @property - def stored_mask(self): - """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` as it is stored - (i.e. with no blurring/centering applied). """ - dims = (self.stored_size, self.stored_size, 1) - mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims) - logger.trace("stored mask shape: %s", mask.shape) - return mask - - @property - def original_roi(self): - """ :class: `numpy.ndarray`: The original region of interest of the mask in the - source frame. """ - points = np.array([[0, 0], - [0, self.stored_size - 1], - [self.stored_size - 1, self.stored_size - 1], - [self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2)) - matrix = cv2.invertAffineTransform(self._affine_matrix) - roi = cv2.transform(points, matrix).reshape((4, 2)) - logger.trace("Returning: %s", roi) - return roi - - @property - def affine_matrix(self): - """ :class: `numpy.ndarray`: The affine matrix to transpose the mask to a full frame. """ - return self._affine_matrix - - @property - def interpolator(self): - """ int: The cv2 interpolator required to transpose the mask to a full frame. """ - return self._interpolator - - def get_full_frame_mask(self, width, height): - """ Return the stored mask in a full size frame of the given dimensions - - Parameters - ---------- - width: int - The width of the original frame that the mask was extracted from - height: int - The height of the original frame that the mask was extracted from - - Returns - ------- - numpy.ndarray: The mask affined to the original full frame of the given dimensions - """ - frame = np.zeros((width, height, 1), dtype="uint8") - mask = cv2.warpAffine(self.mask, - self._affine_matrix, - (width, height), - frame, - flags=cv2.WARP_INVERSE_MAP | self._interpolator, - borderMode=cv2.BORDER_CONSTANT) - logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, mask max: %s", - mask.shape, mask.dtype, mask.min(), mask.max()) - return mask - - def add(self, mask, affine_matrix, interpolator): - """ Add a Faceswap mask to this :class:`Mask`. - - The mask should be the original output from :mod:`plugins.extract.mask` - - Parameters - ---------- - mask: numpy.ndarray - The mask that is to be added as output from :mod:`plugins.extract.mask` - It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` - affine_matrix: numpy.ndarray - The transformation matrix required to transform the mask to the original frame. - interpolator, int: - The CV2 interpolator required to transform this mask to it's original frame - """ - logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, mask max: %s, " - "affine_matrix: %s, interpolator: %s)", mask.shape, mask.dtype, mask.min(), - affine_matrix, mask.max(), interpolator) - self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix) - self._interpolator = interpolator - self.replace_mask(mask) - - def replace_mask(self, mask): - """ Replace the existing :attr:`_mask` with the given mask. - - Parameters - ---------- - mask: numpy.ndarray - The mask that is to be added as output from :mod:`plugins.extract.mask`. - It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` - """ - mask = (cv2.resize(mask, - (self.stored_size, self.stored_size), - interpolation=cv2.INTER_AREA) * 255.0).astype("uint8") - self._mask = compress(mask) - - def set_blur_and_threshold(self, - blur_kernel=0, blur_type="gaussian", blur_passes=1, threshold=0): - """ Set the internal blur kernel and threshold amount for returned masks - - Parameters - ---------- - blur_kernel: int, optional - The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no - blurring. Should be odd, if an even number is passed in (outside of 0) then it is - rounded up to the next odd number. Default: 0 - blur_type: ["gaussian", "normalized"], optional - The blur type to use. ``gaussian`` or ``normalized`` box filter. Default: ``gaussian`` - blur_passes: int, optional - The number of passed to perform when blurring. Default: 1 - threshold: int, optional - The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value. - Default: 0 - """ - logger.trace("blur_kernel: %s, threshold: %s", blur_kernel, threshold) - if blur_type is not None: - blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1 - self._blur["kernel"] = blur_kernel - self._blur["type"] = blur_type - self._blur["passes"] = blur_passes - self._threshold = (threshold / 100.0) * 255.0 - - def set_sub_crop(self, offset, centering): - """ Set the internal crop area of the mask to be returned. - - This impacts the returned mask from :attr:`mask` if the requested mask is required for - different face centering than what has been stored. - - Parameters - ---------- - offset: :class:`numpy.ndarray` - The (x, y) offset from the center point to return the mask for - centering: str - The centering to set the sub crop area for. One of `"legacy"`, `"face"`. `"head"` - """ - if centering == self.stored_centering: - return - - src_size = self.stored_size - (self.stored_size * _EXTRACT_RATIOS[self.stored_centering]) - offset *= ((self.stored_size - (src_size / 2)) / 2) - center = np.rint(offset + self.stored_size / 2).astype("int32") - - crop_size = get_centered_size(self.stored_centering, centering, self.stored_size) - roi = np.array([center - crop_size // 2, center + crop_size // 2]).ravel() - - self._sub_crop["size"] = crop_size - self._sub_crop["slice_in"] = [slice(max(roi[1], 0), max(roi[3], 0)), - slice(max(roi[0], 0), max(roi[2], 0))] - self._sub_crop["slice_out"] = [ - slice(max(roi[1] * -1, 0), - crop_size - min(crop_size, max(0, roi[3] - self.stored_size))), - slice(max(roi[0] * -1, 0), - crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))] - - logger.trace("src_size: %s, roi: %s, sub_crop: %s", src_size, roi, self._sub_crop) - - def _adjust_affine_matrix(self, mask_size, affine_matrix): - """ Adjust the affine matrix for the mask's storage size - - Parameters - ---------- - mask_size: int - The original size of the mask. - affine_matrix: numpy.ndarray - The affine matrix to transform the mask at original size to the parent frame. - - Returns - ------- - affine_matrix: numpy,ndarray - The affine matrix adjusted for the mask at its stored dimensions. - """ - zoom = self.stored_size / mask_size - zoom_mat = np.array([[zoom, 0, 0.], [0, zoom, 0.]]) - adjust_mat = np.dot(zoom_mat, np.concatenate((affine_matrix, np.array([[0., 0., 1.]])))) - logger.trace("storage_size: %s, mask_size: %s, zoom: %s, original matrix: %s, " - "adjusted_matrix: %s", self.stored_size, mask_size, zoom, affine_matrix.shape, - adjust_mat.shape) - return adjust_mat - - def to_dict(self): - """ Convert the mask to a dictionary for saving to an alignments file - - Returns - ------- - dict: - The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - retval = dict() - for key in ("mask", "affine_matrix", "interpolator", "stored_size", "stored_centering"): - retval[key] = getattr(self, self._attr_name(key)) - logger.trace({k: v if k != "mask" else type(v) for k, v in retval.items()}) - return retval - - def to_png_meta(self): - """ Convert the mask to a dictionary supported by png itxt headers. - - Returns - ------- - dict: - The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - retval = dict() - for key in ("mask", "affine_matrix", "interpolator", "stored_size", "stored_centering"): - val = getattr(self, self._attr_name(key)) - if isinstance(val, np.ndarray): - retval[key] = val.tolist() - else: - retval[key] = val - logger.trace({k: v if k != "mask" else type(v) for k, v in retval.items()}) - return retval - - def from_dict(self, mask_dict): - """ Populates the :class:`Mask` from a dictionary loaded from an alignments file. - - Parameters - ---------- - mask_dict: dict - A dictionary stored in an alignments file containing the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - for key in ("mask", "affine_matrix", "interpolator", "stored_size", "stored_centering"): - val = mask_dict.get(key) - val = "face" if key == "stored_centering" and val is None else val - if key == "affine_matrix" and not isinstance(val, np.ndarray): - val = np.array(val, dtype="float64") - setattr(self, self._attr_name(key), val) - logger.trace("%s - %s", key, val if key != "mask" else type(val)) - - @staticmethod - def _attr_name(dict_key): - """ The :class:`Mask` attribute name for the given dictionary key - - Parameters - ---------- - dict_key: str - The key name from an alignments dictionary - - Returns - ------- - attribute_name: str - The attribute name for the given key for :class:`Mask` - """ - retval = "_{}".format(dict_key) if not dict_key.startswith("stored") else dict_key - logger.trace("dict_key: %s, attribute_name: %s", dict_key, retval) - return retval - - -class BlurMask(): # pylint:disable=too-few-public-methods - """ Factory class to return the correct blur object for requested blur type. - - Works for square images only. Currently supports Gaussian and Normalized Box Filters. - - Parameters - ---------- - blur_type: ["gaussian", "normalized"] - The type of blur to use - mask: :class:`numpy.ndarray` - The mask to apply the blur to - kernel: int or float - Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size - is_ratio: bool, optional - Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the - actual kernel size will be calculated from the given ratio and the mask size. If - ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. - Default: ``False`` - passes: int, optional - The number of passes to perform when blurring. Default: ``1`` - - Example - ------- - >>> print(mask.shape) - (128, 128, 1) - >>> new_mask = BlurMask("gaussian", mask, 3, is_ratio=False, passes=1).blurred - >>> print(new_mask.shape) - (128, 128, 1) - """ - def __init__(self, blur_type, mask, kernel, is_ratio=False, passes=1): - logger.trace("Initializing %s: (blur_type: '%s', mask_shape: %s, kernel: %s, " - "is_ratio: %s, passes: %s)", self.__class__.__name__, blur_type, mask.shape, - kernel, is_ratio, passes) - self._blur_type = blur_type.lower() - self._mask = mask - self._passes = passes - kernel_size = self._get_kernel_size(kernel, is_ratio) - self._kernel_size = self._get_kernel_tuple(kernel_size) - logger.trace("Initialized %s", self.__class__.__name__) - - @property - def blurred(self): - """ :class:`numpy.ndarray`: The final mask with blurring applied. """ - func = self._func_mapping[self._blur_type] - kwargs = self._get_kwargs() - blurred = self._mask - for i in range(self._passes): - ksize = int(kwargs["ksize"][0]) - logger.trace("Pass: %s, kernel_size: %s", i + 1, (ksize, ksize)) - blurred = func(blurred, **kwargs) - ksize = int(round(ksize * self._multipass_factor)) - kwargs["ksize"] = self._get_kernel_tuple(ksize) - blurred = blurred[..., None] - logger.trace("Returning blurred mask. Shape: %s", blurred.shape) - return blurred - - @property - def _multipass_factor(self): - """ For multiple passes the kernel must be scaled down. This value is - different for box filter and gaussian """ - factor = dict(gaussian=0.8, normalized=0.5) - return factor[self._blur_type] - - @property - def _sigma(self): - """ int: The Sigma for Gaussian Blur. Returns 0 to force calculation from kernel size. """ - return 0 - - @property - def _func_mapping(self): - """ dict: :attr:`_blur_type` mapped to cv2 Function name. """ - return dict(gaussian=cv2.GaussianBlur, # pylint: disable = no-member - normalized=cv2.blur) # pylint: disable = no-member - - @property - def _kwarg_requirements(self): - """ dict: :attr:`_blur_type` mapped to cv2 Function required keyword arguments. """ - return dict(gaussian=["ksize", "sigmaX"], - normalized=["ksize"]) - - @property - def _kwarg_mapping(self): - """ dict: cv2 function keyword arguments mapped to their parameters. """ - return dict(ksize=self._kernel_size, - sigmaX=self._sigma) - - def _get_kernel_size(self, kernel, is_ratio): - """ Set the kernel size to absolute value. - - If :attr:`is_ratio` is ``True`` then the kernel size is calculated from the given ratio and - the :attr:`_mask` size, otherwise the given kernel size is just returned. - - Parameters - ---------- - kernel: int or float - Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size - is_ratio: bool, optional - Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the - actual kernel size will be calculated from the given ratio and the mask size. If - ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. - - Returns - ------- - int - The size (in pixels) of the blur kernel - """ - if not is_ratio: - return kernel - - mask_diameter = np.sqrt(np.sum(self._mask)) - radius = round(max(1., mask_diameter * kernel / 100.)) - kernel_size = int(radius * 2 + 1) - logger.trace("kernel_size: %s", kernel_size) - return kernel_size - - @staticmethod - def _get_kernel_tuple(kernel_size): - """ Make sure kernel_size is odd and return it as a tuple. - - Parameters - ---------- - kernel_size: int - The size in pixels of the blur kernel - - Returns - ------- - tuple - The kernel size as a tuple of ('int', 'int') - """ - kernel_size += 1 if kernel_size % 2 == 0 else 0 - retval = (kernel_size, kernel_size) - logger.trace(retval) - return retval - - def _get_kwargs(self): - """ dict: the valid keyword arguments for the requested :attr:`_blur_type` """ - retval = {kword: self._kwarg_mapping[kword] - for kword in self._kwarg_requirements[self._blur_type]} - logger.trace("BlurMask kwargs: %s", retval) - return retval + logger.trace("Loading aligned face: (size: %s, " # type:ignore[attr-defined] + "dtype: %s)", size, dtype) + self._aligned = AlignedFace(self.landmarks_xy, + image=image, + centering=centering, + size=size, + coverage_ratio=coverage_ratio, + dtype=dtype, + is_aligned=is_aligned, + is_legacy=is_aligned and is_legacy) -_HASHES_SEEN = dict() +_HASHES_SEEN: dict[str, dict[str, int]] = {} -def update_legacy_png_header(filename, alignments): +def update_legacy_png_header(filename: str, alignments: Alignments + ) -> PNGHeaderDict | None: """ Update a legacy extracted face from pre v2.1 alignments by placing the alignment data for the face in the png exif header for the given filename with the given alignment data. @@ -911,7 +523,7 @@ def update_legacy_png_header(filename, alignments): # effective enough folder = os.path.dirname(filename) if folder not in _HASHES_SEEN: - _HASHES_SEEN[folder] = dict() + _HASHES_SEEN[folder] = {} hashes_seen = _HASHES_SEEN[folder] in_image = read_image(filename, raise_error=True) @@ -927,13 +539,15 @@ def update_legacy_png_header(filename, alignments): detected_face.from_alignment(alignment) # For dupe hash handling, make sure we get a different filename for repeat hashes src_fname, face_idx = list(alignments.hashes_to_frame[in_hash].items())[hashes_seen[in_hash]] - orig_filename = "{}_{}.png".format(os.path.splitext(src_fname)[0], face_idx) - meta = dict(alignments=detected_face.to_png_meta(), - source=dict(alignments_version=alignments.version, + orig_filename = f"{os.path.splitext(src_fname)[0]}_{face_idx}.png" + meta = PNGHeaderDict(alignments=detected_face.to_png_meta(), + source=PNGHeaderSourceDict( + alignments_version=alignments.version, original_filename=orig_filename, face_index=face_idx, source_filename=src_fname, - source_is_video=False)) # Can't check so set false + source_is_video=False, # Can't check so set false + source_frame_dims=None)) out_filename = f"{os.path.splitext(filename)[0]}.png" # Make sure saved file is png out_image = encode_image(in_image, ".png", metadata=meta) diff --git a/lib/align/pose.py b/lib/align/pose.py new file mode 100644 index 0000000000..fe186a3f5c --- /dev/null +++ b/lib/align/pose.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +""" Holds estimated pose information for a faceswap aligned face """ +from __future__ import annotations + +import logging +import typing as T + +import cv2 +import numpy as np + +from lib.logger import parse_class_init + +from .constants import _MEAN_FACE, LandmarkType + +logger = logging.getLogger(__name__) + +if T.TYPE_CHECKING: + from .constants import CenteringType + + +class PoseEstimate(): + """ Estimates pose from a generic 3D head model for the given 2D face landmarks. + + Parameters + ---------- + landmarks: :class:`numpy.ndarry` + The original 68 point landmarks aligned to 0.0 - 1.0 range + landmarks_type: :class:`~LandmarksType` + The type of landmarks that are generating this face + + References + ---------- + Head Pose Estimation using OpenCV and Dlib - https://www.learnopencv.com/tag/solvepnp/ + 3D Model points - http://aifi.isr.uc.pt/Downloads/OpenGL/glAnthropometric3DModel.cpp + """ + _logged_once = False + + def __init__(self, landmarks: np.ndarray, landmarks_type: LandmarkType) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self._distortion_coefficients = np.zeros((4, 1)) # Assuming no lens distortion + self._xyz_2d: np.ndarray | None = None + + if landmarks_type != LandmarkType.LM_2D_68: + self._log_once("Pose estimation is not available for non-68 point landmarks. Pose and " + "offset data will all be returned as the incorrect value of '0'") + self._landmarks_type = landmarks_type + self._camera_matrix = self._get_camera_matrix() + self._rotation, self._translation = self._solve_pnp(landmarks) + self._offset = self._get_offset() + self._pitch_yaw_roll: tuple[float, float, float] = (0, 0, 0) + logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] + + @property + def xyz_2d(self) -> np.ndarray: + """ :class:`numpy.ndarray` projected (x, y) coordinates for each x, y, z point at a + constant distance from adjusted center of the skull (0.5, 0.5) in the 2D space. """ + if self._xyz_2d is None: + xyz = cv2.projectPoints(np.array([[6., 0., -2.3], + [0., 6., -2.3], + [0., 0., 3.7]]).astype("float32"), + self._rotation, + self._translation, + self._camera_matrix, + self._distortion_coefficients)[0].squeeze() + self._xyz_2d = xyz - self._offset["head"] + return self._xyz_2d + + @property + def offset(self) -> dict[CenteringType, np.ndarray]: + """ dict: The amount to offset a standard 0.0 - 1.0 umeyama transformation matrix for a + from the center of the face (between the eyes) or center of the head (middle of skull) + rather than the nose area. """ + return self._offset + + @property + def pitch(self) -> float: + """ float: The pitch of the aligned face in eular angles """ + if not any(self._pitch_yaw_roll): + self._get_pitch_yaw_roll() + return self._pitch_yaw_roll[0] + + @property + def yaw(self) -> float: + """ float: The yaw of the aligned face in eular angles """ + if not any(self._pitch_yaw_roll): + self._get_pitch_yaw_roll() + return self._pitch_yaw_roll[1] + + @property + def roll(self) -> float: + """ float: The roll of the aligned face in eular angles """ + if not any(self._pitch_yaw_roll): + self._get_pitch_yaw_roll() + return self._pitch_yaw_roll[2] + + @classmethod + def _log_once(cls, message: str) -> None: + """ Log a warning about unsupported landmarks if a message has not already been logged """ + if cls._logged_once: + return + logger.warning(message) + cls._logged_once = True + + def _get_pitch_yaw_roll(self) -> None: + """ Obtain the yaw, roll and pitch from the :attr:`_rotation` in eular angles. """ + proj_matrix = np.zeros((3, 4), dtype="float32") + proj_matrix[:3, :3] = cv2.Rodrigues(self._rotation)[0] + euler = cv2.decomposeProjectionMatrix(proj_matrix)[-1] + self._pitch_yaw_roll = T.cast(tuple[float, float, float], tuple(euler.squeeze())) + logger.trace("yaw_pitch: %s", self._pitch_yaw_roll) # type:ignore[attr-defined] + + @classmethod + def _get_camera_matrix(cls) -> np.ndarray: + """ Obtain an estimate of the camera matrix based off the original frame dimensions. + + Returns + ------- + :class:`numpy.ndarray` + An estimated camera matrix + """ + focal_length = 4 + camera_matrix = np.array([[focal_length, 0, 0.5], + [0, focal_length, 0.5], + [0, 0, 1]], dtype="double") + logger.trace("camera_matrix: %s", camera_matrix) # type:ignore[attr-defined] + return camera_matrix + + def _solve_pnp(self, landmarks: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ Solve the Perspective-n-Point for the given landmarks. + + Takes 2D landmarks in world space and estimates the rotation and translation vectors + in 3D space. + + Parameters + ---------- + landmarks: :class:`numpy.ndarry` + The original 68 point landmark co-ordinates relating to the original frame + + Returns + ------- + rotation: :class:`numpy.ndarray` + The solved rotation vector + translation: :class:`numpy.ndarray` + The solved translation vector + """ + if self._landmarks_type != LandmarkType.LM_2D_68: + points: np.ndarray = np.empty([]) + rotation = np.array([[0.0], [0.0], [0.0]]) + translation = rotation.copy() + else: + points = landmarks[[6, 7, 8, 9, 10, 17, 21, 22, 26, 31, 32, 33, 34, + 35, 36, 39, 42, 45, 48, 50, 51, 52, 54, 56, 57, 58]] + _, rotation, translation = cv2.solvePnP(_MEAN_FACE[LandmarkType.LM_3D_26], + points, + self._camera_matrix, + self._distortion_coefficients, + flags=cv2.SOLVEPNP_ITERATIVE) + logger.trace("points: %s, rotation: %s, translation: %s", # type:ignore[attr-defined] + points, rotation, translation) + return rotation, translation + + def _get_offset(self) -> dict[CenteringType, np.ndarray]: + """ Obtain the offset between the original center of the extracted face to the new center + of the head in 2D space. + + Returns + ------- + :class:`numpy.ndarray` + The x, y offset of the new center from the old center. + """ + offset: dict[CenteringType, np.ndarray] = {"legacy": np.array([0.0, 0.0])} + if self._landmarks_type != LandmarkType.LM_2D_68: + offset["face"] = np.array([0.0, 0.0]) + offset["head"] = np.array([0.0, 0.0]) + else: + points: dict[T.Literal["face", "head"], tuple[float, ...]] = {"head": (0.0, 0.0, -2.3), + "face": (0.0, -1.5, 4.2)} + for key, pnts in points.items(): + center = cv2.projectPoints(np.array([pnts]).astype("float32"), + self._rotation, + self._translation, + self._camera_matrix, + self._distortion_coefficients)[0].squeeze() + logger.trace("center %s: %s", key, center) # type:ignore[attr-defined] + offset[key] = center - np.array([0.5, 0.5]) + logger.trace("offset: %s", offset) # type:ignore[attr-defined] + return offset diff --git a/lib/align/thumbnails.py b/lib/align/thumbnails.py new file mode 100644 index 0000000000..d97801d1d9 --- /dev/null +++ b/lib/align/thumbnails.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" Handles the generation of thumbnail jpgs for storing inside an alignments file/png header """ +from __future__ import annotations + +import logging +import typing as T + +import numpy as np + +from lib.logger import parse_class_init + +if T.TYPE_CHECKING: + from .alignments import Alignments + +logger = logging.getLogger(__name__) + + +class Thumbnails(): + """ Thumbnail images stored in the alignments file. + + The thumbnails are stored as low resolution (64px), low quality jpg in the alignments file + and are used for the Manual Alignments tool. + + Parameters + ---------- + alignments: :class:'~lib.align.alignments.Alignments` + The parent alignments class that these thumbs belong to + """ + def __init__(self, alignments: Alignments) -> None: + logger.debug(parse_class_init(locals())) + self._alignments_dict = alignments.data + self._frame_list = list(sorted(self._alignments_dict)) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def has_thumbnails(self) -> bool: + """ bool: ``True`` if all faces in the alignments file contain thumbnail images + otherwise ``False``. """ + retval = all(np.any(T.cast(np.ndarray, face.get("thumb"))) + for frame in self._alignments_dict.values() + for face in frame["faces"]) + logger.trace(retval) # type:ignore[attr-defined] + return retval + + def get_thumbnail_by_index(self, frame_index: int, face_index: int) -> np.ndarray: + """ Obtain a jpg thumbnail from the given frame index for the given face index + + Parameters + ---------- + frame_index: int + The frame index that contains the thumbnail + face_index: int + The face index within the frame to retrieve the thumbnail for + + Returns + ------- + :class:`numpy.ndarray` + The encoded jpg thumbnail + """ + retval = self._alignments_dict[self._frame_list[frame_index]]["faces"][face_index]["thumb"] + assert retval is not None + logger.trace( # type:ignore[attr-defined] + "frame index: %s, face_index: %s, thumb shape: %s", + frame_index, face_index, retval.shape) + return retval + + def add_thumbnail(self, frame: str, face_index: int, thumb: np.ndarray) -> None: + """ Add a thumbnail for the given face index for the given frame. + + Parameters + ---------- + frame: str + The name of the frame to add the thumbnail for + face_index: int + The face index within the given frame to add the thumbnail for + thumb: :class:`numpy.ndarray` + The encoded jpg thumbnail at 64px to add to the alignments file + """ + logger.debug("frame: %s, face_index: %s, thumb shape: %s thumb dtype: %s", + frame, face_index, thumb.shape, thumb.dtype) + self._alignments_dict[frame]["faces"][face_index]["thumb"] = thumb diff --git a/lib/align/updater.py b/lib/align/updater.py new file mode 100644 index 0000000000..7a98e3c8fc --- /dev/null +++ b/lib/align/updater.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 +""" Handles updating of an alignments file from an older version to the current version. """ +from __future__ import annotations + +import logging +import os +import typing as T + +import numpy as np + +from lib.logger import parse_class_init +from lib.utils import VIDEO_EXTENSIONS + +logger = logging.getLogger(__name__) + +if T.TYPE_CHECKING: + from .alignments import Alignments, AlignmentFileDict + + +class _Updater(): + """ Base class for inheriting to test for and update of an alignments file property + + Parameters + ---------- + alignments: :class:`~Alignments` + The alignments object that is being tested and updated + """ + def __init__(self, alignments: Alignments) -> None: + logger.debug(parse_class_init(locals())) + self._alignments = alignments + self._needs_update = self._test() + if self._needs_update: + self._update() + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def is_updated(self) -> bool: + """ bool. ``True`` if this updater has been run otherwise ``False`` """ + return self._needs_update + + def _test(self) -> bool: + """ Calls the child's :func:`test` method and logs output + + Returns + ------- + bool + ``True`` if the test condition is met otherwise ``False`` + """ + logger.debug("checking %s", self.__class__.__name__) + retval = self.test() + logger.debug("legacy %s: %s", self.__class__.__name__, retval) + return retval + + def test(self) -> bool: + """ Override to set the condition to test for. + + Returns + ------- + bool + ``True`` if the test condition is met otherwise ``False`` + """ + raise NotImplementedError() + + def _update(self) -> int: + """ Calls the child's :func:`update` method, logs output and sets the + :attr:`is_updated` flag + + Returns + ------- + int + The number of items that were updated + """ + retval = self.update() + logger.debug("Updated %s: %s", self.__class__.__name__, retval) + return retval + + def update(self) -> int: + """ Override to set the action to perform on the alignments object if the test has + passed + + Returns + ------- + int + The number of items that were updated + """ + raise NotImplementedError() + + +class VideoExtension(_Updater): + """ Alignments files from video files used to have a dummy '.png' extension for each of the + keys. This has been changed to be file extension of the original input video (for better) + identification of alignments files generated from video files + + Parameters + ---------- + alignments: :class:`~Alignments` + The alignments object that is being tested and updated + video_filename: str + The video filename that holds these alignments + """ + def __init__(self, alignments: Alignments, video_filename: str) -> None: + self._video_name, self._extension = os.path.splitext(video_filename) + super().__init__(alignments) + + def test(self) -> bool: + """ Requires update if the extension of the key in the alignment file is not the same + as for the input video file + + Returns + ------- + bool + ``True`` if the key extensions need updating otherwise ``False`` + """ + # Note: Don't check on alignments file version. It's possible that the file gets updated to + # a newer version before this check is run + if self._extension.lower() not in VIDEO_EXTENSIONS: + return False + + exts = set(os.path.splitext(k)[-1] for k in self._alignments.data) + if len(exts) != 1: + logger.debug("Alignments file has multiple key extensions. Skipping") + return False + + if self._extension in exts: + logger.debug("Alignments file contains correct key extensions. Skipping") + return False + + logger.debug("Needs update for video extension (version: %s, extension: %s)", + self._alignments.version, self._extension) + return True + + def update(self) -> int: + """ Update alignments files that have been extracted from videos to have the key end in the + video file extension rather than ',png' (the old way) + + Parameters + ---------- + video_filename: str + The filename of the video file that created these alignments + """ + updated = 0 + for key in list(self._alignments.data): + fname = os.path.splitext(key)[0] + if fname.rsplit("_", maxsplit=1)[0] != self._video_name: + continue # Key is from a different source + + val = self._alignments.data[key] + new_key = f"{fname}{self._extension}" + + del self._alignments.data[key] + self._alignments.data[new_key] = val + + updated += 1 + + logger.debug("Updated alignment keys for video extension: %s", updated) + return updated + + +class FileStructure(_Updater): + """ Alignments were structured: {frame_name: }. We need to be able to store + information at the frame level, so new structure is: {frame_name: {faces: }} + """ + def test(self) -> bool: + """ Test whether the alignments file is laid out in the old structure of + `{frame_name: [faces]}` + + Returns + ------- + bool + ``True`` if the file has legacy structure otherwise ``False`` + """ + return any(isinstance(val, list) for val in self._alignments.data.values()) + + def update(self) -> int: + """ Update legacy alignments files from the format `{frame_name: [faces}` to the + format `{frame_name: {faces: [faces]}`. + + Returns + ------- + int + The number of items that were updated + """ + updated = 0 + for key, val in self._alignments.data.items(): + if not isinstance(val, list): + continue + self._alignments.data[key] = {"faces": val} + updated += 1 + return updated + + +class LandmarkRename(_Updater): + """ Landmarks renamed from landmarksXY to landmarks_xy for PEP compliance """ + def test(self) -> bool: + """ check for legacy landmarksXY keys. + + Returns + ------- + bool + ``True`` if the alignments file contains legacy `landmarksXY` keys otherwise ``False`` + """ + return (any(key == "landmarksXY" + for val in self._alignments.data.values() + for alignment in val["faces"] + for key in alignment)) + + def update(self) -> int: + """ Update legacy `landmarksXY` keys to PEP compliant `landmarks_xy` keys. + + Returns + ------- + int + The number of landmarks keys that were changed + """ + update_count = 0 + for val in self._alignments.data.values(): + for alignment in val["faces"]: + if "landmarksXY" in alignment: + alignment["landmarks_xy"] = alignment.pop("landmarksXY") # type:ignore + update_count += 1 + return update_count + + +class ListToNumpy(_Updater): + """ Landmarks stored as list instead of numpy array """ + def test(self) -> bool: + """ check for legacy landmarks stored as `list` rather than :class:`numpy.ndarray`. + + Returns + ------- + bool + ``True`` if not all landmarks are :class:`numpy.ndarray` otherwise ``False`` + """ + return not all(isinstance(face["landmarks_xy"], np.ndarray) + for val in self._alignments.data.values() + for face in val["faces"]) + + def update(self) -> int: + """ Update landmarks stored as `list` to :class:`numpy.ndarray`. + + Returns + ------- + int + The number of landmarks keys that were changed + """ + update_count = 0 + for val in self._alignments.data.values(): + for alignment in val["faces"]: + test = alignment["landmarks_xy"] + if not isinstance(test, np.ndarray): + alignment["landmarks_xy"] = np.array(test, dtype="float32") + update_count += 1 + return update_count + + +class MaskCentering(_Updater): + """ Masks not containing the stored_centering parameters. Prior to this implementation all + masks were stored with face centering """ + + def test(self) -> bool: + """ Mask centering was introduced in alignments version 2.2 + + Returns + ------- + bool + ``True`` mask centering requires updating otherwise ``False`` + """ + return self._alignments.version < 2.2 + + def update(self) -> int: + """ Add the mask key to the alignment file and update the centering of existing masks + + Returns + ------- + int + The number of masks that were updated + """ + update_count = 0 + for val in self._alignments.data.values(): + for alignment in val["faces"]: + if "mask" not in alignment: + alignment["mask"] = {} + for mask in alignment["mask"].values(): + mask["stored_centering"] = "face" + update_count += 1 + return update_count + + +class IdentityAndVideoMeta(_Updater): + """ Prior to version 2.3 the identity key did not exist and the video_meta key was not + compulsory. These should now both always appear, but do not need to be populated. """ + + def test(self) -> bool: + """ Identity Key was introduced in alignments version 2.3 + + Returns + ------- + bool + ``True`` identity key needs inserting otherwise ``False`` + """ + return self._alignments.version < 2.3 + + # Identity information was not previously stored in the alignments file. + def update(self) -> int: + """ Add the video_meta and identity keys to the alignment file and leave empty + + Returns + ------- + int + The number of keys inserted + """ + update_count = 0 + for val in self._alignments.data.values(): + this_update = 0 + if "video_meta" not in val: + val["video_meta"] = {} + this_update = 1 + for alignment in val["faces"]: + if "identity" not in alignment: + alignment["identity"] = {} + this_update = 1 + update_count += this_update + return update_count + + +class Legacy(): + """ Legacy alignments properties that are no longer used, but are still required for backwards + compatibility/upgrading reasons. + + Parameters + ---------- + alignments: :class:`~Alignments` + The alignments object that requires these legacy properties + """ + def __init__(self, alignments: Alignments) -> None: + self._alignments = alignments + self._hashes_to_frame: dict[str, dict[str, int]] = {} + self._hashes_to_alignment: dict[str, AlignmentFileDict] = {} + + @property + def hashes_to_frame(self) -> dict[str, dict[str, int]]: + """ dict: The SHA1 hash of the face mapped to the frame(s) and face index within the frame + that the hash corresponds to. The structure of the dictionary is: + + {**SHA1_hash** (`str`): {**filename** (`str`): **face_index** (`int`)}}. + + Notes + ----- + This method is deprecated and exists purely for updating legacy hash based alignments + to new png header storage in :class:`lib.align.update_legacy_png_header`. + + The first time this property is referenced, the dictionary will be created and cached. + Subsequent references will be made to this cached dictionary. + """ + if not self._hashes_to_frame: + logger.debug("Generating hashes to frame") + for frame_name, val in self._alignments.data.items(): + for idx, face in enumerate(val["faces"]): + self._hashes_to_frame.setdefault( + face["hash"], {})[frame_name] = idx # type:ignore + return self._hashes_to_frame + + @property + def hashes_to_alignment(self) -> dict[str, AlignmentFileDict]: + """ dict: The SHA1 hash of the face mapped to the alignment for the face that the hash + corresponds to. The structure of the dictionary is: + + Notes + ----- + This method is deprecated and exists purely for updating legacy hash based alignments + to new png header storage in :class:`lib.align.update_legacy_png_header`. + + The first time this property is referenced, the dictionary will be created and cached. + Subsequent references will be made to this cached dictionary. + """ + if not self._hashes_to_alignment: + logger.debug("Generating hashes to alignment") + self._hashes_to_alignment = {face["hash"]: face # type:ignore + for val in self._alignments.data.values() + for face in val["faces"]} + return self._hashes_to_alignment diff --git a/lib/cli/actions.py b/lib/cli/actions.py index c5599d2b04..e3f36ae48c 100644 --- a/lib/cli/actions.py +++ b/lib/cli/actions.py @@ -7,18 +7,19 @@ import argparse import os +import typing as T # << FILE HANDLING >> -class _FullPaths(argparse.Action): # pylint: disable=too-few-public-methods +class _FullPaths(argparse.Action): """ Parent class for various file type and file path handling classes. Expands out given paths to their full absolute paths. This class should not be called directly. It is the base class for the various different file handling methods. """ - def __call__(self, parser, namespace, values, option_string=None): + def __call__(self, parser, namespace, values, option_string=None) -> None: if isinstance(values, (list, tuple)): vals = [os.path.abspath(os.path.expanduser(val)) for val in values] else: @@ -41,8 +42,7 @@ class DirFullPaths(_FullPaths): >>> opts=("-f", "--folder_location"), >>> action=DirFullPaths)), """ - # pylint: disable=too-few-public-methods,unnecessary-pass - pass + pass # pylint:disable=unnecessary-pass class FileFullPaths(_FullPaths): @@ -67,8 +67,7 @@ class FileFullPaths(_FullPaths): >>> action=FileFullPaths, >>> filetypes="video))" """ - # pylint: disable=too-few-public-methods - def __init__(self, *args, filetypes=None, **kwargs): + def __init__(self, *args, filetypes: str | None = None, **kwargs) -> None: super().__init__(*args, **kwargs) self.filetypes = filetypes @@ -86,7 +85,7 @@ def _get_kwargs(self): return [(name, getattr(self, name)) for name in names] -class FilesFullPaths(FileFullPaths): # pylint: disable=too-few-public-methods +class FilesFullPaths(FileFullPaths): """ Adds support for a File browser to select multiple files in the GUI. This extends the standard :class:`argparse.Action` and adds an additional parameter @@ -110,14 +109,14 @@ class FilesFullPaths(FileFullPaths): # pylint: disable=too-few-public-methods >>> filetypes="image", >>> nargs="+")) """ - def __init__(self, *args, filetypes=None, **kwargs): + def __init__(self, *args, filetypes: str | None = None, **kwargs) -> None: if kwargs.get("nargs", None) is None: opt = kwargs["option_strings"] - raise ValueError("nargs must be provided for FilesFullPaths: {}".format(opt)) + raise ValueError(f"nargs must be provided for FilesFullPaths: {opt}") super().__init__(*args, **kwargs) -class DirOrFileFullPaths(FileFullPaths): # pylint: disable=too-few-public-methods +class DirOrFileFullPaths(FileFullPaths): """ Adds support to the GUI to launch either a file browser or a folder browser. Some inputs (for example source frames) can come from a folder of images or from a @@ -144,7 +143,50 @@ class DirOrFileFullPaths(FileFullPaths): # pylint: disable=too-few-public-metho >>> action=DirOrFileFullPaths, >>> filetypes="video))" """ - pass # pylint: disable=unnecessary-pass + + +class DirOrFilesFullPaths(FileFullPaths): + """ Adds support to the GUI to launch either a file browser for selecting multiple files + or a folder browser. + + Some inputs (for example face filter) can come from a folder of images or from multiple + image file. This indicates to the GUI that it should place 2 buttons (one for a folder + browser, one for a multi-file browser) for file/folder browsing. + + The standard :class:`argparse.Action` is extended with the additional parameter + :attr:`filetypes`, indicating to the GUI that it should pop a file browser, and limit + the results to the file types listed. As well as the standard parameters, the following + parameter is required: + + Parameters + ---------- + filetypes: str + The accepted file types for this option. This is the key for the GUIs lookup table which + can be found in :class:`lib.gui.utils.FileHandler`. NB: This parameter is only used for + the file browser and not the folder browser + + Example + ------- + >>> argument_list = [] + >>> argument_list.append(dict( + >>> opts=("-f", "--input_frames"), + >>> action=DirOrFileFullPaths, + >>> filetypes="video))" + """ + def __call__(self, parser, namespace, values, option_string=None) -> None: + """ Override :class:`_FullPaths` __call__ function. + + The input for this option can be a space separated list of files or a single folder. + Folders can have spaces in them, so we don't want to blindly expand the paths. + + We check whether the input can be resolved to a folder first before expanding. + """ + assert isinstance(values, (list, tuple)) + folder = os.path.abspath(os.path.expanduser(" ".join(values))) + if os.path.isdir(folder): + setattr(namespace, self.dest, [folder]) + else: # file list so call parent method + super().__call__(parser, namespace, values, option_string) class SaveFileFullPaths(FileFullPaths): @@ -169,8 +211,7 @@ class SaveFileFullPaths(FileFullPaths): >>> action=SaveFileFullPaths, >>> filetypes="video")) """ - # pylint: disable=too-few-public-methods,unnecessary-pass - pass + pass # pylint:disable=unnecessary-pass class ContextFullPaths(FileFullPaths): @@ -203,19 +244,23 @@ class ContextFullPaths(FileFullPaths): >>> filetypes="video", >>> action_option="-a")) """ - # pylint: disable=too-few-public-methods, too-many-arguments - def __init__(self, *args, filetypes=None, action_option=None, **kwargs): + # pylint:disable=too-many-arguments + def __init__(self, + *args, + filetypes: str | None = None, + action_option: str | None = None, + **kwargs) -> None: opt = kwargs["option_strings"] if kwargs.get("nargs", None) is not None: - raise ValueError("nargs not allowed for ContextFullPaths: {}".format(opt)) + raise ValueError(f"nargs not allowed for ContextFullPaths: {opt}") if filetypes is None: - raise ValueError("filetypes is required for ContextFullPaths: {}".format(opt)) + raise ValueError(f"filetypes is required for ContextFullPaths: {opt}") if action_option is None: - raise ValueError("action_option is required for ContextFullPaths: {}".format(opt)) + raise ValueError(f"action_option is required for ContextFullPaths: {opt}") super().__init__(*args, filetypes=filetypes, **kwargs) self.action_option = action_option - def _get_kwargs(self): + def _get_kwargs(self) -> list[tuple[str, T.Any]]: names = ["option_strings", "dest", "nargs", @@ -232,7 +277,7 @@ def _get_kwargs(self): # << GUI DISPLAY OBJECTS >> -class Radio(argparse.Action): # pylint: disable=too-few-public-methods +class Radio(argparse.Action): """ Adds support for a GUI Radio options box. This is a standard :class:`argparse.Action` (with stock parameters) which indicates to the GUI @@ -249,19 +294,19 @@ class Radio(argparse.Action): # pylint: disable=too-few-public-methods >>> action=Radio, >>> choices=["foo", "bar")) """ - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: opt = kwargs["option_strings"] if kwargs.get("nargs", None) is not None: - raise ValueError("nargs not allowed for Radio buttons: {}".format(opt)) + raise ValueError(f"nargs not allowed for Radio buttons: {opt}") if not kwargs.get("choices", []): - raise ValueError("Choices must be provided for Radio buttons: {}".format(opt)) + raise ValueError(f"Choices must be provided for Radio buttons: {opt}") super().__init__(*args, **kwargs) - def __call__(self, parser, namespace, values, option_string=None): + def __call__(self, parser, namespace, values, option_string=None) -> None: setattr(namespace, self.dest, values) -class MultiOption(argparse.Action): # pylint: disable=too-few-public-methods +class MultiOption(argparse.Action): """ Adds support for multiple option checkboxes in the GUI. This is a standard :class:`argparse.Action` (with stock parameters) which indicates to the GUI @@ -277,19 +322,19 @@ class MultiOption(argparse.Action): # pylint: disable=too-few-public-methods >>> action=MultiOption, >>> choices=["foo", "bar")) """ - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: opt = kwargs["option_strings"] if not kwargs.get("nargs", []): - raise ValueError("nargs must be provided for MultiOption: {}".format(opt)) + raise ValueError(f"nargs must be provided for MultiOption: {opt}") if not kwargs.get("choices", []): - raise ValueError("Choices must be provided for MultiOption: {}".format(opt)) + raise ValueError(f"Choices must be provided for MultiOption: {opt}") super().__init__(*args, **kwargs) - def __call__(self, parser, namespace, values, option_string=None): + def __call__(self, parser, namespace, values, option_string=None) -> None: setattr(namespace, self.dest, values) -class Slider(argparse.Action): # pylint: disable=too-few-public-methods +class Slider(argparse.Action): """ Adds support for a slider in the GUI. The standard :class:`argparse.Action` is extended with the additional parameters listed below. @@ -332,24 +377,28 @@ class Slider(argparse.Action): # pylint: disable=too-few-public-methods >>> type=float, >>> default=5.00)) """ - def __init__(self, *args, min_max=None, rounding=None, **kwargs): + def __init__(self, + *args, + min_max: tuple[int, int] | tuple[float, float] | None = None, + rounding: int | None = None, + **kwargs) -> None: opt = kwargs["option_strings"] if kwargs.get("nargs", None) is not None: - raise ValueError("nargs not allowed for Slider: {}".format(opt)) + raise ValueError(f"nargs not allowed for Slider: {opt}") if kwargs.get("default", None) is None: - raise ValueError("A default value must be supplied for Slider: {}".format(opt)) + raise ValueError(f"A default value must be supplied for Slider: {opt}") if kwargs.get("type", None) not in (int, float): - raise ValueError("Sliders only accept int and float data types: {}".format(opt)) + raise ValueError(f"Sliders only accept int and float data types: {opt}") if min_max is None: - raise ValueError("min_max must be provided for Sliders: {}".format(opt)) + raise ValueError(f"min_max must be provided for Sliders: {opt}") if rounding is None: - raise ValueError("rounding must be provided for Sliders: {}".format(opt)) + raise ValueError(f"rounding must be provided for Sliders: {opt}") super().__init__(*args, **kwargs) self.min_max = min_max self.rounding = rounding - def _get_kwargs(self): + def _get_kwargs(self) -> list[tuple[str, T.Any]]: names = ["option_strings", "dest", "nargs", @@ -363,5 +412,5 @@ def _get_kwargs(self): "rounding"] # Decimal places to round floats to or step interval for ints return [(name, getattr(self, name)) for name in names] - def __call__(self, parser, namespace, values, option_string=None): + def __call__(self, parser, namespace, values, option_string=None) -> None: setattr(namespace, self.dest, values) diff --git a/lib/cli/args.py b/lib/cli/args.py index 3f9e62fa5c..29aa74ee40 100644 --- a/lib/cli/args.py +++ b/lib/cli/args.py @@ -1,24 +1,21 @@ #!/usr/bin/env python3 -""" The Command Line Argument options for faceswap.py """ +""" The global and GUI Command Line Argument options for faceswap.py """ -# pylint:disable=too-many-lines import argparse import gettext import logging import re import sys import textwrap +import typing as T from lib.utils import get_backend from lib.gpu_stats import GPUStats -from plugins.plugin_loader import PluginLoader - -from .actions import (DirFullPaths, DirOrFileFullPaths, FileFullPaths, FilesFullPaths, MultiOption, - Radio, SaveFileFullPaths, Slider) +from .actions import FileFullPaths, MultiOption, SaveFileFullPaths from .launcher import ScriptExecutor -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) _GPUS = GPUStats().cli_devices # LOCALES @@ -28,9 +25,9 @@ class FullHelpArgumentParser(argparse.ArgumentParser): """ Extends :class:`argparse.ArgumentParser` to output full help on bad arguments. """ - def error(self, message): + def error(self, message: str) -> T.NoReturn: self.print_help(sys.stderr) - self.exit(2, "{}: error: {}\n".format(self.prog, message)) + self.exit(2, f"{self.prog}: error: {message}\n") class SmartFormatter(argparse.HelpFormatter): @@ -45,11 +42,15 @@ class SmartFormatter(argparse.HelpFormatter): Prefixing a new line within the help text with "L|" will turn that line into a list item in both the cli help text and the GUI. """ - def __init__(self, prog, indent_increment=2, max_help_position=24, width=None): + def __init__(self, + prog: str, + indent_increment: int = 2, + max_help_position: int = 24, + width: int | None = None) -> None: super().__init__(prog, indent_increment, max_help_position, width) self._whitespace_matcher_limited = re.compile(r'[ \r\f\v]+', re.ASCII) - def _split_lines(self, text, width): + def _split_lines(self, text: str, width: int) -> list[str]: """ Split the given text by the given display width. If the text is not prefixed with "R|" then the standard @@ -62,18 +63,25 @@ def _split_lines(self, text, width): The help text that is to be formatted for display width: int The display width, in characters, for the help text + + Returns + ------- + list + A list of split strings """ if text.startswith("R|"): text = self._whitespace_matcher_limited.sub(' ', text).strip()[2:] - output = list() + output = [] for txt in text.splitlines(): indent = "" if txt.startswith("L|"): indent = " " - txt = " - {}".format(txt[2:]) + txt = f" - {txt[2:]}" output.extend(textwrap.wrap(txt, width, subsequent_indent=indent)) return output - return argparse.HelpFormatter._split_lines(self, text, width) + return argparse.HelpFormatter._split_lines(self, # pylint:disable=protected-access + text, + width) class FaceSwapArgs(): @@ -87,16 +95,20 @@ class FaceSwapArgs(): Parameters ---------- - subparser: :class:`argparse._SubParsersAction` - The subparser for the given command + subparser: :class:`argparse._SubParsersAction` | None + The subparser for the given command. ``None`` if the class is being called for reading + rather than processing command: str The faceswap command that is to be executed description: str, optional The description for the given command. Default: "default" """ - def __init__(self, subparser, command, description="default"): + def __init__(self, + subparser: argparse._SubParsersAction | None, + command: str, + description: str = "default") -> None: self.global_arguments = self._get_global_arguments() - self.info = self.get_info() + self.info: str = self.get_info() self.argument_list = self.get_argument_list() self.optional_arguments = self.get_optional_arguments() self._process_suppressions() @@ -108,7 +120,7 @@ def __init__(self, subparser, command, description="default"): self.parser.set_defaults(func=script.execute_script) @staticmethod - def get_info(): + def get_info() -> str: """ Returns the information text for the current command. This function should be overridden with the actual command help text for each @@ -119,10 +131,10 @@ def get_info(): str The information text for this command. """ - return None + return "" @staticmethod - def get_argument_list(): + def get_argument_list() -> list[dict[str, T.Any]]: """ Returns the argument list for the current command. The argument list should be a list of dictionaries pertaining to each option for a command. @@ -136,11 +148,11 @@ def get_argument_list(): list The list of command line options for the given command """ - argument_list = [] + argument_list: list[dict[str, T.Any]] = [] return argument_list @staticmethod - def get_optional_arguments(): + def get_optional_arguments() -> list[dict[str, T.Any]]: """ Returns the optional argument list for the current command. The optional arguments list is not always required, but is used when there are shared @@ -151,11 +163,11 @@ def get_optional_arguments(): list The list of optional command line options for the given command """ - argument_list = [] + argument_list: list[dict[str, T.Any]] = [] return argument_list @staticmethod - def _get_global_arguments(): + def _get_global_arguments() -> list[dict[str, T.Any]]: """ Returns the global Arguments list that are required for ALL commands in Faceswap. This method should NOT be overridden. @@ -165,71 +177,82 @@ def _get_global_arguments(): list The list of global command line options for all Faceswap commands. """ - global_args = list() + global_args: list[dict[str, T.Any]] = [] if _GPUS: - global_args.append(dict( - opts=("-X", "--exclude-gpus"), - dest="exclude_gpus", - action=MultiOption, - type=str.lower, - nargs="+", - choices=[str(idx) for idx in range(len(_GPUS))], - group=_("Global Options"), - help=_("R|Exclude GPUs from use by Faceswap. Select the number(s) which " - "correspond to any GPU(s) that you do not wish to be made available to " - "Faceswap. Selecting all GPUs here will force Faceswap into CPU mode." - "\nL|{}").format(" \nL|".join(_GPUS)))) - global_args.append(dict( - opts=("-C", "--configfile"), - action=FileFullPaths, - filetypes="ini", - type=str, - group=_("Global Options"), - help=_("Optionally overide the saved config with the path to a custom config file."))) - global_args.append(dict( - opts=("-L", "--loglevel"), - type=str.upper, - dest="loglevel", - default="INFO", - choices=("INFO", "VERBOSE", "DEBUG", "TRACE"), - group=_("Global Options"), - help=_("Log level. Stick with INFO or VERBOSE unless you need to file an error " - "report. Be careful with TRACE as it will generate a lot of data"))) - global_args.append(dict( - opts=("-LF", "--logfile"), - action=SaveFileFullPaths, - filetypes='log', - type=str, - dest="logfile", - default=None, - group=_("Global Options"), - help=_("Path to store the logfile. Leave blank to store in the faceswap folder"))) + global_args.append({ + "opts": ("-X", "--exclude-gpus"), + "dest": "exclude_gpus", + "action": MultiOption, + "type": str.lower, + "nargs": "+", + "choices": [str(idx) for idx in range(len(_GPUS))], + "group": _("Global Options"), + "help": _( + "R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond " + "to any GPU(s) that you do not wish to be made available to Faceswap. " + "Selecting all GPUs here will force Faceswap into CPU mode." + "\nL|{}".format(' \nL|'.join(_GPUS)))}) + global_args.append({ + "opts": ("-C", "--configfile"), + "action": FileFullPaths, + "filetypes": "ini", + "type": str, + "group": _("Global Options"), + "help": _( + "Optionally overide the saved config with the path to a custom config file.")}) + global_args.append({ + "opts": ("-L", "--loglevel"), + "type": str.upper, + "dest": "loglevel", + "default": "INFO", + "choices": ("INFO", "VERBOSE", "DEBUG", "TRACE"), + "group": _("Global Options"), + "help": _( + "Log level. Stick with INFO or VERBOSE unless you need to file an error report. " + "Be careful with TRACE as it will generate a lot of data")}) + global_args.append({ + "opts": ("-F", "--logfile"), + "action": SaveFileFullPaths, + "filetypes": 'log', + "type": str, + "dest": "logfile", + "default": None, + "group": _("Global Options"), + "help": _("Path to store the logfile. Leave blank to store in the faceswap folder")}) # These are hidden arguments to indicate that the GUI/Colab is being used - global_args.append(dict( - opts=("-gui", "--gui"), - action="store_true", - dest="redirect_gui", - default=False, - help=argparse.SUPPRESS)) - global_args.append(dict( - opts=("-colab", "--colab"), - action="store_true", - dest="colab", - default=False, - help=argparse.SUPPRESS)) + global_args.append({ + "opts": ("-gui", "--gui"), + "action": "store_true", + "dest": "redirect_gui", + "default": False, + "help": argparse.SUPPRESS}) + # Deprecated multi-character switches + global_args.append({ + "opts": ("-LF",), + "action": SaveFileFullPaths, + "filetypes": 'log', + "type": str, + "dest": "depr_logfile_LF_F", + "help": argparse.SUPPRESS}) + return global_args @staticmethod - def _create_parser(subparser, command, description): + def _create_parser(subparser: argparse._SubParsersAction, + command: str, + description: str) -> argparse.ArgumentParser: """ Create the parser for the selected command. Parameters ---------- + subparser: :class:`argparse._SubParsersAction` + The subparser for the given command command: str The faceswap command that is to be executed description: str The description for the given command + Returns ------- :class:`~lib.cli.args.FullHelpArgumentParser` @@ -242,7 +265,7 @@ def _create_parser(subparser, command, description): formatter_class=SmartFormatter) return parser - def _add_arguments(self): + def _add_arguments(self) -> None: """ Parse the list of dictionaries containing the command line arguments and convert to argparse parser arguments. """ options = self.global_arguments + self.argument_list + self.optional_arguments @@ -251,7 +274,7 @@ def _add_arguments(self): kwargs = {key: option[key] for key in option.keys() if key not in ("opts", "group")} self.parser.add_argument(*args, **kwargs) - def _process_suppressions(self): + def _process_suppressions(self) -> None: """ Certain options are only available for certain backends. Suppresses command line options that are not available for the running backend. @@ -270,868 +293,11 @@ def _process_suppressions(self): opts["help"] = argparse.SUPPRESS -class ExtractConvertArgs(FaceSwapArgs): - """ Parent class to capture arguments that will be used in both extract and convert processes. - - Extract and Convert share a fair amount of arguments, so arguments that can be used in both of - these processes should be placed here. - - No further processing is done in this class (this is handled by the children), this just - captures the shared arguments. - """ - - @staticmethod - def get_argument_list(): - """ Returns the argument list for shared Extract and Convert arguments. - - Returns - ------- - list - The list of command line options for the given Extract and Convert - """ - argument_list = list() - argument_list.append(dict( - opts=("-i", "--input-dir"), - action=DirOrFileFullPaths, - filetypes="video", - dest="input_dir", - required=True, - group=_("Data"), - help=_("Input directory or video. Either a directory containing the image files you " - "wish to process or path to a video file. NB: This should be the source video/" - "frames NOT the source faces."))) - argument_list.append(dict( - opts=("-o", "--output-dir"), - action=DirFullPaths, - dest="output_dir", - required=True, - group=_("Data"), - help=_("Output directory. This is where the converted files will be saved."))) - argument_list.append(dict( - opts=("-al", "--alignments"), - action=FileFullPaths, - filetypes="alignments", - type=str, - dest="alignments_path", - group=_("Data"), - help=_("Optional path to an alignments file. Leave blank if the alignments file is " - "at the default location."))) - return argument_list - - -class ExtractArgs(ExtractConvertArgs): - """ Creates the command line arguments for extraction. - - This class inherits base options from :class:`ExtractConvertArgs` where arguments that are used - for both Extract and Convert should be placed. - - Commands explicit to Extract should be added in :func:`get_optional_arguments` - """ - - @staticmethod - def get_info(): - """ The information text for the Extract command. - - Returns - ------- - str - The information text for the Extract command. - """ - return _("Extract faces from image or video sources.\n" - "Extraction plugins can be configured in the 'Settings' Menu") - - @staticmethod - def get_optional_arguments(): - """ Returns the argument list unique to the Extract command. - - Returns - ------- - list - The list of optional command line options for the Extract command - """ - if get_backend() == "cpu": - default_detector = default_aligner = "cv2-dnn" - else: - default_detector = "s3fd" - default_aligner = "fan" - - argument_list = [] - argument_list.append(dict( - opts=("-D", "--detector"), - action=Radio, - type=str.lower, - default=default_detector, - choices=PluginLoader.get_available_extractors("detect"), - group=_("Plugins"), - help=_("R|Detector to use. Some of these have configurable settings in " - "'/config/extract.ini' or 'Settings > Configure Extract 'Plugins':" - "\nL|cv2-dnn: A CPU only extractor which is the least reliable and least " - "resource intensive. Use this if not using a GPU and time is important." - "\nL|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " - "than other GPU detectors but can often return more false positives." - "\nL|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces " - "and fewer false positives than other GPU detectors, but is a lot more " - "resource intensive."))) - argument_list.append(dict( - opts=("-A", "--aligner"), - action=Radio, - type=str.lower, - default=default_aligner, - choices=PluginLoader.get_available_extractors("align"), - group=_("Plugins"), - help=_("R|Aligner to use." - "\nL|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " - "but less accurate. Only use this if not using a GPU and time is important." - "\nL|fan: Best aligner. Fast on GPU, slow on CPU."))) - argument_list.append(dict( - opts=("-M", "--masker"), - action=MultiOption, - type=str.lower, - nargs="+", - choices=[mask for mask in PluginLoader.get_available_extractors("mask") - if mask not in ("components", "extended")], - group=_("Plugins"), - help=_("R|Additional Masker(s) to use. The masks generated here will all take up GPU " - "RAM. You can select none, one or multiple masks, but the extraction may take " - "longer the more you select. NB: The Extended and Components (landmark based) " - "masks are automatically generated on extraction." - "\nL|bisenet-fp: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked including full head masking " - "(configurable in mask settings)." - "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " - "faces clear of obstructions. Profile faces and obstructions may result in " - "sub-par performance." - "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " - "frontal faces. The mask model has been specifically trained to recognize " - "some facial obstructions (hands and eyeglasses). Profile faces may result in " - "sub-par performance." - "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " - "faces. The mask model has been trained by community members and will need " - "testing for further description. Profile faces may result in sub-par " - "performance." - "\nThe auto generated masks are as follows:" - "\nL|components: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks to create a mask." - "\nL|extended: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks and the mask is extended upwards onto the " - "forehead." - "\n(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)"))) - argument_list.append(dict( - opts=("-nm", "--normalization"), - action=Radio, - type=str.lower, - dest="normalization", - default="none", - choices=["none", "clahe", "hist", "mean"], - group=_("Plugins"), - help=_("R|Performing normalization can help the aligner better align faces with " - "difficult lighting conditions at an extraction speed cost. Different methods " - "will yield different results on different sets. NB: This does not impact the " - "output face, just the input to the aligner." - "\nL|none: Don't perform normalization on the face." - "\nL|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " - "face." - "\nL|hist: Equalize the histograms on the RGB channels." - "\nL|mean: Normalize the face colors to the mean."))) - argument_list.append(dict( - opts=("-rf", "--re-feed"), - action=Slider, - min_max=(0, 10), - rounding=1, - type=int, - dest="re_feed", - default=0, - group=_("Plugins"), - help=_("The number of times to re-feed the detected face into the aligner. Each time " - "the face is re-fed into the aligner the bounding box is adjusted by a small " - "amount. The final landmarks are then averaged from each iteration. Helps to " - "remove 'micro-jitter' but at the cost of slower extraction speed. The more " - "times the face is re-fed into the aligner, the less micro-jitter should occur " - "but the longer extraction will take."))) - argument_list.append(dict( - opts=("-r", "--rotate-images"), - type=str, - dest="rotate_images", - default=None, - group=_("Plugins"), - help=_("If a face isn't found, rotate the images to try to find a face. Can find " - "more faces at the cost of extraction speed. Pass in a single number to use " - "increments of that size up to 360, or pass in a list of numbers to enumerate " - "exactly what angles to check."))) - argument_list.append(dict( - opts=("-min", "--min-size"), - action=Slider, - min_max=(0, 1080), - rounding=20, - type=int, - dest="min_size", - default=0, - group=_("Face Processing"), - help=_("Filters out faces detected below this size. Length, in pixels across the " - "diagonal of the bounding box. Set to 0 for off"))) - argument_list.append(dict( - opts=("-n", "--nfilter"), - action=FilesFullPaths, - filetypes="image", - dest="nfilter", - default=None, - nargs="+", - group=_("Face Processing"), - help=_("Optionally filter out people who you do not wish to process by passing in an " - "image of that person. Should be a front portrait with a single person in the " - "image. Multiple images can be added space separated. NB: Using face filter " - "will significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-f", "--filter"), - action=FilesFullPaths, - filetypes="image", - dest="filter", - default=None, - nargs="+", - group=_("Face Processing"), - help=_("Optionally select people you wish to process by passing in an image of that " - "person. Should be a front portrait with a single person in the image. " - "Multiple images can be added space separated. NB: Using face filter will " - "significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-l", "--ref_threshold"), - action=Slider, - min_max=(0.01, 0.99), - rounding=2, - type=float, - dest="ref_threshold", - default=0.4, - group=_("Face Processing"), - help=_("For use with the optional nfilter/filter files. Threshold for positive face " - "recognition. Lower values are stricter. NB: Using face filter will " - "significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-sz", "--size"), - action=Slider, - min_max=(256, 1024), - rounding=64, - type=int, - default=512, - group=_("output"), - help=_("The output size of extracted faces. Make sure that the model you intend to " - "train supports your required size. This will only need to be changed for " - "hi-res models."))) - argument_list.append(dict( - opts=("-een", "--extract-every-n"), - action=Slider, - min_max=(1, 100), - rounding=1, - type=int, - dest="extract_every_n", - default=1, - group=_("output"), - help=_("Extract every 'nth' frame. This option will skip frames when extracting " - "faces. For example a value of 1 will extract faces from every frame, a value " - "of 10 will extract faces from every 10th frame."))) - argument_list.append(dict( - opts=("-si", "--save-interval"), - action=Slider, - min_max=(0, 1000), - rounding=10, - type=int, - dest="save_interval", - default=0, - group=_("output"), - help=_("Automatically save the alignments file after a set amount of frames. By " - "default the alignments file is only saved at the end of the extraction " - "process. NB: If extracting in 2 passes then the alignments file will only " - "start to be saved out during the second pass. WARNING: Don't interrupt the " - "script when writing the file because it might get corrupted. Set to 0 to " - "turn off"))) - argument_list.append(dict( - opts=("-dl", "--debug-landmarks"), - action="store_true", - dest="debug_landmarks", - default=False, - group=_("output"), - help=_("Draw landmarks on the ouput faces for debugging purposes."))) - argument_list.append(dict( - opts=("-sp", "--singleprocess"), - action="store_true", - default=False, - backend="nvidia", - group=_("settings"), - help=_("Don't run extraction in parallel. Will run each part of the extraction " - "process separately (one after the other) rather than all at the smae time. " - "Useful if VRAM is at a premium."))) - argument_list.append(dict( - opts=("-s", "--skip-existing"), - action="store_true", - dest="skip_existing", - default=False, - group=_("settings"), - help=_("Skips frames that have already been extracted and exist in the alignments " - "file"))) - argument_list.append(dict( - opts=("-sf", "--skip-existing-faces"), - action="store_true", - dest="skip_faces", - default=False, - group=_("settings"), - help=_("Skip frames that already have detected faces in the alignments file"))) - argument_list.append(dict( - opts=("-ssf", "--skip-saving-faces"), - action="store_true", - dest="skip_saving_faces", - default=False, - group=_("settings"), - help=_("Skip saving the detected faces to disk. Just create an alignments file"))) - return argument_list - - -class ConvertArgs(ExtractConvertArgs): - """ Creates the command line arguments for conversion. - - This class inherits base options from :class:`ExtractConvertArgs` where arguments that are used - for both Extract and Convert should be placed. - - Commands explicit to Convert should be added in :func:`get_optional_arguments` - """ - - @staticmethod - def get_info(): - """ The information text for the Convert command. - - Returns - ------- - str - The information text for the Convert command. - """ - return _("Swap the original faces in a source video/images to your final faces.\n" - "Conversion plugins can be configured in the 'Settings' Menu") - - @staticmethod - def get_optional_arguments(): - """ Returns the argument list unique to the Convert command. - - Returns - ------- - list - The list of optional command line options for the Convert command - """ - - argument_list = [] - argument_list.append(dict( - opts=("-ref", "--reference-video"), - action=FileFullPaths, - filetypes="video", - type=str, - dest="reference_video", - group=_("Data"), - help=_("Only required if converting from images to video. Provide The original video " - "that the source frames were extracted from (for extracting the fps and " - "audio)."))) - argument_list.append(dict( - opts=("-m", "--model-dir"), - action=DirFullPaths, - dest="model_dir", - required=True, - group=_("Data"), - help=_("Model directory. The directory containing the trained model you wish to use " - "for conversion."))) - argument_list.append(dict( - opts=("-c", "--color-adjustment"), - action=Radio, - type=str.lower, - dest="color_adjustment", - default="avg-color", - choices=PluginLoader.get_available_convert_plugins("color", True), - group=_("Plugins"), - help=_("R|Performs color adjustment to the swapped face. Some of these options have " - "configurable settings in '/config/convert.ini' or 'Settings > Configure " - "Convert Plugins':" - "\nL|avg-color: Adjust the mean of each color channel in the swapped " - "reconstruction to equal the mean of the masked area in the original image." - "\nL|color-transfer: Transfers the color distribution from the source to the " - "target image using the mean and standard deviations of the L*a*b* " - "color space." - "\nL|manual-balance: Manually adjust the balance of the image in a variety of " - "color spaces. Best used with the Preview tool to set correct values." - "\nL|match-hist: Adjust the histogram of each color channel in the swapped " - "reconstruction to equal the histogram of the masked area in the original " - "image." - "\nL|seamless-clone: Use cv2's seamless clone function to remove extreme " - "gradients at the mask seam by smoothing colors. Generally does not give " - "very satisfactory results." - "\nL|none: Don't perform color adjustment."))) - argument_list.append(dict( - opts=("-M", "--mask-type"), - action=Radio, - type=str.lower, - dest="mask_type", - default="extended", - choices=PluginLoader.get_available_extractors("mask", - add_none=True, - extend_plugin=True) + ["predicted"], - group=_("Plugins"), - help=_("R|Masker to use. NB: The mask you require must exist within the alignments " - "file. You can add additional masks with the Mask Tool." - "\nL|none: Don't use a mask." - "\nL|bisenet-fp-face: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked (configurable in mask settings). " - "Use this version of bisenet-fp if your model is trained with 'face' or " - "'legacy' centering." - "\nL|bisenet-fp-head: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked (configurable in mask settings). " - "Use this version of bisenet-fp if your model is trained with 'head' centering." - "\nL|components: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks to create a mask." - "\nL|extended: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks and the mask is extended upwards onto the forehead." - "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " - "faces clear of obstructions. Profile faces and obstructions may result in " - "sub-par performance." - "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " - "frontal faces. The mask model has been specifically trained to recognize " - "some facial obstructions (hands and eyeglasses). Profile faces may result in " - "sub-par performance." - "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " - "faces. The mask model has been trained by community members and will need " - "testing for further description. Profile faces may result in sub-par " - "performance." - "\nL|predicted: If the 'Learn Mask' option was enabled during training, this " - "will use the mask that was created by the trained model."))) - argument_list.append(dict( - opts=("-w", "--writer"), - action=Radio, - type=str, - default="opencv", - choices=PluginLoader.get_available_convert_plugins("writer", False), - group=_("Plugins"), - help=_("R|The plugin to use to output the converted images. The writers are " - "configurable in '/config/convert.ini' or 'Settings > Configure Convert " - "Plugins:'" - "\nL|ffmpeg: [video] Writes out the convert straight to video. When the input " - "is a series of images then the '-ref' (--reference-video) parameter must be " - "set." - "\nL|gif: [animated image] Create an animated gif." - "\nL|opencv: [images] The fastest image writer, but less options and formats " - "than other plugins." - "\nL|pillow: [images] Slower than opencv, but has more options and supports " - "more formats."))) - argument_list.append(dict( - opts=("-osc", "--output-scale"), - action=Slider, - min_max=(25, 400), - rounding=1, - type=int, - dest="output_scale", - default=100, - group=_("Frame Processing"), - help=_("Scale the final output frames by this amount. 100%% will output the frames " - "at source dimensions. 50%% at half size 200%% at double size"))) - argument_list.append(dict( - opts=("-fr", "--frame-ranges"), - type=str, - nargs="+", - group=_("Frame Processing"), - help=_("Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " - "--frame-ranges 10-50 90-100. Frames falling outside of the selected range " - "will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " - "converting from images, then the filenames must end with the frame-number!"))) - argument_list.append(dict( - opts=("-a", "--input-aligned-dir"), - action=DirFullPaths, - dest="input_aligned_dir", - default=None, - group=_("Face Processing"), - help=_("If you have not cleansed your alignments file, then you can filter out faces " - "by defining a folder here that contains the faces extracted from your input " - "files/video. If this folder is defined, then only faces that exist within " - "your alignments file and also exist within the specified folder will be " - "converted. Leaving this blank will convert all faces that exist within the " - "alignments file."))) - argument_list.append(dict( - opts=("-n", "--nfilter"), - action=FilesFullPaths, - filetypes="image", - dest="nfilter", - default=None, - nargs="+", - group=_("Face Processing"), - help=_("Optionally filter out people who you do not wish to process by passing in an " - "image of that person. Should be a front portrait with a single person in the " - "image. Multiple images can be added space separated. NB: Using face filter " - "will significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-f", "--filter"), - action=FilesFullPaths, - filetypes="image", - dest="filter", - default=None, - nargs="+", - group=_("Face Processing"), - help=_("Optionally select people you wish to process by passing in an image of that " - "person. Should be a front portrait with a single person in the image. " - "Multiple images can be added space separated. NB: Using face filter will " - "significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-l", "--ref_threshold"), - action=Slider, - min_max=(0.01, 0.99), - rounding=2, - type=float, - dest="ref_threshold", - default=0.4, - group=_("Face Processing"), - help=_("For use with the optional nfilter/filter files. Threshold for positive face " - "recognition. Lower values are stricter. NB: Using face filter will " - "significantly decrease extraction speed and its accuracy cannot be " - "guaranteed."))) - argument_list.append(dict( - opts=("-j", "--jobs"), - action=Slider, - min_max=(0, 40), - rounding=1, - type=int, - dest="jobs", - default=0, - group=_("settings"), - help=_("The maximum number of parallel processes for performing conversion. " - "Converting images is system RAM heavy so it is possible to run out of memory " - "if you have a lot of processes and not enough RAM to accommodate them all. " - "Setting this to 0 will use the maximum available. No matter what you set " - "this to, it will never attempt to use more processes than are available on " - "your system. If singleprocess is enabled this setting will be ignored."))) - argument_list.append(dict( - opts=("-t", "--trainer"), - type=str.lower, - choices=PluginLoader.get_available_models(), - group=_("settings"), - help=_("[LEGACY] This only needs to be selected if a legacy model is being loaded or " - "if there are multiple models in the model folder"))) - argument_list.append(dict( - opts=("-otf", "--on-the-fly"), - action="store_true", - dest="on_the_fly", - default=False, - group=_("settings"), - help=_("Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " - "alignments file for your destination video. However, if you wish you can " - "generate the alignments on-the-fly by enabling this option. This will use " - "an inferior extraction pipeline and will lead to substandard results. If an " - "alignments file is found, this option will be ignored."))) - argument_list.append(dict( - opts=("-k", "--keep-unchanged"), - action="store_true", - dest="keep_unchanged", - default=False, - group=_("Frame Processing"), - help=_("When used with --frame-ranges outputs the unchanged frames that are not " - "processed instead of discarding them."))) - argument_list.append(dict( - opts=("-s", "--swap-model"), - action="store_true", - dest="swap_model", - default=False, - group=_("settings"), - help=_("Swap the model. Instead converting from of A -> B, converts B -> A"))) - argument_list.append(dict( - opts=("-sp", "--singleprocess"), - action="store_true", - default=False, - group=_("settings"), - help=_("Disable multiprocessing. Slower but less resource intensive."))) - return argument_list - - -class TrainArgs(FaceSwapArgs): - """ Creates the command line arguments for training. """ - - @staticmethod - def get_info(): - """ The information text for the Train command. - - Returns - ------- - str - The information text for the Train command. - """ - return _("Train a model on extracted original (A) and swap (B) faces.\n" - "Training models can take a long time. Anything from 24hrs to over a week\n" - "Model plugins can be configured in the 'Settings' Menu") - - @staticmethod - def get_argument_list(): - """ Returns the argument list for Train arguments. - - Returns - ------- - list - The list of command line options for training - """ - argument_list = list() - argument_list.append(dict( - opts=("-A", "--input-A"), - action=DirFullPaths, - dest="input_a", - required=True, - group=_("faces"), - help=_("Input directory. A directory containing training images for face A. This is " - "the original face, i.e. the face that you want to remove and replace with " - "face B."))) - argument_list.append(dict( - opts=("-B", "--input-B"), - action=DirFullPaths, - dest="input_b", - required=True, - group=_("faces"), - help=_("Input directory. A directory containing training images for face B. This is " - "the swap face, i.e. the face that you want to place onto the head of person " - "A."))) - argument_list.append(dict( - opts=("-m", "--model-dir"), - action=DirFullPaths, - dest="model_dir", - required=True, - group=_("model"), - help=_("Model directory. This is where the training data will be stored. You should " - "always specify a new folder for new models. If starting a new model, select " - "either an empty folder, or a folder which does not exist (which will be " - "created). If continuing to train an existing model, specify the location of " - "the existing model."))) - argument_list.append(dict( - opts=("-l", "--load-weights"), - action=FileFullPaths, - filetypes="model", - dest="load_weights", - required=False, - group=_("model"), - help=_("R|Load the weights from a pre-existing model into a newly created model. " - "For most models this will load weights from the Encoder of the given model " - "into the encoder of the newly created model. Some plugins may have specific " - "configuration options allowing you to load weights from other layers. Weights " - "will only be loaded when creating a new model. This option will be ignored if " - "you are resuming an existing model. Generally you will also want to 'freeze-" - "weights' whilst the rest of your model catches up with your Encoder.\n" - "NB: Weights can only be loaded from models of the same plugin as you intend " - "to train."))) - argument_list.append(dict( - opts=("-t", "--trainer"), - action=Radio, - type=str.lower, - default=PluginLoader.get_default_model(), - choices=PluginLoader.get_available_models(), - group=_("model"), - help=_("R|Select which trainer to use. Trainers can be configured from the Settings " - "menu or the config folder." - "\nL|original: The original model created by /u/deepfakes." - "\nL|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " - "for full dfaker method." - "\nL|dfl-h128: 128px in/out model from deepfacelab" - "\nL|dfl-sae: Adaptable model from deepfacelab" - "\nL|dlight: A lightweight, high resolution DFaker variant." - "\nL|iae: A model that uses intermediate layers to try to get better details" - "\nL|lightweight: A lightweight model for low-end cards. Don't expect great " - "results. Can train as low as 1.6GB with batch size 8." - "\nL|realface: A high detail, dual density model based on DFaker, with " - "customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " - "won't work so well. By andenixa et al. Very configurable." - "\nL|unbalanced: 128px in/out model from andenixa. The autoencoders are " - "unbalanced so B>A swaps won't work so well. Very configurable." - "\nL|villain: 128px in/out model from villainguy. Very resource hungry (You " - "will require a GPU with a fair amount of VRAM). Good for details, but more " - "susceptible to color differences."))) - argument_list.append(dict( - opts=("-su", "--summary"), - action="store_true", - dest="summary", - default=False, - group=_("model"), - help=_("Output a summary of the model and exit. If a model folder is provided then a " - "summary of the saved model is displayed. Otherwise a summary of the model " - "that would be created by the chosen plugin and configuration settings is " - "displayed."))) - argument_list.append(dict( - opts=("-f", "--freeze-weights"), - action="store_true", - dest="freeze_weights", - default=False, - group=_("model"), - help=_("Freeze the weights of the model. Freezing weights means that some of the " - "parameters in the model will no longer continue to learn, but those that are " - "not frozen will continue to learn. For most models, this will freeze the " - "encoder, but some models may have configuration options for freezing other " - "layers."))) - argument_list.append(dict( - opts=("-bs", "--batch-size"), - action=Slider, - min_max=(1, 256), - rounding=1, - type=int, - dest="batch_size", - default=16, - group=_("training"), - help=_("Batch size. This is the number of images processed through the model for each " - "side per iteration. NB: As the model is fed 2 sides at a time, the actual " - "number of images within the model at any one time is double the number that " - "you set here. Larger batches require more GPU RAM."))) - argument_list.append(dict( - opts=("-it", "--iterations"), - action=Slider, - min_max=(0, 5000000), - rounding=20000, - type=int, - default=1000000, - group=_("training"), - help=_("Length of training in iterations. This is only really used for automation. " - "There is no 'correct' number of iterations a model should be trained for. " - "You should stop training when you are happy with the previews. However, if " - "you want the model to stop automatically at a set number of iterations, you " - "can set that value here."))) - argument_list.append(dict( - opts=("-d", "--distributed"), - action="store_true", - default=False, - backend="nvidia", - group=_("training"), - help=_("Use the Tensorflow Mirrored Distrubution Strategy to train on multiple " - "GPUs."))) - argument_list.append(dict( - opts=("-s", "--save-interval"), - action=Slider, - min_max=(10, 1000), - rounding=10, - type=int, - dest="save_interval", - default=250, - group=_("Saving"), - help=_("Sets the number of iterations between each model save."))) - argument_list.append(dict( - opts=("-ss", "--snapshot-interval"), - action=Slider, - min_max=(0, 100000), - rounding=5000, - type=int, - dest="snapshot_interval", - default=25000, - group=_("Saving"), - help=_("Sets the number of iterations before saving a backup snapshot of the model " - "in it's current state. Set to 0 for off."))) - argument_list.append(dict( - opts=("-tia", "--timelapse-input-A"), - action=DirFullPaths, - dest="timelapse_input_a", - default=None, - group=_("timelapse"), - help=_("Optional for creating a timelapse. Timelapse will save an image of your " - "selected faces into the timelapse-output folder at every save iteration. " - "This should be the input folder of 'A' faces that you would like to use for " - "creating the timelapse. You must also supply a --timelapse-output and a " - "--timelapse-input-B parameter."))) - argument_list.append(dict( - opts=("-tib", "--timelapse-input-B"), - action=DirFullPaths, - dest="timelapse_input_b", - default=None, - group=_("timelapse"), - help=_("Optional for creating a timelapse. Timelapse will save an image of your " - "selected faces into the timelapse-output folder at every save iteration. " - "This should be the input folder of 'B' faces that you would like to use for " - "creating the timelapse. You must also supply a --timelapse-output and a " - "--timelapse-input-A parameter."))) - argument_list.append(dict( - opts=("-to", "--timelapse-output"), - action=DirFullPaths, - dest="timelapse_output", - default=None, - group=_("timelapse"), - help=_("Optional for creating a timelapse. Timelapse will save an image of your " - "selected faces into the timelapse-output folder at every save iteration. If " - "the input folders are supplied but no output folder, it will default to your " - "model folder /timelapse/"))) - argument_list.append(dict( - opts=("-ps", "--preview-scale"), - action=Slider, - min_max=(25, 200), - rounding=25, - type=int, - dest="preview_scale", - default=100, - group=_("preview"), - help=_("Percentage amount to scale the preview by. 100%% is the model output size."))) - argument_list.append(dict( - opts=("-p", "--preview"), - action="store_true", - dest="preview", - default=False, - group=_("preview"), - help=_("Show training preview output. in a separate window."))) - argument_list.append(dict( - opts=("-w", "--write-image"), - action="store_true", - dest="write_image", - default=False, - group=_("preview"), - help=_("Writes the training result to a file. The image will be stored in the root " - "of your FaceSwap folder."))) - argument_list.append(dict( - opts=("-nl", "--no-logs"), - action="store_true", - dest="no_logs", - default=False, - group=_("training"), - help=_("Disables TensorBoard logging. NB: Disabling logs means that you will not be " - "able to use the graph or analysis for this session in the GUI."))) - argument_list.append(dict( - opts=("-wl", "--warp-to-landmarks"), - action="store_true", - dest="warp_to_landmarks", - default=False, - group=_("augmentation"), - help=_("Warps training faces to closely matched Landmarks from the opposite face-set " - "rather than randomly warping the face. This is the 'dfaker' way of doing " - "warping."))) - argument_list.append(dict( - opts=("-nf", "--no-flip"), - action="store_true", - dest="no_flip", - default=False, - group=_("augmentation"), - help=_("To effectively learn, a random set of images are flipped horizontally. " - "Sometimes it is desirable for this not to occur. Generally this should be " - "left off except for during 'fit training'."))) - argument_list.append(dict( - opts=("-nac", "--no-augment-color"), - action="store_true", - dest="no_augment_color", - default=False, - group=_("augmentation"), - help=_("Color augmentation helps make the model less susceptible to color " - "differences between the A and B sets, at an increased training time cost. " - "Enable this option to disable color augmentation."))) - argument_list.append(dict( - opts=("-nw", "--no-warp"), - action="store_true", - dest="no_warp", - default=False, - group=_("augmentation"), - help=_("Warping is integral to training the Neural Network. This option should only " - "be enabled towards the very end of training to try to bring out more detail. " - "Think of it as 'fine-tuning'. Enabling this option from the beginning is " - "likely to kill a model and lead to terrible results."))) - return argument_list - - class GuiArgs(FaceSwapArgs): """ Creates the command line arguments for the GUI. """ @staticmethod - def get_argument_list(): + def get_argument_list() -> list[dict[str, T.Any]]: """ Returns the argument list for GUI arguments. Returns @@ -1139,11 +305,11 @@ def get_argument_list(): list The list of command line options for the GUI """ - argument_list = [] - argument_list.append(dict( - opts=("-d", "--debug"), - action="store_true", - dest="debug", - default=False, - help=_("Output to Shell console instead of GUI console"))) + argument_list: list[dict[str, T.Any]] = [] + argument_list.append({ + "opts": ("-d", "--debug"), + "action": "store_true", + "dest": "debug", + "default": False, + "help": _("Output to Shell console instead of GUI console")}) return argument_list diff --git a/lib/cli/args_extract_convert.py b/lib/cli/args_extract_convert.py new file mode 100644 index 0000000000..ad3b4da9de --- /dev/null +++ b/lib/cli/args_extract_convert.py @@ -0,0 +1,747 @@ +#!/usr/bin/env python3 +""" The Command Line Argument options for extracting and converting with faceswap.py """ +import argparse +import gettext +import typing as T + +from lib.utils import get_backend +from plugins.plugin_loader import PluginLoader + +from .actions import (DirFullPaths, DirOrFileFullPaths, DirOrFilesFullPaths, FileFullPaths, + FilesFullPaths, MultiOption, Radio, Slider) +from .args import FaceSwapArgs + + +# LOCALES +_LANG = gettext.translation("lib.cli.args_extract_convert", localedir="locales", fallback=True) +_ = _LANG.gettext + + +class ExtractConvertArgs(FaceSwapArgs): + """ Parent class to capture arguments that will be used in both extract and convert processes. + + Extract and Convert share a fair amount of arguments, so arguments that can be used in both of + these processes should be placed here. + + No further processing is done in this class (this is handled by the children), this just + captures the shared arguments. + """ + + @staticmethod + def get_argument_list() -> list[dict[str, T.Any]]: + """ Returns the argument list for shared Extract and Convert arguments. + + Returns + ------- + list + The list of command line options for the given Extract and Convert + """ + argument_list: list[dict[str, T.Any]] = [] + argument_list.append({ + "opts": ("-i", "--input-dir"), + "action": DirOrFileFullPaths, + "filetypes": "video", + "dest": "input_dir", + "required": True, + "group": _("Data"), + "help": _( + "Input directory or video. Either a directory containing the image files you wish " + "to process or path to a video file. NB: This should be the source video/frames " + "NOT the source faces.")}) + argument_list.append({ + "opts": ("-o", "--output-dir"), + "action": DirFullPaths, + "dest": "output_dir", + "required": True, + "group": _("Data"), + "help": _("Output directory. This is where the converted files will be saved.")}) + argument_list.append({ + "opts": ("-p", "--alignments"), + "action": FileFullPaths, + "filetypes": "alignments", + "type": str, + "dest": "alignments_path", + "group": _("Data"), + "help": _( + "Optional path to an alignments file. Leave blank if the alignments file is at " + "the default location.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-al", ), + "action": FileFullPaths, + "filetypes": "alignments", + "type": str, + "dest": "depr_alignments_al_p", + "help": argparse.SUPPRESS}) + return argument_list + + +class ExtractArgs(ExtractConvertArgs): + """ Creates the command line arguments for extraction. + + This class inherits base options from :class:`ExtractConvertArgs` where arguments that are used + for both Extract and Convert should be placed. + + Commands explicit to Extract should be added in :func:`get_optional_arguments` + """ + + @staticmethod + def get_info() -> str: + """ The information text for the Extract command. + + Returns + ------- + str + The information text for the Extract command. + """ + return _("Extract faces from image or video sources.\n" + "Extraction plugins can be configured in the 'Settings' Menu") + + @staticmethod + def get_optional_arguments() -> list[dict[str, T.Any]]: + """ Returns the argument list unique to the Extract command. + + Returns + ------- + list + The list of optional command line options for the Extract command + """ + if get_backend() == "cpu": + default_detector = "mtcnn" + default_aligner = "cv2-dnn" + else: + default_detector = "s3fd" + default_aligner = "fan" + + argument_list: list[dict[str, T.Any]] = [] + argument_list.append({ + "opts": ("-b", "--batch-mode"), + "action": "store_true", + "dest": "batch_mode", + "default": False, + "group": _("Data"), + "help": _( + "R|If selected then the input_dir should be a parent folder containing multiple " + "videos and/or folders of images you wish to extract from. The faces will be " + "output to separate sub-folders in the output_dir.")}) + argument_list.append({ + "opts": ("-D", "--detector"), + "action": Radio, + "type": str.lower, + "default": default_detector, + "choices": PluginLoader.get_available_extractors("detect"), + "group": _("Plugins"), + "help": _( + "R|Detector to use. Some of these have configurable settings in " + "'/config/extract.ini' or 'Settings > Configure Extract 'Plugins':" + "\nL|cv2-dnn: A CPU only extractor which is the least reliable and least resource " + "intensive. Use this if not using a GPU and time is important." + "\nL|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources than " + "other GPU detectors but can often return more false positives." + "\nL|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " + "fewer false positives than other GPU detectors, but is a lot more resource " + "intensive." + "\nL|external: Import a face detection bounding box from a json file. (" + "configurable in Detect settings)")}) + argument_list.append({ + "opts": ("-A", "--aligner"), + "action": Radio, + "type": str.lower, + "default": default_aligner, + "choices": PluginLoader.get_available_extractors("align"), + "group": _("Plugins"), + "help": _( + "R|Aligner to use." + "\nL|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, but " + "less accurate. Only use this if not using a GPU and time is important." + "\nL|fan: Best aligner. Fast on GPU, slow on CPU." + "\nL|external: Import 68 point 2D landmarks or an aligned bounding box from a " + "json file. (configurable in Align settings)")}) + argument_list.append({ + "opts": ("-M", "--masker"), + "action": MultiOption, + "type": str.lower, + "nargs": "+", + "choices": [mask for mask in PluginLoader.get_available_extractors("mask") + if mask not in ("components", "extended")], + "group": _("Plugins"), + "help": _( + "R|Additional Masker(s) to use. The masks generated here will all take up GPU " + "RAM. You can select none, one or multiple masks, but the extraction may take " + "longer the more you select. NB: The Extended and Components (landmark based) " + "masks are automatically generated on extraction." + "\nL|bisenet-fp: Relatively lightweight NN based mask that provides more refined " + "control over the area to be masked including full head masking (configurable in " + "mask settings)." + "\nL|custom: A dummy mask that fills the mask area with all 1s or 0s (" + "configurable in settings). This is only required if you intend to manually edit " + "the custom masks yourself in the manual tool. This mask does not use the GPU so " + "will not use any additional VRAM." + "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " + "faces clear of obstructions. Profile faces and obstructions may result in " + "sub-par performance." + "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " + "frontal faces. The mask model has been specifically trained to recognize some " + "facial obstructions (hands and eyeglasses). Profile faces may result in sub-par " + "performance." + "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " + "faces. The mask model has been trained by community members and will need " + "testing for further description. Profile faces may result in sub-par " + "performance." + "\nThe auto generated masks are as follows:" + "\nL|components: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks to create a mask." + "\nL|extended: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks and the mask is extended upwards onto the forehead." + "\n(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)")}) + argument_list.append({ + "opts": ("-O", "--normalization"), + "action": Radio, + "type": str.lower, + "dest": "normalization", + "default": "none", + "choices": ["none", "clahe", "hist", "mean"], + "group": _("Plugins"), + "help": _( + "R|Performing normalization can help the aligner better align faces with " + "difficult lighting conditions at an extraction speed cost. Different methods " + "will yield different results on different sets. NB: This does not impact the " + "output face, just the input to the aligner." + "\nL|none: Don't perform normalization on the face." + "\nL|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the face." + "\nL|hist: Equalize the histograms on the RGB channels." + "\nL|mean: Normalize the face colors to the mean.")}) + argument_list.append({ + "opts": ("-R", "--re-feed"), + "action": Slider, + "min_max": (0, 10), + "rounding": 1, + "type": int, + "dest": "re_feed", + "default": 0, + "group": _("Plugins"), + "help": _( + "The number of times to re-feed the detected face into the aligner. Each time the " + "face is re-fed into the aligner the bounding box is adjusted by a small amount. " + "The final landmarks are then averaged from each iteration. Helps to remove " + "'micro-jitter' but at the cost of slower extraction speed. The more times the " + "face is re-fed into the aligner, the less micro-jitter should occur but the " + "longer extraction will take.")}) + argument_list.append({ + "opts": ("-a", "--re-align"), + "action": "store_true", + "dest": "re_align", + "default": False, + "group": _("Plugins"), + "help": _( + "Re-feed the initially found aligned face through the aligner. Can help produce " + "better alignments for faces that are rotated beyond 45 degrees in the frame or " + "are at extreme angles. Slows down extraction.")}) + argument_list.append({ + "opts": ("-r", "--rotate-images"), + "type": str, + "dest": "rotate_images", + "default": None, + "group": _("Plugins"), + "help": _( + "If a face isn't found, rotate the images to try to find a face. Can find more " + "faces at the cost of extraction speed. Pass in a single number to use increments " + "of that size up to 360, or pass in a list of numbers to enumerate exactly what " + "angles to check.")}) + argument_list.append({ + "opts": ("-I", "--identity"), + "action": "store_true", + "default": False, + "group": _("Plugins"), + "help": _( + "Obtain and store face identity encodings from VGGFace2. Slows down extract a " + "little, but will save time if using 'sort by face'")}) + argument_list.append({ + "opts": ("-m", "--min-size"), + "action": Slider, + "min_max": (0, 1080), + "rounding": 20, + "type": int, + "dest": "min_size", + "default": 0, + "group": _("Face Processing"), + "help": _( + "Filters out faces detected below this size. Length, in pixels across the " + "diagonal of the bounding box. Set to 0 for off")}) + argument_list.append({ + "opts": ("-n", "--nfilter"), + "action": DirOrFilesFullPaths, + "filetypes": "image", + "dest": "nfilter", + "default": None, + "nargs": "+", + "group": _("Face Processing"), + "help": _( + "Optionally filter out people who you do not wish to extract by passing in images " + "of those people. Should be a small variety of images at different angles and in " + "different conditions. A folder containing the required images or multiple image " + "files, space separated, can be selected.")}) + argument_list.append({ + "opts": ("-f", "--filter"), + "action": DirOrFilesFullPaths, + "filetypes": "image", + "dest": "filter", + "default": None, + "nargs": "+", + "group": _("Face Processing"), + "help": _( + "Optionally select people you wish to extract by passing in images of that " + "person. Should be a small variety of images at different angles and in different " + "conditions A folder containing the required images or multiple image files, " + "space separated, can be selected.")}) + argument_list.append({ + "opts": ("-l", "--ref_threshold"), + "action": Slider, + "min_max": (0.01, 0.99), + "rounding": 2, + "type": float, + "dest": "ref_threshold", + "default": 0.60, + "group": _("Face Processing"), + "help": _( + "For use with the optional nfilter/filter files. Threshold for positive face " + "recognition. Higher values are stricter.")}) + argument_list.append({ + "opts": ("-z", "--size"), + "action": Slider, + "min_max": (256, 1024), + "rounding": 64, + "type": int, + "default": 512, + "group": _("output"), + "help": _( + "The output size of extracted faces. Make sure that the model you intend to train " + "supports your required size. This will only need to be changed for hi-res " + "models.")}) + argument_list.append({ + "opts": ("-N", "--extract-every-n"), + "action": Slider, + "min_max": (1, 100), + "rounding": 1, + "type": int, + "dest": "extract_every_n", + "default": 1, + "group": _("output"), + "help": _( + "Extract every 'nth' frame. This option will skip frames when extracting faces. " + "For example a value of 1 will extract faces from every frame, a value of 10 will " + "extract faces from every 10th frame.")}) + argument_list.append({ + "opts": ("-v", "--save-interval"), + "action": Slider, + "min_max": (0, 1000), + "rounding": 10, + "type": int, + "dest": "save_interval", + "default": 0, + "group": _("output"), + "help": _( + "Automatically save the alignments file after a set amount of frames. By default " + "the alignments file is only saved at the end of the extraction process. NB: If " + "extracting in 2 passes then the alignments file will only start to be saved out " + "during the second pass. WARNING: Don't interrupt the script when writing the " + "file because it might get corrupted. Set to 0 to turn off")}) + argument_list.append({ + "opts": ("-B", "--debug-landmarks"), + "action": "store_true", + "dest": "debug_landmarks", + "default": False, + "group": _("output"), + "help": _("Draw landmarks on the ouput faces for debugging purposes.")}) + argument_list.append({ + "opts": ("-P", "--singleprocess"), + "action": "store_true", + "default": False, + "backend": ("nvidia", "directml", "rocm", "apple_silicon"), + "group": _("settings"), + "help": _( + "Don't run extraction in parallel. Will run each part of the extraction process " + "separately (one after the other) rather than all at the same time. Useful if " + "VRAM is at a premium.")}) + argument_list.append({ + "opts": ("-s", "--skip-existing"), + "action": "store_true", + "dest": "skip_existing", + "default": False, + "group": _("settings"), + "help": _( + "Skips frames that have already been extracted and exist in the alignments file")}) + argument_list.append({ + "opts": ("-e", "--skip-existing-faces"), + "action": "store_true", + "dest": "skip_faces", + "default": False, + "group": _("settings"), + "help": _("Skip frames that already have detected faces in the alignments file")}) + argument_list.append({ + "opts": ("-K", "--skip-saving-faces"), + "action": "store_true", + "dest": "skip_saving_faces", + "default": False, + "group": _("settings"), + "help": _("Skip saving the detected faces to disk. Just create an alignments file")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-min", ), + "type": int, + "dest": "depr_min-size_min_m", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-een", ), + "type": int, + "dest": "depr_extract-every-n_een_N", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-nm",), + "type": str.lower, + "dest": "depr_normalization_nm_O", + "choices": ["none", "clahe", "hist", "mean"], + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-rf", ), + "type": int, + "dest": "depr_re-feed_rf_R", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-sz", ), + "type": int, + "dest": "depr_size_sz_z", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-si", ), + "type": int, + "dest": "depr_save-interval_si_v", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-dl", ), + "action": "store_true", + "dest": "depr_debug-landmarks_dl_B", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-sp", ), + "dest": "depr_singleprocess_sp_P", + "action": "store_true", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-sf", ), + "action": "store_true", + "dest": "depr_skip-existing-faces_sf_e", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-ssf", ), + "action": "store_true", + "dest": "depr_skip-saving-faces_ssf_K", + "help": argparse.SUPPRESS}) + return argument_list + + +class ConvertArgs(ExtractConvertArgs): + """ Creates the command line arguments for conversion. + + This class inherits base options from :class:`ExtractConvertArgs` where arguments that are used + for both Extract and Convert should be placed. + + Commands explicit to Convert should be added in :func:`get_optional_arguments` + """ + + @staticmethod + def get_info() -> str: + """ The information text for the Convert command. + + Returns + ------- + str + The information text for the Convert command. + """ + return _("Swap the original faces in a source video/images to your final faces.\n" + "Conversion plugins can be configured in the 'Settings' Menu") + + @staticmethod + def get_optional_arguments() -> list[dict[str, T.Any]]: + """ Returns the argument list unique to the Convert command. + + Returns + ------- + list + The list of optional command line options for the Convert command + """ + + argument_list: list[dict[str, T.Any]] = [] + argument_list.append({ + "opts": ("-r", "--reference-video"), + "action": FileFullPaths, + "filetypes": "video", + "type": str, + "dest": "reference_video", + "group": _("Data"), + "help": _( + "Only required if converting from images to video. Provide The original video " + "that the source frames were extracted from (for extracting the fps and audio).")}) + argument_list.append({ + "opts": ("-m", "--model-dir"), + "action": DirFullPaths, + "dest": "model_dir", + "required": True, + "group": _("Data"), + "help": _( + "Model directory. The directory containing the trained model you wish to use for " + "conversion.")}) + argument_list.append({ + "opts": ("-c", "--color-adjustment"), + "action": Radio, + "type": str.lower, + "dest": "color_adjustment", + "default": "avg-color", + "choices": PluginLoader.get_available_convert_plugins("color", True), + "group": _("Plugins"), + "help": _( + "R|Performs color adjustment to the swapped face. Some of these options have " + "configurable settings in '/config/convert.ini' or 'Settings > Configure Convert " + "Plugins':" + "\nL|avg-color: Adjust the mean of each color channel in the swapped " + "reconstruction to equal the mean of the masked area in the original image." + "\nL|color-transfer: Transfers the color distribution from the source to the " + "target image using the mean and standard deviations of the L*a*b* color space." + "\nL|manual-balance: Manually adjust the balance of the image in a variety of " + "color spaces. Best used with the Preview tool to set correct values." + "\nL|match-hist: Adjust the histogram of each color channel in the swapped " + "reconstruction to equal the histogram of the masked area in the original image." + "\nL|seamless-clone: Use cv2's seamless clone function to remove extreme " + "gradients at the mask seam by smoothing colors. Generally does not give very " + "satisfactory results." + "\nL|none: Don't perform color adjustment.")}) + argument_list.append({ + "opts": ("-M", "--mask-type"), + "action": Radio, + "type": str.lower, + "dest": "mask_type", + "default": "extended", + "choices": PluginLoader.get_available_extractors("mask", + add_none=True, + extend_plugin=True) + ["predicted"], + "group": _("Plugins"), + "help": _( + "R|Masker to use. NB: The mask you require must exist within the alignments file. " + "You can add additional masks with the Mask Tool." + "\nL|none: Don't use a mask." + "\nL|bisenet-fp_face: Relatively lightweight NN based mask that provides more " + "refined control over the area to be masked (configurable in mask settings). Use " + "this version of bisenet-fp if your model is trained with 'face' or " + "'legacy' centering." + "\nL|bisenet-fp_head: Relatively lightweight NN based mask that provides more " + "refined control over the area to be masked (configurable in mask settings). Use " + "this version of bisenet-fp if your model is trained with 'head' centering." + "\nL|custom_face: Custom user created, face centered mask." + "\nL|custom_head: Custom user created, head centered mask." + "\nL|components: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks to create a mask." + "\nL|extended: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks and the mask is extended upwards onto the forehead." + "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " + "faces clear of obstructions. Profile faces and obstructions may result in sub-" + "par performance." + "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " + "frontal faces. The mask model has been specifically trained to recognize some " + "facial obstructions (hands and eyeglasses). Profile faces may result in sub-par " + "performance." + "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " + "faces. The mask model has been trained by community members and will need " + "testing for further description. Profile faces may result in sub-par " + "performance." + "\nL|predicted: If the 'Learn Mask' option was enabled during training, this will " + "use the mask that was created by the trained model.")}) + argument_list.append({ + "opts": ("-w", "--writer"), + "action": Radio, + "type": str, + "default": "opencv", + "choices": PluginLoader.get_available_convert_plugins("writer", False), + "group": _("Plugins"), + "help": _( + "R|The plugin to use to output the converted images. The writers are configurable " + "in '/config/convert.ini' or 'Settings > Configure Convert Plugins:'" + "\nL|ffmpeg: [video] Writes out the convert straight to video. When the input is " + "a series of images then the '-ref' (--reference-video) parameter must be set." + "\nL|gif: [animated image] Create an animated gif." + "\nL|opencv: [images] The fastest image writer, but less options and formats than " + "other plugins." + "\nL|patch: [images] Outputs the raw swapped face patch, along with the " + "transformation matrix required to re-insert the face back into the original " + "frame. Use this option if you wish to post-process and composite the final face " + "within external tools." + "\nL|pillow: [images] Slower than opencv, but has more options and supports more " + "formats.")}) + argument_list.append({ + "opts": ("-O", "--output-scale"), + "action": Slider, + "min_max": (25, 400), + "rounding": 1, + "type": int, + "dest": "output_scale", + "default": 100, + "group": _("Frame Processing"), + "help": _( + "Scale the final output frames by this amount. 100%% will output the frames at " + "source dimensions. 50%% at half size 200%% at double size")}) + argument_list.append({ + "opts": ("-R", "--frame-ranges"), + "type": str, + "nargs": "+", + "dest": "frame_ranges", + "group": _("Frame Processing"), + "help": _( + "Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " + "--frame-ranges 10-50 90-100. Frames falling outside of the selected range will " + "be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " + "converting from images, then the filenames must end with the frame-number!")}) + argument_list.append({ + "opts": ("-S", "--face-scale"), + "action": Slider, + "min_max": (-10.0, 10.0), + "rounding": 2, + "dest": "face_scale", + "type": float, + "default": 0.0, + "group": _("Face Processing"), + "help": _( + "Scale the swapped face by this percentage. Positive values will enlarge the " + "face, Negative values will shrink the face.")}) + argument_list.append({ + "opts": ("-a", "--input-aligned-dir"), + "action": DirFullPaths, + "dest": "input_aligned_dir", + "default": None, + "group": _("Face Processing"), + "help": _( + "If you have not cleansed your alignments file, then you can filter out faces by " + "defining a folder here that contains the faces extracted from your input files/" + "video. If this folder is defined, then only faces that exist within your " + "alignments file and also exist within the specified folder will be converted. " + "Leaving this blank will convert all faces that exist within the alignments " + "file.")}) + argument_list.append({ + "opts": ("-n", "--nfilter"), + "action": FilesFullPaths, + "filetypes": "image", + "dest": "nfilter", + "default": None, + "nargs": "+", + "group": _("Face Processing"), + "help": _( + "Optionally filter out people who you do not wish to process by passing in an " + "image of that person. Should be a front portrait with a single person in the " + "image. Multiple images can be added space separated. NB: Using face filter will " + "significantly decrease extraction speed and its accuracy cannot be guaranteed.")}) + argument_list.append({ + "opts": ("-f", "--filter"), + "action": FilesFullPaths, + "filetypes": "image", + "dest": "filter", + "default": None, + "nargs": "+", + "group": _("Face Processing"), + "help": _( + "Optionally select people you wish to process by passing in an image of that " + "person. Should be a front portrait with a single person in the image. Multiple " + "images can be added space separated. NB: Using face filter will significantly " + "decrease extraction speed and its accuracy cannot be guaranteed.")}) + argument_list.append({ + "opts": ("-l", "--ref_threshold"), + "action": Slider, + "min_max": (0.01, 0.99), + "rounding": 2, + "type": float, + "dest": "ref_threshold", + "default": 0.4, + "group": _("Face Processing"), + "help": _( + "For use with the optional nfilter/filter files. Threshold for positive face " + "recognition. Lower values are stricter. NB: Using face filter will significantly " + "decrease extraction speed and its accuracy cannot be guaranteed.")}) + argument_list.append({ + "opts": ("-j", "--jobs"), + "action": Slider, + "min_max": (0, 40), + "rounding": 1, + "type": int, + "dest": "jobs", + "default": 0, + "group": _("settings"), + "help": _( + "The maximum number of parallel processes for performing conversion. Converting " + "images is system RAM heavy so it is possible to run out of memory if you have a " + "lot of processes and not enough RAM to accommodate them all. Setting this to 0 " + "will use the maximum available. No matter what you set this to, it will never " + "attempt to use more processes than are available on your system. If " + "singleprocess is enabled this setting will be ignored.")}) + argument_list.append({ + "opts": ("-T", "--on-the-fly"), + "action": "store_true", + "dest": "on_the_fly", + "default": False, + "group": _("settings"), + "help": _( + "Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " + "alignments file for your destination video. However, if you wish you can " + "generate the alignments on-the-fly by enabling this option. This will use an " + "inferior extraction pipeline and will lead to substandard results. If an " + "alignments file is found, this option will be ignored.")}) + argument_list.append({ + "opts": ("-k", "--keep-unchanged"), + "action": "store_true", + "dest": "keep_unchanged", + "default": False, + "group": _("Frame Processing"), + "help": _( + "When used with --frame-ranges outputs the unchanged frames that are not " + "processed instead of discarding them.")}) + argument_list.append({ + "opts": ("-s", "--swap-model"), + "action": "store_true", + "dest": "swap_model", + "default": False, + "group": _("settings"), + "help": _("Swap the model. Instead converting from of A -> B, converts B -> A")}) + argument_list.append({ + "opts": ("-P", "--singleprocess"), + "action": "store_true", + "default": False, + "group": _("settings"), + "help": _("Disable multiprocessing. Slower but less resource intensive.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-sp", ), + "action": "store_true", + "dest": "depr_singleprocess_sp_P", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-ref", ), + "type": str, + "dest": "depr_reference-video_ref_r", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-fr", ), + "type": str, + "nargs": "+", + "dest": "depr_frame-ranges_fr_R", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-osc", ), + "type": int, + "dest": "depr_output-scale_osc_O", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-otf", ), + "action": "store_true", + "dest": "depr_on-the-fly_otf_T", + "help": argparse.SUPPRESS}) + return argument_list diff --git a/lib/cli/args_train.py b/lib/cli/args_train.py new file mode 100644 index 0000000000..efbaa93c5b --- /dev/null +++ b/lib/cli/args_train.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python3 +""" The Command Line Argument options for training with faceswap.py """ +import argparse +import gettext +import typing as T + +from plugins.plugin_loader import PluginLoader + +from .actions import DirFullPaths, FileFullPaths, Radio, Slider +from .args import FaceSwapArgs + + +# LOCALES +_LANG = gettext.translation("lib.cli.args_train", localedir="locales", fallback=True) +_ = _LANG.gettext + + +class TrainArgs(FaceSwapArgs): + """ Creates the command line arguments for training. """ + + @staticmethod + def get_info() -> str: + """ The information text for the Train command. + + Returns + ------- + str + The information text for the Train command. + """ + return _("Train a model on extracted original (A) and swap (B) faces.\n" + "Training models can take a long time. Anything from 24hrs to over a week\n" + "Model plugins can be configured in the 'Settings' Menu") + + @staticmethod + def get_argument_list() -> list[dict[str, T.Any]]: + """ Returns the argument list for Train arguments. + + Returns + ------- + list + The list of command line options for training + """ + argument_list: list[dict[str, T.Any]] = [] + argument_list.append({ + "opts": ("-A", "--input-A"), + "action": DirFullPaths, + "dest": "input_a", + "required": True, + "group": _("faces"), + "help": _( + "Input directory. A directory containing training images for face A. This is the " + "original face, i.e. the face that you want to remove and replace with face B.")}) + argument_list.append({ + "opts": ("-B", "--input-B"), + "action": DirFullPaths, + "dest": "input_b", + "required": True, + "group": _("faces"), + "help": _( + "Input directory. A directory containing training images for face B. This is the " + "swap face, i.e. the face that you want to place onto the head of person A.")}) + argument_list.append({ + "opts": ("-m", "--model-dir"), + "action": DirFullPaths, + "dest": "model_dir", + "required": True, + "group": _("model"), + "help": _( + "Model directory. This is where the training data will be stored. You should " + "always specify a new folder for new models. If starting a new model, select " + "either an empty folder, or a folder which does not exist (which will be " + "created). If continuing to train an existing model, specify the location of the " + "existing model.")}) + argument_list.append({ + "opts": ("-l", "--load-weights"), + "action": FileFullPaths, + "filetypes": "model", + "dest": "load_weights", + "required": False, + "group": _("model"), + "help": _( + "R|Load the weights from a pre-existing model into a newly created model. For " + "most models this will load weights from the Encoder of the given model into the " + "encoder of the newly created model. Some plugins may have specific configuration " + "options allowing you to load weights from other layers. Weights will only be " + "loaded when creating a new model. This option will be ignored if you are " + "resuming an existing model. Generally you will also want to 'freeze-weights' " + "whilst the rest of your model catches up with your Encoder.\n" + "NB: Weights can only be loaded from models of the same plugin as you intend to " + "train.")}) + argument_list.append({ + "opts": ("-t", "--trainer"), + "action": Radio, + "type": str.lower, + "default": PluginLoader.get_default_model(), + "choices": PluginLoader.get_available_models(), + "group": _("model"), + "help": _( + "R|Select which trainer to use. Trainers can be configured from the Settings menu " + "or the config folder." + "\nL|original: The original model created by /u/deepfakes." + "\nL|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' for " + "full dfaker method." + "\nL|dfl-h128: 128px in/out model from deepfacelab" + "\nL|dfl-sae: Adaptable model from deepfacelab" + "\nL|dlight: A lightweight, high resolution DFaker variant." + "\nL|iae: A model that uses intermediate layers to try to get better details" + "\nL|lightweight: A lightweight model for low-end cards. Don't expect great " + "results. Can train as low as 1.6GB with batch size 8." + "\nL|realface: A high detail, dual density model based on DFaker, with " + "customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " + "won't work so well. By andenixa et al. Very configurable." + "\nL|unbalanced: 128px in/out model from andenixa. The autoencoders are " + "unbalanced so B>A swaps won't work so well. Very configurable." + "\nL|villain: 128px in/out model from villainguy. Very resource hungry (You will " + "require a GPU with a fair amount of VRAM). Good for details, but more " + "susceptible to color differences.")}) + argument_list.append({ + "opts": ("-u", "--summary"), + "action": "store_true", + "dest": "summary", + "default": False, + "group": _("model"), + "help": _( + "Output a summary of the model and exit. If a model folder is provided then a " + "summary of the saved model is displayed. Otherwise a summary of the model that " + "would be created by the chosen plugin and configuration settings is displayed.")}) + argument_list.append({ + "opts": ("-f", "--freeze-weights"), + "action": "store_true", + "dest": "freeze_weights", + "default": False, + "group": _("model"), + "help": _( + "Freeze the weights of the model. Freezing weights means that some of the " + "parameters in the model will no longer continue to learn, but those that are not " + "frozen will continue to learn. For most models, this will freeze the encoder, " + "but some models may have configuration options for freezing other layers.")}) + argument_list.append({ + "opts": ("-b", "--batch-size"), + "action": Slider, + "min_max": (1, 256), + "rounding": 1, + "type": int, + "dest": "batch_size", + "default": 16, + "group": _("training"), + "help": _( + "Batch size. This is the number of images processed through the model for each " + "side per iteration. NB: As the model is fed 2 sides at a time, the actual number " + "of images within the model at any one time is double the number that you set " + "here. Larger batches require more GPU RAM.")}) + argument_list.append({ + "opts": ("-i", "--iterations"), + "action": Slider, + "min_max": (0, 5000000), + "rounding": 20000, + "type": int, + "default": 1000000, + "group": _("training"), + "help": _( + "Length of training in iterations. This is only really used for automation. There " + "is no 'correct' number of iterations a model should be trained for. You should " + "stop training when you are happy with the previews. However, if you want the " + "model to stop automatically at a set number of iterations, you can set that " + "value here.")}) + argument_list.append({ + "opts": ("-D", "--distribution-strategy"), + "dest": "distribution_strategy", + "action": Radio, + "type": str.lower, + "choices": ["default", "central-storage", "mirrored"], + "default": "default", + "backend": ("nvidia", "directml", "rocm", "apple_silicon"), + "group": _("training"), + "help": _( + "R|Select the distribution stategy to use." + "\nL|default: Use Tensorflow's default distribution strategy." + "\nL|central-storage: Centralizes variables on the CPU whilst operations are " + "performed on 1 or more local GPUs. This can help save some VRAM at the cost of " + "some speed by not storing variables on the GPU. Note: Mixed-Precision is not " + "supported on multi-GPU setups." + "\nL|mirrored: Supports synchronous distributed training across multiple local " + "GPUs. A copy of the model and all variables are loaded onto each GPU with " + "batches distributed to each GPU at each iteration.")}) + argument_list.append({ + "opts": ("-n", "--no-logs"), + "action": "store_true", + "dest": "no_logs", + "default": False, + "group": _("training"), + "help": _( + "Disables TensorBoard logging. NB: Disabling logs means that you will not be able " + "to use the graph or analysis for this session in the GUI.")}) + argument_list.append({ + "opts": ("-r", "--use-lr-finder"), + "action": "store_true", + "dest": "use_lr_finder", + "default": False, + "group": _("training"), + "help": _( + "Use the Learning Rate Finder to discover the optimal learning rate for training. " + "For new models, this will calculate the optimal learning rate for the model. For " + "existing models this will use the optimal learning rate that was discovered when " + "initializing the model. Setting this option will ignore the manually configured " + "learning rate (configurable in train settings).")}) + argument_list.append({ + "opts": ("-s", "--save-interval"), + "action": Slider, + "min_max": (10, 1000), + "rounding": 10, + "type": int, + "dest": "save_interval", + "default": 250, + "group": _("Saving"), + "help": _("Sets the number of iterations between each model save.")}) + argument_list.append({ + "opts": ("-I", "--snapshot-interval"), + "action": Slider, + "min_max": (0, 100000), + "rounding": 5000, + "type": int, + "dest": "snapshot_interval", + "default": 25000, + "group": _("Saving"), + "help": _( + "Sets the number of iterations before saving a backup snapshot of the model in " + "it's current state. Set to 0 for off.")}) + argument_list.append({ + "opts": ("-x", "--timelapse-input-A"), + "action": DirFullPaths, + "dest": "timelapse_input_a", + "default": None, + "group": _("timelapse"), + "help": _( + "Optional for creating a timelapse. Timelapse will save an image of your selected " + "faces into the timelapse-output folder at every save iteration. This should be " + "the input folder of 'A' faces that you would like to use for creating the " + "timelapse. You must also supply a --timelapse-output and a --timelapse-input-B " + "parameter.")}) + argument_list.append({ + "opts": ("-y", "--timelapse-input-B"), + "action": DirFullPaths, + "dest": "timelapse_input_b", + "default": None, + "group": _("timelapse"), + "help": _( + "Optional for creating a timelapse. Timelapse will save an image of your selected " + "faces into the timelapse-output folder at every save iteration. This should be " + "the input folder of 'B' faces that you would like to use for creating the " + "timelapse. You must also supply a --timelapse-output and a --timelapse-input-A " + "parameter.")}) + argument_list.append({ + "opts": ("-z", "--timelapse-output"), + "action": DirFullPaths, + "dest": "timelapse_output", + "default": None, + "group": _("timelapse"), + "help": _( + "Optional for creating a timelapse. Timelapse will save an image of your selected " + "faces into the timelapse-output folder at every save iteration. If the input " + "folders are supplied but no output folder, it will default to your model folder/" + "timelapse/")}) + argument_list.append({ + "opts": ("-p", "--preview"), + "action": "store_true", + "dest": "preview", + "default": False, + "group": _("preview"), + "help": _("Show training preview output. in a separate window.")}) + argument_list.append({ + "opts": ("-w", "--write-image"), + "action": "store_true", + "dest": "write_image", + "default": False, + "group": _("preview"), + "help": _( + "Writes the training result to a file. The image will be stored in the root of " + "your FaceSwap folder.")}) + argument_list.append({ + "opts": ("-M", "--warp-to-landmarks"), + "action": "store_true", + "dest": "warp_to_landmarks", + "default": False, + "group": _("augmentation"), + "help": _( + "Warps training faces to closely matched Landmarks from the opposite face-set " + "rather than randomly warping the face. This is the 'dfaker' way of doing " + "warping.")}) + argument_list.append({ + "opts": ("-P", "--no-flip"), + "action": "store_true", + "dest": "no_flip", + "default": False, + "group": _("augmentation"), + "help": _( + "To effectively learn, a random set of images are flipped horizontally. Sometimes " + "it is desirable for this not to occur. Generally this should be left off except " + "for during 'fit training'.")}) + argument_list.append({ + "opts": ("-c", "--no-augment-color"), + "action": "store_true", + "dest": "no_augment_color", + "default": False, + "group": _("augmentation"), + "help": _( + "Color augmentation helps make the model less susceptible to color differences " + "between the A and B sets, at an increased training time cost. Enable this option " + "to disable color augmentation.")}) + argument_list.append({ + "opts": ("-W", "--no-warp"), + "action": "store_true", + "dest": "no_warp", + "default": False, + "group": _("augmentation"), + "help": _( + "Warping is integral to training the Neural Network. This option should only be " + "enabled towards the very end of training to try to bring out more detail. Think " + "of it as 'fine-tuning'. Enabling this option from the beginning is likely to " + "kill a model and lead to terrible results.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-su", ), + "action": "store_true", + "dest": "depr_summary_su_u", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-bs", ), + "type": int, + "dest": "depr_batch-size_bs_b", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-it", ), + "type": int, + "dest": "depr_iterations_it_i", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-nl", ), + "action": "store_true", + "dest": "depr_no-logs_nl_n", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-ss", ), + "type": int, + "dest": "depr_snapshot-interval_ss_I", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-tia", ), + "type": str, + "dest": "depr_timelapse-input-A_tia_x", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-tib", ), + "type": str, + "dest": "depr_timelapse-input-B_tib_y", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-to", ), + "type": str, + "dest": "depr_timelapse-output_to_z", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-wl", ), + "action": "store_true", + "dest": "depr_warp-to-landmarks_wl_M", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-nf", ), + "action": "store_true", + "dest": "depr_no-flip_nf_P", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-nac", ), + "action": "store_true", + "dest": "depr_no-augment-color_nac_c", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-nw", ), + "action": "store_true", + "dest": "depr_no-warp_nw_W", + "help": argparse.SUPPRESS}) + return argument_list diff --git a/lib/cli/launcher.py b/lib/cli/launcher.py index feb352435b..8c15a25815 100644 --- a/lib/cli/launcher.py +++ b/lib/cli/launcher.py @@ -1,21 +1,27 @@ #!/usr/bin/env python3 """ Launches the correct script with the given Command Line Arguments """ +from __future__ import annotations import logging import os import platform import sys +import typing as T from importlib import import_module from lib.gpu_stats import set_exclude_devices, GPUStats from lib.logger import crash_log, log_setup -from lib.utils import (FaceswapError, get_backend, get_tf_version, safe_shutdown, - set_backend, set_system_verbosity) +from lib.utils import (FaceswapError, get_backend, get_tf_version, + safe_shutdown, set_backend, set_system_verbosity) -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + import argparse + from collections.abc import Callable +logger = logging.getLogger(__name__) -class ScriptExecutor(): # pylint:disable=too-few-public-methods + +class ScriptExecutor(): """ Loads the relevant script modules and executes the script. This class is initialized in each of the argparsers for the relevant @@ -27,10 +33,10 @@ class ScriptExecutor(): # pylint:disable=too-few-public-methods command: str The faceswap command that is being executed """ - def __init__(self, command): + def __init__(self, command: str) -> None: self._command = command.lower() - def _import_script(self): + def _import_script(self) -> Callable: """ Imports the relevant script as indicated by :attr:`_command` from the scripts folder. Returns @@ -38,6 +44,7 @@ def _import_script(self): class: Faceswap Script The uninitialized script from the faceswap scripts folder. """ + self._set_environment_variables() self._test_for_tf_version() self._test_for_gui() cmd = os.path.basename(sys.argv[0]) @@ -47,21 +54,56 @@ def _import_script(self): script = getattr(module, self._command.title()) return script - def _test_for_tf_version(self): + def _set_environment_variables(self) -> None: + """ Set the number of threads that numexpr can use and TF environment variables. """ + # Allocate a decent number of threads to numexpr to suppress warnings + cpu_count = os.cpu_count() + allocate = cpu_count - cpu_count // 3 if cpu_count is not None else 1 + if "OMP_NUM_THREADS" in os.environ: + # If this is set above NUMEXPR_MAX_THREADS, numexpr will error. + # ref: https://github.com/pydata/numexpr/issues/322 + os.environ.pop("OMP_NUM_THREADS") + os.environ["NUMEXPR_MAX_THREADS"] = str(max(1, allocate)) + + # Ensure tensorflow doesn't pin all threads to one core when using Math Kernel Library + os.environ["TF_MIN_GPU_MULTIPROCESSOR_COUNT"] = "4" + os.environ["KMP_AFFINITY"] = "disabled" + + # If running under CPU on Windows, the following error can be encountered: + # OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5 already initialized. + # OMP: Hint This means that multiple copies of the OpenMP runtime have been linked into + # the program. That is dangerous, since it can degrade performance or cause incorrect + # results. The best thing to do is to ensure that only a single OpenMP runtime is linked + # into the process, e.g. by avoiding static linking of the OpenMP runtime in any library. + # As an unsafe, unsupported, undocumented workaround you can set the environment variable + # KMP_DUPLICATE_LIB_OK=TRUE to allow the program to continue to execute, but that may cause + # crashes or silently produce incorrect results. For more information, + # please see http://www.intel.com/software/products/support/. + # + # TODO find a better way than just allowing multiple libs + if get_backend() == "cpu" and platform.system() == "Windows": + logger.debug("Setting `KMP_DUPLICATE_LIB_OK` environment variable to `TRUE`") + os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" + + # There is a memory leak in TF2.10+ predict function. This fix will work for tf2.10 but not + # for later versions. This issue has been patched recently, but we'll probably need to + # skip some TF versions + # ref: https://github.com/tensorflow/tensorflow/issues/58676 + # TODO remove this fix post TF2.10 and check memleak is fixed + logger.debug("Setting TF_RUN_EAGER_OP_AS_FUNCTION env var to False") + os.environ["TF_RUN_EAGER_OP_AS_FUNCTION"] = "false" + + def _test_for_tf_version(self) -> None: """ Check that the required Tensorflow version is installed. Raises ------ FaceswapError - If Tensorflow is not found, or is not between versions 2.4 and 2.8 + If Tensorflow is not found, or is not between versions 2.4 and 2.9 """ - amd_ver = 2.2 - min_ver = 2.4 - max_ver = 2.8 + min_ver = (2, 10) + max_ver = (2, 10) try: - # Ensure tensorflow doesn't pin all threads to one core when using Math Kernel Library - os.environ["TF_MIN_GPU_MULTIPROCESSOR_COUNT"] = "4" - os.environ["KMP_AFFINITY"] = "disabled" import tensorflow as tf # noqa pylint:disable=import-outside-toplevel,unused-import except ImportError as err: if "DLL load failed while importing" in str(err): @@ -79,23 +121,18 @@ def _test_for_tf_version(self): self._handle_import_error(msg) tf_ver = get_tf_version() - backend = get_backend() - if backend != "amd" and tf_ver < min_ver: + if tf_ver < min_ver: msg = (f"The minimum supported Tensorflow is version {min_ver} but you have version " f"{tf_ver} installed. Please upgrade Tensorflow.") self._handle_import_error(msg) - if backend != "amd" and tf_ver > max_ver: + if tf_ver > max_ver: msg = (f"The maximum supported Tensorflow is version {max_ver} but you have version " f"{tf_ver} installed. Please downgrade Tensorflow.") self._handle_import_error(msg) - if backend == "amd" and tf_ver != amd_ver: - msg = (f"The supported Tensorflow version for AMD cards is {amd_ver} but you have " - "version {tf_ver} installed. Please install the correct version.") - self._handle_import_error(msg) logger.debug("Installed Tensorflow Version: %s", tf_ver) @classmethod - def _handle_import_error(cls, message): + def _handle_import_error(cls, message: str) -> None: """ Display the error message to the console and wait for user input to dismiss it, if running GUI under Windows, otherwise use standard error handling. @@ -112,15 +149,15 @@ def _handle_import_error(cls, message): else: raise FaceswapError(message) - def _test_for_gui(self): + def _test_for_gui(self) -> None: """ If running the gui, performs check to ensure necessary prerequisites are present. """ if self._command != "gui": return self._test_tkinter() self._check_display() - @staticmethod - def _test_tkinter(): + @classmethod + def _test_tkinter(cls) -> None: """ If the user is running the GUI, test whether the tkinter app is available on their machine. If not exit gracefully. @@ -133,7 +170,7 @@ def _test_tkinter(): If tkinter cannot be imported """ try: - import tkinter # noqa pylint: disable=unused-import,import-outside-toplevel + import tkinter # noqa pylint:disable=unused-import,import-outside-toplevel except ImportError as err: logger.error("It looks like TkInter isn't installed for your OS, so the GUI has been " "disabled. To enable the GUI please install the TkInter application. You " @@ -147,8 +184,8 @@ def _test_tkinter(): logger.info("Fedora: sudo dnf install python3-tkinter") raise FaceswapError("TkInter not found") from err - @staticmethod - def _check_display(): + @classmethod + def _check_display(cls) -> None: """ Check whether there is a display to output the GUI to. If running on Windows then it is assumed that we are not running in headless mode @@ -164,7 +201,7 @@ def _check_display(): "See https://support.apple.com/en-gb/HT201341") raise FaceswapError("No display detected. GUI mode has been disabled.") - def execute_script(self, arguments): + def execute_script(self, arguments: argparse.Namespace) -> None: """ Performs final set up and launches the requested :attr:`_command` with the given command line arguments. @@ -190,11 +227,11 @@ def execute_script(self, arguments): except FaceswapError as err: for line in str(err).splitlines(): logger.error(line) - except KeyboardInterrupt: # pylint: disable=try-except-raise + except KeyboardInterrupt: # pylint:disable=try-except-raise raise except SystemExit: pass - except Exception: # pylint: disable=broad-except + except Exception: # pylint:disable=broad-except crash_file = crash_log() logger.exception("Got Exception on main handler:") logger.critical("An unexpected crash has occurred. Crash report written to '%s'. " @@ -205,7 +242,7 @@ def execute_script(self, arguments): finally: safe_shutdown(got_error=not success) - def _configure_backend(self, arguments): + def _configure_backend(self, arguments: argparse.Namespace) -> None: """ Configure the backend. Exclude any GPUs for use by Faceswap when requested. @@ -217,8 +254,8 @@ def _configure_backend(self, arguments): arguments: :class:`argparse.Namespace` The command line arguments passed to Faceswap. """ - if get_backend() == "cpu": - # Cpu backends will not have this attribute + if not hasattr(arguments, "exclude_gpus"): + # CPU backends and systems where no GPU was detected will not have this attribute logger.debug("Adding missing exclude gpus argument to namespace") setattr(arguments, "exclude_gpus", None) return @@ -233,33 +270,7 @@ def _configure_backend(self, arguments): if GPUStats().exclude_all_devices: msg = "Switching backend to CPU" - if get_backend() == "amd": - msg += (". Using Tensorflow for CPU operations.") - os.environ["KERAS_BACKEND"] = "tensorflow" set_backend("cpu") logger.info(msg) logger.debug("Executing: %s. PID: %s", self._command, os.getpid()) - - if get_backend() == "amd" and not self._setup_amd(arguments): - safe_shutdown(got_error=True) - - @classmethod - def _setup_amd(cls, arguments): - """ Test for plaidml and perform setup for AMD. - - Parameters - ---------- - arguments: :class:`argparse.Namespace` - The command line arguments passed to Faceswap. - """ - logger.debug("Setting up for AMD") - try: - import plaidml # noqa pylint:disable=unused-import,import-outside-toplevel - except ImportError: - logger.error("PlaidML not found. Run `pip install plaidml-keras` for AMD support") - return False - from lib.gpu_stats import setup_plaidml # pylint:disable=import-outside-toplevel - setup_plaidml(arguments.loglevel, arguments.exclude_gpus) - logger.debug("setup up for PlaidML") - return True diff --git a/lib/config.py b/lib/config.py index b5981797e6..d4a2f1298b 100644 --- a/lib/config.py +++ b/lib/config.py @@ -3,52 +3,121 @@ Extends out :class:`configparser.ConfigParser` functionality by checking for default configuration updates and returning data in it's correct format """ +import gettext import logging import os import sys import textwrap + from collections import OrderedDict from configparser import ConfigParser +from dataclasses import dataclass from importlib import import_module from lib.utils import full_path_split -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +# LOCALES +_LANG = gettext.translation("lib.config", localedir="locales", fallback=True) +_ = _LANG.gettext + +OrderedDictSectionType = OrderedDict[str, "ConfigSection"] +OrderedDictItemType = OrderedDict[str, "ConfigItem"] + +logger = logging.getLogger(__name__) +ConfigValueType = bool | int | float | list[str] | str | None + + +@dataclass +class ConfigItem: + """ Dataclass for holding information about configuration items + + Parameters + ---------- + default: any + The default value for the configuration item + helptext: str + The helptext to be displayed for the configuration item + datatype: type + The type of the configuration item + rounding: int + The decimal places for floats or the step interval for ints for slider updates + min_max: tuple + The minumum and maximum value for the GUI slider for the configuration item + gui_radio: bool + ``True`` to display the configuration item in a Radio Box + fixed: bool + ``True`` if the item cannot be changed for existing models (training only) + group: str + The group that this configuration item belongs to in the GUI + """ + default: ConfigValueType + helptext: str + datatype: type + rounding: int + min_max: tuple[int, int] | tuple[float, float] | None + choices: str | list[str] + gui_radio: bool + fixed: bool + group: str | None + + +@dataclass +class ConfigSection: + """ Dataclass for holding information about configuration sections + + Parameters + ---------- + helptext: str + The helptext to be displayed for the configuration section + items: :class:`collections.OrderedDict` + Dictionary of configuration items for the section + """ + helptext: str + items: OrderedDictItemType class FaceswapConfig(): """ Config Items """ - def __init__(self, section, configfile=None): - """ Init Configuration """ + def __init__(self, section: str | None, configfile: str | None = None) -> None: + """ Init Configuration + + Parameters + ---------- + section: str or ``None`` + The configuration section. ``None`` for all sections + configfile: str, optional + Optional path to a config file. ``None`` for default location. Default: ``None`` + """ logger.debug("Initializing: %s", self.__class__.__name__) - self.configfile = self.get_config_file(configfile) + self.configfile = self._get_config_file(configfile) self.config = ConfigParser(allow_no_value=True) - self.defaults = OrderedDict() - self.config.optionxform = str + self.defaults: OrderedDictSectionType = OrderedDict() + self.config.optionxform = str # type:ignore self.section = section self.set_defaults() - self.handle_config() + self._handle_config() logger.debug("Initialized: %s", self.__class__.__name__) @property - def changeable_items(self): + def changeable_items(self) -> dict[str, ConfigValueType]: """ Training only. Return a dict of config items with their set values for items that can be altered after the model has been created """ - retval = dict() + retval: dict[str, ConfigValueType] = {} sections = [sect for sect in self.config.sections() if sect.startswith("global")] - for sect in sections + [self.section]: + all_sections = sections if self.section is None else sections + [self.section] + for sect in all_sections: if sect not in self.defaults: continue - for key, val in self.defaults[sect].items(): - if key == "helptext" or val["fixed"]: + for key, val in self.defaults[sect].items.items(): + if val.fixed: continue retval[key] = self.get(sect, key) logger.debug("Alterable for existing models: %s", retval) return retval - def set_defaults(self): + def set_defaults(self) -> None: """ Override for plugin specific config defaults Should be a series of self.add_section() and self.add_item() calls @@ -56,8 +125,8 @@ def set_defaults(self): e.g: section = "sect_1" - self.add_section(title=section, - info="Section 1 Information") + self.add_section(section, + "Section 1 Information") self.add_item(section=section, title="option_1", @@ -67,7 +136,7 @@ def set_defaults(self): """ raise NotImplementedError - def _defaults_from_plugin(self, plugin_folder): + def _defaults_from_plugin(self, plugin_folder: str) -> None: """ Scan the given plugins folder for config defaults.py files and update the default configuration. @@ -83,11 +152,14 @@ def _defaults_from_plugin(self, plugin_folder): base_path = os.path.dirname(os.path.realpath(sys.argv[0])) # Can't use replace as there is a bug on some Windows installs that lowers some paths import_path = ".".join(full_path_split(dirpath[len(base_path):])[1:]) - plugin_type = import_path.split(".")[-1] + plugin_type = import_path.rsplit(".", maxsplit=1)[-1] for filename in default_files: self._load_defaults_from_module(filename, import_path, plugin_type) - def _load_defaults_from_module(self, filename, module_path, plugin_type): + def _load_defaults_from_module(self, + filename: str, + module_path: str, + plugin_type: str) -> None: """ Load the plugin's defaults module, extract defaults and add to default configuration. Parameters @@ -104,19 +176,20 @@ def _load_defaults_from_module(self, filename, module_path, plugin_type): module = os.path.splitext(filename)[0] section = ".".join((plugin_type, module.replace("_defaults", ""))) logger.debug("Importing defaults module: %s.%s", module_path, module) - mod = import_module("{}.{}".format(module_path, module)) - self.add_section(title=section, info=mod._HELPTEXT) # pylint:disable=protected-access - for key, val in mod._DEFAULTS.items(): # pylint:disable=protected-access + mod = import_module(f"{module_path}.{module}") + self.add_section(section, mod._HELPTEXT) # type:ignore[attr-defined] # pylint:disable=protected-access # noqa:E501 + for key, val in mod._DEFAULTS.items(): # type:ignore[attr-defined] # pylint:disable=protected-access # noqa:E501 self.add_item(section=section, title=key, **val) logger.debug("Added defaults: %s", section) @property - def config_dict(self): - """ Collate global options and requested section into a dictionary with the correct + def config_dict(self) -> dict[str, ConfigValueType]: + """ dict: Collate global options and requested section into a dictionary with the correct data types """ - conf = dict() + conf: dict[str, ConfigValueType] = {} sections = [sect for sect in self.config.sections() if sect.startswith("global")] - sections.append(self.section) + if self.section is not None: + sections.append(self.section) for sect in sections: if sect not in self.config.sections(): continue @@ -126,7 +199,7 @@ def config_dict(self): conf[key] = self.get(sect, key) return conf - def get(self, section, option): + def get(self, section: str, option: str) -> ConfigValueType: """ Return a config item in it's correct format. Parameters @@ -142,24 +215,26 @@ def get(self, section, option): The selected configuration option in the correct data format """ logger.debug("Getting config item: (section: '%s', option: '%s')", section, option) - datatype = self.defaults[section][option]["type"] + datatype = self.defaults[section].items[option].datatype + + retval: ConfigValueType if datatype == bool: - func = self.config.getboolean + retval = self.config.getboolean(section, option) elif datatype == int: - func = self.config.getint + retval = self.config.getint(section, option) elif datatype == float: - func = self.config.getfloat + retval = self.config.getfloat(section, option) elif datatype == list: - func = self._parse_list + retval = self._parse_list(section, option) else: - func = self.config.get - retval = func(section, option) + retval = self.config.get(section, option) + if isinstance(retval, str) and retval.lower() == "none": retval = None logger.debug("Returning item: (type: %s, value: %s)", datatype, retval) return retval - def _parse_list(self, section, option): + def _parse_list(self, section: str, option: str) -> list[str]: """ Parse options that are stored as lists in the config file. These can be space or comma-separated items in the config file. They will be returned as a list of strings, regardless of what the final data type should be, so conversion from strings to other @@ -187,32 +262,58 @@ def _parse_list(self, section, option): raw_option, retval, section, option) return retval - def get_config_file(self, configfile): - """ Return the config file from the calling folder or the provided file """ + def _get_config_file(self, configfile: str | None) -> str: + """ Return the config file from the calling folder or the provided file + + Parameters + ---------- + configfile: str or ``None`` + Path to a config file. ``None`` for default location. + + Returns + ------- + str + The full path to the configuration file + """ if configfile is not None: if not os.path.isfile(configfile): - err = "Config file does not exist at: {}".format(configfile) + err = f"Config file does not exist at: {configfile}" logger.error(err) raise ValueError(err) return configfile - dirname = os.path.dirname(sys.modules[self.__module__].__file__) + filepath = sys.modules[self.__module__].__file__ + assert filepath is not None + dirname = os.path.dirname(filepath) folder, fname = os.path.split(dirname) - retval = os.path.join(os.path.dirname(folder), "config", "{}.ini".format(fname)) + retval = os.path.join(os.path.dirname(folder), "config", f"{fname}.ini") logger.debug("Config File location: '%s'", retval) return retval - def add_section(self, title=None, info=None): - """ Add a default section to config file """ + def add_section(self, title: str, info: str) -> None: + """ Add a default section to config file + + Parameters + ---------- + title: str + The title for the section + info: str + The helptext for the section + """ logger.debug("Add section: (title: '%s', info: '%s')", title, info) - if None in (title, info): - raise ValueError("Default config sections must have a title and " - "information text") - self.defaults[title] = OrderedDict() - self.defaults[title]["helptext"] = info - - def add_item(self, section=None, title=None, datatype=str, default=None, info=None, - rounding=None, min_max=None, choices=None, gui_radio=False, fixed=True, - group=None): + self.defaults[title] = ConfigSection(helptext=info, items=OrderedDict()) + + def add_item(self, + section: str | None = None, + title: str | None = None, + datatype: type = str, + default: ConfigValueType = None, + info: str | None = None, + rounding: int | None = None, + min_max: tuple[int, int] | tuple[float, float] | None = None, + choices: str | list[str] | None = None, + gui_radio: bool = False, + fixed: bool = True, + group: str | None = None) -> None: """ Add a default item to a config section For int or float values, rounding and min_max must be set @@ -243,107 +344,164 @@ def add_item(self, section=None, title=None, datatype=str, default=None, info=No "fixed: %s, group: %s)", section, title, datatype, default, info, rounding, min_max, choices, gui_radio, fixed, group) - choices = list() if not choices else choices + choices = [] if not choices else choices - if None in (section, title, default, info): - raise ValueError("Default config items must have a section, title, defult and " - "information text") + assert (section is not None and + title is not None and + default is not None and + info is not None), ("Default config items must have a section, title, defult and " + "information text") if not self.defaults.get(section, None): - raise ValueError("Section does not exist: {}".format(section)) - if datatype not in (str, bool, float, int, list): - raise ValueError("'datatype' must be one of str, bool, float or " - "int: {} - {}".format(section, title)) + raise ValueError(f"Section does not exist: {section}") + assert datatype in (str, bool, float, int, list), ( + f"'datatype' must be one of str, bool, float or int: {section} - {title}") if datatype in (float, int) and (rounding is None or min_max is None): raise ValueError("'rounding' and 'min_max' must be set for numerical options") if isinstance(datatype, list) and not choices: raise ValueError("'choices' must be defined for list based configuration items") - if not isinstance(choices, (list, tuple)): - raise ValueError("'choices' must be a list or tuple") - - info = self.expand_helptext(info, choices, default, datatype, min_max, fixed) - self.defaults[section][title] = {"default": default, - "helptext": info, - "type": datatype, - "rounding": rounding, - "min_max": min_max, - "choices": choices, - "gui_radio": gui_radio, - "fixed": fixed, - "group": group} - - @staticmethod - def expand_helptext(helptext, choices, default, datatype, min_max, fixed): + if choices != "colorchooser" and not isinstance(choices, (list, tuple)): + raise ValueError("'choices' must be a list or tuple or 'colorchooser") + + info = self._expand_helptext(info, choices, default, datatype, min_max, fixed) + self.defaults[section].items[title] = ConfigItem(default=default, + helptext=info, + datatype=datatype, + rounding=rounding or 0, + min_max=min_max, + choices=choices, + gui_radio=gui_radio, + fixed=fixed, + group=group) + + @classmethod + def _expand_helptext(cls, + helptext: str, + choices: str | list[str], + default: ConfigValueType, + datatype: type, + min_max: tuple[int, int] | tuple[float, float] | None, + fixed: bool) -> str: """ Add extra helptext info from parameters """ helptext += "\n" if not fixed: - helptext += "\nThis option can be updated for existing models.\n" + helptext += _("\nThis option can be updated for existing models.\n") if datatype == list: - helptext += ("\nIf selecting multiple options then each option should be separated " - "by a space or a comma (e.g. item1, item2, item3)\n") - if choices: - helptext += "\nChoose from: {}".format(choices) + helptext += _("\nIf selecting multiple options then each option should be separated " + "by a space or a comma (e.g. item1, item2, item3)\n") + if choices and choices != "colorchooser": + helptext += _("\nChoose from: {}").format(choices) elif datatype == bool: - helptext += "\nChoose from: True, False" + helptext += _("\nChoose from: True, False") elif datatype == int: + assert min_max is not None cmin, cmax = min_max - helptext += "\nSelect an integer between {} and {}".format(cmin, cmax) + helptext += _("\nSelect an integer between {} and {}").format(cmin, cmax) elif datatype == float: + assert min_max is not None cmin, cmax = min_max - helptext += "\nSelect a decimal number between {} and {}".format(cmin, cmax) - helptext += "\n[Default: {}]".format(default) + helptext += _("\nSelect a decimal number between {} and {}").format(cmin, cmax) + helptext += _("\n[Default: {}]").format(default) return helptext - def check_exists(self): - """ Check that a config file exists """ + def _check_exists(self) -> bool: + """ Check that a config file exists + + Returns + ------- + bool + ``True`` if the given configuration file exists + """ if not os.path.isfile(self.configfile): logger.debug("Config file does not exist: '%s'", self.configfile) return False logger.debug("Config file exists: '%s'", self.configfile) return True - def create_default(self): + def _create_default(self) -> None: """ Generate a default config if it does not exist """ logger.debug("Creating default Config") - for section, items in self.defaults.items(): - logger.debug("Adding section: '%s')", section) - self.insert_config_section(section, items["helptext"]) - for item, opt in items.items(): + for name, section in self.defaults.items(): + logger.debug("Adding section: '%s')", name) + self.insert_config_section(name, section.helptext) + for item, opt in section.items.items(): logger.debug("Adding option: (item: '%s', opt: '%s')", item, opt) - if item == "helptext": - continue - self.insert_config_item(section, - item, - opt["default"], - opt) + self._insert_config_item(name, item, opt.default, opt) self.save_config() - def insert_config_section(self, section, helptext, config=None): - """ Insert a section into the config """ + def insert_config_section(self, + section: str, + helptext: str, + config: ConfigParser | None = None) -> None: + """ Insert a section into the config + + Parameters + ---------- + section: str + The section title to insert + helptext: str + The help text for the config section + config: :class:`configparser.ConfigParser`, optional + The config parser object to insert the section into. ``None`` to insert it into the + default config. Default: ``None`` + """ logger.debug("Inserting section: (section: '%s', helptext: '%s', config: '%s')", section, helptext, config) config = self.config if config is None else config - config.optionxform = str + config.optionxform = str # type:ignore helptext = self.format_help(helptext, is_section=True) config.add_section(section) config.set(section, helptext) logger.debug("Inserted section: '%s'", section) - def insert_config_item(self, section, item, default, option, - config=None): - """ Insert an item into a config section """ + def _insert_config_item(self, + section: str, + item: str, + default: ConfigValueType, + option: ConfigItem, + config: ConfigParser | None = None) -> None: + """ Insert an item into a config section + + Parameters + ---------- + section: str + The section to insert the item into + item: str + The name of the item to insert + default: ConfigValueType + The default value for the item + option: :class:`ConfigItem` + The configuration option to insert + config: :class:`configparser.ConfigParser`, optional + The config parser object to insert the section into. ``None`` to insert it into the + default config. Default: ``None`` + """ logger.debug("Inserting item: (section: '%s', item: '%s', default: '%s', helptext: '%s', " - "config: '%s')", section, item, default, option["helptext"], config) + "config: '%s')", section, item, default, option.helptext, config) config = self.config if config is None else config - config.optionxform = str - helptext = option["helptext"] + config.optionxform = str # type:ignore + helptext = option.helptext helptext = self.format_help(helptext, is_section=False) config.set(section, helptext) config.set(section, item, str(default)) logger.debug("Inserted item: '%s'", item) - @staticmethod - def format_help(helptext, is_section=False): - """ Format comments for default ini file """ + @classmethod + def format_help(cls, helptext: str, is_section: bool = False) -> str: + """ Format comments for default ini file + + Parameters + ---------- + helptext: str + The help text to be formatted + is_section: bool, optional + ``True`` if the help text pertains to a section. ``False`` if it pertains to an item. + Default: ``True`` + + Returns + ------- + str + The formatted help text + """ logger.debug("Formatting help: (helptext: '%s', is_section: '%s')", helptext, is_section) formatted = "" for hlp in helptext.split("\n"): @@ -357,94 +515,101 @@ def format_help(helptext, is_section=False): if is_section: helptext = helptext.upper() else: - helptext = "\n{}".format(helptext) + helptext = f"\n{helptext}" logger.debug("formatted help: '%s'", helptext) return helptext - def load_config(self): + def _load_config(self) -> None: """ Load values from config """ - logger.verbose("Loading config: '%s'", self.configfile) - self.config.read(self.configfile) + logger.verbose("Loading config: '%s'", self.configfile) # type:ignore[attr-defined] + self.config.read(self.configfile, encoding="utf-8") - def save_config(self): + def save_config(self) -> None: """ Save a config file """ logger.info("Updating config at: '%s'", self.configfile) - with open(self.configfile, "w") as f_cfgfile: + with open(self.configfile, "w", encoding="utf-8", errors="replace") as f_cfgfile: self.config.write(f_cfgfile) logger.debug("Updated config at: '%s'", self.configfile) - def validate_config(self): + def _validate_config(self) -> None: """ Check for options in default config against saved config and add/remove as appropriate """ logger.debug("Validating config") - if self.check_config_change(): - self.add_new_config_items() - self.check_config_choices() + if self._check_config_change(): + self._add_new_config_items() + self._check_config_choices() logger.debug("Validated config") - def add_new_config_items(self): + def _add_new_config_items(self) -> None: """ Add new items to the config file """ logger.debug("Updating config") new_config = ConfigParser(allow_no_value=True) - for section, items in self.defaults.items(): - self.insert_config_section(section, items["helptext"], new_config) - for item, opt in items.items(): - if item == "helptext": - continue - if section not in self.config.sections(): - logger.debug("Adding new config section: '%s'", section) - opt_value = opt["default"] + for section_name, section in self.defaults.items(): + self.insert_config_section(section_name, section.helptext, new_config) + for item, opt in section.items.items(): + if section_name not in self.config.sections(): + logger.debug("Adding new config section: '%s'", section_name) + opt_value = opt.default else: - opt_value = self.config[section].get(item, opt["default"]) - self.insert_config_item(section, - item, - opt_value, - opt, - new_config) + opt_value = self.config[section_name].get(item, str(opt.default)) + self._insert_config_item(section_name, + item, + opt_value, + opt, + new_config) self.config = new_config - self.config.optionxform = str + self.config.optionxform = str # type:ignore self.save_config() logger.debug("Updated config") - def check_config_choices(self): + def _check_config_choices(self) -> None: """ Check that config items are valid choices """ logger.debug("Checking config choices") - for section, items in self.defaults.items(): - for item, opt in items.items(): - if item == "helptext" or not opt["choices"]: + for section_name, section in self.defaults.items(): + for item, opt in section.items.items(): + if not opt.choices: continue - if opt["type"] == list: # Multi-select items - opt_value = self._parse_list(section, item) - if not opt_value: # No option selected + if opt.datatype == list: # Multi-select items + opt_values = self._parse_list(section_name, item) + if not opt_values: # No option selected continue - if not all(val in opt["choices"] for val in opt_value): - invalid = [val for val in opt_value if val not in opt["choices"]] - valid = ", ".join(val for val in opt_value if val in opt["choices"]) + if not all(val in opt.choices for val in opt_values): + invalid = [val for val in opt_values if val not in opt.choices] + valid = ", ".join(val for val in opt_values if val in opt.choices) logger.warning("The option(s) %s are not valid selections for '%s': '%s'. " - "setting to: '%s'", invalid, section, item, valid) - self.config.set(section, item, valid) + "setting to: '%s'", invalid, section_name, item, valid) + self.config.set(section_name, item, valid) else: # Single-select items - opt_value = self.config.get(section, item) + if opt.choices == "colorchooser": + continue + opt_value = self.config.get(section_name, item) if opt_value.lower() == "none" and any(choice.lower() == "none" - for choice in opt["choices"]): + for choice in opt.choices): continue - if opt_value not in opt["choices"]: - default = str(opt["default"]) + if opt_value not in opt.choices: + default = str(opt.default) logger.warning("'%s' is not a valid config choice for '%s': '%s'. " - "Defaulting to: '%s'", opt_value, section, item, default) - self.config.set(section, item, default) + "Defaulting to: '%s'", + opt_value, section_name, item, default) + self.config.set(section_name, item, default) logger.debug("Checked config choices") - def check_config_change(self): - """ Check whether new default items have been added or removed - from the config file compared to saved version """ + def _check_config_change(self) -> bool: + """ Check whether new default items have been added or removed from the config file + compared to saved version + + Returns + ------- + bool + ``True`` if a config option has been added or removed + """ if set(self.config.sections()) != set(self.defaults.keys()): logger.debug("Default config has new section(s)") return True - for section, items in self.defaults.items(): - opts = [opt for opt in items.keys() if opt != "helptext"] - exists = [opt for opt in self.config[section].keys() + for section_name, section in self.defaults.items(): + opts = list(section.items) + exists = [opt for opt in self.config[section_name].keys() if not opt.startswith(("# ", "\n# "))] if set(exists) != set(opts): logger.debug("Default config has new item(s)") @@ -452,7 +617,7 @@ def check_config_change(self): logger.debug("Default config has not changed") return False - def handle_config(self): + def _handle_config(self) -> None: """ Handle the config. Checks whether a config file exists for this section. If not then a default is created. @@ -461,27 +626,26 @@ def handle_config(self): """ logger.debug("Handling config: (section: %s, configfile: '%s')", self.section, self.configfile) - if not self.check_exists(): - self.create_default() - self.load_config() - self.validate_config() + if not self._check_exists(): + self._create_default() + self._load_config() + self._validate_config() logger.debug("Handled config") -def generate_configs(): +def generate_configs() -> None: """ Generate config files if they don't exist. This script is run prior to anything being set up, so don't use logging Generates the default config files for plugins in the faceswap config folder """ - base_path = os.path.realpath(os.path.dirname(sys.argv[0])) plugins_path = os.path.join(base_path, "plugins") configs_path = os.path.join(base_path, "config") for dirpath, _, filenames in os.walk(plugins_path): if "_config.py" in filenames: section = os.path.split(dirpath)[-1] - config_file = os.path.join(configs_path, "{}.ini".format(section)) + config_file = os.path.join(configs_path, f"{section}.ini") if not os.path.exists(config_file): - mod = import_module("plugins.{}.{}".format(section, "_config")) - mod.Config(None) + mod = import_module(f"plugins.{section}._config") + mod.Config(None) # type:ignore[attr-defined] diff --git a/lib/convert.py b/lib/convert.py index 224d6daca2..c96b759213 100644 --- a/lib/convert.py +++ b/lib/convert.py @@ -1,14 +1,50 @@ #!/usr/bin/env python3 """ Converter for Faceswap """ - +from __future__ import annotations import logging +import typing as T +from dataclasses import dataclass import cv2 import numpy as np from plugins.plugin_loader import PluginLoader -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Callable + from lib.align.aligned_face import AlignedFace, CenteringType + from lib.align.detected_face import DetectedFace + from lib.config import FaceswapConfig + from lib.queue_manager import EventQueue + from scripts.convert import ConvertItem + from plugins.convert.color._base import Adjustment as ColorAdjust + from plugins.convert.color.seamless_clone import Color as SeamlessAdjust + from plugins.convert.mask.mask_blend import Mask as MaskAdjust + from plugins.convert.scaling._base import Adjustment as ScalingAdjust + +logger = logging.getLogger(__name__) + + +@dataclass +class Adjustments: + """ Dataclass to hold the optional processing plugins + + Parameters + ---------- + color: :class:`~plugins.color._base.Adjustment`, Optional + The selected color processing plugin. Default: `None` + mask: :class:`~plugins.mask_blend.Mask`, Optional + The selected mask processing plugin. Default: `None` + seamless: :class:`~plugins.color.seamless_clone.Color`, Optional + The selected mask processing plugin. Default: `None` + sharpening: :class:`~plugins.scaling._base.Adjustment`, Optional + The selected mask processing plugin. Default: `None` + """ + color: ColorAdjust | None = None + mask: MaskAdjust | None = None + seamless: SeamlessAdjust | None = None + sharpening: ScalingAdjust | None = None class Converter(): @@ -37,8 +73,14 @@ class Converter(): Optional location of custom configuration ``ini`` file. If ``None`` then use the default config location. Default: ``None`` """ - def __init__(self, output_size, coverage_ratio, centering, draw_transparent, pre_encode, - arguments, configfile=None): + def __init__(self, + output_size: int, + coverage_ratio: float, + centering: CenteringType, + draw_transparent: bool, + pre_encode: Callable | None, + arguments: Namespace, + configfile: str | None = None) -> None: logger.debug("Initializing %s: (output_size: %s, coverage_ratio: %s, centering: %s, " "draw_transparent: %s, pre_encode: %s, arguments: %s, configfile: %s)", self.__class__.__name__, output_size, coverage_ratio, centering, @@ -52,18 +94,20 @@ def __init__(self, output_size, coverage_ratio, centering, draw_transparent, pre self._configfile = configfile self._scale = arguments.output_scale / 100 - self._adjustments = dict(mask=None, color=None, seamless=None, sharpening=None) + self._face_scale = 1.0 - arguments.face_scale / 100. + self._adjustments = Adjustments() + self._full_frame_output: bool = arguments.writer != "patch" self._load_plugins() logger.debug("Initialized %s", self.__class__.__name__) @property - def cli_arguments(self): + def cli_arguments(self) -> Namespace: """:class:`argparse.Namespace`: The command line arguments passed to the convert process """ return self._args - def reinitialize(self, config): + def reinitialize(self, config: FaceswapConfig) -> None: """ Reinitialize this :class:`Converter`. Called as part of the :mod:`~tools.preview` tool. Resets all adjustments then loads the @@ -75,11 +119,14 @@ def reinitialize(self, config): Pre-loaded :class:`lib.config.FaceswapConfig`. used over any configuration on disk. """ logger.debug("Reinitializing converter") - self._adjustments = dict(mask=None, color=None, seamless=None, sharpening=None) + self._face_scale = 1.0 - self._args.face_scale / 100. + self._adjustments = Adjustments() self._load_plugins(config=config, disable_logging=True) logger.debug("Reinitialized converter") - def _load_plugins(self, config=None, disable_logging=False): + def _load_plugins(self, + config: FaceswapConfig | None = None, + disable_logging: bool = False) -> None: """ Load the requested adjustment plugins. Loads the :mod:`plugins.converter` plugins that have been requested for this conversion @@ -95,30 +142,32 @@ def _load_plugins(self, config=None, disable_logging=False): suppress these messages otherwise ``False``. Default: ``False`` """ logger.debug("Loading plugins. config: %s", config) - self._adjustments["mask"] = PluginLoader.get_converter( - "mask", - "mask_blend", - disable_logging=disable_logging)(self._args.mask_type, - self._output_size, - self._coverage_ratio, - configfile=self._configfile, - config=config) + self._adjustments.mask = PluginLoader.get_converter("mask", + "mask_blend", + disable_logging=disable_logging)( + self._args.mask_type, + self._output_size, + self._coverage_ratio, + configfile=self._configfile, + config=config) if self._args.color_adjustment != "none" and self._args.color_adjustment is not None: - self._adjustments["color"] = PluginLoader.get_converter( - "color", - self._args.color_adjustment, - disable_logging=disable_logging)(configfile=self._configfile, config=config) - - sharpening = PluginLoader.get_converter( - "scaling", - "sharpen", - disable_logging=disable_logging)(configfile=self._configfile, config=config) - if sharpening.config.get("method", None) is not None: - self._adjustments["sharpening"] = sharpening + self._adjustments.color = PluginLoader.get_converter("color", + self._args.color_adjustment, + disable_logging=disable_logging)( + configfile=self._configfile, + config=config) + + sharpening = PluginLoader.get_converter("scaling", + "sharpen", + disable_logging=disable_logging)( + configfile=self._configfile, + config=config) + if sharpening.config.get("method") is not None: + self._adjustments.sharpening = sharpening logger.debug("Loaded plugins: %s", self._adjustments) - def process(self, in_queue, out_queue): + def process(self, in_queue: EventQueue, out_queue: EventQueue): """ Main convert process. Takes items from the in queue, runs the relevant adjustments, patches faces to final frame @@ -126,56 +175,82 @@ def process(self, in_queue, out_queue): Parameters ---------- - in_queue: :class:`queue.Queue` + in_queue: :class:`~lib.queue_manager.EventQueue` The output from :class:`scripts.convert.Predictor`. Contains detected faces from the Faceswap model as well as the frame to be patched. - out_queue: :class:`queue.Queue` + out_queue: :class:`~lib.queue_manager.EventQueue` The queue to place patched frames into for writing by one of Faceswap's :mod:`plugins.convert.writer` plugins. """ logger.debug("Starting convert process. (in_queue: %s, out_queue: %s)", in_queue, out_queue) - log_once = False + logged = False while True: - items = in_queue.get() - if items == "EOF": + inbound: T.Literal["EOF"] | ConvertItem | list[ConvertItem] = in_queue.get() + if inbound == "EOF": logger.debug("EOF Received") logger.debug("Patch queue finished") # Signal EOF to other processes in pool logger.debug("Putting EOF back to in_queue") - in_queue.put(items) + in_queue.put(inbound) break - if isinstance(items, dict): - items = [items] + items = inbound if isinstance(inbound, list) else [inbound] for item in items: - logger.trace("Patch queue got: '%s'", item["filename"]) + logger.trace("Patch queue got: '%s'", # type: ignore[attr-defined] + item.inbound.filename) try: image = self._patch_image(item) - except Exception as err: # pylint: disable=broad-except + except Exception as err: # pylint:disable=broad-except # Log error and output original frame logger.error("Failed to convert image: '%s'. Reason: %s", - item["filename"], str(err)) - image = item["image"] + item.inbound.filename, str(err)) + image = item.inbound.image - loglevel = logger.trace if log_once else logger.warning - loglevel("Convert error traceback:", exc_info=True) - log_once = True + lvl = logger.trace if logged else logger.warning # type: ignore[attr-defined] + lvl("Convert error traceback:", exc_info=True) + logged = True # UNCOMMENT THIS CODE BLOCK TO PRINT TRACEBACK ERRORS - import sys ; import traceback - exc_info = sys.exc_info() ; traceback.print_exception(*exc_info) - logger.trace("Out queue put: %s", item["filename"]) - out_queue.put((item["filename"], image)) + # import sys; import traceback + # exc_info = sys.exc_info(); traceback.print_exception(*exc_info) + logger.trace("Out queue put: %s", # type: ignore[attr-defined] + item.inbound.filename) + out_queue.put((item.inbound.filename, image)) logger.debug("Completed convert process") - def _patch_image(self, predicted): + def _get_warp_matrix(self, matrix: np.ndarray, size: int) -> np.ndarray: + """ Obtain the final scaled warp transformation matrix based on face scaling from the + original transformation matrix + + Parameters + ---------- + matrix: :class:`numpy.ndarray` + The transformation for patching the swapped face back onto the output frame + size: int + The size of the face patch, in pixels + + Returns + ------- + :class:`numpy.ndarray` + The final transformation matrix with any scaling applied + """ + if self._face_scale == 1.0: + mat = matrix + else: + mat = matrix * self._face_scale + patch_center = (size / 2, size / 2) + mat[..., 2] += (1 - self._face_scale) * np.array(patch_center) + + return mat + + def _patch_image(self, predicted: ConvertItem) -> np.ndarray | list[bytes]: """ Patch a swapped face onto a frame. Run selected adjustments and swap the faces in a frame. Parameters ---------- - predicted: dict + predicted: :class:`~scripts.convert.ConvertItem` The output from :class:`scripts.convert.Predictor`. Returns @@ -186,21 +261,70 @@ def _patch_image(self, predicted): function (if it has one) """ - logger.trace("Patching image: '%s'", predicted["filename"]) - frame_size = (predicted["image"].shape[1], predicted["image"].shape[0]) + logger.trace("Patching image: '%s'", # type: ignore[attr-defined] + predicted.inbound.filename) + frame_size = (predicted.inbound.image.shape[1], predicted.inbound.image.shape[0]) new_image, background = self._get_new_image(predicted, frame_size) - patched_face = self._post_warp_adjustments(background, new_image) - patched_face = self._scale_image(patched_face) - patched_face *= 255.0 - patched_face = np.rint(patched_face, - out=np.empty(patched_face.shape, dtype="uint8"), - casting='unsafe') - if self._writer_pre_encode is not None: - patched_face = self._writer_pre_encode(patched_face) - logger.trace("Patched image: '%s'", predicted["filename"]) - return patched_face - - def _get_new_image(self, predicted, frame_size): + + if self._full_frame_output: + patched_face = self._post_warp_adjustments(background, new_image) + patched_face = self._scale_image(patched_face) + patched_face *= 255.0 + patched_face = np.rint(patched_face, + out=np.empty(patched_face.shape, dtype="uint8"), + casting='unsafe') + else: + patched_face = new_image + + if self._writer_pre_encode is None: + retval: np.ndarray | list[bytes] = patched_face + else: + kwargs: dict[str, T.Any] = {} + if self.cli_arguments.writer == "patch": + kwargs["canvas_size"] = (background.shape[1], background.shape[0]) + kwargs["matrices"] = np.array([self._get_warp_matrix(face.adjusted_matrix, + patched_face.shape[1]) + for face in predicted.reference_faces], + dtype="float32") + retval = self._writer_pre_encode(patched_face, **kwargs) + logger.trace("Patched image: '%s'", # type: ignore[attr-defined] + predicted.inbound.filename) + return retval + + def _warp_to_frame(self, + reference: AlignedFace, + face: np.ndarray, + frame: np.ndarray, + multiple_faces: bool) -> None: + """ Perform affine transformation to place a face patch onto the given frame. + + Affine is done in place on the `frame` array, so this function does not return a value + + Parameters + ---------- + reference: :class:`lib.align.AlignedFace` + The object holding the original aligned face + face: :class:`numpy.ndarray` + The swapped face patch + frame: :class:`numpy.ndarray` + The frame to affine the face onto + multiple_faces: bool + Controls the border mode to use. Uses BORDER_CONSTANT if there is only 1 face in + the image, otherwise uses the inferior BORDER_TRANSPARENT + """ + # Warp face with the mask + mat = self._get_warp_matrix(reference.adjusted_matrix, face.shape[0]) + border = cv2.BORDER_TRANSPARENT if multiple_faces else cv2.BORDER_CONSTANT + cv2.warpAffine(face, + mat, + (frame.shape[1], frame.shape[0]), + frame, + flags=cv2.WARP_INVERSE_MAP | reference.interpolators[1], + borderMode=border) + + def _get_new_image(self, + predicted: ConvertItem, + frame_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]: """ Get the new face from the predictor and apply pre-warp manipulations. Applies any requested adjustments to the raw output of the Faceswap model @@ -208,7 +332,7 @@ def _get_new_image(self, predicted, frame_size): Parameters ---------- - predicted: dict + predicted: :class:`~scripts.convert.ConvertItem` The output from :class:`scripts.convert.Predictor`. frame_size: tuple The (`width`, `height`) of the final frame in pixels @@ -220,39 +344,47 @@ def _get_new_image(self, predicted, frame_size): background: :class: `numpy.ndarray` The original frame """ - logger.trace("Getting: (filename: '%s', faces: %s)", - predicted["filename"], len(predicted["swapped_faces"])) + logger.trace("Getting: (filename: '%s', faces: %s)", # type: ignore[attr-defined] + predicted.inbound.filename, len(predicted.swapped_faces)) placeholder = np.zeros((frame_size[1], frame_size[0], 4), dtype="float32") - background = predicted["image"] / np.array(255.0, dtype="float32") - placeholder[:, :, :3] = background + if self._full_frame_output: + background = predicted.inbound.image / np.array(255.0, dtype="float32") + placeholder[:, :, :3] = background + else: + faces = [] # Collect the faces into final array + background = placeholder # Used for obtaining original frame dimensions - for new_face, detected_face, reference_face in zip(predicted["swapped_faces"], - predicted["detected_faces"], - predicted["reference_faces"]): + for new_face, detected_face, reference_face in zip(predicted.swapped_faces, + predicted.inbound.detected_faces, + predicted.reference_faces): predicted_mask = new_face[:, :, -1] if new_face.shape[2] == 4 else None new_face = new_face[:, :, :3] - interpolator = reference_face.interpolators[1] - new_face = self._pre_warp_adjustments(new_face, detected_face, reference_face, predicted_mask) - # Warp face with the mask - cv2.warpAffine(new_face, - reference_face.adjusted_matrix, - frame_size, - placeholder, - flags=cv2.WARP_INVERSE_MAP | interpolator, - borderMode=cv2.BORDER_CONSTANT) + if self._full_frame_output: + self._warp_to_frame(reference_face, + new_face, placeholder, + len(predicted.swapped_faces) > 1) + else: + faces.append(new_face) - logger.trace("Got filename: '%s'. (placeholders: %s)", - predicted["filename"], placeholder.shape) + if not self._full_frame_output: + placeholder = np.array(faces, dtype="float32") + + logger.trace("Got filename: '%s'. (placeholders: %s)", # type: ignore[attr-defined] + predicted.inbound.filename, placeholder.shape) return placeholder, background - def _pre_warp_adjustments(self, new_face, detected_face, reference_face, predicted_mask): + def _pre_warp_adjustments(self, + new_face: np.ndarray, + detected_face: DetectedFace, + reference_face: AlignedFace, + predicted_mask: np.ndarray | None) -> np.ndarray: """ Run any requested adjustments that can be performed on the raw output from the Faceswap model. @@ -277,21 +409,25 @@ def _pre_warp_adjustments(self, new_face, detected_face, reference_face, predict The face output from the Faceswap Model with any requested pre-warp adjustments performed. """ - logger.trace("new_face shape: %s, predicted_mask shape: %s", new_face.shape, - predicted_mask.shape if predicted_mask is not None else None) - old_face = reference_face.face[..., :3] / 255.0 + logger.trace("new_face shape: %s, predicted_mask shape: %s", # type: ignore[attr-defined] + new_face.shape, predicted_mask.shape if predicted_mask is not None else None) + old_face = T.cast(np.ndarray, reference_face.face)[..., :3] / 255.0 new_face, raw_mask = self._get_image_mask(new_face, detected_face, predicted_mask, reference_face) - if self._adjustments["color"] is not None: - new_face = self._adjustments["color"].run(old_face, new_face, raw_mask) - if self._adjustments["seamless"] is not None: - new_face = self._adjustments["seamless"].run(old_face, new_face, raw_mask) - logger.trace("returning: new_face shape %s", new_face.shape) + if self._adjustments.color is not None: + new_face = self._adjustments.color.run(old_face, new_face, raw_mask) + if self._adjustments.seamless is not None: + new_face = self._adjustments.seamless.run(old_face, new_face, raw_mask) + logger.trace("returning: new_face shape %s", new_face.shape) # type: ignore[attr-defined] return new_face - def _get_image_mask(self, new_face, detected_face, predicted_mask, reference_face): + def _get_image_mask(self, + new_face: np.ndarray, + detected_face: DetectedFace, + predicted_mask: np.ndarray | None, + reference_face: AlignedFace) -> tuple[np.ndarray, np.ndarray]: """ Return any selected image mask Places the requested mask into the new face's Alpha channel. @@ -312,22 +448,26 @@ def _get_image_mask(self, new_face, detected_face, predicted_mask, reference_fac ------- :class:`numpy.ndarray` The swapped face with the requested mask added to the Alpha channel + :class:`numpy.ndarray` + The raw mask with no erosion or blurring applied """ - logger.trace("Getting mask. Image shape: %s", new_face.shape) + logger.trace("Getting mask. Image shape: %s", new_face.shape) # type: ignore[attr-defined] if self._args.mask_type not in ("none", "predicted"): mask_centering = detected_face.mask[self._args.mask_type].stored_centering else: mask_centering = "face" # Unused but requires a valid value - crop_offset = (reference_face.pose.offset[self._centering] - - reference_face.pose.offset[mask_centering]) - mask, raw_mask = self._adjustments["mask"].run(detected_face, crop_offset, self._centering, - predicted_mask=predicted_mask) - logger.trace("Adding mask to alpha channel") + assert self._adjustments.mask is not None + mask, raw_mask = self._adjustments.mask.run(detected_face, + reference_face.pose.offset[mask_centering], + reference_face.pose.offset[self._centering], + self._centering, + predicted_mask=predicted_mask) + logger.trace("Adding mask to alpha channel") # type: ignore[attr-defined] new_face = np.concatenate((new_face, mask), -1) - logger.trace("Got mask. Image shape: %s", new_face.shape) + logger.trace("Got mask. Image shape: %s", new_face.shape) # type: ignore[attr-defined] return new_face, raw_mask - def _post_warp_adjustments(self, background, new_image): + def _post_warp_adjustments(self, background: np.ndarray, new_image: np.ndarray) -> np.ndarray: """ Perform any requested adjustments to the swapped faces after they have been transformed into the final frame. @@ -343,8 +483,8 @@ def _post_warp_adjustments(self, background, new_image): :class:`numpy.ndarray` The final merged and swapped frame with any requested post-warp adjustments applied """ - if self._adjustments["sharpening"] is not None: - new_image = self._adjustments["sharpening"].run(new_image) + if self._adjustments.sharpening is not None: + new_image = self._adjustments.sharpening.run(new_image) if self._draw_transparent: frame = new_image @@ -359,7 +499,7 @@ def _post_warp_adjustments(self, background, new_image): np.clip(frame, 0.0, 1.0, out=frame) return frame - def _scale_image(self, frame): + def _scale_image(self, frame: np.ndarray) -> np.ndarray: """ Scale the final image if requested. If output scale has been requested in command line arguments, scale the output @@ -377,11 +517,11 @@ def _scale_image(self, frame): """ if self._scale == 1: return frame - logger.trace("source frame: %s", frame.shape) + logger.trace("source frame: %s", frame.shape) # type: ignore[attr-defined] interp = cv2.INTER_CUBIC if self._scale > 1 else cv2.INTER_AREA dims = (round((frame.shape[1] / 2 * self._scale) * 2), round((frame.shape[0] / 2 * self._scale) * 2)) frame = cv2.resize(frame, dims, interpolation=interp) - logger.trace("resized frame: %s", frame.shape) + logger.trace("resized frame: %s", frame.shape) # type: ignore[attr-defined] np.clip(frame, 0.0, 1.0, out=frame) return frame diff --git a/lib/face_filter.py b/lib/face_filter.py deleted file mode 100644 index c6589a386f..0000000000 --- a/lib/face_filter.py +++ /dev/null @@ -1,181 +0,0 @@ -#!/usr/bin python3 -""" Face Filterer for extraction in faceswap.py """ - -import logging - -from lib.align import AlignedFace -from lib.vgg_face import VGGFace -from lib.image import read_image -from plugins.extract.pipeline import Extractor, ExtractMedia - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - - -def avg(arr): - """ Return an average """ - return sum(arr) * 1.0 / len(arr) - - -class FaceFilter(): - """ Face filter for extraction - NB: we take only first face, so the reference file should only contain one face. """ - - def __init__(self, reference_file_paths, nreference_file_paths, detector, aligner, - multiprocess=False, threshold=0.4): - logger.debug("Initializing %s: (reference_file_paths: %s, nreference_file_paths: %s, " - "detector: %s, aligner: %s, multiprocess: %s, threshold: %s)", - self.__class__.__name__, reference_file_paths, nreference_file_paths, - detector, aligner, multiprocess, threshold) - self.vgg_face = VGGFace() - self.filters = self.load_images(reference_file_paths, nreference_file_paths) - # TODO Revert face-filter to use the selected detector and aligner. - # Currently Tensorflow does not release vram after it has been allocated - # Whilst this vram can still be used, the pipeline for the extraction process can't see - # it so thinks there is not enough vram available. - # Either the pipeline will need to be changed to be re-usable by face-filter and extraction - # Or another vram measurement technique will need to be implemented to for when tensorflow - # has already performed allocation. For now we force CPU detectors. - - # self.align_faces(detector, aligner, multiprocess) - self.align_faces("cv2-dnn", "cv2-dnn", "none", multiprocess) - - self.get_filter_encodings() - self.threshold = threshold - logger.debug("Initialized %s", self.__class__.__name__) - - @staticmethod - def load_images(reference_file_paths, nreference_file_paths): - """ Load the images """ - retval = dict() - for fpath in reference_file_paths: - retval[fpath] = {"image": read_image(fpath, raise_error=True), - "type": "filter"} - for fpath in nreference_file_paths: - retval[fpath] = {"image": read_image(fpath, raise_error=True), - "type": "nfilter"} - logger.debug("Loaded filter images: %s", {k: v["type"] for k, v in retval.items()}) - return retval - - # Extraction pipeline - def align_faces(self, detector_name, aligner_name, masker_name, multiprocess): - """ Use the requested detectors to retrieve landmarks for filter images """ - extractor = Extractor(detector_name, - aligner_name, - masker_name, - multiprocess=multiprocess) - self.run_extractor(extractor) - del extractor - self.load_aligned_face() - - def run_extractor(self, extractor): - """ Run extractor to get faces """ - for _ in range(extractor.passes): - extractor.launch() - self.queue_images(extractor) - for faces in extractor.detected_faces(): - filename = faces.filename - detected_faces = faces.detected_faces - if len(detected_faces) > 1: - logger.warning("Multiple faces found in %s file: '%s'. Using first detected " - "face.", self.filters[filename]["type"], filename) - self.filters[filename]["detected_face"] = detected_faces[0] - - def queue_images(self, extractor): - """ queue images for detection and alignment """ - in_queue = extractor.input_queue - for fname, img in self.filters.items(): - logger.debug("Adding to filter queue: '%s' (%s)", fname, img["type"]) - feed_dict = ExtractMedia(fname, img["image"], detected_faces=img.get("detected_faces")) - logger.debug("Queueing filename: '%s' items: %s", fname, feed_dict) - in_queue.put(feed_dict) - logger.debug("Sending EOF to filter queue") - in_queue.put("EOF") - - def load_aligned_face(self): - """ Align the faces for vgg_face input """ - for filename, face in self.filters.items(): - logger.debug("Loading aligned face: '%s'", filename) - image = face["image"] - detected_face = face["detected_face"] - detected_face.load_aligned(image, centering="legacy", size=224) - face["face"] = detected_face.aligned.face - del face["image"] - logger.debug("Loaded aligned face: ('%s', shape: %s)", - filename, face["face"].shape) - - def get_filter_encodings(self): - """ Return filter face encodings from Keras VGG Face """ - for filename, face in self.filters.items(): - logger.debug("Getting encodings for: '%s'", filename) - encodings = self.vgg_face.predict(face["face"]) - logger.debug("Filter Filename: %s, encoding shape: %s", filename, encodings.shape) - face["encoding"] = encodings - del face["face"] - - def check(self, image, detected_face): - """ Check the extracted Face - - Parameters - ---------- - image: :class:`numpy.ndarray` - The original frame that contains the face to be checked - detected_face: :class:`lib.align.DetectedFace` - The detected face object that contains the face to be checked - - Returns - ------- - bool - ``True`` if the face matches a filter otherwise ``False`` - """ - logger.trace("Checking face with FaceFilter") - distances = {"filter": list(), "nfilter": list()} - feed = AlignedFace(detected_face.landmarks_xy, image=image, size=224, centering="legacy") - encodings = self.vgg_face.predict(feed.face) - for filt in self.filters.values(): - similarity = self.vgg_face.find_cosine_similiarity(filt["encoding"], encodings) - distances[filt["type"]].append(similarity) - - avgs = {key: avg(val) if val else None for key, val in distances.items()} - mins = {key: min(val) if val else None for key, val in distances.items()} - # Filter - if distances["filter"] and avgs["filter"] > self.threshold: - msg = "Rejecting filter face: {} > {}".format(round(avgs["filter"], 2), self.threshold) - retval = False - # nFilter no Filter - elif not distances["filter"] and avgs["nfilter"] < self.threshold: - msg = "Rejecting nFilter face: {} < {}".format(round(avgs["nfilter"], 2), - self.threshold) - retval = False - # Filter with nFilter - elif distances["filter"] and distances["nfilter"] and mins["filter"] > mins["nfilter"]: - msg = ("Rejecting face as distance from nfilter sample is smaller: (filter: {}, " - "nfilter: {})".format(round(mins["filter"], 2), round(mins["nfilter"], 2))) - retval = False - elif distances["filter"] and distances["nfilter"] and avgs["filter"] > avgs["nfilter"]: - msg = ("Rejecting face as average distance from nfilter sample is smaller: (filter: " - "{}, nfilter: {})".format(round(mins["filter"], 2), round(mins["nfilter"], 2))) - retval = False - elif distances["filter"] and distances["nfilter"]: - # k-nearest-neighbor classifier - var_k = min(5, min(len(distances["filter"]), len(distances["nfilter"])) + 1) - var_n = sum(list(map(lambda x: x[0], - list(sorted([(1, d) for d in distances["filter"]] + - [(0, d) for d in distances["nfilter"]], - key=lambda x: x[1]))[:var_k]))) - ratio = var_n/var_k - if ratio < 0.5: - msg = ("Rejecting face as k-nearest neighbors classification is less than " - "0.5: {}".format(round(ratio, 2))) - retval = False - else: - msg = None - retval = True - else: - msg = None - retval = True - if msg: - logger.verbose(msg) - else: - logger.trace("Accepted face: (similarity: %s, threshold: %s)", - distances, self.threshold) - return retval diff --git a/lib/git.py b/lib/git.py new file mode 100644 index 0000000000..90cba0af3c --- /dev/null +++ b/lib/git.py @@ -0,0 +1,157 @@ +#!/usr/bin python3 +""" Handles command line calls to git """ +import logging +import os +import sys + +from subprocess import PIPE, Popen + +logger = logging.getLogger(__name__) + + +class Git(): + """ Handles calls to github """ + def __init__(self) -> None: + logger.debug("Initializing: %s", self.__class__.__name__) + self._working_dir = os.path.dirname(os.path.realpath(sys.argv[0])) + self._available = self._check_available() + logger.debug("Initialized: %s", self.__class__.__name__) + + def _from_git(self, command: str) -> tuple[bool, list[str]]: + """ Execute a git command + + Parameters + ---------- + command : str + The command to send to git + + Returns + ------- + success: bool + ``True`` if the command succesfully executed otherwise ``False`` + list[str] + The output lines from stdout if there was no error, otherwise from stderr + """ + logger.debug("command: '%s'", command) + cmd = f"git {command}" + with Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE, cwd=self._working_dir) as proc: + stdout, stderr = proc.communicate() + retcode = proc.returncode + success = retcode == 0 + lines = stdout.decode("utf-8", errors="replace").splitlines() + if not lines: + lines = stderr.decode("utf-8", errors="replace").splitlines() + logger.debug("command: '%s', returncode: %s, success: %s, lines: %s", + cmd, retcode, success, lines) + return success, lines + + def _check_available(self) -> bool: + """ Check if git is available. Does a call to git status. If the process errors due to + folder ownership, attempts to add the folder to github safe folders list and tries + again + + Returns + ------- + bool + ``True`` if git is available otherwise ``False`` + + """ + success, msg = self._from_git("status") + if success: + return True + config = next((line.strip() for line in msg if "add safe.directory" in line), None) + if not config: + return False + success, _ = self._from_git(config.split("git ", 1)[-1]) + return True + + @property + def status(self) -> list[str]: + """ Obtain the output of git status for tracked files only """ + if not self._available: + return [] + success, status = self._from_git("status -uno") + if not success or not status: + return [] + return status + + @property + def branch(self) -> str: + """ str: The git branch that is currently being used to execute Faceswap. """ + status = next((line.strip() for line in self.status if "On branch" in line), "Not Found") + return status.replace("On branch ", "") + + @property + def branches(self) -> list[str]: + """ list[str]: List of all available branches. """ + if not self._available: + return [] + success, branches = self._from_git("branch -a") + if not success or not branches: + return [] + return branches + + def update_remote(self) -> bool: + """ Update all branches to track remote + + Returns + ------- + bool + ``True`` if update was succesful otherwise ``False`` + """ + if not self._available: + return False + return self._from_git("remote update")[0] + + def pull(self) -> bool: + """ Pull the current branch + + Returns + ------- + bool + ``True`` if pull is successful otherwise ``False`` + """ + if not self._available: + return False + return self._from_git("pull")[0] + + def checkout(self, branch: str) -> bool: + """ Checkout the requested branch + + Parameters + ---------- + branch : str + The branch to checkout + + Returns + ------- + bool + ``True`` if the branch was succesfully checkout out otherwise ``False`` + """ + if not self._available: + return False + return self._from_git(f"checkout {branch}")[0] + + def get_commits(self, count: int) -> list[str]: + """ Obtain the last commits to the repo + + Parameters + ---------- + count : int + The last number of commits to obtain + + Returns + ------- + list[str] + list of commits, or empty list if none found + """ + if not self._available: + return [] + success, commits = self._from_git(f"log --pretty=oneline --abbrev-commit -n {count}") + if not success or not commits: + return [] + return commits + + +git = Git() +""" :class:`Git`: Handles calls to github """ diff --git a/lib/gpu_stats/__init__.py b/lib/gpu_stats/__init__.py index e10cac5c3e..070a53d057 100644 --- a/lib/gpu_stats/__init__.py +++ b/lib/gpu_stats/__init__.py @@ -6,17 +6,19 @@ from lib.utils import get_backend -from ._base import set_exclude_devices # noqa +from ._base import set_exclude_devices, GPUInfo backend = get_backend() if backend == "nvidia" and platform.system().lower() == "darwin": - from .nvidia_apple import NvidiaAppleStats as GPUStats # noqa + from .nvidia_apple import NvidiaAppleStats as GPUStats # type:ignore elif backend == "nvidia": - from .nvidia import NvidiaStats as GPUStats # noqa -elif backend == "amd": - from .amd import AMDStats as GPUStats, setup_plaidml # noqa + from .nvidia import NvidiaStats as GPUStats # type:ignore elif backend == "apple_silicon": - from .apple_silicon import AppleSiliconStats as GPUStats # noqa -elif backend == "cpu": - from .cpu import CPUStats as GPUStats # noqa + from .apple_silicon import AppleSiliconStats as GPUStats # type:ignore +elif backend == "directml": + from .directml import DirectML as GPUStats # type:ignore +elif backend == "rocm": + from .rocm import ROCm as GPUStats # type:ignore +else: + from .cpu import CPUStats as GPUStats # type:ignore diff --git a/lib/gpu_stats/_base.py b/lib/gpu_stats/_base.py index 286de83160..8953f134ae 100644 --- a/lib/gpu_stats/_base.py +++ b/lib/gpu_stats/_base.py @@ -1,49 +1,73 @@ #!/usr/bin/env python3 """ Parent class for obtaining Stats for various GPU/TPU backends. All GPU Stats should inherit -from the :class:`GPUStats` class contained here. """ +from the :class:`_GPUStats` class contained here. """ import logging -import os -import sys -from typing import List, Optional +from dataclasses import dataclass from lib.utils import get_backend -if sys.version_info < (3, 8): - from typing_extensions import TypedDict -else: - from typing import TypedDict - +_EXCLUDE_DEVICES: list[int] = [] -_EXCLUDE_DEVICES: List[int] = [] +@dataclass +class GPUInfo(): + """Dataclass for storing information about the available GPUs on the system. - -class GPUInfo(TypedDict): - """ Typed Dictionary for returning Full GPU Information. """ - vram: List[int] + Attributes: + ---------- + vram: list[int] + List of integers representing the total VRAM available on each GPU, in MB. + vram_free: list[int] + List of integers representing the free VRAM available on each GPU, in MB. + driver: str + String representing the driver version being used for the GPUs. + devices: list[str] + List of strings representing the names of each GPU device. + devices_active: list[int] + List of integers representing the indices of the active GPU devices. + """ + vram: list[int] + vram_free: list[int] driver: str - devices: List[str] - active_devices: List[str] + devices: list[str] + devices_active: list[int] -class BiggestGPUInfo(TypedDict): - """ Typed Dictionary for returning GPU Information about the card with most available VRAM. """ +@dataclass +class BiggestGPUInfo(): + """ Dataclass for holding GPU Information about the card with most available VRAM. + + Attributes + ---------- + card_id: int + Integer representing the index of the GPU device. + device: str + The name of the device + free: float + The amount of available VRAM on the GPU + total: float + the total amount of VRAM on the GPU + """ card_id: int device: str free: float total: float -def set_exclude_devices(devices: List[int]) -> None: +def set_exclude_devices(devices: list[int]) -> None: """ Add any explicitly selected GPU devices to the global list of devices to be excluded from use by Faceswap. Parameters ---------- - devices: list - list of indices corresponding to the GPU devices connected to the computer + devices: list[int] + list of GPU device indices to exclude + + Example + ------- + >>> set_exclude_devices([0, 1]) # Exclude the first two GPU devices """ logger = logging.getLogger(__name__) logger.debug("Excluding GPU indicies: %s", devices) @@ -52,25 +76,31 @@ def set_exclude_devices(devices: List[int]) -> None: _EXCLUDE_DEVICES.extend(devices) -class GPUStats(): - """ Parent class for returning information of GPUs used. """ +class _GPUStats(): + """ Parent class for collecting GPU device information. + + Parameters: + ----------- + log : bool, optional + Flag indicating whether or not to log debug messages. Default: `True`. + """ def __init__(self, log: bool = True) -> None: # Logger is held internally, as we don't want to log when obtaining system stats on crash # or when querying the backend for command line options - self._logger: Optional[logging.Logger] = logging.getLogger(__name__) if log else None + self._logger: logging.Logger | None = logging.getLogger(__name__) if log else None self._log("debug", f"Initializing {self.__class__.__name__}") self._is_initialized = False self._initialize() self._device_count: int = self._get_device_count() - self._active_devices: List[int] = self._get_active_devices() + self._active_devices: list[int] = self._get_active_devices() self._handles: list = self._get_handles() self._driver: str = self._get_driver() - self._device_names: List[str] = self._get_device_names() - self._vram: List[float] = self._get_vram() - self._vram_free: List[float] = self._get_free_vram() + self._device_names: list[str] = self._get_device_names() + self._vram: list[int] = self._get_vram() + self._vram_free: list[int] = self._get_free_vram() if get_backend() != "cpu" and not self._active_devices: self._log("warning", "No GPU detected") @@ -84,8 +114,8 @@ def device_count(self) -> int: return self._device_count @property - def cli_devices(self) -> List[str]: - """ list: List of available devices for use in faceswap's command line arguments. """ + def cli_devices(self) -> list[str]: + """ list[str]: Formatted index: name text string for each GPU """ return [f"{idx}: {device}" for idx, device in enumerate(self._device_names)] @property @@ -95,22 +125,9 @@ def exclude_all_devices(self) -> bool: @property def sys_info(self) -> GPUInfo: - """ dict: GPU Stats that are required for system information logging. - - The dictionary contains the following data: - - **vram** (`list`): the total amount of VRAM in Megabytes for each GPU as pertaining to - :attr:`_handles` - - **driver** (`str`): The GPU driver version that is installed on the OS - - **devices** (`list`): The device name of each GPU on the system as pertaining - to :attr:`_handles` - - **devices_active** (`list`): The device name of each active GPU on the system as - pertaining to :attr:`_handles` - """ + """ :class:`GPUInfo`: The GPU Stats that are required for system information logging """ return GPUInfo(vram=self._vram, + vram_free=self._get_free_vram(), driver=self._driver, devices=self._device_names, devices_active=self._active_devices) @@ -131,16 +148,16 @@ def _log(self, level: str, message: str) -> None: logger = getattr(self._logger, level.lower()) logger(message) - def _initialize(self): - """ Override for GPU specific initialization code. """ + def _initialize(self) -> None: + """ Override to initialize the GPU device handles and any other necessary resources. """ self._is_initialized = True - def _shutdown(self): - """ Override for GPU specific shutdown code. """ + def _shutdown(self) -> None: + """ Override to shutdown the GPU device handles and any other necessary resources. """ self._is_initialized = False def _get_device_count(self) -> int: - """ Override to obtain GPU specific device count + """ Override to obtain the number of GPU devices Returns ------- @@ -149,14 +166,13 @@ def _get_device_count(self) -> int: """ raise NotImplementedError() - def _get_active_devices(self) -> List[int]: - """ Obtain the indices of active GPUs (those that have not been explicitly excluded by - CUDA_VISIBLE_DEVICES environment variable or explicitly excluded in the command line - arguments). + def _get_active_devices(self) -> list[int]: + """ Obtain the indices of active GPUs (those that have not been explicitly excluded in + the command line arguments). Notes ----- - Override for GPUs that do not use CUDA + Override for GPU specific checking Returns ------- @@ -164,10 +180,6 @@ def _get_active_devices(self) -> List[int]: The list of device indices that are available for Faceswap to use """ devices = [idx for idx in range(self._device_count) if idx not in _EXCLUDE_DEVICES] - env_devices = os.environ.get("CUDA_VISIBLE_DEVICES") - if env_devices: - env_devices = [int(i) for i in env_devices.split(",")] - devices = [idx for idx in devices if idx in env_devices] self._log("debug", f"Active GPU Devices: {devices}") return devices @@ -191,7 +203,7 @@ def _get_driver(self) -> str: """ raise NotImplementedError() - def _get_device_names(self) -> List[str]: + def _get_device_names(self) -> list[str]: """ Override to obtain the names of all connected GPUs. The quality of this information depends on the backend and OS being used, but it should be sufficient for identifying cards. @@ -204,7 +216,7 @@ def _get_device_names(self) -> List[str]: """ raise NotImplementedError() - def _get_vram(self) -> List[float]: + def _get_vram(self) -> list[int]: """ Override to obtain the total VRAM in Megabytes for each connected GPU. Returns @@ -215,8 +227,8 @@ def _get_vram(self) -> List[float]: """ raise NotImplementedError() - def _get_free_vram(self) -> List[float]: - """ Override to obrain the amount of VRAM that is available, in Megabytes, for each + def _get_free_vram(self) -> list[int]: + """ Override to obtain the amount of VRAM that is available, in Megabytes, for each connected GPU. Returns @@ -232,17 +244,7 @@ def get_card_most_free(self) -> BiggestGPUInfo: Returns ------- - dict - The dictionary contains the following data: - - **card_id** (`int`): The index of the card as pertaining to :attr:`_handles` - - **device** (`str`): The name of the device - - **free** (`float`): The amount of available VRAM on the GPU - - **total** (`float`): the total amount of VRAM on the GPU - + :class:`BiggestGpuInfo` If a GPU is not detected then the **card_id** is returned as ``-1`` and the amount of free and total RAM available is fixed to 2048 Megabytes. """ diff --git a/lib/gpu_stats/amd.py b/lib/gpu_stats/amd.py deleted file mode 100644 index 6d81d17852..0000000000 --- a/lib/gpu_stats/amd.py +++ /dev/null @@ -1,364 +0,0 @@ -#!/usr/bin/env python3 -""" Collects and returns Information on available AMD GPUs. """ -import json -import logging -import os -import sys - -from typing import List, Optional - -import plaidml - -from ._base import GPUStats, _EXCLUDE_DEVICES - - -_PLAIDML_INITIALIZED: bool = False - - -def setup_plaidml(log_level: str, exclude_devices: List[int]) -> None: - """ Setup PlaidML for AMD Cards. - - Sets the Keras backend to PlaidML, loads the plaidML backend and makes GPU Device information - from PlaidML available to :class:`AMDStats`. - - Parameters - ---------- - log_level: str - Faceswap's log level. Used for setting the log level inside PlaidML - exclude_devices: list - A list of integers of device IDs that should not be used by Faceswap - """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name - logger.info("Setting up for PlaidML") - logger.verbose("Setting Keras Backend to PlaidML") - # Add explicitly excluded devices to list. The contents are checked in AMDstats - if exclude_devices: - _EXCLUDE_DEVICES.extend(int(idx) for idx in exclude_devices) - os.environ["KERAS_BACKEND"] = "plaidml.keras.backend" - stats = AMDStats(log_level) - logger.info("Using GPU(s): %s", [stats.names[i] for i in stats.active_devices]) - logger.info("Successfully set up for PlaidML") - - -class AMDStats(GPUStats): - """ Holds information and statistics about AMD GPU(s) available on the currently - running system. - - Notes - ----- - The quality of data that returns is very much dependent on the OpenCL implementation used - for a particular OS. Some data is just not available at all, so assumptions and substitutions - are made where required. PlaidML is used as an interface into OpenCL to obtain the required - information. - - PlaidML is explicitly initialized inside this class, as it can be called from the command line - arguments to list available GPUs. PlaidML needs to be set up and configured to obtain reliable - information. As the function :func:`setup_plaidml` is called very early within the Faceswap - and launch process and it references this class, initial PlaidML setup can all be handled here. - - Parameters - ---------- - log: bool, optional - Whether the class should output information to the logger. There may be occasions where the - logger has not yet been set up when this class is queried. Attempting to log in these - instances will raise an error. If GPU stats are being queried prior to the logger being - available then this parameter should be set to ``False``. Otherwise set to ``True``. - Default: ``True`` - """ - def __init__(self, log: bool = True, log_level: str = "INFO") -> None: - - self._log_level: str = log_level.upper() - - # Following attributes are set in :func:``_initialize`` - self._ctx: Optional(plaidml.Context) = None - self._supported_devices: Optional(List[plaidml._DeviceConfig]) = None - self._all_devices: Optional(List[plaidml._DeviceConfig]) = None - self._device_details: Optional(List[dict]) = None - - super().__init__(log=log) - - @property - def active_devices(self) -> List[int]: - """ list: The active device ids in use. """ - return self._active_devices - - @property - def _plaid_ids(self) -> List[str]: - """ list: The device identification for each GPU device that PlaidML has discovered. """ - return [device.id.decode("utf-8") for device in self._all_devices] - - @property - def _experimental_indices(self) -> List[int]: - """ list: The indices corresponding to :attr:`_ids` of GPU devices marked as - "experimental". """ - retval = [idx for idx, device in enumerate(self._all_devices) - if device not in self._supported_indices] - return retval - - @property - def _supported_indices(self) -> List[int]: - """ list: The indices corresponding to :attr:`_ids` of GPU devices marked as - "supported". """ - retval = [idx for idx, device in enumerate(self._all_devices) - if device in self._supported_devices] - return retval - - @property - def _all_vram(self) -> List[float]: - """ list: The VRAM of each GPU device that PlaidML has discovered. """ - return [int(device.get("globalMemSize", 0)) / (1024 * 1024) - for device in self._device_details] - - @property - def names(self) -> List[str]: - """ list: The name of each GPU device that PlaidML has discovered. """ - return [f"{device.get('vendor', 'unknown')} - {device.get('name', 'unknown')} " - f"({ 'supported' if idx in self._supported_indices else 'experimental'})" - for idx, device in enumerate(self._device_details)] - - def _initialize(self) -> None: - """ Initialize PlaidML for AMD GPUs. - - If :attr:`_is_initialized` is ``True`` then this function just returns performing no - action. - - if ``False`` then PlaidML is setup, if not already, and GPU information is extracted - from the PlaidML context. - """ - if self._is_initialized: - return - self._log("debug", "Initializing PlaidML for AMD GPU.") - - self._initialize_plaidml() - - self._ctx = plaidml.Context() - self._supported_devices = self._get_supported_devices() - self._all_devices = self._get_all_devices() - self._device_details = self._get_device_details() - self._select_device() - - super()._initialize() - - def _initialize_plaidml(self) -> None: - """ Initialize PlaidML on first call to this class and set global - :attr:``_PLAIDML_INITIALIZED`` to ``True``. If PlaidML has already been initialized then - return performing no action. """ - global _PLAIDML_INITIALIZED # pylint:disable=global-statement - - if _PLAIDML_INITIALIZED: - return - - self._log("debug", "Performing first time PlaidML setup.") - self._set_plaidml_logger() - - _PLAIDML_INITIALIZED = True - - def _set_plaidml_logger(self) -> None: - """ Set PlaidMLs default logger to Faceswap Logger, prevent propagation and set the correct - log level. """ - self._log("debug", "Setting PlaidML Default Logger") - - plaidml.DEFAULT_LOG_HANDLER = logging.getLogger("plaidml_root") - plaidml.DEFAULT_LOG_HANDLER.propagate = 0 - - numeric_level = getattr(logging, self._log_level, None) - if numeric_level < 10: # DEBUG Logging - plaidml._internal_set_vlog(1) # pylint:disable=protected-access - elif numeric_level < 20: # INFO Logging - plaidml._internal_set_vlog(0) # pylint:disable=protected-access - else: # WARNING LOGGING - plaidml.quiet() - - def _get_supported_devices(self) -> List[plaidml._DeviceConfig]: - """ Obtain GPU devices from PlaidML that are marked as "supported". - - Returns - ------- - list_LOGGER. - The :class:`plaidml._DeviceConfig` objects for all supported GPUs that PlaidML has - discovered. - """ - experimental_setting = plaidml.settings.experimental - - plaidml.settings.experimental = False - devices = plaidml.devices(self._ctx, limit=100, return_all=True)[0] - - plaidml.settings.experimental = experimental_setting - - supported = [d for d in devices - if d.details - and json.loads(d.details.decode("utf-8")).get("type", "cpu").lower() == "gpu"] - - self._log("debug", f"Obtained supported devices: {supported}") - return supported - - def _get_all_devices(self) -> List[plaidml._DeviceConfig]: - """ Obtain all available (experimental and supported) GPU devices from PlaidML. - - Returns - ------- - list - The :class:`pladml._DeviceConfig` objects for GPUs that PlaidML has discovered. - """ - experimental_setting = plaidml.settings.experimental - - plaidml.settings.experimental = True - devices = plaidml.devices(self._ctx, limit=100, return_all=True)[0] - - plaidml.settings.experimental = experimental_setting - - experi = [d for d in devices - if d.details - and json.loads(d.details.decode("utf-8")).get("type", "cpu").lower() == "gpu"] - - self._log("debug", f"Obtained experimental Devices: {experi}") - - all_devices = experi + self._supported_devices - - self._log("debug", f"Obtained all Devices: {all_devices}") - return all_devices - - def _get_device_details(self) -> List[dict]: - """ Obtain the device details for all connected AMD GPUS. - - Returns - ------- - list - The `dict` device detail for all GPUs that PlaidML has discovered. - """ - details = [json.loads(d.details.decode("utf-8")) - for d in self._all_devices if d.details] - self._log("debug", f"Obtained Device details: {details}") - return details - - def _select_device(self) -> None: - """ - If the plaidml user configuration settings exist, then set the default GPU from the - settings file, Otherwise set the GPU to be the one with most VRAM. """ - if os.path.exists(plaidml.settings.user_settings): # pylint:disable=no-member - self._log("debug", "Setting PlaidML devices from user_settings") - else: - self._select_largest_gpu() - - def _select_largest_gpu(self) -> None: - """ Set the default GPU to be a supported device with the most available VRAM. If no - supported device is available, then set the GPU to be an experimental device with the - most VRAM available. """ - category = "supported" if self._supported_devices else "experimental" - self._log("debug", f"Obtaining largest {category} device") - - indices = getattr(self, f"_{category}_indices") - if not indices: - self._log("error", "Failed to automatically detect your GPU.") - self._log("error", "Please run `plaidml-setup` to set up your GPU.") - sys.exit(1) - - max_vram = max([self._all_vram[idx] for idx in indices]) - self._log("debug", f"Max VRAM: {max_vram}") - - gpu_idx = min([idx for idx, vram in enumerate(self._all_vram) - if vram == max_vram and idx in indices]) - self._log("debug", f"GPU IDX: {gpu_idx}") - - selected_gpu = self._plaid_ids[gpu_idx] - self._log("info", f"Setting GPU to largest available {category} device. If you want to " - "override this selection, run `plaidml-setup` from the command line.") - - plaidml.settings.experimental = category == "experimental" - plaidml.settings.device_ids = [selected_gpu] - - def _get_device_count(self) -> int: - """ Detect the number of AMD GPUs available from PlaidML. - - Returns - ------- - int - The total number of AMD GPUs available - """ - retval = len(self._all_devices) - self._log("debug", f"GPU Device count: {retval}") - return retval - - def _get_active_devices(self) -> List[int]: - """ Obtain the indices of active GPUs (those that have not been explicitly excluded by - PlaidML or explicitly excluded in the command line arguments). - - Returns - ------- - list - The list of device indices that are available for Faceswap to use - """ - devices = [idx for idx, d_id in enumerate(self._plaid_ids) - if d_id in plaidml.settings.device_ids and idx not in _EXCLUDE_DEVICES] - self._log("debug", f"Active GPU Devices: {devices}") - return devices - - def _get_handles(self) -> list: - """ AMD Doesn't really use device handles, so we just return the all devices list - - Returns - ------- - list - The list of all AMD discovered GPUs - """ - handles = self._all_devices - self._log("debug", f"AMD GPU Handles found: {handles}") - return handles - - def _get_driver(self) -> str: - """ Obtain the AMD driver version currently in use. - - Returns - ------- - str - The current AMD GPU driver versions - """ - drivers = [device.get("driverVersion", "No Driver Found") - for device in self._device_details] - self._log("debug", f"GPU Drivers: {drivers}") - return drivers - - def _get_device_names(self) -> List[str]: - """ Obtain the list of names of connected AMD GPUs as identified in :attr:`_handles`. - - Returns - ------- - list - The list of connected Nvidia GPU names - """ - names = self.names - self._log("debug", f"GPU Devices: {names}") - return names - - def _get_vram(self) -> List[float]: - """ Obtain the VRAM in Megabytes for each connected AMD GPU as identified in - :attr:`_handles`. - - Returns - ------- - list - The VRAM in Megabytes for each connected Nvidia GPU - """ - vram = self._all_vram - self._log("debug", f"GPU VRAM: {vram}") - return vram - - def _get_free_vram(self) -> List[float]: - """ Obtain the amount of VRAM that is available, in Megabytes, for each connected AMD - GPU. - - Notes - ----- - There is no useful way to get free VRAM on PlaidML. OpenCL loads and unloads VRAM as - required, so this returns the total memory available per card for AMD GPUs, which is - not particularly useful. - - Returns - ------- - list - List of `float`s containing the amount of VRAM available, in Megabytes, for each - connected GPU as corresponding to the values in :attr:`_handles - """ - vram = self._all_vram - self._log("debug", f"GPU VRAM free: {vram}") - return vram diff --git a/lib/gpu_stats/apple_silicon.py b/lib/gpu_stats/apple_silicon.py index 787eb7bc6b..a8b0815015 100644 --- a/lib/gpu_stats/apple_silicon.py +++ b/lib/gpu_stats/apple_silicon.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """ Collects and returns Information on available Apple Silicon SoCs in Apple Macs. """ -from typing import List, Optional +import typing as T import os import psutil @@ -8,13 +8,13 @@ from lib.utils import FaceswapError -from ._base import GPUStats +from ._base import _GPUStats _METAL_INITIALIZED: bool = False -class AppleSiliconStats(GPUStats): +class AppleSiliconStats(_GPUStats): """ Holds information and statistics about Apple Silicon SoC(s) available on the currently running Apple system. @@ -35,7 +35,7 @@ class AppleSiliconStats(GPUStats): """ def __init__(self, log: bool = True) -> None: # Following attribute set in :func:``_initialize`` - self._tf_devices: Optional(List[str]) = None + self._tf_devices: list[T.Any] = [] super().__init__(log=log) @@ -142,7 +142,7 @@ def _get_driver(self) -> str: self._log("debug", f"GPU Driver: {driver}") return driver - def _get_device_names(self) -> List[str]: + def _get_device_names(self) -> list[str]: """ Obtain the list of names of available Apple Silicon SoC(s) as identified in :attr:`_handles`. @@ -155,7 +155,7 @@ def _get_device_names(self) -> List[str]: self._log("debug", f"GPU Devices: {names}") return names - def _get_vram(self) -> List[float]: + def _get_vram(self) -> list[int]: """ Obtain the VRAM in Megabytes for each available Apple Silicon SoC(s) as identified in :attr:`_handles`. @@ -170,12 +170,12 @@ def _get_vram(self) -> List[float]: list The RAM in Megabytes for each available Apple Silicon SoC """ - vram = [(psutil.virtual_memory().total / self._device_count) / (1024 * 1024) + vram = [int((psutil.virtual_memory().total / self._device_count) / (1024 * 1024)) for _ in range(self._device_count)] self._log("debug", f"SoC RAM: {vram}") return vram - def _get_free_vram(self) -> List[float]: + def _get_free_vram(self) -> list[int]: """ Obtain the amount of VRAM that is available, in Megabytes, for each available Apple Silicon SoC. @@ -185,7 +185,7 @@ def _get_free_vram(self) -> List[float]: List of `float`s containing the amount of RAM available, in Megabytes, for each available SoC as corresponding to the values in :attr:`_handles """ - vram = [(psutil.virtual_memory().available / self._device_count) / (1024 * 1024) + vram = [int((psutil.virtual_memory().available / self._device_count) / (1024 * 1024)) for _ in range(self._device_count)] self._log("debug", f"SoC RAM free: {vram}") return vram diff --git a/lib/gpu_stats/cpu.py b/lib/gpu_stats/cpu.py index bbc9ef4d45..ae20c96c09 100644 --- a/lib/gpu_stats/cpu.py +++ b/lib/gpu_stats/cpu.py @@ -1,22 +1,18 @@ #!/usr/bin/env python3 """ Dummy functions for running faceswap on CPU. """ +from ._base import _GPUStats -from typing import List - -from ._base import GPUStats - - -class CPUStats(GPUStats): +class CPUStats(_GPUStats): """ Holds information and statistics about the CPU on the currently running system. Notes ----- - The information held here is not useful, but GPUStats is dynamically imported depending on the - backend used, so we need to make sure this class is available for Faceswap run on the CPU + The information held here is not useful, but _GPUStats is dynamically imported depending on + the backend used, so we need to make sure this class is available for Faceswap run on the CPU Backend. - The base :class:`GPUStats` handles the dummying in of information when no GPU is detected. + The base :class:`_GPUStats` handles the dummying in of information when no GPU is detected. Parameters ---------- @@ -49,7 +45,7 @@ def _get_handles(self) -> list: list An empty list for CPU Backends """ - handles = [] + handles: list = [] self._log("debug", f"GPU Handles found: {len(handles)}") return handles @@ -65,7 +61,7 @@ def _get_driver(self) -> str: self._log("debug", f"GPU Driver: {driver}") return driver - def _get_device_names(self) -> List[str]: + def _get_device_names(self) -> list[str]: """ Obtain the list of names of connected GPUs as identified in :attr:`_handles`. Returns @@ -73,11 +69,11 @@ def _get_device_names(self) -> List[str]: list An empty list for CPU backends """ - names = [] + names: list[str] = [] self._log("debug", f"GPU Devices: {names}") return names - def _get_vram(self) -> List[float]: + def _get_vram(self) -> list[int]: """ Obtain the RAM in Megabytes for the running system. Returns @@ -85,11 +81,11 @@ def _get_vram(self) -> List[float]: list An empty list for CPU backends """ - vram = [] + vram: list[int] = [] self._log("debug", f"GPU VRAM: {vram}") return vram - def _get_free_vram(self) -> List[float]: + def _get_free_vram(self) -> list[int]: """ Obtain the amount of RAM that is available, in Megabytes, for the running system. Returns @@ -97,6 +93,6 @@ def _get_free_vram(self) -> List[float]: list An empty list for CPU backends """ - vram = [] + vram: list[int] = [] self._log("debug", f"GPU VRAM free: {vram}") return vram diff --git a/lib/gpu_stats/directml.py b/lib/gpu_stats/directml.py new file mode 100644 index 0000000000..10bb435b93 --- /dev/null +++ b/lib/gpu_stats/directml.py @@ -0,0 +1,630 @@ +#!/usr/bin/env python3 +""" Collects and returns Information on DirectX 12 hardware devices for DirectML. """ +from __future__ import annotations +import os +import sys +import typing as T +assert sys.platform == "win32" + +import ctypes +from ctypes import POINTER, Structure, windll +from dataclasses import dataclass +from enum import Enum, IntEnum + +from comtypes import COMError, IUnknown, GUID, STDMETHOD, HRESULT # pylint:disable=import-error + +from ._base import _GPUStats + +if T.TYPE_CHECKING: + from collections.abc import Callable + +# Monkey patch default ctypes.c_uint32 value to Enum ctypes property for easier tracking of types +# We can't just subclass as the attribute will be assumed to be part of the Enumeration, so we +# attach it directly and suck up the typing errors. +setattr(Enum, "ctype", ctypes.c_uint32) + + +############################# +# CTYPES SUPPORTING OBJECTS # +############################# +# GUIDs +@dataclass +class LookupGUID: + """ GUIDs that are required for creating COM objects which are used and discarded. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/d3d12/nn-d3d12-id3d12device2 + """ + IDXGIDevice = GUID("{54ec77fa-1377-44e6-8c32-88fd5f44c84c}") + ID3D12Device = GUID("{189819f1-1db6-4b57-be54-1821339b85f7}") + + +# ENUMS +class DXGIGpuPreference(IntEnum): + """ The preference of GPU for the app to run on. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi1_6/ne-dxgi1_6-dxgi_gpu_preference + """ + DXGI_GPU_PREFERENCE_UNSPECIFIED = 0 + DXGI_GPU_PREFERENCE_MINIMUM_POWER = 1 + DXGI_GPU_PREFERENCE_HIGH_PERFORMANCE = 2 + + +class DXGIAdapterFlag(IntEnum): + """ Identifies the type of DXGI adapter. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi/ne-dxgi-dxgi_adapter_flag + """ + DXGI_ADAPTER_FLAG_NONE = 0 + DXGI_ADAPTER_FLAG_REMOTE = 1 + DXGI_ADAPTER_FLAG_SOFTWARE = 2 + DXGI_ADAPTER_FLAG_FORCE_DWORD = 0xffffffff + + +class DXGIMemorySegmentGroup(IntEnum): + """ Constants that specify an adapter's memory segment grouping. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi1_4/ne-dxgi1_4-dxgi_memory_segment_group + """ + DXGI_MEMORY_SEGMENT_GROUP_LOCAL = 0 + DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL = 1 + + +class D3DFeatureLevel(Enum): + """ Describes the set of features targeted by a Direct3D device. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/d3dcommon/ne-d3dcommon-d3d_feature_level + """ + D3D_FEATURE_LEVEL_1_0_CORE = 0x1000 + D3D_FEATURE_LEVEL_9_1 = 0x9100 + D3D_FEATURE_LEVEL_9_2 = 0x9200 + D3D_FEATURE_LEVEL_9_3 = 0x9300 + D3D_FEATURE_LEVEL_10_0 = 0xa000 + D3D_FEATURE_LEVEL_10_1 = 0xa100 + D3D_FEATURE_LEVEL_11_0 = 0xb000 + D3D_FEATURE_LEVEL_11_1 = 0xb100 + D3D_FEATURE_LEVEL_12_0 = 0xc000 + D3D_FEATURE_LEVEL_12_1 = 0xc100 + D3D_FEATURE_LEVEL_12_2 = 0xc200 + + +class VendorID(Enum): + """ DirectX VendorID Enum """ + AMD = 0x1002 + NVIDIA = 0x10DE + MICROSOFT = 0x1414 + QUALCOMM = 0x4D4F4351 + INTEL = 0x8086 + + +# STRUCTS +class StructureRepr(Structure): + """ Override the standard structure class to add a useful __repr__ for logging """ + def __repr__(self) -> str: + """ Output the class name and the structure contents """ + content = ["=".join([field[0], str(getattr(self, field[0]))]) + for field in self._fields_] + if self.__dict__: # Add manually added parameters + content.extend("=".join([key, str(val)]) for key, val in self.__dict__.items()) + return f"{self.__class__.__name__}({', '.join(content)})" + + +class LUID(StructureRepr): # pylint:disable=too-few-public-methods + """ Local Identifier for an adaptor + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-luid """ + _fields_ = [("LowPart", ctypes.c_ulong), ("HighPart", ctypes.c_long)] + + +class DriverVersion(StructureRepr): # pylint:disable=too-few-public-methods + """ Stucture (based off LARGE_INTEGER) to hold the driver version + + Reference + --------- + https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-large_integer-r1""" + _fields_ = [("parts_a", ctypes.c_uint16), + ("parts_b", ctypes.c_uint16), + ("parts_c", ctypes.c_uint16), + ("parts_d", ctypes.c_uint16)] + + +class DXGIAdapterDesc1(StructureRepr): # pylint:disable=too-few-public-methods + """ Describes an adapter (or video card) using DXGI 1.1 + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi/ns-dxgi-DXGIAdapterDesc1 """ + _fields_ = [ + ("Description", ctypes.c_wchar * 128), + ("VendorId", ctypes.c_uint), + ("DeviceId", ctypes.c_uint), + ("SubSysId", ctypes.c_uint), + ("Revision", ctypes.c_uint), + ("DedicatedVideoMemory", ctypes.c_size_t), + ("DedicatedSystemMemory", ctypes.c_size_t), + ("SharedSystemMemory", ctypes.c_size_t), + ("AdapterLuid", LUID), + ("Flags", DXGIAdapterFlag.ctype)] # type:ignore[attr-defined] # pylint:disable=no-member + + +class DXGIQueryVideoMemoryInfo(StructureRepr): # pylint:disable=too-few-public-methods + """ Describes the current video memory budgeting parameters. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi1_4/ns-dxgi1_4-dxgi_query_video_memory_info + """ + _fields_ = [("Budget", ctypes.c_uint64), + ("CurrentUsage", ctypes.c_uint64), + ("AvailableForReservation", ctypes.c_uint64), + ("CurrentReservation", ctypes.c_uint64)] + + +# COM OBjects +class IDXObject(IUnknown): # pylint:disable=too-few-public-methods + """ Base interface for all DXGI objects. + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi/nn-dxgi-idxgiobject + """ + _iid_ = GUID("{aec22fb8-76f3-4639-9be0-28eb43a67a2e}") + _methods_ = [STDMETHOD(HRESULT, "SetPrivateData", + [GUID, ctypes.c_uint, POINTER(ctypes.c_void_p)]), + STDMETHOD(HRESULT, "SetPrivateDataInterface", [GUID, POINTER(IUnknown)]), + STDMETHOD(HRESULT, "GetPrivateData", + [GUID, POINTER(ctypes.c_uint), POINTER(ctypes.c_void_p)]), + STDMETHOD(HRESULT, "GetParent", [GUID, POINTER(POINTER(ctypes.c_void_p))])] + + +class IDXGIFactory6(IDXObject): # pylint:disable=too-few-public-methods + """ Implements methods for generating DXGI objects + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi/nn-dxgi-idxgifactory + """ + _iid_ = GUID("{c1b6694f-ff09-44a9-b03c-77900a0a1d17}") + + _methods_ = [STDMETHOD(HRESULT, "EnumAdapters"), # IDXGIFactory + STDMETHOD(HRESULT, "MakeWindowAssociation"), + STDMETHOD(HRESULT, "GetWindowAssociation"), + STDMETHOD(HRESULT, "CreateSwapChain"), + STDMETHOD(HRESULT, "CreateSoftwareAdapter"), + STDMETHOD(HRESULT, "EnumAdapters1"), # IDXGIFactory1 + STDMETHOD(ctypes.c_bool, "IsCurrent"), + STDMETHOD(ctypes.c_bool, "IsWindowedStereoEnabled"), # IDXGIFactory2 + STDMETHOD(HRESULT, "CreateSwapChainForHwnd"), + STDMETHOD(HRESULT, "CreateSwapChainForCoreWindow"), + STDMETHOD(HRESULT, "GetSharedResourceAdapterLuid"), + STDMETHOD(HRESULT, "RegisterStereoStatusWindow"), + STDMETHOD(HRESULT, "RegisterStereoStatusEvent"), + STDMETHOD(None, "UnregisterStereoStatus"), + STDMETHOD(HRESULT, "RegisterOcclusionStatusWindow"), + STDMETHOD(HRESULT, "RegisterOcclusionStatusEvent"), + STDMETHOD(None, "UnregisterOcclusionStatus"), + STDMETHOD(HRESULT, "CreateSwapChainForComposition"), + STDMETHOD(ctypes.c_uint, "GetCreationFlags"), # IDXGIFactory3 + STDMETHOD(HRESULT, "EnumAdapterByLuid", # IDXGIFactory4 + [LUID, GUID, POINTER(POINTER(ctypes.c_void_p))]), + STDMETHOD(HRESULT, "EnumWarpAdapter"), + STDMETHOD(HRESULT, "CheckFeatureSupport"), # IDXGIFactory5 + STDMETHOD(HRESULT, # IDXGIFactory6 + "EnumAdapterByGpuPreference", + [ctypes.c_uint, + DXGIGpuPreference.ctype, # type:ignore[attr-defined] # pylint:disable=no-member # noqa:E501 + GUID, + POINTER(ctypes.c_void_p)])] + + +class IDXGIAdapter3(IDXObject): # pylint:disable=too-few-public-methods + """ Represents a display sub-system (including one or more GPU's, DACs and video memory). + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi1_4/nn-dxgi1_4-idxgiadapter3 + """ + _iid_ = GUID("{645967a4-1392-4310-a798-8053ce3e93fd}") + _methods_ = [STDMETHOD(HRESULT, "EnumOutputs"), # v1.0 Methods + STDMETHOD(HRESULT, "GetDesc"), + STDMETHOD(HRESULT, "CheckInterfaceSupport", # v1.1 Methods + [GUID, POINTER(DriverVersion)]), + STDMETHOD(HRESULT, "GetDesc1", [POINTER(DXGIAdapterDesc1)]), + STDMETHOD(HRESULT, "GetDesc2"), # v1.2 Methods + STDMETHOD(HRESULT, # v1.3 Methods + "RegisterHardwareContentProtectionTeardownStatusEvent"), + STDMETHOD(None, "UnregisterHardwareContentProtectionTeardownStatus"), + STDMETHOD(HRESULT, + "QueryVideoMemoryInfo", + [ctypes.c_uint, + DXGIMemorySegmentGroup.ctype, # type:ignore[attr-defined] # pylint:disable=no-member # noqa:E501 + POINTER(DXGIQueryVideoMemoryInfo)]), + STDMETHOD(HRESULT, "SetVideoMemoryReservation"), + STDMETHOD(HRESULT, "RegisterVideoMemoryBudgetChangeNotificationEvent"), + STDMETHOD(None, "UnregisterVideoMemoryBudgetChangeNotification")] + + +########################### +# PYTHON COLLATED OBJECTS # +########################### +@dataclass +class Device: + """ Holds information about a device attached to an adapter. + + Parameters + ---------- + description: :class:`DXGIAdapterDesc1` + The information returned from DXGI.dll about the device + driver_version: str + The driver version of the device + local_mem: :class:`DXGIQueryVideoMemoryInfo` + The amount of local memory currently available + non_local_mem: :class:`DXGIQueryVideoMemoryInfo` + The amount of non-local memory currently available + is_d3d12: bool + ``True`` if the device supports DirectX12 + is_compute_only: bool + ``True`` if the device is only compute (no graphics) + """ + description: DXGIAdapterDesc1 + driver_version: str + local_mem: DXGIQueryVideoMemoryInfo + non_local_mem: DXGIQueryVideoMemoryInfo + is_d3d12: bool + is_compute_only: bool = False + + @property + def is_software_adapter(self) -> bool: + """ bool: ``True`` if this is a software adapter. """ + return self.description.Flags == DXGIAdapterFlag.DXGI_ADAPTER_FLAG_SOFTWARE.value + + @property + def is_valid(self) -> bool: + """ bool: ``True`` if this adapter is a hardware adaptor and is not the basic renderer """ + if self.is_software_adapter: + return False + + if (self.description.VendorId == VendorID.MICROSOFT.value and + self.description.DeviceId == 0x8c): + return False + + return True + + +class Adapters(): # pylint:disable=too-few-public-methods + """ Wrapper to obtain connected DirectX Graphics interface adapters from Windows + + Parameters + ---------- + log_func: :func:`~lib.gpu_stats._base._log` + The logging function to use from the parent GPUStats class + """ + def __init__(self, log_func: Callable[[str, str], None]) -> None: + self._log = log_func + self._log("debug", f"Initializing {self.__class__.__name__}: (log_func: {log_func})") + + self._factory = self._get_factory() + self._adapters = self._get_adapters() + self._devices = self._process_adapters() + + self._valid_adaptors: list[Device] = [] + self._log("debug", f"Initialized {self.__class__.__name__}") + + def _get_factory(self) -> ctypes._Pointer: + """ Get a DXGI 1.1 Factory object + + Reference + --------- + https://learn.microsoft.com/en-us/windows/win32/api/dxgi/nf-dxgi-createdxgifactory1 + + Returns + ------- + :class:`ctypes._Pointer` + A pointer to a :class:`IDXGIFactory6` COM instance + """ + factory_func = windll.dxgi.CreateDXGIFactory + factory_func.argtypes = (GUID, POINTER(ctypes.c_void_p)) + factory_func.restype = HRESULT + handle = ctypes.c_void_p(0) + factory_func(IDXGIFactory6._iid_, ctypes.byref(handle)) # pylint:disable=protected-access + retval = ctypes.POINTER(IDXGIFactory6)(T.cast(IDXGIFactory6, handle.value)) + self._log("debug", f"factory: {retval}") + return retval + + @property + def valid_adapters(self) -> list[Device]: + """ list[:class:`Device`]: DirectX 12 compatible hardware :class:`Device` objects """ + if self._valid_adaptors: + return self._valid_adaptors + + for device in self._devices: + if not device.is_valid: + # Sorted by most performant so everything after first basic adapter is skipped + break + if not device.is_d3d12: + continue + self._valid_adaptors.append(device) + self._log("debug", f"valid_adaptors: {self._valid_adaptors}") + return self._valid_adaptors + + def _get_adapters(self) -> list[ctypes._Pointer]: + """ Obtain DirectX 12 supporting hardware adapter objects and add a Device class for + obtaining details + + Returns + ------- + list + List of :class:`ctypes._Pointer` objects + """ + idx = 0 + retval = [] + while True: + try: + handle = ctypes.c_void_p(0) + success = self._factory.EnumAdapterByGpuPreference( # type:ignore[attr-defined] + idx, + DXGIGpuPreference.DXGI_GPU_PREFERENCE_HIGH_PERFORMANCE.value, + IDXGIAdapter3._iid_, # pylint:disable=protected-access + ctypes.byref(handle)) + if success != 0: + raise AttributeError("Error calling EnumAdapterByGpuPreference. Result: " + f"{hex(ctypes.c_ulong(success).value)}") + adapter = POINTER(IDXGIAdapter3)(T.cast(IDXGIAdapter3, handle.value)) + self._log("debug", f"found adapter: {adapter}") + retval.append(adapter) + except COMError as err: + err_code = hex(ctypes.c_ulong(err.hresult).value) # pylint:disable=no-member + self._log( + "debug", + "COM Error. Breaking: " + f"{err.text}({err_code})") # pylint:disable=no-member + break + finally: + idx += 1 + + self._log("debug", f"adapters: {retval}") + return retval + + def _query_adapter(self, func: Callable[[T.Any], T.Any], *args: T.Any) -> None: + """ Query an adapter function, logging if the HRESULT is not a success + + Parameters + ---------- + func: Callable[[Any], Any] + The adaptor function to call + args: Any + The arguments to pass to the adaptor function + """ + check = func(*args) + if check: + self._log("debug", f"Failed HRESULT for func {func}({args}): " + f"{hex(ctypes.c_ulong(check).value)}") + + def _test_d3d12(self, adapter: ctypes._Pointer) -> bool: + """ Test whether the given adapter supports DirectX 12 + + Parameters + ---------- + adapter: :class:`ctypes._Pointer` + A pointer to an adapter instance + + Returns + ------- + bool + ``True`` if the given adapter supports DirectX 12 + """ + factory_func = windll.d3d12.D3D12CreateDevice + factory_func.argtypes = ( + POINTER(IUnknown), + D3DFeatureLevel.ctype, # type:ignore[attr-defined] # pylint:disable=no-member + GUID, + POINTER(ctypes.c_void_p)) + handle = ctypes.c_void_p(0) + factory_func.restype = HRESULT + success = factory_func(adapter, + D3DFeatureLevel.D3D_FEATURE_LEVEL_11_0.value, + LookupGUID.ID3D12Device, + ctypes.byref(handle)) + return success in (0, 1) + + def _process_adapters(self) -> list[Device]: + """ Process the adapters to add discovered information. + + Returns + ------- + list[:class:`Device`] + List of device of objects found in the adapters + """ + retval = [] + for adapter in self._adapters: + # Description + desc = DXGIAdapterDesc1() + self._query_adapter(adapter.GetDesc1, ctypes.byref(desc)) # type:ignore[attr-defined] + + # Driver Version + driver = DriverVersion() + self._query_adapter(adapter.CheckInterfaceSupport, # type:ignore[attr-defined] + LookupGUID.IDXGIDevice, + ctypes.byref(driver)) + driver_version = f"{driver.parts_d}.{driver.parts_c}.{driver.parts_b}.{driver.parts_a}" + + # Current Memory + local_mem = DXGIQueryVideoMemoryInfo() + self._query_adapter(adapter.QueryVideoMemoryInfo, # type:ignore[attr-defined] + 0, + DXGIMemorySegmentGroup.DXGI_MEMORY_SEGMENT_GROUP_LOCAL.value, + local_mem) + non_local_mem = DXGIQueryVideoMemoryInfo() + self._query_adapter( + adapter.QueryVideoMemoryInfo, # type:ignore[attr-defined] + 0, + DXGIMemorySegmentGroup.DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL.value, + non_local_mem) + + # is_d3d12 + is_d3d12 = self._test_d3d12(adapter) + + retval.append(Device(desc, driver_version, local_mem, non_local_mem, is_d3d12)) + + return retval + + +class DirectML(_GPUStats): + """ Holds information and statistics about GPUs connected using Windows API + + Parameters + ---------- + log: bool, optional + Whether the class should output information to the logger. There may be occasions where the + logger has not yet been set up when this class is queried. Attempting to log in these + instances will raise an error. If GPU stats are being queried prior to the logger being + available then this parameter should be set to ``False``. Otherwise set to ``True``. + Default: ``True`` + """ + def __init__(self, log: bool = True) -> None: + self._devices: list[Device] = [] + super().__init__(log=log) + + @property + def _all_vram(self) -> list[int]: + """ list: The VRAM of each GPU device that the DX API has discovered. """ + return [int(device.description.DedicatedVideoMemory / (1024 * 1024)) + for device in self._devices] + + @property + def names(self) -> list[str]: + """ list: The name of each GPU device that the DX API has discovered. """ + return [device.description.Description for device in self._devices] + + def _get_active_devices(self) -> list[int]: + """ Obtain the indices of active GPUs (those that have not been explicitly excluded by + DML_VISIBLE_DEVICES environment variable or explicitly excluded in the command line + arguments). + + Returns + ------- + list + The list of device indices that are available for Faceswap to use + """ + devices = super()._get_active_devices() + env_devices = os.environ.get("DML_VISIBLE_DEVICES") + if env_devices: + new_devices = [int(i) for i in env_devices.split(",")] + devices = [idx for idx in devices if idx in new_devices] + self._log("debug", f"Active GPU Devices: {devices}") + return devices + + def _get_devices(self) -> list[Device]: + """ Obtain all detected DX API devices. + + Returns + ------- + list + The :class:`~dx_lib.Device` objects for GPUs that the DX API has discovered. + """ + adapters = Adapters(log_func=self._log) + devices = adapters.valid_adapters + self._log("debug", f"Obtained Devices: {devices}") + return devices + + def _initialize(self) -> None: + """ Initialize DX Core for DirectML backend. + + If :attr:`_is_initialized` is ``True`` then this function just returns performing no + action. + + if ``False`` then DirectML is setup, if not already, and GPU information is extracted + from the DirectML context. + """ + if self._is_initialized: + return + self._log("debug", "Initializing Win DX API for DirectML.") + self._devices = self._get_devices() + super()._initialize() + + def _get_device_count(self) -> int: + """ Detect the number of GPUs available from the DX API. + + Returns + ------- + int + The total number of GPUs available + """ + retval = len(self._devices) + self._log("debug", f"GPU Device count: {retval}") + return retval + + def _get_handles(self) -> list: + """ The DX API doesn't really use device handles, so we just return the all devices list + + Returns + ------- + list + The list of all discovered GPUs + """ + handles = self._devices + self._log("debug", f"DirectML GPU Handles found: {handles}") + return handles + + def _get_driver(self) -> str: + """ Obtain the driver versions currently in use. + + Returns + ------- + str + The current DirectX 12 GPU driver versions + """ + drivers = "|".join([device.driver_version if device.driver_version else "No Driver Found" + for device in self._devices]) + self._log("debug", f"GPU Drivers: {drivers}") + return drivers + + def _get_device_names(self) -> list[str]: + """ Obtain the list of names of connected GPUs as identified in :attr:`_handles`. + + Returns + ------- + list + The list of connected Nvidia GPU names + """ + names = self.names + self._log("debug", f"GPU Devices: {names}") + return names + + def _get_vram(self) -> list[int]: + """ Obtain the VRAM in Megabytes for each connected DirectML GPU as identified in + :attr:`_handles`. + + Returns + ------- + list + The VRAM in Megabytes for each connected Nvidia GPU + """ + vram = self._all_vram + self._log("debug", f"GPU VRAM: {vram}") + return vram + + def _get_free_vram(self) -> list[int]: + """ Obtain the amount of VRAM that is available, in Megabytes, for each connected DirectX + 12 supporting GPU. + + Returns + ------- + list + List of `float`s containing the amount of VRAM available, in Megabytes, for each + connected GPU as corresponding to the values in :attr:`_handles + """ + vram = [int(device.local_mem.Budget / (1024 * 1024)) for device in self._devices] + self._log("debug", f"GPU VRAM free: {vram}") + return vram diff --git a/lib/gpu_stats/nvidia.py b/lib/gpu_stats/nvidia.py index a1d9d328ef..3347399d33 100644 --- a/lib/gpu_stats/nvidia.py +++ b/lib/gpu_stats/nvidia.py @@ -1,16 +1,15 @@ #!/usr/bin/env python3 """ Collects and returns Information on available Nvidia GPUs. """ - -from typing import List +import os import pynvml from lib.utils import FaceswapError -from ._base import GPUStats +from ._base import _GPUStats -class NvidiaStats(GPUStats): +class NvidiaStats(_GPUStats): """ Holds information and statistics about Nvidia GPU(s) available on the currently running system. @@ -54,7 +53,7 @@ def _initialize(self) -> None: "remove and reinstall your Nvidia drivers before reporting. Original " f"Error: {str(err)}") raise FaceswapError(msg) from err - except Exception as err: # pylint: disable=broad-except + except Exception as err: # pylint:disable=broad-except msg = ("An unhandled exception occured reading from the Nvidia Machine Learning " f"Library. Original error: {str(err)}") raise FaceswapError(msg) from err @@ -83,6 +82,24 @@ def _get_device_count(self) -> int: self._log("debug", f"GPU Device count: {retval}") return retval + def _get_active_devices(self) -> list[int]: + """ Obtain the indices of active GPUs (those that have not been explicitly excluded by + CUDA_VISIBLE_DEVICES environment variable or explicitly excluded in the command line + arguments). + + Returns + ------- + list + The list of device indices that are available for Faceswap to use + """ + devices = super()._get_active_devices() + env_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if env_devices: + new_devices = [int(i) for i in env_devices.split(",")] + devices = [idx for idx in devices if idx in new_devices] + self._log("debug", f"Active GPU Devices: {devices}") + return devices + def _get_handles(self) -> list: """ Obtain the device handles for all connected Nvidia GPUs. @@ -105,14 +122,14 @@ def _get_driver(self) -> str: The current GPU driver version """ try: - driver = pynvml.nvmlSystemGetDriverVersion().decode("utf-8") + driver = pynvml.nvmlSystemGetDriverVersion() except pynvml.NVMLError as err: self._log("debug", f"Unable to obtain driver. Original error: {str(err)}") driver = "No Nvidia driver found" self._log("debug", f"GPU Driver: {driver}") return driver - def _get_device_names(self) -> List[str]: + def _get_device_names(self) -> list[str]: """ Obtain the list of names of connected Nvidia GPUs as identified in :attr:`_handles`. Returns @@ -120,12 +137,12 @@ def _get_device_names(self) -> List[str]: list The list of connected Nvidia GPU names """ - names = [pynvml.nvmlDeviceGetName(handle).decode("utf-8") + names = [pynvml.nvmlDeviceGetName(handle) for handle in self._handles] self._log("debug", f"GPU Devices: {names}") return names - def _get_vram(self) -> List[float]: + def _get_vram(self) -> list[int]: """ Obtain the VRAM in Megabytes for each connected Nvidia GPU as identified in :attr:`_handles`. @@ -139,7 +156,7 @@ def _get_vram(self) -> List[float]: self._log("debug", f"GPU VRAM: {vram}") return vram - def _get_free_vram(self) -> List[float]: + def _get_free_vram(self) -> list[int]: """ Obtain the amount of VRAM that is available, in Megabytes, for each connected Nvidia GPU. @@ -149,7 +166,15 @@ def _get_free_vram(self) -> List[float]: List of `float`s containing the amount of VRAM available, in Megabytes, for each connected GPU as corresponding to the values in :attr:`_handles """ + is_initialized = self._is_initialized + if not is_initialized: + self._initialize() + self._handles = self._get_handles() + vram = [pynvml.nvmlDeviceGetMemoryInfo(handle).free / (1024 * 1024) for handle in self._handles] + if not is_initialized: + self._shutdown() + self._log("debug", f"GPU VRAM free: {vram}") return vram diff --git a/lib/gpu_stats/nvidia_apple.py b/lib/gpu_stats/nvidia_apple.py index 9c5c6c616e..acbcd93f58 100644 --- a/lib/gpu_stats/nvidia_apple.py +++ b/lib/gpu_stats/nvidia_apple.py @@ -1,15 +1,13 @@ #!/usr/bin/env python3 """ Collects and returns Information on available Nvidia GPUs connected to Apple Macs. """ -from typing import List - import pynvx from lib.utils import FaceswapError -from ._base import GPUStats +from ._base import _GPUStats -class NvidiaAppleStats(GPUStats): +class NvidiaAppleStats(_GPUStats): """ Holds information and statistics about Nvidia GPU(s) available on the currently running Apple system. @@ -92,7 +90,7 @@ def _get_driver(self) -> str: self._log("debug", f"GPU Driver: {driver}") return driver - def _get_device_names(self) -> List[str]: + def _get_device_names(self) -> list[str]: """ Obtain the list of names of connected Nvidia GPUs as identified in :attr:`_handles`. Returns @@ -105,7 +103,7 @@ def _get_device_names(self) -> List[str]: self._log("debug", f"GPU Devices: {names}") return names - def _get_vram(self) -> List[float]: + def _get_vram(self) -> list[int]: """ Obtain the VRAM in Megabytes for each connected Nvidia GPU as identified in :attr:`_handles`. @@ -120,7 +118,7 @@ def _get_vram(self) -> List[float]: self._log("debug", f"GPU VRAM: {vram}") return vram - def _get_free_vram(self) -> List[float]: + def _get_free_vram(self) -> list[int]: """ Obtain the amount of VRAM that is available, in Megabytes, for each connected Nvidia GPU. diff --git a/lib/gpu_stats/rocm.py b/lib/gpu_stats/rocm.py new file mode 100644 index 0000000000..dca43b3818 --- /dev/null +++ b/lib/gpu_stats/rocm.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +""" Collects and returns Information about connected AMD GPUs for ROCm using sysfs and from +modinfo + +As no ROCm compatible hardware was available for testing, this just returns information on all AMD +GPUs discovered on the system regardless of ROCm compatibility. + +It is a good starting point but may need to be refined over time +""" +import os +import re +from subprocess import run + +from ._base import _GPUStats + +_DEVICE_LOOKUP = { # ref: https://gist.github.com/roalercon/51f13a387f3754615cce + int("0x130F", 0): "AMD Radeon(TM) R7 Graphics", + int("0x1313", 0): "AMD Radeon(TM) R7 Graphics", + int("0x1316", 0): "AMD Radeon(TM) R5 Graphics", + int("0x6600", 0): "AMD Radeon HD 8600/8700M", + int("0x6601", 0): "AMD Radeon (TM) HD 8500M/8700M", + int("0x6604", 0): "AMD Radeon R7 M265 Series", + int("0x6605", 0): "AMD Radeon R7 M260 Series", + int("0x6606", 0): "AMD Radeon HD 8790M", + int("0x6607", 0): "AMD Radeon (TM) HD8530M", + int("0x6610", 0): "AMD Radeon HD 8670 Graphics", + int("0x6611", 0): "AMD Radeon HD 8570 Graphics", + int("0x6613", 0): "AMD Radeon R7 200 Series", + int("0x6640", 0): "AMD Radeon HD 8950", + int("0x6658", 0): "AMD Radeon R7 200 Series", + int("0x665C", 0): "AMD Radeon HD 7700 Series", + int("0x665D", 0): "AMD Radeon R7 200 Series", + int("0x6660", 0): "AMD Radeon HD 8600M Series", + int("0x6663", 0): "AMD Radeon HD 8500M Series", + int("0x6664", 0): "AMD Radeon R5 M200 Series", + int("0x6665", 0): "AMD Radeon R5 M230 Series", + int("0x6667", 0): "AMD Radeon R5 M200 Series", + int("0x666F", 0): "AMD Radeon HD 8500M", + int("0x6704", 0): "AMD FirePro V7900 (FireGL V)", + int("0x6707", 0): "AMD FirePro V5900 (FireGL V)", + int("0x6718", 0): "AMD Radeon HD 6900 Series", + int("0x6719", 0): "AMD Radeon HD 6900 Series", + int("0x671D", 0): "AMD Radeon HD 6900 Series", + int("0x671F", 0): "AMD Radeon HD 6900 Series", + int("0x6720", 0): "AMD Radeon HD 6900M Series", + int("0x6738", 0): "AMD Radeon HD 6800 Series", + int("0x6739", 0): "AMD Radeon HD 6800 Series", + int("0x673E", 0): "AMD Radeon HD 6700 Series", + int("0x6740", 0): "AMD Radeon HD 6700M Series", + int("0x6741", 0): "AMD Radeon 6600M and 6700M Series", + int("0x6742", 0): "AMD Radeon HD 5570", + int("0x6743", 0): "AMD Radeon E6760", + int("0x6749", 0): "AMD FirePro V4900 (FireGL V)", + int("0x674A", 0): "AMD FirePro V3900 (ATI FireGL)", + int("0x6750", 0): "AMD Radeon HD 6500 series", + int("0x6751", 0): "AMD Radeon HD 7600A Series", + int("0x6758", 0): "AMD Radeon HD 6670", + int("0x6759", 0): "AMD Radeon HD 6570 Graphics", + int("0x675B", 0): "AMD Radeon HD 7600 Series", + int("0x675D", 0): "AMD Radeon HD 7500 Series", + int("0x675F", 0): "AMD Radeon HD 5500 Series", + int("0x6760", 0): "AMD Radeon HD 6400M Series", + int("0x6761", 0): "AMD Radeon HD 6430M", + int("0x6763", 0): "AMD Radeon E6460", + int("0x6770", 0): "AMD Radeon HD 6400 Series", + int("0x6771", 0): "AMD Radeon R5 235X", + int("0x6772", 0): "AMD Radeon HD 7400A Series", + int("0x6778", 0): "AMD Radeon HD 7000 series", + int("0x6779", 0): "AMD Radeon HD 6450", + int("0x677B", 0): "AMD Radeon HD 7400 Series", + int("0x6780", 0): "AMD FirePro W9000 (FireGL V)", + int("0x678A", 0): "AMD FirePro S10000 (FireGL V)", + int("0x6798", 0): "AMD Radeon HD 7900 Series", + int("0x679A", 0): "AMD Radeon HD 7900 Series", + int("0x679B", 0): "AMD Radeon HD 7900 Series", + int("0x679E", 0): "AMD Radeon HD 7800 Series", + int("0x67B0", 0): "AMD Radeon R9 200 Series", + int("0x67B1", 0): "AMD Radeon R9 200 Series", + int("0x6800", 0): "AMD Radeon HD 7970M", + int("0x6801", 0): "AMD Radeon(TM) HD8970M", + int("0x6808", 0): "AMD FirePro S7000 (FireGL V)", + int("0x6809", 0): "AMD FirePro R5000 (FireGL V)", + int("0x6810", 0): "AMD Radeon R9 200 Series", + int("0x6811", 0): "AMD Radeon R9 200 Series", + int("0x6818", 0): "AMD Radeon HD 7800 Series", + int("0x6819", 0): "AMD Radeon HD 7800 Series", + int("0x6820", 0): "AMD Radeon HD 8800M Series", + int("0x6821", 0): "AMD Radeon HD 8800M Series", + int("0x6822", 0): "AMD Radeon E8860", + int("0x6823", 0): "AMD Radeon HD 8800M Series", + int("0x6825", 0): "AMD Radeon HD 7800M Series", + int("0x6827", 0): "AMD Radeon HD 7800M Series", + int("0x6828", 0): "AMD FirePro W600", + int("0x682B", 0): "AMD Radeon HD 8800M Series", + int("0x682D", 0): "AMD Radeon HD 7700M Series", + int("0x682F", 0): "AMD Radeon HD 7700M Series", + int("0x6835", 0): "AMD Radeon R7 Series / HD 9000 Series", + int("0x6837", 0): "AMD Radeon HD 6570", + int("0x683D", 0): "AMD Radeon HD 7700 Series", + int("0x683F", 0): "AMD Radeon HD 7700 Series", + int("0x6840", 0): "AMD Radeon HD 7600M Series", + int("0x6841", 0): "AMD Radeon HD 7500M/7600M Series", + int("0x6842", 0): "AMD Radeon HD 7000M Series", + int("0x6843", 0): "AMD Radeon HD 7670M", + int("0x6858", 0): "AMD Radeon HD 7400 Series", + int("0x6859", 0): "AMD Radeon HD 7400 Series", + int("0x6888", 0): "ATI FirePro V8800 (FireGL V)", + int("0x6889", 0): "ATI FirePro V7800 (FireGL V)", + int("0x688A", 0): "ATI FirePro V9800 (FireGL V)", + int("0x688C", 0): "AMD FireStream 9370", + int("0x688D", 0): "AMD FireStream 9350", + int("0x6898", 0): "AMD Radeon HD 5800 Series", + int("0x6899", 0): "AMD Radeon HD 5800 Series", + int("0x689B", 0): "AMD Radeon HD 6800 Series", + int("0x689C", 0): "AMD Radeon HD 5900 Series", + int("0x689E", 0): "AMD Radeon HD 5800 Series", + int("0x68A0", 0): "AMD Mobility Radeon HD 5800 Series", + int("0x68A1", 0): "AMD Mobility Radeon HD 5800 Series", + int("0x68A8", 0): "AMD Radeon HD 6800M Series", + int("0x68A9", 0): "ATI FirePro V5800 (FireGL V)", + int("0x68B8", 0): "AMD Radeon HD 5700 Series", + int("0x68B9", 0): "AMD Radeon HD 5600/5700", + int("0x68BA", 0): "AMD Radeon HD 6700 Series", + int("0x68BE", 0): "AMD Radeon HD 5700 Series", + int("0x68BF", 0): "AMD Radeon HD 6700 Green Edition", + int("0x68C0", 0): "AMD Mobility Radeon HD 5000", + int("0x68C1", 0): "AMD Mobility Radeon HD 5000 Series", + int("0x68C7", 0): "AMD Mobility Radeon HD 5570", + int("0x68C8", 0): "ATI FirePro V4800 (FireGL V)", + int("0x68C9", 0): "ATI FirePro 3800 (FireGL) Graphics Adapter", + int("0x68D8", 0): "AMD Radeon HD 5670", + int("0x68D9", 0): "AMD Radeon HD 5570", + int("0x68DA", 0): "AMD Radeon HD 5500 Series", + int("0x68E0", 0): "AMD Mobility Radeon HD 5000 Series", + int("0x68E1", 0): "AMD Mobility Radeon HD 5000 Series", + int("0x68E4", 0): "AMD Radeon HD 5450", + int("0x68E5", 0): "AMD Radeon HD 6300M Series", + int("0x68F1", 0): "AMD FirePro 2460", + int("0x68F2", 0): "AMD FirePro 2270 (ATI FireGL)", + int("0x68F9", 0): "AMD Radeon HD 5450", + int("0x68FA", 0): "AMD Radeon HD 7300 Series", + int("0x9640", 0): "AMD Radeon HD 6550D", + int("0x9641", 0): "AMD Radeon HD 6620G", + int("0x9642", 0): "AMD Radeon HD 6370D", + int("0x9643", 0): "AMD Radeon HD 6380G", + int("0x9644", 0): "AMD Radeon HD 6410D", + int("0x9645", 0): "AMD Radeon HD 6410D", + int("0x9647", 0): "AMD Radeon HD 6520G", + int("0x9648", 0): "AMD Radeon HD 6480G", + int("0x9649", 0): "AMD Radeon(TM) HD 6480G", + int("0x964A", 0): "AMD Radeon HD 6530D", + int("0x9802", 0): "AMD Radeon HD 6310 Graphics", + int("0x9803", 0): "AMD Radeon HD 6250 Graphics", + int("0x9804", 0): "AMD Radeon HD 6250 Graphics", + int("0x9805", 0): "AMD Radeon HD 6250 Graphics", + int("0x9806", 0): "AMD Radeon HD 6320 Graphics", + int("0x9807", 0): "AMD Radeon HD 6290 Graphics", + int("0x9808", 0): "AMD Radeon HD 7340 Graphics", + int("0x9809", 0): "AMD Radeon HD 7310 Graphics", + int("0x980A", 0): "AMD Radeon HD 7290 Graphics", + int("0x9830", 0): "AMD Radeon HD 8400", + int("0x9831", 0): "AMD Radeon(TM) HD 8400E", + int("0x9832", 0): "AMD Radeon HD 8330", + int("0x9833", 0): "AMD Radeon(TM) HD 8330E", + int("0x9834", 0): "AMD Radeon HD 8210", + int("0x9835", 0): "AMD Radeon(TM) HD 8210E", + int("0x9836", 0): "AMD Radeon HD 8280", + int("0x9837", 0): "AMD Radeon(TM) HD 8280E", + int("0x9838", 0): "AMD Radeon HD 8240", + int("0x9839", 0): "AMD Radeon HD 8180", + int("0x983D", 0): "AMD Radeon HD 8250", + int("0x9900", 0): "AMD Radeon HD 7660G", + int("0x9901", 0): "AMD Radeon HD 7660D", + int("0x9903", 0): "AMD Radeon HD 7640G", + int("0x9904", 0): "AMD Radeon HD 7560D", + int("0x9906", 0): "AMD FirePro A300 Series (FireGL V) Graphics Adapter", + int("0x9907", 0): "AMD Radeon HD 7620G", + int("0x9908", 0): "AMD Radeon HD 7600G", + int("0x990A", 0): "AMD Radeon HD 7500G", + int("0x990B", 0): "AMD Radeon HD 8650G", + int("0x990C", 0): "AMD Radeon HD 8670D", + int("0x990D", 0): "AMD Radeon HD 8550G", + int("0x990E", 0): "AMD Radeon HD 8570D", + int("0x990F", 0): "AMD Radeon HD 8610G", + int("0x9910", 0): "AMD Radeon HD 7660G", + int("0x9913", 0): "AMD Radeon HD 7640G", + int("0x9917", 0): "AMD Radeon HD 7620G", + int("0x9918", 0): "AMD Radeon HD 7600G", + int("0x9919", 0): "AMD Radeon HD 7500G", + int("0x9990", 0): "AMD Radeon HD 7520G", + int("0x9991", 0): "AMD Radeon HD 7540D", + int("0x9992", 0): "AMD Radeon HD 7420G", + int("0x9993", 0): "AMD Radeon HD 7480D", + int("0x9994", 0): "AMD Radeon HD 7400G", + int("0x9995", 0): "AMD Radeon HD 8450G", + int("0x9996", 0): "AMD Radeon HD 8470D", + int("0x9997", 0): "AMD Radeon HD 8350G", + int("0x9998", 0): "AMD Radeon HD 8370D", + int("0x9999", 0): "AMD Radeon HD 8510G", + int("0x999A", 0): "AMD Radeon HD 8410G", + int("0x999B", 0): "AMD Radeon HD 8310G", + int("0x999C", 0): "AMD Radeon HD 8650D", + int("0x999D", 0): "AMD Radeon HD 8550D", + int("0x99A0", 0): "AMD Radeon HD 7520G", + int("0x99A2", 0): "AMD Radeon HD 7420G", + int("0x99A4", 0): "AMD Radeon HD 7400G"} + + +class ROCm(_GPUStats): + """ Holds information and statistics about GPUs connected using sysfs + + Parameters + ---------- + log: bool, optional + Whether the class should output information to the logger. There may be occasions where the + logger has not yet been set up when this class is queried. Attempting to log in these + instances will raise an error. If GPU stats are being queried prior to the logger being + available then this parameter should be set to ``False``. Otherwise set to ``True``. + Default: ``True`` + """ + def __init__(self, log: bool = True) -> None: + self._vendor_id = "0x1002" # AMD VendorID + self._sysfs_paths: list[str] = [] + super().__init__(log=log) + + def _from_sysfs_file(self, path: str) -> str: + """ Obtain the value from a sysfs file. On permission error or file doesn't exist, log and + return empty value + + Parameters + ---------- + path: str + The path to a sysfs file to obtain the value from + + Returns + ------- + str + The obtained value from the given path + """ + if not os.path.isfile(path): + self._log("debug", f"File '{path}' does not exist. Returning empty string") + return "" + try: + with open(path, "r", encoding="utf-8", errors="ignore") as sysfile: + val = sysfile.read().strip() + except PermissionError: + self._log("debug", f"Permission error accessing file '{path}'. Returning empty string") + val = "" + return val + + def _get_sysfs_paths(self) -> list[str]: + """ Obtain a list of sysfs paths to AMD branded GPUs connected to the system + + Returns + ------- + list[str] + List of full paths to the sysfs entries for connected AMD GPUs + """ + base_dir = "/sys/class/drm/" + + retval: list[str] = [] + if not os.path.exists(base_dir): + self._log("warning", f"sysfs not found at '{base_dir}'") + return retval + + for folder in sorted(os.listdir(base_dir)): + folder_path = os.path.join(base_dir, folder, "device") + vendor_path = os.path.join(folder_path, "vendor") + if not os.path.isdir(vendor_path) and not re.match(r"^card\d+$", folder): + self._log("debug", f"skipping path '{folder_path}'") + continue + + vendor_id = self._from_sysfs_file(vendor_path) + if vendor_id != self._vendor_id: + self._log("debug", f"Skipping non AMD Vendor '{vendor_id}' for device: '{folder}'") + continue + + retval.append(folder_path) + + self._log("debug", f"sysfs AMD devices: {retval}") + return retval + + def _initialize(self) -> None: + """ Initialize sysfs for ROCm backend. + + If :attr:`_is_initialized` is ``True`` then this function just returns performing no + action. + + if ``False`` then the location of AMD cards within sysfs is collected + """ + if self._is_initialized: + return + self._log("debug", "Initializing sysfs for AMDGPU (ROCm).") + self._sysfs_paths = self._get_sysfs_paths() + super()._initialize() + + def _get_device_count(self) -> int: + """ The number of AMD cards found in sysfs + + Returns + ------- + int + The total number of GPUs available + """ + retval = len(self._sysfs_paths) + self._log("debug", f"GPU Device count: {retval}") + return retval + + def _get_handles(self) -> list: + """ The sysfs doesn't use device handles, so we just return the list of the sysfs locations + per card + + Returns + ------- + list + The list of all discovered GPUs + """ + handles = self._sysfs_paths + self._log("debug", f"sysfs GPU Handles found: {handles}") + return handles + + def _get_driver(self) -> str: + """ Obtain the driver versions currently in use from modinfo + + Returns + ------- + str + The current AMDGPU driver versions + """ + retval = "" + cmd = ["modinfo", "amdgpu"] + try: + proc = run(cmd, + check=True, + timeout=5, + capture_output=True, + encoding="utf-8", + errors="ignore") + for line in proc.stdout.split("\n"): + if line.startswith("version:"): + retval = line.split()[-1] + break + except Exception as err: # pylint:disable=broad-except + self._log("debug", f"Error reading modinfo: '{str(err)}'") + + self._log("debug", f"GPU Drivers: {retval}") + return retval + + def _get_device_names(self) -> list[str]: + """ Obtain the list of names of connected GPUs as identified in :attr:`_handles`. + + Returns + ------- + list + The list of connected AMD GPU names + """ + retval = [] + for device in self._sysfs_paths: + name = self._from_sysfs_file(os.path.join(device, "product_name")) + number = self._from_sysfs_file(os.path.join(device, "product_number")) + if name or number: # product_name or product_number populated + self._log("debug", f"Got name from product_name: '{name}', product_number: " + f"'{number}'") + retval.append(f"{name + ' ' if name else ''}{number}") + continue + + device_id = self._from_sysfs_file(os.path.join(device, "device")) + self._log("debug", f"Got device_id: '{device_id}'") + + if not device_id: # Can't get device name + retval.append("Not found") + continue + try: + lookup = int(device_id, 0) + except ValueError: + retval.append(device_id) + continue + + device_name = _DEVICE_LOOKUP.get(lookup, device_id) + retval.append(device_name) + + self._log("debug", f"Device names: {retval}") + return retval + + def _get_active_devices(self) -> list[int]: + """ Obtain the indices of active GPUs (those that have not been explicitly excluded by + HIP_VISIBLE_DEVICES environment variable or explicitly excluded in the command line + arguments). + + Returns + ------- + list + The list of device indices that are available for Faceswap to use + """ + devices = super()._get_active_devices() + env_devices = os.environ.get("HIP_VISIBLE_DEVICES ") + if env_devices: + new_devices = [int(i) for i in env_devices.split(",")] + devices = [idx for idx in devices if idx in new_devices] + self._log("debug", f"Active GPU Devices: {devices}") + return devices + + def _get_vram(self) -> list[int]: + """ Obtain the VRAM in Megabytes for each connected AMD GPU as identified in + :attr:`_handles`. + + Returns + ------- + list + The VRAM in Megabytes for each connected Nvidia GPU + """ + retval = [] + for device in self._sysfs_paths: + query = self._from_sysfs_file(os.path.join(device, "mem_info_vram_total")) + try: + vram = int(query) + except ValueError: + self._log("debug", f"Couldn't extract VRAM from string: '{query}'", ) + vram = 0 + retval.append(int(vram / (1024 * 1024))) + + self._log("debug", f"GPU VRAM: {retval}") + return retval + + def _get_free_vram(self) -> list[int]: + """ Obtain the amount of VRAM that is available, in Megabytes, for each connected AMD + GPU. + + Returns + ------- + list + List of `float`s containing the amount of VRAM available, in Megabytes, for each + connected GPU as corresponding to the values in :attr:`_handles + """ + retval = [] + total_vram = self._get_vram() + for device, vram in zip(self._sysfs_paths, total_vram): + if not vram: + retval.append(0) + continue + query = self._from_sysfs_file(os.path.join(device, "mem_info_vram_used")) + try: + used = int(query) + except ValueError: + self._log("debug", f"Couldn't extract used VRAM from string: '{query}'") + used = 0 + + retval.append(vram - int(used / (1024 * 1024))) + self._log("debug", f"GPU VRAM free: {retval}") + return retval diff --git a/lib/gui/.cache/presets/train/model_phaze_a_clipfaker128_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker128_preset.json new file mode 100644 index 0000000000..0afede8da7 --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker128_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 128, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "clipv_farl-b-16-64", + "enc_scaling": 29, + "enc_load_weights": true, + "bottleneck_type": "flatten", + "bottleneck_norm": "none", + "bottleneck_size": 1024, + "bottleneck_in_encoder": true, + "fc_depth": 1, + "fc_min_filters": 1024, + "fc_max_filters": 1024, + "fc_dimensions": 4, + "fc_filter_slope": -0.5, + "fc_dropout": 0.0, + "fc_upsampler": "subpixel", + "fc_upsamples": 1, + "fc_upsample_filters": 512, + "fc_gblock_depth": 3, + "fc_gblock_min_nodes": 512, + "fc_gblock_max_nodes": 512, + "fc_gblock_filter_slope": -0.5, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "subpixel", + "dec_upscales_in_fc": 0, + "dec_norm": "none", + "dec_min_filters": 64, + "dec_max_filters": 512, + "dec_slope_mode": "cap_min", + "dec_filter_slope": 0.5, + "dec_res_blocks": 1, + "dec_output_kernel": 5, + "dec_gaussian": false, + "dec_skip_last_residual": true, + "freeze_layers": "keras_encoder", + "load_layers": "encoder", + "fs_original_depth": 4, + "fs_original_min_filters": 128, + "fs_original_max_filters": 1024, + "fs_original_use_alt": false, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_clipfaker256_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker256_preset.json new file mode 100644 index 0000000000..974b614b97 --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker256_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 256, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "clipv_farl-b-16-64", + "enc_scaling": 58, + "enc_load_weights": true, + "bottleneck_type": "flatten", + "bottleneck_norm": "none", + "bottleneck_size": 1024, + "bottleneck_in_encoder": true, + "fc_depth": 1, + "fc_min_filters": 1024, + "fc_max_filters": 1024, + "fc_dimensions": 4, + "fc_filter_slope": -0.5, + "fc_dropout": 0.0, + "fc_upsampler": "subpixel", + "fc_upsamples": 1, + "fc_upsample_filters": 512, + "fc_gblock_depth": 3, + "fc_gblock_min_nodes": 512, + "fc_gblock_max_nodes": 512, + "fc_gblock_filter_slope": -0.5, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "subpixel", + "dec_upscales_in_fc": 0, + "dec_norm": "none", + "dec_min_filters": 64, + "dec_max_filters": 1024, + "dec_slope_mode": "cap_min", + "dec_filter_slope": 0.5, + "dec_res_blocks": 1, + "dec_output_kernel": 5, + "dec_gaussian": false, + "dec_skip_last_residual": true, + "freeze_layers": "keras_encoder", + "load_layers": "encoder", + "fs_original_depth": 4, + "fs_original_min_filters": 128, + "fs_original_max_filters": 1024, + "fs_original_use_alt": false, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_clipfaker448_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker448_preset.json new file mode 100644 index 0000000000..59bfedce6b --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_clipfaker448_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 448, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "clipv_farl-b-16-64", + "enc_scaling": 100, + "enc_load_weights": true, + "bottleneck_type": "flatten", + "bottleneck_norm": "none", + "bottleneck_size": 1024, + "bottleneck_in_encoder": true, + "fc_depth": 1, + "fc_min_filters": 384, + "fc_max_filters": 384, + "fc_dimensions": 7, + "fc_filter_slope": -0.5, + "fc_dropout": 0.0, + "fc_upsampler": "subpixel", + "fc_upsamples": 1, + "fc_upsample_filters": 1024, + "fc_gblock_depth": 3, + "fc_gblock_min_nodes": 512, + "fc_gblock_max_nodes": 512, + "fc_gblock_filter_slope": -0.5, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "subpixel", + "dec_upscales_in_fc": 0, + "dec_norm": "none", + "dec_min_filters": 64, + "dec_max_filters": 1024, + "dec_slope_mode": "cap_min", + "dec_filter_slope": 0.5, + "dec_res_blocks": 1, + "dec_output_kernel": 5, + "dec_gaussian": false, + "dec_skip_last_residual": true, + "freeze_layers": "keras_encoder", + "load_layers": "encoder", + "fs_original_depth": 4, + "fs_original_min_filters": 128, + "fs_original_max_filters": 1024, + "fs_original_use_alt": false, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_dny1024_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_dny1024_preset.json new file mode 100644 index 0000000000..161d53336c --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_dny1024_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 1024, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "fs_original", + "enc_scaling": 100, + "enc_load_weights": false, + "bottleneck_type": "dense", + "bottleneck_norm": "none", + "bottleneck_size": 512, + "bottleneck_in_encoder": true, + "fc_depth": 0, + "fc_min_filters": 512, + "fc_max_filters": 512, + "fc_dimensions": 1, + "fc_filter_slope": 0.0, + "fc_dropout": 0.0, + "fc_upsampler": "upsample2d", + "fc_upsamples": 2, + "fc_upsample_filters": 128, + "fc_gblock_depth": 1, + "fc_gblock_min_nodes": 128, + "fc_gblock_max_nodes": 128, + "fc_gblock_filter_slope": 0.0, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "upscale_dny", + "dec_upscales_in_fc": 2, + "dec_norm": "none", + "dec_min_filters": 16, + "dec_max_filters": 512, + "dec_slope_mode": "cap_max", + "dec_filter_slope": 0.5, + "dec_res_blocks": 0, + "dec_output_kernel": 1, + "dec_gaussian": false, + "dec_skip_last_residual": false, + "freeze_layers": "encoder", + "load_layers": "encoder", + "fs_original_depth": 9, + "fs_original_min_filters": 16, + "fs_original_max_filters": 512, + "fs_original_use_alt": true, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_dny256_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_dny256_preset.json new file mode 100644 index 0000000000..e19e61dcf7 --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_dny256_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 256, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "fs_original", + "enc_scaling": 25, + "enc_load_weights": false, + "bottleneck_type": "dense", + "bottleneck_norm": "none", + "bottleneck_size": 512, + "bottleneck_in_encoder": true, + "fc_depth": 0, + "fc_min_filters": 512, + "fc_max_filters": 512, + "fc_dimensions": 1, + "fc_filter_slope": 0.0, + "fc_dropout": 0.0, + "fc_upsampler": "upsample2d", + "fc_upsamples": 2, + "fc_upsample_filters": 128, + "fc_gblock_depth": 1, + "fc_gblock_min_nodes": 128, + "fc_gblock_max_nodes": 128, + "fc_gblock_filter_slope": 0.0, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "upscale_dny", + "dec_upscales_in_fc": 1, + "dec_norm": "none", + "dec_min_filters": 16, + "dec_max_filters": 512, + "dec_slope_mode": "cap_max", + "dec_filter_slope": 0.5, + "dec_res_blocks": 0, + "dec_output_kernel": 1, + "dec_gaussian": false, + "dec_skip_last_residual": false, + "freeze_layers": "encoder", + "load_layers": "encoder", + "fs_original_depth": 7, + "fs_original_min_filters": 16, + "fs_original_max_filters": 512, + "fs_original_use_alt": true, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_dny512_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_dny512_preset.json new file mode 100644 index 0000000000..9e0534d5f9 --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_dny512_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 512, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "fs_original", + "enc_scaling": 50, + "enc_load_weights": false, + "bottleneck_type": "dense", + "bottleneck_norm": "none", + "bottleneck_size": 512, + "bottleneck_in_encoder": true, + "fc_depth": 0, + "fc_min_filters": 512, + "fc_max_filters": 512, + "fc_dimensions": 1, + "fc_filter_slope": 0.0, + "fc_dropout": 0.0, + "fc_upsampler": "upsample2d", + "fc_upsamples": 2, + "fc_upsample_filters": 128, + "fc_gblock_depth": 1, + "fc_gblock_min_nodes": 128, + "fc_gblock_max_nodes": 128, + "fc_gblock_filter_slope": 0.0, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "upscale_dny", + "dec_upscales_in_fc": 2, + "dec_norm": "none", + "dec_min_filters": 16, + "dec_max_filters": 512, + "dec_slope_mode": "cap_max", + "dec_filter_slope": 0.5, + "dec_res_blocks": 0, + "dec_output_kernel": 1, + "dec_gaussian": false, + "dec_skip_last_residual": false, + "freeze_layers": "encoder", + "load_layers": "encoder", + "fs_original_depth": 8, + "fs_original_min_filters": 16, + "fs_original_max_filters": 512, + "fs_original_use_alt": true, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} \ No newline at end of file diff --git a/lib/gui/.cache/presets/train/model_phaze_a_sym384_preset.json b/lib/gui/.cache/presets/train/model_phaze_a_sym384_preset.json new file mode 100644 index 0000000000..e836a856f8 --- /dev/null +++ b/lib/gui/.cache/presets/train/model_phaze_a_sym384_preset.json @@ -0,0 +1,52 @@ +{ + "output_size": 384, + "shared_fc": "none", + "enable_gblock": false, + "split_fc": false, + "split_gblock": false, + "split_decoders": true, + "enc_architecture": "efficientnet_v2_s", + "enc_scaling": 100, + "enc_load_weights": true, + "bottleneck_type": "max_pooling", + "bottleneck_norm": "none", + "bottleneck_size": 1280, + "bottleneck_in_encoder": true, + "fc_depth": 1, + "fc_min_filters": 1536, + "fc_max_filters": 1536, + "fc_dimensions": 3, + "fc_filter_slope": -0.5, + "fc_dropout": 0.0, + "fc_upsampler": "subpixel", + "fc_upsamples": 0, + "fc_upsample_filters": 1280, + "fc_gblock_depth": 3, + "fc_gblock_min_nodes": 512, + "fc_gblock_max_nodes": 512, + "fc_gblock_filter_slope": -0.5, + "fc_gblock_dropout": 0.0, + "dec_upscale_method": "upscale_dny", + "dec_upscales_in_fc": 2, + "dec_norm": "none", + "dec_min_filters": 24, + "dec_max_filters": 1536, + "dec_slope_mode": "cap_max", + "dec_filter_slope": 0.5, + "dec_res_blocks": 1, + "dec_output_kernel": 3, + "dec_gaussian": true, + "dec_skip_last_residual": true, + "freeze_layers": "keras_encoder", + "load_layers": "encoder", + "fs_original_depth": 4, + "fs_original_min_filters": 128, + "fs_original_max_filters": 1024, + "fs_original_use_alt": false, + "mobilenet_width": 1.0, + "mobilenet_depth": 1, + "mobilenet_dropout": 0.001, + "mobilenet_minimalistic": false, + "__filetype": "faceswap_preset", + "__section": "train|model|phaze_a" +} diff --git a/lib/gui/_config.py b/lib/gui/_config.py index 1c26a238f0..9f87e3b775 100644 --- a/lib/gui/_config.py +++ b/lib/gui/_config.py @@ -9,12 +9,12 @@ from lib.config import FaceswapConfig -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Config(FaceswapConfig): """ Config File for GUI """ - # pylint: disable=too-many-statements + # pylint:disable=too-many-statements def set_defaults(self): """ Set the default values for config """ logger.debug("Setting defaults") @@ -26,9 +26,9 @@ def set_globals(self): """ logger.debug("Setting global config") section = "global" - self.add_section(title=section, - info="Faceswap GUI Options.\nConfigure the appearance and behaviour of " - "the GUI") + self.add_section(section, + "Faceswap GUI Options.\nConfigure the appearance and behaviour of " + "the GUI") self.add_item( section=section, title="fullscreen", datatype=bool, default=False, group="startup", info="Start Faceswap maximized.") @@ -102,12 +102,12 @@ def get_clean_fonts(): A list of valid fonts for the system """ fmanager = font_manager.FontManager() - fonts = dict() + fonts = {} for font in fmanager.ttflist: if str(font.weight) in ("400", "normal", "regular"): - fonts.setdefault(font.name, dict())["regular"] = True + fonts.setdefault(font.name, {})["regular"] = True if str(font.weight) in ("700", "bold"): - fonts.setdefault(font.name, dict())["bold"] = True + fonts.setdefault(font.name, {})["bold"] = True valid_fonts = {key for key, val in fonts.items() if len(val) == 2} retval = sorted(list(valid_fonts.intersection(tk_font.families()))) if not retval: @@ -115,5 +115,5 @@ def get_clean_fonts(): # prefixed logger.debug("No bold/regular fonts found. Running simple filter") retval = sorted([fnt for fnt in tk_font.families() - if not fnt.startswith("@") and not any([ord(c) > 127 for c in fnt])]) + if not fnt.startswith("@") and not any(ord(c) > 127 for c in fnt)]) return ["default"] + retval diff --git a/lib/gui/analysis/event_reader.py b/lib/gui/analysis/event_reader.py index 820a7ea24c..60aa45e6e3 100644 --- a/lib/gui/analysis/event_reader.py +++ b/lib/gui/analysis/event_reader.py @@ -1,20 +1,42 @@ #!/usr/bin/env python3 """ Handles the loading and collation of events from Tensorflow event log files. """ - +from __future__ import annotations import logging import os +import re +import typing as T import zlib +from dataclasses import dataclass, field + import numpy as np import tensorflow as tf from tensorflow.core.util import event_pb2 # pylint:disable=no-name-in-module from tensorflow.python.framework import ( # pylint:disable=no-name-in-module errors_impl as tf_errors) +from lib.logger import parse_class_init from lib.serializer import get_serializer -from lib.utils import get_backend -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from collections.abc import Generator, Iterator + +logger = logging.getLogger(__name__) + + +@dataclass +class EventData: + """ Holds data collected from Tensorflow Event Files + + Parameters + ---------- + timestamp: float + The timestamp of the event step (iteration) + loss: list[float] + The loss values collected for A and B sides for the event step + """ + timestamp: float = 0.0 + loss: list[float] = field(default_factory=list) class _LogFiles(): @@ -25,27 +47,27 @@ class _LogFiles(): logs_folder: str The folder that contains the Tensorboard log files """ - def __init__(self, logs_folder): - logger.debug("Initializing: %s: (logs_folder: '%s')", self.__class__.__name__, logs_folder) + def __init__(self, logs_folder: str) -> None: + logger.debug(parse_class_init(locals())) self._logs_folder = logs_folder self._filenames = self._get_log_filenames() logger.debug("Initialized: %s", self.__class__.__name__) @property - def session_ids(self): - """ list: Sorted list of `ints` of available session ids. """ + def session_ids(self) -> list[int]: + """ list[int]: Sorted list of `ints` of available session ids. """ return list(sorted(self._filenames)) - def _get_log_filenames(self): + def _get_log_filenames(self) -> dict[int, str]: """ Get the Tensorflow event filenames for all existing sessions. Returns ------- - dict + dict[int, str] The full path of each log file for each training session id that has been run """ logger.debug("Loading log filenames. base_dir: '%s'", self._logs_folder) - retval = {} + retval: dict[int, str] = {} for dirpath, _, filenames in os.walk(self._logs_folder): if not any(filename.startswith("events.out.tfevents") for filename in filenames): continue @@ -58,7 +80,7 @@ def _get_log_filenames(self): return retval @classmethod - def _get_session_id(cls, folder): + def _get_session_id(cls, folder: str) -> int | None: """ Obtain the session id for the given folder. Parameters @@ -68,7 +90,7 @@ def _get_session_id(cls, folder): Returns ------- - int + int or ``None`` The session ID for the given folder. If no session id can be determined, return ``None`` """ @@ -79,7 +101,7 @@ def _get_session_id(cls, folder): return retval @classmethod - def _get_log_filename(cls, folder, filenames): + def _get_log_filename(cls, folder: str, filenames: list[str]) -> str: """ Obtain the session log file for the given folder. If multiple log files exist for the given folder, then the most recent log file is used, as earlier files are assumed to be obsolete. @@ -88,25 +110,25 @@ def _get_log_filename(cls, folder, filenames): ---------- folder: str The full path to the folder that contains the session's Tensorflow Event Log - filenames: list + filenames: list[str] List of filenames that exist within the given folder Returns ------- str - The full path the the selected log file + The full path of the selected log file """ logfiles = [fname for fname in filenames if fname.startswith("events.out.tfevents")] retval = os.path.join(folder, sorted(logfiles)[-1]) # Take last item if multi matches logger.debug("logfiles: %s, selected: '%s'", logfiles, retval) return retval - def refresh(self): + def refresh(self) -> None: """ Refresh the list of log filenames. """ logger.debug("Refreshing log filenames") self._filenames = self._get_log_filenames() - def get(self, session_id): + def get(self, session_id: int) -> str: """ Obtain the log filename for the given session id. Parameters @@ -119,40 +141,117 @@ def get(self, session_id): str The full path to the log file for the requested session id """ - retval = self._filenames.get(session_id) + retval = self._filenames.get(session_id, "") logger.debug("session_id: %s, log_filename: '%s'", session_id, retval) return retval -class _Cache(): - """ Holds parsed Tensorflow log event data in a compressed cache in memory. +class _CacheData(): + """ Holds cached data that has been retrieved from Tensorflow Event Files and is compressed + in memory for a single or live training session Parameters ---------- - session_ids: list - List of `ints` pertaining to the session ids that exist in the Tensorflow events folder + labels: list[str] + The labels for the loss values + timestamps: :class:`np.ndarray` + The timestamp of the event step (iteration) + loss: :class:`np.ndarray` + The loss values collected for A and B sides for the session """ - def __init__(self, session_ids): - logger.debug("Initializing: %s: (session_ids: %s)", self.__class__.__name__, session_ids) - self._data = {idx: None for idx in session_ids} - self._carry_over = {} - self._loss_labels = [] + def __init__(self, labels: list[str], timestamps: np.ndarray, loss: np.ndarray) -> None: + self.labels = labels + self._loss = zlib.compress(T.cast(bytes, loss)) + self._timestamps = zlib.compress(T.cast(bytes, timestamps)) + self._timestamps_shape = timestamps.shape + self._loss_shape = loss.shape + + @property + def loss(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The loss values for this session """ + retval: np.ndarray = np.frombuffer(zlib.decompress(self._loss), dtype="float32") + if len(self._loss_shape) > 1: + retval = retval.reshape(-1, *self._loss_shape[1:]) + return retval + + @property + def timestamps(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The timestamps for this session """ + retval: np.ndarray = np.frombuffer(zlib.decompress(self._timestamps), dtype="float64") + if len(self._timestamps_shape) > 1: + retval = retval.reshape(-1, *self._timestamps_shape[1:]) + return retval + + def add_live_data(self, timestamps: np.ndarray, loss: np.ndarray) -> None: + """ Add live data to the end of the stored data + + loss: :class:`numpy.ndarray` + The latest loss values to add to the cache + timestamps: :class:`numpy.ndarray` + The latest timestamps to add to the cache + """ + new_buffer: list[bytes] = [] + new_shapes: list[tuple[int, ...]] = [] + for data, buffer, dtype, shape in zip([timestamps, loss], + [self._timestamps, self._loss], + ["float64", "float32"], + [self._timestamps_shape, self._loss_shape]): + + old = np.frombuffer(zlib.decompress(buffer), dtype=dtype) + if data.ndim > 1: + old = old.reshape(-1, *data.shape[1:]) + + new = np.concatenate((old, data)) + + logger.debug("old_shape: %s new_shape: %s", shape, new.shape) + new_buffer.append(zlib.compress(new)) + new_shapes.append(new.shape) + del old + + self._timestamps = new_buffer[0] + self._loss = new_buffer[1] + self._timestamps_shape = new_shapes[0] + self._loss_shape = new_shapes[1] + + +class _Cache(): + """ Holds parsed Tensorflow log event data in a compressed cache in memory. """ + def __init__(self) -> None: + logger.debug(parse_class_init(locals())) + self._data: dict[int, _CacheData] = {} + self._carry_over: dict[int, EventData] = {} + self._loss_labels: list[str] = [] logger.debug("Initialized: %s", self.__class__.__name__) - def is_cached(self, session_id): - """ bool: ``True`` if the data already exists in the cache otherwise ``False``. """ + def is_cached(self, session_id: int) -> bool: + """ Check if the given session_id's data is already cached + + Parameters + ---------- + session_id: int + The session ID to check + + Returns + ------- + bool + ``True`` if the data already exists in the cache otherwise ``False``. + """ return self._data.get(session_id) is not None - def cache_data(self, session_id, data, labels, is_live=False): + def cache_data(self, + session_id: int, + data: dict[int, EventData], + labels: list[str], + is_live: bool = False) -> None: """ Add a full session's worth of event data to :attr:`_data`. Parameters ---------- session_id: int The session id to add the data for - data: dict + data[int, :class:`EventData`] The extracted event data dictionary generated from :class:`_EventParser` - labels: list + labels: list[str] List of `str` for the labels of each loss value output is_live: bool, optional ``True`` if the data to be cached is from a live training session otherwise ``False``. @@ -171,16 +270,14 @@ def cache_data(self, session_id, data, labels, is_live=False): timestamps, loss = self._to_numpy(data, is_live) - if not is_live or (is_live and not self._data.get(session_id, None)): - self._data[session_id] = dict(labels=self._loss_labels, - loss=zlib.compress(loss), - loss_shape=loss.shape, - timestamps=zlib.compress(timestamps), - timestamps_shape=timestamps.shape) + if not is_live or (is_live and not self._data.get(session_id)): + self._data[session_id] = _CacheData(self._loss_labels, timestamps, loss) else: self._add_latest_live(session_id, loss, timestamps) - def _to_numpy(self, data, is_live): + def _to_numpy(self, + data: dict[int, EventData], + is_live: bool) -> tuple[np.ndarray, np.ndarray]: """ Extract each individual step data into separate numpy arrays for loss and timestamps. Timestamps are stored float64 as the extra accuracy is needed for correct timings. Arrays @@ -206,9 +303,7 @@ def _to_numpy(self, data, is_live): logger.debug("Processing carry over: %s", self._carry_over) self._collect_carry_over(data) - times, loss = zip(*[(data[idx].get("timestamp"), data[idx].get("loss", [])) - for idx in sorted(data)]) - times, loss = self._process_data(data, times, loss, is_live) + times, loss = self._process_data(data, is_live) if is_live and not all(len(val) == len(self._loss_labels) for val in loss): # TODO Many attempts have been made to fix this for live graph logging, and the issue @@ -230,19 +325,19 @@ def _to_numpy(self, data, is_live): del loss[idx] del times[idx] - times, loss = (np.array(times, dtype="float64"), np.array(loss, dtype="float32")) + n_times, n_loss = (np.array(times, dtype="float64"), np.array(loss, dtype="float32")) logger.debug("Converted to numpy: (data points: %s, timestamps shape: %s, loss shape: %s)", - len(data), times.shape, loss.shape) + len(data), n_times.shape, n_loss.shape) - return times, loss + return n_times, n_loss - def _collect_carry_over(self, data): + def _collect_carry_over(self, data: dict[int, EventData]) -> None: """ For live data, collect carried over data from the previous update and merge into the current data dictionary. Parameters ---------- - data: dict + data: dict[int, :class:`EventData`] The latest raw data dictionary """ logger.debug("Carry over keys: %s, data keys: %s", list(self._carry_over), list(data)) @@ -254,12 +349,14 @@ def _collect_carry_over(self, data): carry_over = self._carry_over.pop(key) update = data[key] logger.debug("Merging carry over data: %s in to %s", carry_over, update) - timestamp = update.get("timestamp") - update["timestamp"] = carry_over["timestamp"] if timestamp is None else timestamp - update.setdefault("loss", []).extend(carry_over.get("loss", [])) + timestamp = update.timestamp + update.timestamp = carry_over.timestamp if not timestamp else timestamp + update.loss = carry_over.loss + update.loss logger.debug("Merged carry over data: %s", update) - def _process_data(self, data, timestamps, loss, is_live): + def _process_data(self, + data: dict[int, EventData], + is_live: bool) -> tuple[list[float], list[list[float]]]: """ Process live update data. Live data requires different processing as often we will only have partial data for the @@ -271,10 +368,6 @@ def _process_data(self, data, timestamps, loss, is_live): ---------- data: dict The incoming tensorflow event data in dictionary form per step - timestamps: tuple - The raw timestamps for for the latest live query, including any partial reads - loss: tuple - The raw loss for for the latest live query, including any partial reads is_live: bool ``True`` if the data to be cached is from a live training session otherwise ``False``. @@ -285,23 +378,26 @@ def _process_data(self, data, timestamps, loss, is_live): loss: list Cleaned list of complete loss for the latest live query """ - loss = list(loss) - timestamps = list(timestamps) + timestamps, loss = zip(*[(data[idx].timestamp, data[idx].loss) + for idx in sorted(data)]) + + l_loss: list[list[float]] = list(loss) + l_timestamps: list[float] = list(timestamps) - if len(loss[-1]) != len(self._loss_labels): - logger.debug("Truncated loss found. loss count: %s", len(loss)) + if len(l_loss[-1]) != len(self._loss_labels): + logger.debug("Truncated loss found. loss count: %s", len(l_loss)) idx = sorted(data)[-1] if is_live: logger.debug("Setting carried over data: %s", data[idx]) self._carry_over[idx] = data[idx] logger.debug("Removing truncated loss: (timestamp: %s, loss: %s)", - timestamps[-1], loss[-1]) - del loss[-1] - del timestamps[-1] + l_timestamps[-1], loss[-1]) + del l_loss[-1] + del l_timestamps[-1] - return timestamps, loss + return l_timestamps, l_loss - def _add_latest_live(self, session_id, loss, timestamps): + def _add_latest_live(self, session_id: int, loss: np.ndarray, timestamps: np.ndarray) -> None: """ Append the latest received live training data to the cached data. Parameters @@ -318,19 +414,10 @@ def _add_latest_live(self, session_id, loss, timestamps): if not np.any(loss) and not np.any(timestamps): return - cache = self._data[session_id] - for metric in ("loss", "timestamps"): - data = locals()[metric] - old_shape = cache[f"{metric}_shape"] - dtype = "float32" if metric == "loss" else "float64" - old = np.frombuffer(zlib.decompress(cache[metric]), dtype=dtype).reshape(old_shape) - new = np.concatenate((old, data)) - logger.debug("'%s' old_shape: %s new_shape: %s", metric, old_shape, new.shape) - cache[f"{metric}_shape"] = new.shape - cache[metric] = zlib.compress(new) - del old + self._data[session_id].add_live_data(timestamps, loss) - def get_data(self, session_id, metric): + def get_data(self, session_id: int, metric: T.Literal["loss", "timestamps"] + ) -> dict[int, dict[str, np.ndarray | list[str]]] | None: """ Retrieve the decompressed cached data from the cache for the given session id. Parameters @@ -343,9 +430,10 @@ def get_data(self, session_id, metric): Returns ------- - dict - The `session_id` (s) as key, the values are a dictionary containing the requested - metric information for each session returned + dict or ``None`` + The `session_id`(s) as key, the values are a dictionary containing the requested + metric information for each session returned. ``None`` if no data is stored for the + given session_id """ if session_id is None: raw = self._data @@ -355,14 +443,12 @@ def get_data(self, session_id, metric): return None raw = {session_id: data} - dtype = "float32" if metric == "loss" else "float64" - - retval = {} + retval: dict[int, dict[str, np.ndarray | list[str]]] = {} for idx, data in raw.items(): - val = {metric: np.frombuffer(zlib.decompress(data[metric]), - dtype=dtype).reshape(data[f"{metric}_shape"])} + array = data.loss if metric == "loss" else data.timestamps + val: dict[str, np.ndarray | list[str]] = {str(metric): array} if metric == "loss": - val["labels"] = data["labels"] + val["labels"] = data.labels retval[idx] = val logger.debug("Obtained cached data: %s", @@ -386,25 +472,24 @@ class TensorBoardLogs(): is_training: bool ``True`` if the events are being read whilst Faceswap is training otherwise ``False`` """ - def __init__(self, logs_folder, is_training): - logger.debug("Initializing: %s: (logs_folder: %s, is_training: %s)", - self.__class__.__name__, logs_folder, is_training) + def __init__(self, logs_folder: str, is_training: bool) -> None: + logger.debug(parse_class_init(locals())) self._is_training = False self._training_iterator = None self._log_files = _LogFiles(logs_folder) self.set_training(is_training) - self._cache = _Cache(self.session_ids) + self._cache = _Cache() logger.debug("Initialized: %s", self.__class__.__name__) @property - def session_ids(self): - """ list: Sorted list of integers of available session ids. """ + def session_ids(self) -> list[int]: + """ list[int]: Sorted list of integers of available session ids. """ return self._log_files.session_ids - def set_training(self, is_training): + def set_training(self, is_training: bool) -> None: """ Set the internal training flag to the given `is_training` value. If a new training session is being instigated, refresh the log filenames @@ -431,7 +516,7 @@ def set_training(self, is_training): del self._training_iterator self._training_iterator = None - def _cache_data(self, session_id): + def _cache_data(self, session_id: int) -> None: """ Cache TensorBoard logs for the given session ID on first access. Populates :attr:`_cache` with timestamps and loss data. @@ -447,10 +532,11 @@ def _cache_data(self, session_id): live_data = self._is_training and session_id == max(self.session_ids) iterator = self._training_iterator if live_data else tf.compat.v1.io.tf_record_iterator( self._log_files.get(session_id)) + assert iterator is not None parser = _EventParser(iterator, self._cache, live_data) parser.cache_events(session_id) - def _check_cache(self, session_id=None): + def _check_cache(self, session_id: int | None = None) -> None: """ Check if the given session_id has been cached and if not, cache it. Parameters @@ -468,7 +554,7 @@ def _check_cache(self, session_id=None): if not self._cache.is_cached(idx): self._cache_data(idx) - def get_loss(self, session_id=None): + def get_loss(self, session_id: int | None = None) -> dict[int, dict[str, np.ndarray]]: """ Read the loss from the TensorBoard event logs Parameters @@ -484,19 +570,22 @@ def get_loss(self, session_id=None): and list of loss values for each step """ logger.debug("Getting loss: (session_id: %s)", session_id) - retval = {} + retval: dict[int, dict[str, np.ndarray]] = {} for idx in [session_id] if session_id else self.session_ids: self._check_cache(idx) - data = self._cache.get_data(idx, "loss") - if not data: + full_data = self._cache.get_data(idx, "loss") + if not full_data: continue - data = data[idx] - retval[idx] = {title: data["loss"][:, idx] for idx, title in enumerate(data["labels"])} + data = full_data[idx] + loss = data["loss"] + assert isinstance(loss, np.ndarray) + retval[idx] = {title: loss[:, idx] for idx, title in enumerate(data["labels"])} + logger.debug({key: {k: v.shape for k, v in val.items()} for key, val in retval.items()}) return retval - def get_timestamps(self, session_id=None): + def get_timestamps(self, session_id: int | None = None) -> dict[int, np.ndarray]: """ Read the timestamps from the TensorBoard logs. As loss timestamps are slightly different for each loss, we collect the timestamp from the @@ -516,18 +605,20 @@ def get_timestamps(self, session_id=None): logger.debug("Getting timestamps: (session_id: %s, is_training: %s)", session_id, self._is_training) - retval = {} + retval: dict[int, np.ndarray] = {} for idx in [session_id] if session_id else self.session_ids: self._check_cache(idx) data = self._cache.get_data(idx, "timestamps") if not data: continue - retval[idx] = data[idx]["timestamps"] + timestamps = data[idx]["timestamps"] + assert isinstance(timestamps, np.ndarray) + retval[idx] = timestamps logger.debug({k: v.shape for k, v in retval.items()}) return retval -class _EventParser(): # pylint:disable=too-few-public-methods +class _EventParser(): """ Parses Tensorflow event and populates data to :class:`_Cache`. Parameters @@ -540,17 +631,17 @@ class _EventParser(): # pylint:disable=too-few-public-methods ``True`` if the iterator to be loaded is a training iterator for reading live data otherwise ``False`` """ - def __init__(self, iterator, cache, live_data): - logger.debug("Initializing: %s: (iterator: %s, cache: %s, live_data: %s)", - self.__class__.__name__, iterator, cache, live_data) + def __init__(self, iterator: Iterator[bytes], cache: _Cache, live_data: bool) -> None: + logger.debug(parse_class_init(locals())) self._live_data = live_data self._cache = cache self._iterator = self._get_latest_live(iterator) if live_data else iterator - self._loss_labels = [] + self._loss_labels: list[str] = [] + self._num_strip = re.compile(r"_\d+$") logger.debug("Initialized: %s", self.__class__.__name__) @classmethod - def _get_latest_live(cls, iterator): + def _get_latest_live(cls, iterator: Iterator[bytes]) -> Generator[bytes, None, None]: """ Obtain the latest event logs for live training data. The live data iterator remains open so that it can be re-queried @@ -580,7 +671,7 @@ def _get_latest_live(cls, iterator): break logger.debug("Collected %s records from live log file", i) - def cache_events(self, session_id): + def cache_events(self, session_id: int) -> None: """ Parse the Tensorflow events logs and add to :attr:`_cache`. Parameters @@ -588,7 +679,8 @@ def cache_events(self, session_id): session_id: int The session id that the data is being cached for """ - data = {} + assert self._iterator is not None + data: dict[int, EventData] = {} try: for record in self._iterator: event = event_pb2.Event.FromString(record) # pylint:disable=no-member @@ -596,11 +688,9 @@ def cache_events(self, session_id): continue if event.summary.value[0].tag == "keras": self._parse_outputs(event) - if get_backend() == "amd": - # No model is logged for AMD so need to get loss labels from state file - self._add_amd_loss_labels(session_id) if event.summary.value[0].tag.startswith("batch_"): - data[event.step] = self._process_event(event, data.get(event.step, {})) + data[event.step] = self._process_event(event, + data.get(event.step, EventData())) except tf_errors.DataLossError as err: logger.warning("The logs for Session %s are corrupted and cannot be displayed. " @@ -609,7 +699,7 @@ def cache_events(self, session_id): self._cache.cache_data(session_id, data, self._loss_labels, is_live=self._live_data) - def _parse_outputs(self, event): + def _parse_outputs(self, event: event_pb2.Event) -> None: """ Parse the outputs from the stored model structure for mapping loss names to model outputs. @@ -639,7 +729,7 @@ def _parse_outputs(self, event): if layer["name"] == layer_name)["config"] layer_outputs = self._get_outputs(output_config) for output in layer_outputs: # Drill into sub-model to get the actual output names - loss_name = output[0][0] + loss_name = self._num_strip.sub("", output[0][0]) # strip trailing numbers if loss_name[-2:] not in ("_a", "_b"): # Rename losses to reflect the side output new_name = f"{loss_name.replace('_both', '')}_{side}" logger.debug("Renaming loss output from '%s' to '%s'", loss_name, new_name) @@ -650,7 +740,7 @@ def _parse_outputs(self, event): logger.debug("Collated loss labels: %s", self._loss_labels) @classmethod - def _get_outputs(cls, model_config): + def _get_outputs(cls, model_config: dict[str, T.Any]) -> np.ndarray: """ Obtain the output names, instance index and output index for the given model. If there is only a single output, the shape of the array is expanded to remain consistent @@ -674,30 +764,8 @@ def _get_outputs(cls, model_config): outputs, outputs.shape) return outputs - def _add_amd_loss_labels(self, session_id): - """ It is not possible to store the model config in the Tensorboard logs for AMD so we - need to obtain the loss labels from the model's state file. This is called now so we know - event data is being written, and therefore the most current loss label data is available - in the state file. - - Loss names are added to :attr:`_loss_labels` - - Parameters - ---------- - session_id: int - The session id that the data is being cached for - - """ - if self._cache._loss_labels: # pylint:disable=protected-access - return - # Import global session here to prevent circular import - from . import Session # pylint:disable=import-outside-toplevel - loss_labels = sorted(Session.get_loss_keys(session_id=session_id)) - self._loss_labels = loss_labels - logger.debug("Collated loss labels: %s", self._loss_labels) - @classmethod - def _process_event(cls, event, step): + def _process_event(cls, event: event_pb2.Event, step: EventData) -> EventData: """ Process a single Tensorflow event. Adds timestamp to the step `dict` if a total loss value is received, process the labels for @@ -707,28 +775,29 @@ def _process_event(cls, event, step): ---------- event: :class:`tensorflow.core.util.event_pb2` The event data to be processed - step: dict - The dictionary to populated with the extracted data from the tensorflow event + step: :class:`EventData` + The currently processing dictionary to be populated with the extracted data from the + tensorflow event for this step Returns ------- - dict - The given step `dict` with the given event data added to it. + :class:`EventData` + The given step :class:`EventData` with the given event data added to it. """ summary = event.summary.value[0] - if summary.tag in ("batch_loss", "batch_total"): # Pre tf2.3 totals were "batch_total" - step["timestamp"] = event.wall_time + if summary.tag == "batch_total": + step.timestamp = event.wall_time return step loss = summary.simple_value if not loss: # Need to convert a tensor to a float for TF2.8 logged data. This maybe due to change # in logging or may be due to work around put in place in FS training function for the - # following bug in TF 2.8 when writing records: + # following bug in TF 2.8/2.9 when writing records: # https://github.com/keras-team/keras/issues/16173 loss = float(tf.make_ndarray(summary.tensor)) - step.setdefault("loss", []).append(loss) + step.loss.append(loss) return step diff --git a/lib/gui/analysis/stats.py b/lib/gui/analysis/stats.py index 9ab8347d1e..b055cfa332 100644 --- a/lib/gui/analysis/stats.py +++ b/lib/gui/analysis/stats.py @@ -5,25 +5,24 @@ the analysis tab) or the currently training session. """ - +from __future__ import annotations import logging -import time import os +import time +import typing as T import warnings from math import ceil from threading import Event -from typing import List, Optional, Tuple, Union -from typing_extensions import Self import numpy as np +from lib.logger import parse_class_init from lib.serializer import get_serializer -from lib.utils import get_backend from .event_reader import TensorBoardLogs -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class GlobalSession(): @@ -32,13 +31,13 @@ class GlobalSession(): :attr:`lib.gui.analysis.Session` """ def __init__(self) -> None: - logger.debug("Initializing %s", self.__class__.__name__) - self._state = None - self._model_dir = None - self._model_name = None + logger.debug(parse_class_init(locals())) + self._state: dict[str, T.Any] = {} + self._model_dir = "" + self._model_name = "" - self._tb_logs = None - self._summary = None + self._tb_logs: TensorBoardLogs | None = None + self._summary: SessionsSummary | None = None self._is_training = False self._is_querying = Event() @@ -48,7 +47,7 @@ def __init__(self) -> None: @property def is_loaded(self) -> bool: """ bool: ``True`` if session data is loaded otherwise ``False`` """ - return self._model_dir is not None + return bool(self._model_dir) @property def is_training(self) -> bool: @@ -62,29 +61,33 @@ def model_filename(self) -> str: return os.path.join(self._model_dir, self._model_name) @property - def batch_sizes(self) -> dict: + def batch_sizes(self) -> dict[int, int]: """ dict: The batch sizes for each session_id for the model. """ - if self._state is None: + if not self._state: return {} return {int(sess_id): sess["batchsize"] for sess_id, sess in self._state.get("sessions", {}).items()} @property - def full_summary(self) -> List[dict]: + def full_summary(self) -> list[dict]: """ list: List of dictionaries containing summary statistics for each session id. """ + assert self._summary is not None return self._summary.get_summary_stats() @property def logging_disabled(self) -> bool: """ bool: ``True`` if logging is enabled for the currently training session otherwise ``False``. """ - if self._state is None: + if not self._state: return True - return self._state["sessions"][str(self.session_ids[-1])]["no_logs"] + max_id = str(max(int(idx) for idx in self._state["sessions"])) + return self._state["sessions"][max_id]["no_logs"] @property - def session_ids(self) -> List[int]: + def session_ids(self) -> list[int]: """ list: The sorted list of all existing session ids in the state file """ + if self._tb_logs is None: + return [] return self._tb_logs.session_ids def _load_state_file(self) -> None: @@ -96,8 +99,8 @@ def _load_state_file(self) -> None: logger.debug("Loaded state: %s", self._state) def initialize_session(self, - model_folder: Optional[str], - model_name: Optional[str], + model_folder: str, + model_name: str, is_training: bool = False) -> None: """ Initialize a Session. @@ -106,12 +109,14 @@ def initialize_session(self, Parameters ---------- - model_folder: str, optional + model_folder: str, If loading a session manually (e.g. for the analysis tab), then the path to the model - folder must be provided. For training sessions, this should be left at ``None`` + folder must be provided. For training sessions, this should be passed through from the + launcher model_name: str, optional If loading a session manually (e.g. for the analysis tab), then the model filename - must be provided. For training sessions, this should be left at ``None`` + must be provided. For training sessions, this should be passed through from the + launcher is_training: bool, optional ``True`` if the session is being initialized for a training session, otherwise ``False``. Default: ``False`` @@ -120,6 +125,7 @@ def initialize_session(self, if self._model_dir == model_folder and self._model_name == model_name: if is_training: + assert self._tb_logs is not None self._tb_logs.set_training(is_training) self._load_state_file() self._is_training = True @@ -131,8 +137,10 @@ def initialize_session(self, self._model_dir = model_folder self._model_name = model_name self._load_state_file() - self._tb_logs = TensorBoardLogs(os.path.join(self._model_dir, f"{self._model_name}_logs"), - is_training) + if not self.logging_disabled: + self._tb_logs = TensorBoardLogs(os.path.join(self._model_dir, + f"{self._model_name}_logs"), + is_training) self._summary = SessionsSummary(self) logger.debug("Initialized session. Session_IDS: %s", self.session_ids) @@ -146,8 +154,8 @@ def stop_training(self) -> None: def clear(self) -> None: """ Clear the currently loaded session. """ self._state = {} - self._model_dir = None - self._model_name = None + self._model_dir = "" + self._model_name = "" del self._tb_logs self._tb_logs = None @@ -157,7 +165,7 @@ def clear(self) -> None: self._is_training = False - def get_loss(self, session_id: Optional[int]) -> dict: + def get_loss(self, session_id: int | None) -> dict[str, np.ndarray]: """ Obtain the loss values for the given session_id. Parameters @@ -176,13 +184,15 @@ def get_loss(self, session_id: Optional[int]) -> dict: if self._is_training: self._is_querying.set() + assert self._tb_logs is not None loss_dict = self._tb_logs.get_loss(session_id=session_id) if session_id is None: - retval = {} + all_loss: dict[str, list[float]] = {} for key in sorted(loss_dict): for loss_key, loss in loss_dict[key].items(): - retval.setdefault(loss_key, []).extend(loss) - retval = {key: np.array(val, dtype="float32") for key, val in retval.items()} + all_loss.setdefault(loss_key, []).extend(loss) + retval: dict[str, np.ndarray] = {key: np.array(val, dtype="float32") + for key, val in all_loss.items()} else: retval = loss_dict.get(session_id, {}) @@ -190,7 +200,15 @@ def get_loss(self, session_id: Optional[int]) -> dict: self._is_querying.clear() return retval - def get_timestamps(self, session_id: Optional[int]) -> Union[dict, np.ndarray]: + @T.overload + def get_timestamps(self, session_id: None) -> dict[int, np.ndarray]: + ... + + @T.overload + def get_timestamps(self, session_id: int) -> np.ndarray: + ... + + def get_timestamps(self, session_id): """ Obtain the time stamps keys for the given session_id. Parameters @@ -201,7 +219,7 @@ def get_timestamps(self, session_id: Optional[int]) -> Union[dict, np.ndarray]: Returns ------- - dict or :class:`numpy.ndarray` + dict[int] or :class:`numpy.ndarray` If a session ID has been given then a single :class:`numpy.ndarray` will be returned with the session's time stamps. Otherwise a 'dict' will be returned with the session IDs as key with :class:`numpy.ndarray` of timestamps as values @@ -211,6 +229,7 @@ def get_timestamps(self, session_id: Optional[int]) -> Union[dict, np.ndarray]: if self._is_training: self._is_querying.set() + assert self._tb_logs is not None retval = self._tb_logs.get_timestamps(session_id=session_id) if session_id is not None: retval = retval[session_id] @@ -229,7 +248,7 @@ def _wait_for_thread(self) -> None: continue break - def get_loss_keys(self, session_id: Optional[int]) -> List[str]: + def get_loss_keys(self, session_id: int | None) -> list[str]: """ Obtain the loss keys for the given session_id. Parameters @@ -244,28 +263,24 @@ def get_loss_keys(self, session_id: Optional[int]) -> List[str]: The loss keys for the given session. If ``None`` is passed as session_id then a unique list of all loss keys for all sessions is returned """ - if get_backend() == "amd": - # We can't log the graph in Tensorboard logs for AMD so need to obtain from state file - loss_keys = {int(sess_id): [name for name in session["loss_names"] if name != "total"] - for sess_id, session in self._state["sessions"].items()} - else: - loss_keys = {sess_id: list(logs.keys()) - for sess_id, logs - in self._tb_logs.get_loss(session_id=session_id).items()} + assert self._tb_logs is not None + loss_keys = {sess_id: list(logs.keys()) + for sess_id, logs + in self._tb_logs.get_loss(session_id=session_id).items()} if session_id is None: - retval = list(set(loss_key - for session in loss_keys.values() - for loss_key in session)) + retval: list[str] = list(set(loss_key + for session in loss_keys.values() + for loss_key in session)) else: - retval = loss_keys.get(session_id) + retval = loss_keys.get(session_id, []) return retval _SESSION = GlobalSession() -class SessionsSummary(): # pylint:disable=too-few-public-methods +class SessionsSummary(): """ Performs top level summary calculations for each session ID within the loaded or currently training Session for display in the Analysis tree view. @@ -275,15 +290,15 @@ class SessionsSummary(): # pylint:disable=too-few-public-methods The loaded or currently training session """ def __init__(self, session: GlobalSession) -> None: - logger.debug("Initializing %s: (session: %s)", self.__class__.__name__, session) + logger.debug(parse_class_init(locals())) self._session = session self._state = session._state - self._time_stats = None - self._per_session_stats = None + self._time_stats: dict[int, dict[str, float | int]] = {} + self._per_session_stats: list[dict[str, T.Any]] = [] logger.debug("Initialized %s", self.__class__.__name__) - def get_summary_stats(self) -> List[dict]: + def get_summary_stats(self) -> list[dict]: """ Compile the individual session statistics and calculate the total. Format the stats for display @@ -315,25 +330,26 @@ def _get_time_stats(self) -> None: If the main Session is currently training, then the training session ID is updated with the latest stats. """ - if self._time_stats is None: + if not self._time_stats: logger.debug("Collating summary time stamps") self._time_stats = { - sess_id: dict(start_time=np.min(timestamps) if np.any(timestamps) else 0, - end_time=np.max(timestamps) if np.any(timestamps) else 0, - iterations=timestamps.shape[0] if np.any(timestamps) else 0) - for sess_id, timestamps in self._session.get_timestamps(None).items()} + sess_id: {"start_time": np.min(timestamps) if np.any(timestamps) else 0, + "end_time": np.max(timestamps) if np.any(timestamps) else 0, + "iterations": timestamps.shape[0] if np.any(timestamps) else 0} + for sess_id, timestamps in T.cast(dict[int, np.ndarray], + self._session.get_timestamps(None)).items()} elif _SESSION.is_training: logger.debug("Updating summary time stamps for training session") session_id = _SESSION.session_ids[-1] - latest = self._session.get_timestamps(session_id) + latest = T.cast(np.ndarray, self._session.get_timestamps(session_id)) - self._time_stats[session_id] = dict( - start_time=np.min(latest) if np.any(latest) else 0, - end_time=np.max(latest) if np.any(latest) else 0, - iterations=latest.shape[0] if np.any(latest) else 0) + self._time_stats[session_id] = { + "start_time": np.min(latest) if np.any(latest) else 0, + "end_time": np.max(latest) if np.any(latest) else 0, + "iterations": latest.shape[0] if np.any(latest) else 0} logger.debug("time_stats: %s", self._time_stats) @@ -344,12 +360,12 @@ def _get_per_session_stats(self) -> None: If a training session is running, then updates the training sessions stats only. """ - if self._per_session_stats is None: + if not self._per_session_stats: logger.debug("Collating per session stats") compiled = [] for session_id in self._time_stats: logger.debug("Compiling session ID: %s", session_id) - if self._state is None: + if not self._state: logger.debug("Session state dict doesn't exist. Most likely task has been " "terminated during compilation") return @@ -367,15 +383,17 @@ def _get_per_session_stats(self) -> None: stats = self._per_session_stats[-1] - stats["start"] = ts_data["start_time"] - stats["end"] = ts_data["end_time"] - stats["elapsed"] = int(stats["end"] - stats["start"]) + start = np.nan_to_num(ts_data["start_time"]) + end = np.nan_to_num(ts_data["end_time"]) + stats["start"] = start + stats["end"] = end + stats["elapsed"] = int(end - start) stats["iterations"] = ts_data["iterations"] stats["rate"] = (((stats["batch"] * 2) * stats["iterations"]) - / stats["elapsed"] if stats["elapsed"] != 0 else 0) + / stats["elapsed"] if stats["elapsed"] > 0 else 0) logger.debug("per_session_stats: %s", self._per_session_stats) - def _collate_stats(self, session_id: int) -> dict: + def _collate_stats(self, session_id: int) -> dict[str, int | float]: """ Collate the session summary statistics for the given session ID. Parameters @@ -389,27 +407,30 @@ def _collate_stats(self, session_id: int) -> dict: The collated session summary statistics """ timestamps = self._time_stats[session_id] - elapsed = int(timestamps["end_time"] - timestamps["start_time"]) + start = np.nan_to_num(timestamps["start_time"]) + end = np.nan_to_num(timestamps["end_time"]) + elapsed = int(end - start) batchsize = self._session.batch_sizes.get(session_id, 0) - retval = dict( - session=session_id, - start=timestamps["start_time"], - end=timestamps["end_time"], - elapsed=elapsed, - rate=(((batchsize * 2) * timestamps["iterations"]) / elapsed if elapsed != 0 else 0), - batch=batchsize, - iterations=timestamps["iterations"]) + retval = { + "session": session_id, + "start": start, + "end": end, + "elapsed": elapsed, + "rate": (((batchsize * 2) * timestamps["iterations"]) / elapsed + if elapsed != 0 else 0), + "batch": batchsize, + "iterations": timestamps["iterations"]} logger.debug(retval) return retval - def _total_stats(self) -> dict: + def _total_stats(self) -> dict[str, str | int | float]: """ Compile the Totals stats. Totals are fully calculated each time as they will change on the basis of the training session. Returns ------- - dict: + dict The Session name, start time, end time, elapsed time, rate, batch size and number of iterations for all session ids within the loaded data. """ @@ -439,7 +460,7 @@ def _total_stats(self) -> dict: logger.debug(totals) return totals - def _format_stats(self, compiled_stats: List[dict]) -> List[dict]: + def _format_stats(self, compiled_stats: list[dict]) -> list[dict]: """ Format for the incoming list of statistics for display. Parameters @@ -469,7 +490,7 @@ def _format_stats(self, compiled_stats: List[dict]) -> List[dict]: return retval @classmethod - def _convert_time(cls, timestamp: float) -> Tuple[str, str, str]: + def _convert_time(cls, timestamp: float) -> tuple[str, str, str]: """ Convert time stamp to total hours, minutes and seconds. Parameters @@ -482,8 +503,8 @@ def _convert_time(cls, timestamp: float) -> Tuple[str, str, str]: tuple (`hours`, `minutes`, `seconds`) as strings """ - hrs = int(timestamp // 3600) - hrs = f"{hrs:02d}" if hrs < 10 else str(hrs) + ihrs = int(timestamp // 3600) + hrs = f"{ihrs:02d}" if ihrs < 10 else str(ihrs) mins = f"{(int(timestamp % 3600) // 60):02d}" secs = f"{(int(timestamp % 3600) % 60):02d}" return hrs, mins, secs @@ -514,16 +535,12 @@ class Calculations(): """ def __init__(self, session_id, display: str = "loss", - loss_keys: str = "loss", - selections: str = "raw", + loss_keys: list[str] | str = "loss", + selections: list[str] | str = "raw", avg_samples: int = 500, smooth_amount: float = 0.90, flatten_outliers: bool = False) -> None: - logger.debug("Initializing %s: (session_id: %s, display: %s, loss_keys: %s, " - "selections: %s, avg_samples: %s, smooth_amount: %s, flatten_outliers: %s)", - self.__class__.__name__, session_id, display, loss_keys, selections, - avg_samples, smooth_amount, flatten_outliers) - + logger.debug(parse_class_init(locals())) warnings.simplefilter("ignore", np.RankWarning) self._session_id = session_id @@ -532,13 +549,13 @@ def __init__(self, session_id, self._loss_keys = loss_keys if isinstance(loss_keys, list) else [loss_keys] self._selections = selections if isinstance(selections, list) else [selections] self._is_totals = session_id is None - self._args = dict(avg_samples=avg_samples, - smooth_amount=smooth_amount, - flatten_outliers=flatten_outliers) + self._args: dict[str, int | float] = {"avg_samples": avg_samples, + "smooth_amount": smooth_amount, + "flatten_outliers": flatten_outliers} self._iterations = 0 self._limit = 0 self._start_iteration = 0 - self._stats = {} + self._stats: dict[str, np.ndarray] = {} self.refresh() logger.debug("Initialized %s", self.__class__.__name__) @@ -553,11 +570,11 @@ def start_iteration(self) -> int: return self._start_iteration @property - def stats(self) -> dict: + def stats(self) -> dict[str, np.ndarray]: """ dict: The final calculated statistics """ return self._stats - def refresh(self) -> Optional[Self]: + def refresh(self) -> Calculations | None: """ Refresh the stats """ logger.debug("Refreshing") if not _SESSION.is_loaded: @@ -654,11 +671,11 @@ def _get_raw(self) -> None: if len(iterations) > 1: # Crop all losses to the same number of items if self._iterations == 0: - self.stats = {lossname: np.array([], dtype=loss.dtype) - for lossname, loss in self.stats.items()} + self._stats = {lossname: np.array([], dtype=loss.dtype) + for lossname, loss in self.stats.items()} else: - self.stats = {lossname: loss[:self._iterations] - for lossname, loss in self.stats.items()} + self._stats = {lossname: loss[:self._iterations] + for lossname, loss in self.stats.items()} else: # Rate calculation data = self._calc_rate_total() if self._is_totals else self._calc_rate() @@ -715,8 +732,9 @@ def _calc_rate(self) -> np.ndarray: The training rate for each iteration of the selected session """ logger.debug("Calculating rate") - retval = (_SESSION.batch_sizes[self._session_id] * 2) / np.diff(_SESSION.get_timestamps( - self._session_id)) + batch_size = _SESSION.batch_sizes[self._session_id] * 2 + retval = batch_size / np.diff(T.cast(np.ndarray, + _SESSION.get_timestamps(self._session_id))) logger.debug("Calculated rate: Item_count: %s", len(retval)) return retval @@ -737,7 +755,7 @@ def _calc_rate_total(cls) -> np.ndarray: logger.debug("Calculating totals rate") batchsizes = _SESSION.batch_sizes total_timestamps = _SESSION.get_timestamps(None) - rate = [] + rate: list[float] = [] for sess_id in sorted(total_timestamps.keys()): batchsize = batchsizes[sess_id] timestamps = total_timestamps[sess_id] @@ -777,7 +795,7 @@ def _calc_avg(self, data: np.ndarray) -> np.ndarray: The moving average for the given data """ logger.debug("Calculating Average. Data points: %s", len(data)) - window = self._args["avg_samples"] + window = T.cast(int, self._args["avg_samples"]) pad = ceil(window / 2) datapoints = data.shape[0] @@ -835,7 +853,7 @@ def _calc_trend(cls, data: np.ndarray) -> np.ndarray: return trend -class _ExponentialMovingAverage(): # pylint:disable=too-few-public-methods +class _ExponentialMovingAverage(): """ Reshapes data before calculating exponential moving average, then iterates once over the rows to calculate the offset without precision issues. @@ -851,6 +869,7 @@ class _ExponentialMovingAverage(): # pylint:disable=too-few-public-methods Adapted from: https://stackoverflow.com/questions/42869495 """ def __init__(self, data: np.ndarray, amount: float) -> None: + logger.debug(parse_class_init(locals())) assert data.ndim == 1 amount = min(max(amount, 0.001), 0.999) @@ -859,6 +878,7 @@ def __init__(self, data: np.ndarray, amount: float) -> None: self._dtype = "float32" if data.dtype == np.float32 else "float64" self._row_size = self._get_max_row_size() self._out = np.empty_like(data, dtype=self._dtype) + logger.debug("Initialized %s", self.__class__.__name__) def __call__(self) -> np.ndarray: """ Perform the exponential moving average calculation. @@ -933,7 +953,7 @@ def _ewma_vectorized_safe(self) -> None: def _ewma_vectorized(self, data: np.ndarray, out: np.ndarray, - offset: Optional[float] = None) -> None: + offset: float | None = None) -> None: """ Calculates the exponential moving average over a vector. Will fail for large inputs. The result is processed in place into the array passed to the `out` parameter @@ -964,8 +984,8 @@ def _ewma_vectorized(self, out /= scaling_factors[-2::-1] # cumulative sums / scaling if offset != 0: - offset = np.array(offset, copy=False).astype(self._dtype, copy=False) - out += offset * scaling_factors[1:] + noffset = np.array(offset, copy=False).astype(self._dtype, copy=False) + out += noffset * scaling_factors[1:] def _ewma_vectorized_2d(self, data: np.ndarray, out: np.ndarray) -> None: """ Calculates the exponential moving average over the last axis. diff --git a/lib/gui/command.py b/lib/gui/command.py index bac9e8c109..95528be95f 100644 --- a/lib/gui/command.py +++ b/lib/gui/command.py @@ -9,8 +9,9 @@ from .control_helper import ControlPanel from .custom_widgets import Tooltip from .utils import get_images, get_config +from .options import CliOption -logger = logging.getLogger(__name__) # pylint:disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) @@ -22,7 +23,7 @@ class CommandNotebook(ttk.Notebook): # pylint:disable=too-many-ancestors def __init__(self, parent): logger.debug("Initializing %s: (parent: %s)", self.__class__.__name__, parent) - self.actionbtns = dict() + self.actionbtns = {} super().__init__(parent) parent.add(self) @@ -50,7 +51,7 @@ def set_running_task_trace(self): to change the action buttons text and command """ logger.debug("Set running trace") tk_vars = get_config().tk_vars - tk_vars["runningtask"].trace("w", self.change_action_button) + tk_vars.running_task.trace("w", self.change_action_button) def build_tabs(self): """ Build the tabs for the relevant command """ @@ -73,14 +74,14 @@ def change_action_button(self, *args): for cmd, action in self.actionbtns.items(): btnact = action - if tk_vars["runningtask"].get(): + if tk_vars.running_task.get(): ttl = " Stop" img = get_images().icons["stop"] hlp = "Exit the running process" else: - ttl = " {}".format(cmd.title()) + ttl = f" {cmd.title()}" img = get_images().icons["start"] - hlp = "Run the {} script".format(cmd.title()) + hlp = f"Run the {cmd.title()} script" logger.debug("Updated Action Button: '%s'", ttl) btnact.config(text=ttl, image=img) Tooltip(btnact, text=hlp, wrap_length=200) @@ -88,7 +89,7 @@ def change_action_button(self, *args): def _set_modified_vars(self): """ Set the tkinter variable for each tab to indicate whether contents have been modified """ - tkvars = dict() + tkvars = {} for tab in self.tab_names: if tab == "tools": for ttab in self.tools_tab_names: @@ -116,7 +117,7 @@ class CommandTab(ttk.Frame): # pylint:disable=too-many-ancestors def __init__(self, parent, category, command): logger.debug("Initializing %s: (category: '%s', command: '%s')", self.__class__.__name__, category, command) - super().__init__(parent, name="tab_{}".format(command.lower())) + super().__init__(parent, name=f"tab_{command.lower()}") self.category = category self.actionbtns = parent.actionbtns @@ -129,7 +130,7 @@ def build_tab(self): """ Build the tab """ logger.debug("Build Tab: '%s'", self.command) options = get_config().cli_opts.opts[self.command] - cp_opts = [val["cpanel_option"] for key, val in options.items() if key != "helptext"] + cp_opts = [val.cpanel_option for val in options.values() if isinstance(val, CliOption)] ControlPanel(self, cp_opts, label_width=16, @@ -171,14 +172,14 @@ def add_action_button(self, category, actionbtns): actframe.pack(fill=tk.X, side=tk.RIGHT) tk_vars = get_config().tk_vars - var_value = "{},{}".format(category, self.command) + var_value = f"{category},{self.command}" btngen = ttk.Button(actframe, image=get_images().icons["generate"], text=" Generate", compound=tk.LEFT, width=14, - command=lambda: tk_vars["generate"].set(var_value)) + command=lambda: tk_vars.generate_command.set(var_value)) btngen.pack(side=tk.LEFT, padx=5) Tooltip(btngen, text=_("Output command line options to the console"), @@ -186,10 +187,10 @@ def add_action_button(self, category, actionbtns): btnact = ttk.Button(actframe, image=get_images().icons["start"], - text=" {}".format(self.title), + text=f" {self.title}", compound=tk.LEFT, width=14, - command=lambda: tk_vars["action"].set(var_value)) + command=lambda: tk_vars.action_command.set(var_value)) btnact.pack(side=tk.LEFT, fill=tk.X, expand=True) Tooltip(btnact, text=_("Run the {} script").format(self.title), diff --git a/lib/gui/control_helper.py b/lib/gui/control_helper.py index 09524be57b..5179256759 100644 --- a/lib/gui/control_helper.py +++ b/lib/gui/control_helper.py @@ -5,6 +5,7 @@ import re import tkinter as tk +import typing as T from tkinter import colorchooser, ttk from itertools import zip_longest from functools import partial @@ -14,7 +15,7 @@ from .custom_widgets import ContextMenu, MultiOption, ToggledFrame, Tooltip from .utils import FileHandler, get_config, get_images -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) @@ -23,16 +24,30 @@ # We store Tooltips, ContextMenus and Commands globally when they are created # Because we need to add them back to newly cloned widgets (they are not easily accessible from # original config or are prone to getting destroyed when the original widget is destroyed) -_RECREATE_OBJECTS = dict(tooltips=dict(), commands=dict(), contextmenus=dict()) +_RECREATE_OBJECTS: dict[str, dict[str, T.Any]] = {"tooltips": {}, + "commands": {}, + "contextmenus": {}} -def _get_tooltip(widget, text=None, text_variable=None, wrap_length=600): - """ Store the tooltip layout and widget id in _TOOLTIPS and return a tooltip """ +def _get_tooltip(widget, text=None, text_variable=None): + """ Store the tooltip layout and widget id in _TOOLTIPS and return a tooltip. + + Auto adjust tooltip width based on amount of text. + + """ _RECREATE_OBJECTS["tooltips"][str(widget)] = {"text": text, - "text_variable": text_variable, - "wrap_length": wrap_length} - logger.debug("Adding to tooltips dict: (widget: %s. text: '%s', wrap_length: %s)", - widget, text, wrap_length) + "text_variable": text_variable} + logger.debug("Adding to tooltips dict: (widget: %s. text: '%s')", widget, text) + + wrap_length = 400 + if text is not None: + while True: + if len(text) < wrap_length * 5: + break + if wrap_length > 800: + break + wrap_length = int(wrap_length * 1.10) + return Tooltip(widget, text=text, text_variable=text_variable, wrap_length=wrap_length) @@ -141,17 +156,17 @@ def __init__(self, title, dtype, # pylint:disable=too-many-arguments self.dtype = dtype self.sysbrowser = sysbrowser self._command = command - self._options = dict(title=title, - subgroup=subgroup, - group=group, - default=default, - initial_value=initial_value, - choices=choices, - is_radio=is_radio, - is_multi_option=is_multi_option, - rounding=rounding, - min_max=min_max, - helptext=helptext) + self._options = {"title": title, + "subgroup": subgroup, + "group": group, + "default": default, + "initial_value": initial_value, + "choices": choices, + "is_radio": is_radio, + "is_multi_option": is_multi_option, + "rounding": rounding, + "min_max": min_max, + "helptext": helptext} self.control = self.get_control() self.tk_var = self.get_tk_var(initial_value, track_modified) logger.debug("Initialized %s", self.__class__.__name__) @@ -330,12 +345,12 @@ def _model_callback(var): if not config.user_config_dict["auto_load_model_stats"]: logger.debug("Session updating disabled by user config") return - if config.tk_vars["runningtask"].get(): + if config.tk_vars.running_task.get(): logger.debug("Task running. Not updating session") return folder = var.get() logger.debug("Setting analysis model folder callback: '%s'", folder) - get_config().tk_vars["analysis_folder"].set(folder) + get_config().tk_vars.analysis_folder.set(folder) class ControlPanel(ttk.Frame): # pylint:disable=too-many-ancestors @@ -405,10 +420,10 @@ def __init__(self, parent, options, # pylint:disable=too-many-arguments if self._style.startswith("SPanel"): self._theme = {**self._theme, **get_config().user_theme["group_settings"]} - self.group_frames = dict() - self._sub_group_frames = dict() + self.group_frames = {} + self._sub_group_frames = {} - canvas_kwargs = dict(bd=0, highlightthickness=0, bg=self._theme["panel_background"]) + canvas_kwargs = {"bd": 0, "highlightthickness": 0, "bg": self._theme["panel_background"]} self._canvas = tk.Canvas(self, **canvas_kwargs) self._canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True) @@ -512,8 +527,8 @@ def get_group_frame(self, group): group_frame.pack(side=tk.TOP, fill=tk.X, padx=5, pady=5, anchor=tk.NW) - self.group_frames[group] = dict(frame=retval, - chkbtns=self.checkbuttons_frame(retval)) + self.group_frames[group] = {"frame": retval, + "chkbtns": self.checkbuttons_frame(retval)} group_frame = self.group_frames[group] return group_frame @@ -528,7 +543,7 @@ def add_scrollbar(self): self.mainframe.bind("", self.update_scrollbar) logger.debug("Added Config Scrollbar") - def update_scrollbar(self, event): # pylint: disable=unused-argument + def update_scrollbar(self, event): # pylint:disable=unused-argument """ Update the options frame scrollbar """ self._canvas.configure(scrollregion=self._canvas.bbox("all")) @@ -630,7 +645,7 @@ def set_subframes(self): """ Set a sub-frame for each possible column """ subframes = [] for idx in range(self.max_columns): - name = "af_subframe_{}".format(idx) + name = f"af_subframe_{idx}" subframe = ttk.Frame(self.parent, name=name, style=f"{self._style}TFrame") if idx < self.columns: # Only pack visible columns @@ -705,14 +720,14 @@ def _custom_kwargs(cls, widget): dict The custom keyword arguments required for recreating the given widget """ - retval = dict() + retval = {} if widget.__class__.__name__ == "MultiOption": - retval = dict(value=widget._value, # pylint:disable=protected-access - variable=widget._master_variable) # pylint:disable=protected-access + retval = {"value": widget._value, # pylint:disable=protected-access + "variable": widget._master_variable} # pylint:disable=protected-access elif widget.__class__.__name__ == "ToggledFrame": # Toggled Frames need to have their variable tracked - retval = dict(text=widget._text, # pylint:disable=protected-access - toggle_var=widget._toggle_var) # pylint:disable=protected-access + retval = {"text": widget._text, # pylint:disable=protected-access + "toggle_var": widget._toggle_var} # pylint:disable=protected-access return retval def get_all_children_config(self, widget, child_list): @@ -773,7 +788,7 @@ def config_cleaner(widget): configuration from a widget We use config() instead of configure() because some items (ttk Scale) do not populate configure()""" - new_config = dict() + new_config = {} for key in widget.config(): if key == "class": continue @@ -893,7 +908,7 @@ class ControlBuilder(): blank_nones: bool Sets selected values to an empty string rather than None if this is true. """ - def __init__(self, parent, option, option_columns, # pylint: disable=too-many-arguments + def __init__(self, parent, option, option_columns, # pylint:disable=too-many-arguments label_width, checkbuttons_frame, style, blank_nones): logger.debug("Initializing %s: (parent: %s, option: %s, option_columns: %s, " "label_width: %s, checkbuttons_frame: %s, style: %s, blank_nones: %s)", @@ -923,7 +938,7 @@ def control_frame(self, parent): """ Frame to hold control and it's label """ logger.debug("Build control frame") frame = ttk.Frame(parent, - name="fr_{}".format(self.option.name), + name=f"fr_{self.option.name}", style=f"{self._style}Group.TFrame") frame.pack(fill=tk.X) logger.debug("Built control frame") @@ -955,7 +970,7 @@ def build_control_label(self): style=f"{self._style}Group.TLabel") lbl.pack(padx=5, pady=5, side=tk.LEFT, anchor=tk.N) if self.option.helptext is not None: - _get_tooltip(lbl, text=self.option.helptext, wrap_length=600) + _get_tooltip(lbl, text=self.option.helptext) logger.debug("Built control label: (widget: '%s', title: '%s'", self.option.name, self.option.title) @@ -975,7 +990,7 @@ def build_one_control(self): if self.option.control != ttk.Checkbutton: ctl.pack(padx=5, pady=5, fill=tk.X, expand=True) if self.option.helptext is not None and not self.helpset: - tooltip_kwargs = dict(text=self.option.helptext, wrap_length=600) + tooltip_kwargs = {"text": self.option.helptext} if self.option.sysbrowser is not None: tooltip_kwargs["text_variable"] = self.option.tk_var _get_tooltip(ctl, **tooltip_kwargs) @@ -996,7 +1011,7 @@ def _multi_option_control(self, option_type): help_intro, help_items = self._get_multi_help_items(self.option.helptext) ctl = ttk.LabelFrame(self.frame, text=self.option.title, - name="{}_labelframe".format(option_type), + name=f"{option_type}_labelframe", style=f"{self._style}Group.TLabelframe") holder = AutoFillContainer(ctl, self.option_columns, @@ -1018,8 +1033,8 @@ def _multi_option_control(self, option_type): if choice.lower() in help_items: self.helpset = True helptext = help_items[choice.lower()] - helptext = "{}\n\n - {}".format(helptext, help_intro) - _get_tooltip(ctl, text=helptext, wrap_length=600) + helptext = f"{helptext}\n\n - {help_intro}" + _get_tooltip(ctl, text=helptext) ctl.pack(anchor=tk.W, fill=tk.X) logger.debug("Added %s option %s", option_type, choice) return holder.parent @@ -1046,7 +1061,7 @@ def _get_multi_help_items(helptext): if any(line.startswith(" - ") for line in all_help): intro = all_help[0] retval = (intro, - {re.sub(r"[^A-Za-z0-9\-\_]+", "", + {re.sub(r"[^\w\-\_]+", "", line.split()[1].lower()): " ".join(line.replace("_", " ").split()[1:]) for line in all_help if line.startswith(" - ")}) logger.debug("help items: %s", retval) @@ -1058,7 +1073,7 @@ def slider_control(self): "rounding: %s, min_max: %s)", self.option.name, self.option.dtype, self.option.rounding, self.option.min_max) validate = self.slider_check_int if self.option.dtype == int else self.slider_check_float - vcmd = (self.frame.register(validate)) + vcmd = self.frame.register(validate) tbox = tk.Entry(self.frame, width=8, textvariable=self.option.tk_var, @@ -1166,34 +1181,35 @@ def _color_control(self): logger.debug("Add control to Options Frame: (widget: '%s', control: %s, choices: %s)", self.option.name, self.option.control, self.option.choices) frame = ttk.Frame(self.frame, style=f"{self._style}Group.TFrame") + lbl = ttk.Label(frame, + text=self.option.title, + width=self.label_width, + anchor=tk.W, + style=f"{self._style}Group.TLabel") ctl = tk.Frame(frame, - bg=self.option.default, + bg=self.option.tk_var.get(), bd=2, cursor="hand2", relief=tk.SUNKEN, width=round(int(20 * get_config().scaling_factor)), - height=round(int(12 * get_config().scaling_factor))) + height=round(int(14 * get_config().scaling_factor))) ctl.bind("", lambda *e, c=ctl, t=self.option.title: self._ask_color(c, t)) - ctl.pack(side=tk.LEFT, anchor=tk.W) - lbl = ttk.Label(frame, - text=self.option.title, - width=self.label_width, - anchor=tk.W, - style=f"{self._style}Group.TLabel") - lbl.pack(padx=2, pady=5, side=tk.RIGHT, anchor=tk.N) - frame.pack(side=tk.LEFT, anchor=tk.W) + lbl.pack(side=tk.LEFT, anchor=tk.N) + ctl.pack(side=tk.RIGHT, anchor=tk.W) + frame.pack(padx=5, side=tk.LEFT, anchor=tk.W) if self.option.helptext is not None: - _get_tooltip(lbl, text=self.option.helptext, wrap_length=600) + _get_tooltip(frame, text=self.option.helptext) + # Callback to set the color chooser background on an update (e.g. reset) + self.option.tk_var.trace("w", lambda *e: ctl.config(bg=self.option.tk_var.get())) logger.debug("Added control to Options Frame: %s", self.option.name) return ctl def _ask_color(self, frame, title): """ Pop ask color dialog set to variable and change frame color """ color = self.option.tk_var.get() - chosen = colorchooser.askcolor(color=color, title="{} Color".format(title))[1] + chosen = colorchooser.askcolor(parent=frame, color=color, title=f"{title} Color")[1] if chosen is None: return - frame.config(bg=chosen) self.option.tk_var.set(chosen) def control_to_checkframe(self): @@ -1205,7 +1221,7 @@ def control_to_checkframe(self): text=self.option.title, name=self.option.name, style=f"{self._style}Group.TCheckbutton") - _get_tooltip(ctl, text=self.option.helptext, wrap_length=600) + _get_tooltip(ctl, text=self.option.helptext) ctl.pack(side=tk.TOP, anchor=tk.W, fill=tk.X) logger.debug("Added control checkframe: '%s'", self.option.name) return ctl @@ -1232,15 +1248,15 @@ def __init__(self, opt_name, tk_var, control_frame, sysbrowser_dict, style): @property def helptext(self): """ Dict containing tooltip text for buttons """ - retval = dict(folder=_("Select a folder..."), - load=_("Select a file..."), - load2=_("Select a file..."), - picture=_("Select a folder of images..."), - video=_("Select a video..."), - model=_("Select a model folder..."), - multi_load=_("Select one or more files..."), - context=_("Select a file or folder..."), - save_as=_("Select a save location...")) + retval = {"folder": _("Select a folder..."), + "load": _("Select a file..."), + "load2": _("Select a file..."), + "picture": _("Select a folder of images..."), + "video": _("Select a video..."), + "model": _("Select a model folder..."), + "multi_load": _("Select one or more files..."), + "context": _("Select a file or folder..."), + "save_as": _("Select a save location...")} return retval @staticmethod @@ -1286,7 +1302,7 @@ def add_browser_buttons(self): cursor="hand2") _add_command(fileopn.cget("command"), cmd) fileopn.pack(padx=1, side=tk.RIGHT) - _get_tooltip(fileopn, text=self.helptext[lbl], wrap_length=600) + _get_tooltip(fileopn, text=self.helptext[lbl]) logger.debug("Added browser buttons: (action: %s, filetypes: %s", action, self.filetypes) @@ -1324,7 +1340,7 @@ def ask_multi_load(filepath, filetypes): """ Pop-up to get path to a file """ filenames = FileHandler("filename_multi", filetypes).return_file if filenames: - final_names = " ".join("\"{}\"".format(fname) for fname in filenames) + final_names = " ".join(f"\"{fname}\"" for fname in filenames) logger.debug(final_names) filepath.set(final_names) diff --git a/lib/gui/custom_widgets.py b/lib/gui/custom_widgets.py index 01e19c09ed..d18d832ed8 100644 --- a/lib/gui/custom_widgets.py +++ b/lib/gui/custom_widgets.py @@ -5,6 +5,7 @@ import platform import re import sys +import typing as T import tkinter as tk from tkinter import ttk, TclError @@ -12,10 +13,10 @@ from .utils import get_config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -class ContextMenu(tk.Menu): # pylint: disable=too-many-ancestors +class ContextMenu(tk.Menu): # pylint:disable=too-many-ancestors """ A Pop up menu to be triggered when right clicking on widgets that this menu has been applied to. @@ -71,7 +72,7 @@ def _select_all(self): self._widget.select_range(0, tk.END) -class RightClickMenu(tk.Menu): # pylint: disable=too-many-ancestors +class RightClickMenu(tk.Menu): # pylint:disable=too-many-ancestors """ A Pop up menu that can be bound to a right click mouse event to bring up a context menu Parameters @@ -101,7 +102,7 @@ def __init__(self, labels, actions, hotkeys=None): def _create_menu(self): """ Create the menu based on :attr:`_labels` and :attr:`_actions`. """ for idx, (label, action) in enumerate(zip(self._labels, self._actions)): - kwargs = dict(label=label, command=action) + kwargs = {"label": label, "command": action} if isinstance(self._hotkeys, (list, tuple)) and self._hotkeys[idx]: kwargs["accelerator"] = self._hotkeys[idx] self.add_command(**kwargs) @@ -117,7 +118,7 @@ def popup(self, event): self.tk_popup(event.x_root, event.y_root) -class ConsoleOut(ttk.Frame): # pylint: disable=too-many-ancestors +class ConsoleOut(ttk.Frame): # pylint:disable=too-many-ancestors """ The Console out section of the GUI. A Read only text box for displaying the output from stdout/stderr. @@ -141,7 +142,7 @@ def __init__(self, parent, debug): self._console = _ReadOnlyText(self, relief=tk.FLAT) rc_menu = ContextMenu(self._console) rc_menu.cm_bind() - self._console_clear = get_config().tk_vars['console_clear'] + self._console_clear = get_config().tk_vars.console_clear self._set_console_clear_var_trace() self._debug = debug self._build_console() @@ -194,7 +195,7 @@ def _redirect_console(self): sys.stderr = _SysOutRouter(self._console, "stderr") logger.debug("Redirected console") - def _clear(self, *args): # pylint: disable=unused-argument + def _clear(self, *args): # pylint:disable=unused-argument """ Clear the console output screen """ logger.debug("Clear console") if not self._console_clear.get(): @@ -205,7 +206,7 @@ def _clear(self, *args): # pylint: disable=unused-argument logger.debug("Cleared console") -class _ReadOnlyText(tk.Text): # pylint: disable=too-many-ancestors +class _ReadOnlyText(tk.Text): # pylint:disable=too-many-ancestors """ A read only text widget. Standard tkinter Text widgets are read/write by default. As we want to make the console @@ -238,6 +239,7 @@ def __init__(self, console, out_type): self._console = console self._out_type = out_type self._recolor = re.compile(r".+?(\s\d+:\d+:\d+\s)(?P[A-Z]+)\s") + self._ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") logger.debug("Initialized %s", self.__class__.__name__) def _get_tag(self, string): @@ -254,6 +256,7 @@ def _get_tag(self, string): def write(self, string): """ Capture stdout/stderr """ + string = self._ansi_escape.sub("", string) self._console.insert(tk.END, string, self._get_tag(string)) self._console.see(tk.END) @@ -315,9 +318,8 @@ def __init__(self, widget): tk_.createcommand(wgt, self.dispatch) def __repr__(self): - return "%s(%s<%s>)" % (self.__class__.__name__, - self.widget.__class__.__name__, - self.widget._w) # pylint:disable=protected-access + return (f"{self.__class__.__name__}({self.widget.__class__.__name__}" + f"<{self.widget._w}>)") # pylint:disable=protected-access def close(self): "de-register operations and revert redirection created by .__init__." @@ -409,14 +411,13 @@ def __init__(self, redirect, operation): self.orig_and_operation = (redirect.orig, operation) def __repr__(self): - return "%s(%r, %r)" % (self.__class__.__name__, - self.redirect, self.operation) + return f"{self.__class__.__name__}({self.redirect}, {self.operation})" def __call__(self, *args): return self.tk_call(self.orig_and_operation + args) -class StatusBar(ttk.Frame): # pylint: disable=too-many-ancestors +class StatusBar(ttk.Frame): # pylint:disable=too-many-ancestors """ Status Bar for displaying the Status Message and Progress Bar at the bottom of the GUI. Parameters @@ -428,12 +429,13 @@ class StatusBar(ttk.Frame): # pylint: disable=too-many-ancestors frame otherwise ``False``. Default: ``False`` """ - def __init__(self, parent, hide_status=False): + def __init__(self, parent: ttk.Frame, hide_status: bool = False) -> None: super().__init__(parent) self._frame = ttk.Frame(self) self._message = tk.StringVar() self._pbar_message = tk.StringVar() self._pbar_position = tk.IntVar() + self._mode: T.Literal["indeterminate", "determinate"] = "determinate" self._message.set("Ready") @@ -443,12 +445,12 @@ def __init__(self, parent, hide_status=False): self._frame.pack(padx=10, pady=2, fill=tk.X, expand=False) @property - def message(self): + def message(self) -> tk.StringVar: """:class:`tkinter.StringVar`: The variable to hold the status bar message on the left hand side of the status bar. """ return self._message - def _status(self, hide_status): + def _status(self, hide_status: bool) -> None: """ Place Status label into left of the status bar. Parameters @@ -472,8 +474,14 @@ def _status(self, hide_status): anchor=tk.W) lblstatus.pack(side=tk.LEFT, anchor=tk.W, fill=tk.X, expand=True) - def _progress_bar(self): - """ Place progress bar into right of the status bar. """ + def _progress_bar(self) -> ttk.Progressbar: + """ Place progress bar into right of the status bar. + + Returns + ------- + :class:`tkinter.ttk.Progressbar` + The progress bar object + """ progressframe = ttk.Frame(self._frame) progressframe.pack(side=tk.RIGHT, anchor=tk.E, fill=tk.X) @@ -484,12 +492,12 @@ def _progress_bar(self): length=200, variable=self._pbar_position, maximum=100, - mode="determinate") + mode=self._mode) pbar.pack(side=tk.LEFT, padx=2, fill=tk.X, expand=True) pbar.pack_forget() return pbar - def start(self, mode): + def start(self, mode: T.Literal["indeterminate", "determinate"]) -> None: """ Set progress bar mode and display, Parameters @@ -500,16 +508,24 @@ def start(self, mode): self._set_mode(mode) self._pbar.pack() - def stop(self): + def stop(self) -> None: """ Reset progress bar and hide """ self._pbar_message.set("") self._pbar_position.set(0) - self._set_mode("determinate") + self._mode = "determinate" + self._set_mode(self._mode) self._pbar.pack_forget() - def _set_mode(self, mode): - """ Set the progress bar mode """ - self._pbar.config(mode=mode) + def _set_mode(self, mode: T.Literal["indeterminate", "determinate"]) -> None: + """ Set the progress bar mode + + Parameters + ---------- + mode: ["indeterminate", "determinate"] + The mode that the progress bar should be executed in + """ + self._mode = mode + self._pbar.config(mode=self._mode) if mode == "indeterminate": self._pbar.config(maximum=100) self._pbar.start() @@ -517,7 +533,23 @@ def _set_mode(self, mode): self._pbar.stop() self._pbar.config(maximum=100) - def progress_update(self, message, position, update_position=True): + def set_mode(self, mode: T.Literal["indeterminate", "determinate"]) -> None: + """ Set the mode of a currently displayed progress bar and reset position to 0. + + If the given mode is the same as the currently configured mode, returns without performing + any action. + + Parameters + ---------- + mode: ["indeterminate", "determinate"] + The mode that the progress bar should be set to + """ + if mode == self._mode: + return + self.stop() + self.start(mode) + + def progress_update(self, message: str, position: int, update_position: bool = True) -> None: """ Update the GUIs progress bar and position. Parameters @@ -619,12 +651,8 @@ def tip_pos_calculator(widget, label, x_1, y_1 = mouse_x + tip_delta[0], mouse_y + tip_delta[1] x_2, y_2 = x_1 + width, y_1 + height - x_delta = x_2 - s_width - if x_delta < 0: - x_delta = 0 - y_delta = y_2 - s_height - if y_delta < 0: - y_delta = 0 + x_delta = max(x_2 - s_width, 0) + y_delta = max(y_2 - s_height, 0) offscreen = (x_delta, y_delta) != (0, 0) @@ -670,7 +698,7 @@ def tip_pos_calculator(widget, label, text = self._text if self._text_variable and self._text_variable.get(): - text += "\n\nCurrent value: '{}'".format(self._text_variable.get()) + text += f"\n\nCurrent value: '{self._text_variable.get()}'" label = tk.Label(win, text=text, justify=tk.LEFT, @@ -687,7 +715,7 @@ def tip_pos_calculator(widget, label, xpos, ypos = tip_pos_calculator(widget, label) - self._topwidget.wm_geometry("+%d+%d" % (xpos, ypos)) + self._topwidget.wm_geometry(f"+{xpos}+{ypos}") def _hide(self): """ Hide the tooltip """ @@ -697,7 +725,7 @@ def _hide(self): self._topwidget = None -class MultiOption(ttk.Checkbutton): # pylint: disable=too-many-ancestors +class MultiOption(ttk.Checkbutton): # pylint:disable=too-many-ancestors """ Similar to the standard :class:`ttk.Radio` widget, but with the ability to select multiple pre-defined options. Selected options are generated as `nargs` for the argument parser to consume. @@ -739,7 +767,7 @@ def _master_needs_update(self): logger.trace(retval) return retval - def _on_update(self, *args): # pylint: disable=unused-argument + def _on_update(self, *args): # pylint:disable=unused-argument """ Update the master variable on a check button change. The value for this checked option is added or removed from the :attr:`_master_variable` @@ -760,7 +788,7 @@ def _on_update(self, *args): # pylint: disable=unused-argument logger.trace("Setting master variable to: %s", val) self._master_variable.set(val) - def _on_master_update(self, *args): # pylint: disable=unused-argument + def _on_master_update(self, *args): # pylint:disable=unused-argument """ Update the check button on a master variable change (e.g. load .fsw file in the GUI). The value for this option is set to ``True`` or ``False`` depending on it's existence in @@ -819,7 +847,7 @@ def __init__(self, title, total): center = np.array(( (self.master.winfo_width() // 2) - (self.winfo_width() // 2), (self.master.winfo_height() // 2) - (self.winfo_height() // 2))) + offset - self.wm_geometry("+{}+{}".format(*center)) + self.wm_geometry(f"+{center[0]}+{center[1]}") get_config().set_cursor_busy() self.grab_set() diff --git a/lib/gui/display.py b/lib/gui/display.py index 86141296cd..3729e35437 100644 --- a/lib/gui/display.py +++ b/lib/gui/display.py @@ -10,18 +10,20 @@ import tkinter as tk from tkinter import ttk +from lib.logger import parse_class_init + from .display_analysis import Analysis from .display_command import GraphDisplay, PreviewExtract, PreviewTrain from .utils import get_config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) _ = _LANG.gettext -class DisplayNotebook(ttk.Notebook): # pylint: disable=too-many-ancestors +class DisplayNotebook(ttk.Notebook): # pylint:disable=too-many-ancestors """ The tkinter Notebook that holds the display items. Parameters @@ -31,12 +33,12 @@ class DisplayNotebook(ttk.Notebook): # pylint: disable=too-many-ancestors """ def __init__(self, parent): - logger.debug("Initializing %s", self.__class__.__name__) + logger.debug(parse_class_init(locals())) super().__init__(parent) parent.add(self) tk_vars = get_config().tk_vars - self._wrapper_var = tk_vars["display"] - self._runningtask = tk_vars["runningtask"] + self._wrapper_var = tk_vars.display + self._running_task = tk_vars.running_task self._set_wrapper_var_trace() self._add_static_tabs() @@ -46,10 +48,10 @@ def __init__(self, parent): logger.debug("Initialized %s", self.__class__.__name__) @property - def runningtask(self): + def running_task(self): """ :class:`tkinter.BooleanVar`: The global tkinter variable that indicates whether a Faceswap task is currently running or not. """ - return self._runningtask + return self._running_task def _set_wrapper_var_trace(self): """ Sets the trigger to update the displayed notebook's pages when the global tkinter @@ -95,7 +97,7 @@ def _command_display(self, command): command: str The Faceswap command that is being executed """ - build_tabs = getattr(self, "_{}_tabs".format(command)) + build_tabs = getattr(self, f"_{command}_tabs") build_tabs() def _extract_tabs(self, command="extract"): @@ -150,7 +152,7 @@ def _remove_tabs(self): child_object.close() # Call the OptionalDisplayPage close() method self.forget(child) - def _update_displaybook(self, *args): # pylint: disable=unused-argument + def _update_displaybook(self, *args): # pylint:disable=unused-argument """ Callback to be executed when the global tkinter variable `display` (:attr:`wrapper_var`) is updated when a Faceswap task is executed. diff --git a/lib/gui/display_analysis.py b/lib/gui/display_analysis.py index a682227622..9dcef89164 100644 --- a/lib/gui/display_analysis.py +++ b/lib/gui/display_analysis.py @@ -8,20 +8,22 @@ import tkinter as tk from tkinter import ttk +from lib.logger import parse_class_init + from .custom_widgets import Tooltip from .display_page import DisplayPage from .popup_session import SessionPopUp from .analysis import Session from .utils import FileHandler, get_config, get_images, LongRunningTask -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) _ = _LANG.gettext -class Analysis(DisplayPage): # pylint: disable=too-many-ancestors +class Analysis(DisplayPage): # pylint:disable=too-many-ancestors """ Session Analysis Tab. The area of the GUI that holds the session summary stats for model training sessions. @@ -36,8 +38,7 @@ class Analysis(DisplayPage): # pylint: disable=too-many-ancestors The help text to display for the summary statistics page """ def __init__(self, parent, tab_name, helptext): - logger.debug("Initializing: %s: (parent, %s, tab_name: '%s', helptext: '%s')", - self.__class__.__name__, parent, tab_name, helptext) + logger.debug(parse_class_init(locals())) super().__init__(parent, tab_name, helptext) self._summary = None @@ -62,10 +63,10 @@ def set_vars(self): dict The dictionary of variable names to tkinter variables """ - return dict(selected_id=tk.StringVar(), - refresh_graph=get_config().tk_vars["refreshgraph"], - is_training=get_config().tk_vars["istraining"], - analysis_folder=get_config().tk_vars["analysis_folder"]) + return {"selected_id": tk.StringVar(), + "refresh_graph": get_config().tk_vars.refresh_graph, + "is_training": get_config().tk_vars.is_training, + "analysis_folder": get_config().tk_vars.analysis_folder} def on_tab_select(self): """ Callback for when the analysis tab is selected. @@ -299,7 +300,7 @@ class _Options(): # pylint:disable=too-few-public-methods The Analysis Display Tab that holds the options buttons """ def __init__(self, parent): - logger.debug("Initializing: %s (parent: %s)", self.__class__.__name__, parent) + logger.debug(parse_class_init(locals())) self._parent = parent self._buttons = self._add_buttons() self._add_training_callback() @@ -365,7 +366,7 @@ def _set_buttons_state(self, *args): # pylint:disable=unused-argument button.state([state]) -class StatsData(ttk.Frame): # pylint: disable=too-many-ancestors +class StatsData(ttk.Frame): # pylint:disable=too-many-ancestors """ Stats frame of analysis tab. Holds the tree-view containing the summarized session statistics in the Analysis tab. @@ -380,8 +381,7 @@ class StatsData(ttk.Frame): # pylint: disable=too-many-ancestors The help text to display for the summary statistics page """ def __init__(self, parent, selected_id, helptext): - logger.debug("Initializing: %s: (parent, %s, selected_id: %s, helptext: '%s')", - self.__class__.__name__, parent, selected_id, helptext) + logger.debug(parse_class_init(locals())) super().__init__(parent) self._selected_id = selected_id diff --git a/lib/gui/display_command.py b/lib/gui/display_command.py index a43b54acab..3cc2e31887 100644 --- a/lib/gui/display_command.py +++ b/lib/gui/display_command.py @@ -5,9 +5,12 @@ import logging import os import tkinter as tk +import typing as T from tkinter import ttk +from lib.logger import parse_class_init +from lib.training.preview_tk import PreviewTk from .display_graph import TrainingGraph from .display_page import DisplayOptionalPage @@ -16,47 +19,56 @@ from .control_helper import set_slider_rounding from .utils import FileHandler, get_config, get_images, preview_trigger -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) _ = _LANG.gettext -class PreviewExtract(DisplayOptionalPage): # pylint: disable=too-many-ancestors +class PreviewExtract(DisplayOptionalPage): # pylint:disable=too-many-ancestors """ Tab to display output preview images for extract and convert """ + def __init__(self, *args, **kwargs) -> None: + logger.debug(parse_class_init(locals())) + self._preview = get_images().preview_extract + super().__init__(*args, **kwargs) + logger.debug("Initialized %s", self.__class__.__name__) - def display_item_set(self): + def display_item_set(self) -> None: """ Load the latest preview if available """ - logger.trace("Loading latest preview") - size = 256 if self.command == "convert" else 128 - get_images().load_latest_preview(thumbnail_size=int(size * get_config().scaling_factor), - frame_dims=(self.winfo_width(), self.winfo_height())) - self.display_item = get_images().previewoutput + logger.trace("Loading latest preview") # type:ignore[attr-defined] + size = int(256 if self.command == "convert" else 128 * get_config().scaling_factor) + if not self._preview.load_latest_preview(thumbnail_size=size, + frame_dims=(self.winfo_width(), + self.winfo_height())): + logger.trace("Preview not updated") # type:ignore[attr-defined] + return + logger.debug("Preview loaded") + self.display_item = True - def display_item_process(self): + def display_item_process(self) -> None: """ Display the preview """ - logger.trace("Displaying preview") + logger.trace("Displaying preview") # type:ignore[attr-defined] if not self.subnotebook.children: self.add_child() else: self.update_child() - def add_child(self): + def add_child(self) -> None: """ Add the preview label child """ logger.debug("Adding child") preview = self.subnotebook_add_page(self.tabname, widget=None) - lblpreview = ttk.Label(preview, image=get_images().previewoutput[1]) + lblpreview = ttk.Label(preview, image=self._preview.image) lblpreview.pack(side=tk.TOP, anchor=tk.NW) Tooltip(lblpreview, text=self.helptext, wrap_length=200) - def update_child(self): + def update_child(self) -> None: """ Update the preview image on the label """ - logger.trace("Updating preview") + logger.trace("Updating preview") # type:ignore[attr-defined] for widget in self.subnotebook_get_widgets(): - widget.configure(image=get_images().previewoutput[1]) + widget.configure(image=self._preview.image) - def save_items(self): + def save_items(self) -> None: """ Open save dialogue and save preview """ location = FileHandler("dir", None).return_file if not location: @@ -64,149 +76,108 @@ def save_items(self): filename = "extract_convert_preview" now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") filename = os.path.join(location, f"{filename}_{now}.png") - get_images().previewoutput[0].save(filename) - logger.debug("Saved preview to %s", filename) + self._preview.save(filename) print(f"Saved preview to {filename}") -class PreviewTrain(DisplayOptionalPage): # pylint: disable=too-many-ancestors +class PreviewTrain(DisplayOptionalPage): # pylint:disable=too-many-ancestors """ Training preview image(s) """ - def __init__(self, *args, **kwargs): - self.update_preview = get_config().tk_vars["updatepreview"] + def __init__(self, *args, **kwargs) -> None: + logger.debug(parse_class_init(locals())) + self._preview = get_images().preview_train + self._display: PreviewTk | None = None super().__init__(*args, **kwargs) + logger.debug("Initialized %s", self.__class__.__name__) - def add_options(self): + def add_options(self) -> None: """ Add the additional options """ self._add_option_refresh() self._add_option_mask_toggle() super().add_options() - def _add_option_refresh(self): + def subnotebook_hide(self) -> None: + """ Override default subnotebook hide action to also remove the embedded option bar + control and reset the training image buffer """ + if self.subnotebook and self.subnotebook.winfo_ismapped(): + logger.debug("Removing preview controls from options bar") + if self._display is not None: + self._display.remove_option_controls() + super().subnotebook_hide() + del self._display + self._display = None + self._preview.reset() + + def _add_option_refresh(self) -> None: """ Add refresh button to refresh preview immediately """ logger.debug("Adding refresh option") btnrefresh = ttk.Button(self.optsframe, image=get_images().icons["reload"], - command=lambda x="update": preview_trigger().set(x)) + command=lambda x="update": preview_trigger().set(x)) # type:ignore btnrefresh.pack(padx=2, side=tk.RIGHT) Tooltip(btnrefresh, text=_("Preview updates at every model save. Click to refresh now."), wrap_length=200) logger.debug("Added refresh option") - def _add_option_mask_toggle(self): + def _add_option_mask_toggle(self) -> None: """ Add button to toggle mask display on and off """ logger.debug("Adding mask toggle option") - btntoggle = ttk.Button(self.optsframe, - image=get_images().icons["mask2"], - command=lambda x="mask_toggle": preview_trigger().set(x)) + btntoggle = ttk.Button( + self.optsframe, + image=get_images().icons["mask2"], + command=lambda x="mask_toggle": preview_trigger().set(x)) # type:ignore btntoggle.pack(padx=2, side=tk.RIGHT) Tooltip(btntoggle, text=_("Click to toggle mask overlay on and off."), wrap_length=200) logger.debug("Added mask toggle option") - def display_item_set(self): + def display_item_set(self) -> None: """ Load the latest preview if available """ - logger.trace("Loading latest preview") - if not self.update_preview.get(): - logger.trace("Preview not updated") + # TODO This seems to be triggering faster than the waittime + logger.trace("Loading latest preview") # type:ignore[attr-defined] + if not self._preview.load(): + logger.trace("Preview not updated") # type:ignore[attr-defined] return - get_images().load_training_preview() - self.display_item = get_images().previewtrain + logger.debug("Preview loaded") + self.display_item = True - def display_item_process(self): + def display_item_process(self) -> None: """ Display the preview(s) resized as appropriate """ - logger.trace("Displaying preview") - sortednames = sorted(list(get_images().previewtrain.keys())) - existing = self.subnotebook_get_titles_ids() - should_update = self.update_preview.get() - - for name in sortednames: - if name not in existing: - self.add_child(name) - elif should_update: - tab_id = existing[name] - self.update_child(tab_id, name) - - if should_update: - self.update_preview.set(False) - - def add_child(self, name): - """ Add the preview canvas child """ - logger.debug("Adding child") - preview = PreviewTrainCanvas(self.subnotebook, name) - preview = self.subnotebook_add_page(name, widget=preview) - Tooltip(preview, text=self.helptext, wrap_length=200) - self.vars["modified"].set(get_images().previewtrain[name][2]) - - def update_child(self, tab_id, name): - """ Update the preview canvas """ - logger.debug("Updating preview") - if self.vars["modified"].get() != get_images().previewtrain[name][2]: - self.vars["modified"].set(get_images().previewtrain[name][2]) - widget = self.subnotebook_page_from_id(tab_id) - widget.reload() - - def save_items(self): + if self.subnotebook.children: + return + + logger.debug("Displaying preview") + self._display = PreviewTk(self._preview.buffer, self.subnotebook, self.optsframe, None) + self.subnotebook_add_page(self.tabname, widget=self._display.master_frame) + + def save_items(self) -> None: """ Open save dialogue and save preview """ + if self._display is None: + return + location = FileHandler("dir", None).return_file if not location: return - for preview in self.subnotebook.children.values(): - preview.save_preview(location) - - -class PreviewTrainCanvas(ttk.Frame): # pylint: disable=too-many-ancestors - """ Canvas to hold a training preview image """ - def __init__(self, parent, previewname): - logger.debug("Initializing %s: (previewname: '%s')", self.__class__.__name__, previewname) - ttk.Frame.__init__(self, parent) - - self.name = previewname - get_images().resize_image(self.name, None) - self.previewimage = get_images().previewtrain[self.name][1] - - self.canvas = tk.Canvas(self, bd=0, highlightthickness=0) - self.canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - self.imgcanvas = self.canvas.create_image(0, - 0, - image=self.previewimage, - anchor=tk.NW) - self.bind("", self.resize) - logger.debug("Initialized %s:", self.__class__.__name__) - - def resize(self, event): - """ Resize the image to fit the frame, maintaining aspect ratio """ - logger.trace("Resizing preview image") - framesize = (event.width, event.height) - # Sometimes image is resized before frame is drawn - framesize = None if framesize == (1, 1) else framesize - get_images().resize_image(self.name, framesize) - self.reload() - - def reload(self): - """ Reload the preview image """ - logger.trace("Reloading preview image") - self.previewimage = get_images().previewtrain[self.name][1] - self.canvas.itemconfig(self.imgcanvas, image=self.previewimage) - - def save_preview(self, location): - """ Save the figure to file """ - filename = self.name - now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - filename = os.path.join(location, f"{filename}_{now}.png") - get_images().previewtrain[self.name][0].save(filename) - logger.debug("Saved preview to %s", filename) - print(f"Saved preview to {filename}") + + self._display.save(location) -class GraphDisplay(DisplayOptionalPage): # pylint: disable=too-many-ancestors +class GraphDisplay(DisplayOptionalPage): # pylint:disable=too-many-ancestors """ The Graph Tab of the Display section """ - def __init__(self, parent, tab_name, helptext, wait_time, command=None): - self._trace_vars = {} + def __init__(self, + parent: ttk.Notebook, + tab_name: str, + helptext: str, + wait_time: int, + command: str | None = None) -> None: + logger.debug(parse_class_init(locals())) + self._trace_vars: dict[T.Literal["smoothgraph", "display_iterations"], + tuple[tk.BooleanVar, str]] = {} super().__init__(parent, tab_name, helptext, wait_time, command) + logger.debug("Initialized %s", self.__class__.__name__) - def set_vars(self): + def set_vars(self) -> None: """ Add graphing specific variables to the default variables. Overrides original method. @@ -237,17 +208,18 @@ def set_vars(self): logger.debug(tk_vars) return tk_vars - def on_tab_select(self): + def on_tab_select(self) -> None: """ Callback for when the graph tab is selected. Pull latest data and run the tab's update code when the tab is selected. """ - logger.debug("Callback received for '%s' tab", self.tabname) + logger.debug("Callback received for '%s' tab (display_item: %s)", + self.tabname, self.display_item) if self.display_item is not None: - get_config().tk_vars["refreshgraph"].set(True) + get_config().tk_vars.refresh_graph.set(True) self._update_page() - def add_options(self): + def add_options(self) -> None: """ Add the additional options """ self._add_option_refresh() super().add_options() @@ -256,10 +228,10 @@ def add_options(self): self._add_option_smoothing() self._add_option_iterations() - def _add_option_refresh(self): + def _add_option_refresh(self) -> None: """ Add refresh button to refresh graph immediately """ logger.debug("Adding refresh option") - tk_var = get_config().tk_vars["refreshgraph"] + tk_var = get_config().tk_vars.refresh_graph btnrefresh = ttk.Button(self.optsframe, image=get_images().icons["reload"], command=lambda: tk_var.set(True)) @@ -269,7 +241,7 @@ def _add_option_refresh(self): wrap_length=200) logger.debug("Added refresh option") - def _add_option_raw(self): + def _add_option_raw(self) -> None: """ Add check-button to hide/display raw data """ logger.debug("Adding display raw option") tk_var = self.vars["raw_data"] @@ -277,11 +249,11 @@ def _add_option_raw(self): self.optsframe, variable=tk_var, text="Raw", - command=lambda v=tk_var: self._display_data_callback("raw", v)) + command=lambda v=tk_var: self._display_data_callback("raw", v)) # type:ignore chkbtn.pack(side=tk.RIGHT, padx=5, anchor=tk.W) Tooltip(chkbtn, text=_("Display the raw loss data"), wrap_length=200) - def _add_option_smoothed(self): + def _add_option_smoothed(self) -> None: """ Add check-button to hide/display smoothed data """ logger.debug("Adding display smoothed option") tk_var = self.vars["smooth_data"] @@ -289,11 +261,11 @@ def _add_option_smoothed(self): self.optsframe, variable=tk_var, text="Smoothed", - command=lambda v=tk_var: self._display_data_callback("smoothed", v)) + command=lambda v=tk_var: self._display_data_callback("smoothed", v)) # type:ignore chkbtn.pack(side=tk.RIGHT, padx=5, anchor=tk.W) Tooltip(chkbtn, text=_("Display the smoothed loss data"), wrap_length=200) - def _add_option_smoothing(self): + def _add_option_smoothing(self) -> None: """ Add a slider to adjust the smoothing amount """ logger.debug("Adding Smoothing Slider") tk_var = self.vars["smoothgraph"] @@ -312,7 +284,7 @@ def _add_option_smoothing(self): ctl = ttk.Scale( ctl_frame, variable=tk_var, - command=lambda val, var=tk_var, dt=float, rn=3, mm=min_max: + command=lambda val, var=tk_var, dt=float, rn=3, mm=min_max: # type:ignore set_slider_rounding(val, var, dt, rn, mm)) ctl["from_"] = min_max[0] ctl["to"] = min_max[1] @@ -323,7 +295,7 @@ def _add_option_smoothing(self): wrap_length=200) logger.debug("Added Smoothing Slider") - def _add_option_iterations(self): + def _add_option_iterations(self) -> None: """ Add a slider to adjust the amount if iterations to display """ logger.debug("Adding Iterations Slider") tk_var = self.vars["display_iterations"] @@ -342,7 +314,7 @@ def _add_option_iterations(self): ctl = ttk.Scale( ctl_frame, variable=tk_var, - command=lambda val, var=tk_var, dt=int, rn=1000, mm=min_max: + command=lambda val, var=tk_var, dt=int, rn=1000, mm=min_max: # type:ignore set_slider_rounding(val, var, dt, rn, mm)) ctl["from_"] = min_max[0] ctl["to"] = min_max[1] @@ -353,32 +325,31 @@ def _add_option_iterations(self): wrap_length=200) logger.debug("Added Iterations Slider") - def display_item_set(self): + def display_item_set(self) -> None: """ Load the graph(s) if available """ if Session.is_training and Session.logging_disabled: - logger.trace("Logs disabled. Hiding graph") + logger.trace("Logs disabled. Hiding graph") # type:ignore[attr-defined] self.set_info("Graph is disabled as 'no-logs' has been selected") self.display_item = None self._clear_trace_variables() elif Session.is_training and self.display_item is None: - logger.trace("Loading graph") + logger.trace("Loading graph") # type:ignore[attr-defined] self.display_item = Session self._add_trace_variables() elif Session.is_training and self.display_item is not None: - logger.trace("Graph already displayed. Nothing to do.") + logger.trace("Graph already displayed. Nothing to do.") # type:ignore[attr-defined] else: - logger.trace("Clearing graph") + logger.trace("Clearing graph") # type:ignore[attr-defined] self.display_item = None self._clear_trace_variables() - def display_item_process(self): + def display_item_process(self) -> None: """ Add a single graph to the graph window """ if not Session.is_training: logger.debug("Waiting for Session Data to become available to graph") self.after(1000, self.display_item_process) return - logger.debug("Adding graph") existing = list(self.subnotebook_get_titles_ids().keys()) loss_keys = self.display_item.get_loss_keys(Session.session_ids[-1]) @@ -395,6 +366,7 @@ def display_item_process(self): tabname = loss_key.replace("_", " ").title() if tabname in existing: continue + logger.debug("Adding graph '%s'", tabname) display_keys = [key for key in loss_keys if key.startswith(loss_key)] data = Calculations(session_id=Session.session_ids[-1], @@ -404,7 +376,7 @@ def display_item_process(self): smooth_amount=self.vars["smoothgraph"].get()) self.add_child(tabname, data) - def _smooth_amount_callback(self, *args): + def _smooth_amount_callback(self, *args) -> None: """ Update each graph's smooth amount on variable change """ try: smooth_amount = self.vars["smoothgraph"].get() @@ -416,7 +388,7 @@ def _smooth_amount_callback(self, *args): for graph in self.subnotebook.children.values(): graph.calcs.set_smooth_amount(smooth_amount) - def _iteration_limit_callback(self, *args): + def _iteration_limit_callback(self, *args) -> None: """ Limit the amount of data displayed in the live graph on a iteration slider variable change. """ try: @@ -429,7 +401,7 @@ def _iteration_limit_callback(self, *args): for graph in self.subnotebook.children.values(): graph.calcs.set_iterations_limit(limit) - def _display_data_callback(self, line, variable): + def _display_data_callback(self, line: str, variable: tk.BooleanVar) -> None: """ Update the displayed graph lines based on option check button selection. Parameters @@ -444,15 +416,23 @@ def _display_data_callback(self, line, variable): for graph in self.subnotebook.children.values(): graph.calcs.update_selections(line, var) - def add_child(self, name, data): - """ Add the graph for the selected keys """ + def add_child(self, name: str, data: Calculations) -> None: + """ Add the graph for the selected keys. + + Parameters + ---------- + name: str + The name of the graph to add to the notebook + data: :class:`~lib.gui.analysis.stats.Calculations` + The object holding the data to be graphed + """ logger.debug("Adding child: %s", name) graph = TrainingGraph(self.subnotebook, data, "Loss") graph.build() graph = self.subnotebook_add_page(name, widget=graph) Tooltip(graph, text=self.helptext, wrap_length=200) - def save_items(self): + def save_items(self) -> None: """ Open save dialogue and save graphs """ graphlocation = FileHandler("dir", None).return_file if not graphlocation: @@ -460,15 +440,15 @@ def save_items(self): for graph in self.subnotebook.children.values(): graph.save_fig(graphlocation) - def _add_trace_variables(self): + def _add_trace_variables(self) -> None: """ Add tracing for when the option sliders are updated, for updating the graph. """ - for name, action in zip(("smoothgraph", "display_iterations"), + for name, action in zip(T.get_args(T.Literal["smoothgraph", "display_iterations"]), (self._smooth_amount_callback, self._iteration_limit_callback)): var = self.vars[name] if name not in self._trace_vars: self._trace_vars[name] = (var, var.trace("w", action)) - def _clear_trace_variables(self): + def _clear_trace_variables(self) -> None: """ Clear all of the trace variables from :attr:`_trace_vars` and reset the dictionary. """ if self._trace_vars: for name, (var, trace) in self._trace_vars.items(): @@ -476,7 +456,7 @@ def _clear_trace_variables(self): var.trace_vdelete("w", trace) self._trace_vars = {} - def close(self): + def close(self) -> None: """ Clear the plots from RAM """ self._clear_trace_variables() if self.subnotebook is None: diff --git a/lib/gui/display_graph.py b/lib/gui/display_graph.py index c2187450e4..c5f1304a81 100755 --- a/lib/gui/display_graph.py +++ b/lib/gui/display_graph.py @@ -1,12 +1,13 @@ #!/usr/bin python3 """ Graph functions for Display Frame area of the Faceswap GUI """ +from __future__ import annotations import datetime import logging import os import tkinter as tk +import typing as T from tkinter import ttk -from typing import Union, List, Tuple from math import ceil, floor import numpy as np @@ -17,15 +18,18 @@ NavigationToolbar2Tk) from matplotlib.backend_bases import NavigationToolbar2 +from lib.logger import parse_class_init + from .custom_widgets import Tooltip from .utils import get_config, get_images, LongRunningTask -matplotlib.use("TkAgg") +if T.TYPE_CHECKING: + from matplotlib.lines import Line2D logger: logging.Logger = logging.getLogger(__name__) -class GraphBase(ttk.Frame): # pylint: disable=too-many-ancestors +class GraphBase(ttk.Frame): # pylint:disable=too-many-ancestors """ Base class for matplotlib line graphs. Parameters @@ -38,16 +42,16 @@ class GraphBase(ttk.Frame): # pylint: disable=too-many-ancestors The data label for the y-axis """ def __init__(self, parent: ttk.Frame, data, ylabel: str) -> None: - logger.debug("Initializing %s", self.__class__.__name__) super().__init__(parent) + matplotlib.use("TkAgg") # Can't be at module level as breaks Github CI style.use("ggplot") self._calcs = data self._ylabel = ylabel self._colourmaps = ["Reds", "Blues", "Greens", "Purples", "Oranges", "Greys", "copper", "summer", "bone", "hot", "cool", "pink", "Wistia", "spring", "winter"] - self._lines = [] - self._toolbar = None + self._lines: list[Line2D] = [] + self._toolbar: "NavigationToolbar" | None = None self._fig = Figure(figsize=(4, 4), dpi=75) self._ax1 = self._fig.add_subplot(1, 1, 1) @@ -55,7 +59,6 @@ def __init__(self, parent: ttk.Frame, data, ylabel: str) -> None: self._initiate_graph() self._update_plot(initiate=True) - logger.debug("Initialized %s", self.__class__.__name__) @property def calcs(self): @@ -84,7 +87,7 @@ def _update_plot(self, initiate: bool = True) -> None: Whether the graph should be initialized for the first time (``True``) or data is being updated for an existing graph (``False``). Default: ``True`` """ - logger.trace("Updating plot") + logger.trace("Updating plot") # type:ignore[attr-defined] if initiate: logger.debug("Initializing plot") self._lines = [] @@ -112,7 +115,7 @@ def _update_plot(self, initiate: bool = True) -> None: if initiate: self._legend_place() - logger.trace("Updated plot") + logger.trace("Updated plot") # type:ignore[attr-defined] def _axes_labels_set(self) -> None: """ Set the X and Y axes labels. """ @@ -126,7 +129,7 @@ def _axes_limits_set_default(self) -> None: self._ax1.set_ylim(0.00, 100.0) self._ax1.set_xlim(0, 1) - def _axes_limits_set(self, data: List[float]) -> None: + def _axes_limits_set(self, data: list[float]) -> None: """ Set the axes limits. Parameters @@ -145,12 +148,13 @@ def _axes_limits_set(self, data: List[float]) -> None: ymin, ymax = self._axes_data_get_min_max(data) self._ax1.set_ylim(ymin, ymax) self._ax1.set_xlim(xmin, xmax) - logger.trace("axes ranges: (y: (%s, %s), x:(0, %s)", ymin, ymax, xmax) + logger.trace("axes ranges: (y: (%s, %s), x:(0, %s)", # type:ignore[attr-defined] + ymin, ymax, xmax) else: self._axes_limits_set_default() @staticmethod - def _axes_data_get_min_max(data: List[float]) -> Tuple[float]: + def _axes_data_get_min_max(data: list[float]) -> tuple[float, float]: """ Obtain the minimum and maximum values for the y-axis from the given data points. Parameters @@ -163,14 +167,14 @@ def _axes_data_get_min_max(data: List[float]) -> Tuple[float]: tuple The minimum and maximum values for the y axis """ - ymin, ymax = [], [] + ymins, ymaxs = [], [] for item in data: # TODO Handle as array not loop - ymin.append(np.nanmin(item) * 1000) - ymax.append(np.nanmax(item) * 1000) - ymin = floor(min(ymin)) / 1000 - ymax = ceil(max(ymax)) / 1000 - logger.trace("ymin: %s, ymax: %s", ymin, ymax) + ymins.append(np.nanmin(item) * 1000) + ymaxs.append(np.nanmax(item) * 1000) + ymin = floor(min(ymins)) / 1000 + ymax = ceil(max(ymaxs)) / 1000 + logger.trace("ymin: %s, ymax: %s", ymin, ymax) # type:ignore[attr-defined] return ymin, ymax def _axes_set_yscale(self, scale: str) -> None: @@ -184,7 +188,7 @@ def _axes_set_yscale(self, scale: str) -> None: logger.debug("yscale: '%s'", scale) self._ax1.set_yscale(scale) - def _lines_sort(self, keys: List[str]) -> List[List[Union[str, int, Tuple[float]]]]: + def _lines_sort(self, keys: list[str]) -> list[list[str | int | tuple[float]]]: """ Sort the data keys into consistent order and set line color map and line width. Parameters @@ -197,9 +201,9 @@ def _lines_sort(self, keys: List[str]) -> List[List[Union[str, int, Tuple[float] list A list of loss keys with their corresponding line formatting and color information """ - logger.trace("Sorting lines") - raw_lines = [] - sorted_lines = [] + logger.trace("Sorting lines") # type:ignore[attr-defined] + raw_lines: list[list[str]] = [] + sorted_lines: list[list[str]] = [] for key in sorted(keys): title = key.replace("_", " ").title() if key.startswith("raw"): @@ -213,7 +217,7 @@ def _lines_sort(self, keys: List[str]) -> List[List[Union[str, int, Tuple[float] return lines @staticmethod - def _lines_groupsize(raw_lines: List[str], sorted_lines: List[str]) -> int: + def _lines_groupsize(raw_lines: list[list[str]], sorted_lines: list[list[str]]) -> int: """ Get the number of items in each group. If raw data isn't selected, then check the length of remaining groups until something is @@ -238,12 +242,12 @@ def _lines_groupsize(raw_lines: List[str], sorted_lines: List[str]) -> int: keys = [key[0][:key[0].find("_")] for key in sorted_lines] distinct_keys = set(keys) groupsize = len(keys) // len(distinct_keys) - logger.trace(groupsize) + logger.trace(groupsize) # type:ignore[attr-defined] return groupsize def _lines_style(self, - lines: List[str], - groupsize: int) -> List[List[Union[str, int, Tuple[float]]]]: + lines: list[list[str]], + groupsize: int) -> list[list[str | int | tuple[float]]]: """ Obtain the color map and line width for each group. Parameters @@ -258,16 +262,17 @@ def _lines_style(self, list A list of loss keys with their corresponding line formatting and color information """ - logger.trace("Setting lines style") + logger.trace("Setting lines style") # type:ignore[attr-defined] groups = int(len(lines) / groupsize) colours = self._lines_create_colors(groupsize, groups) widths = list(range(1, groups + 1)) - for idx, item in enumerate(lines): + retval = T.cast(list[list[str | int | tuple[float]]], lines) + for idx, item in enumerate(retval): linewidth = widths[idx // groupsize] item.extend((linewidth, colours[idx])) - return lines + return retval - def _lines_create_colors(self, groupsize: int, groups: int) -> List[Tuple[float]]: + def _lines_create_colors(self, groupsize: int, groups: int) -> list[tuple[float]]: """ Create the color maps. Parameters @@ -288,7 +293,7 @@ def _lines_create_colors(self, groupsize: int, groups: int) -> List[Tuple[float] cmap = matplotlib.cm.get_cmap(colour) cpoint = 1 - (i / 5) colours.append(cmap(cpoint)) - logger.trace(colours) + logger.trace(colours) # type:ignore[attr-defined] return colours def _legend_place(self) -> None: @@ -316,7 +321,7 @@ def clear(self) -> None: del self._fig -class TrainingGraph(GraphBase): # pylint: disable=too-many-ancestors +class TrainingGraph(GraphBase): # pylint:disable=too-many-ancestors """ Live graph to be displayed during training. Parameters @@ -330,14 +335,16 @@ class TrainingGraph(GraphBase): # pylint: disable=too-many-ancestors """ def __init__(self, parent: ttk.Frame, data, ylabel: str) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent, data, ylabel) - self._thread = None # Thread for LongRunningTask - self._displayed_keys = [] + self._thread: LongRunningTask | None = None # Thread for LongRunningTask + self._displayed_keys: list[str] = [] self._add_callback() + logger.debug("Initialized %s", self.__class__.__name__) def _add_callback(self) -> None: """ Add the variable trace to update graph on refresh button press or save iteration. """ - get_config().tk_vars["refreshgraph"].trace("w", self.refresh) + get_config().tk_vars.refresh_graph.trace("w", self.refresh) # type:ignore def build(self) -> None: """ Build the Training graph. """ @@ -345,9 +352,9 @@ def build(self) -> None: self._plotcanvas.draw() logger.debug("Built training graph") - def refresh(self, *args) -> None: # pylint: disable=unused-argument + def refresh(self, *args) -> None: # pylint:disable=unused-argument """ Read the latest loss data and apply to current graph """ - refresh_var = get_config().tk_vars["refreshgraph"] + refresh_var = T.cast(tk.BooleanVar, get_config().tk_vars.refresh_graph) if not refresh_var.get() and self._thread is None: return @@ -399,15 +406,15 @@ def save_fig(self, location: str) -> None: def _resize_fig(self) -> None: """ Resize the figure to the current canvas size. """ - class Event(): # pylint: disable=too-few-public-methods + class Event(): # pylint:disable=too-few-public-methods """ Event class that needs to be passed to plotcanvas.resize """ - pass # pylint: disable=unnecessary-pass - Event.width = self.winfo_width() - Event.height = self.winfo_height() - self._plotcanvas.resize(Event) # pylint: disable=no-value-for-parameter + pass # pylint:disable=unnecessary-pass + setattr(Event, "width", self.winfo_width()) + setattr(Event, "height", self.winfo_height()) + self._plotcanvas.resize(Event) # pylint:disable=no-value-for-parameter -class SessionGraph(GraphBase): # pylint: disable=too-many-ancestors +class SessionGraph(GraphBase): # pylint:disable=too-many-ancestors """ Session Graph for session pop-up. Parameters @@ -422,8 +429,10 @@ class SessionGraph(GraphBase): # pylint: disable=too-many-ancestors Should be one of ``"log"`` or ``"linear"`` """ def __init__(self, parent: ttk.Frame, data, ylabel: str, scale: str) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent, data, ylabel) self._scale = scale + logger.debug("Initialized %s", self.__class__.__name__) def build(self) -> None: """ Build the session graph """ @@ -458,6 +467,7 @@ def set_yscale_type(self, scale: str) -> None: scale: str Should be one of ``"log"`` or ``"linear"`` """ + scale = scale.lower() logger.debug("Updating scale type: '%s'", scale) self._scale = scale self._update_plot(initiate=True) @@ -466,7 +476,7 @@ def set_yscale_type(self, scale: str) -> None: logger.debug("Updated scale type") -class NavigationToolbar(NavigationToolbar2Tk): # pylint: disable=too-many-ancestors +class NavigationToolbar(NavigationToolbar2Tk): # pylint:disable=too-many-ancestors """ Overrides the default Navigation Toolbar to provide only the buttons we require and to layout the items in a consistent manner with the rest of the GUI for the Analysis Session Graph pop up Window. @@ -483,12 +493,12 @@ class NavigationToolbar(NavigationToolbar2Tk): # pylint: disable=too-many-ances toolitems = [t for t in NavigationToolbar2Tk.toolitems if t[0] in ("Home", "Pan", "Zoom", "Save")] - def __init__(self, # pylint: disable=super-init-not-called + def __init__(self, # pylint:disable=super-init-not-called canvas: FigureCanvasTkAgg, - window: SessionGraph, + window: ttk.Frame, *, pack_toolbar: bool = True) -> None: - + logger.debug(parse_class_init(locals())) # Avoid using self.window (prefer self.canvas.get_tk_widget().master), # so that Tool implementations can reuse the methods. @@ -522,13 +532,14 @@ def __init__(self, # pylint: disable=super-init-not-called NavigationToolbar2.__init__(self, canvas) # pylint:disable=non-parent-init-called if pack_toolbar: self.pack(side=tk.BOTTOM, fill=tk.X) + logger.debug("Initialized %s", self.__class__.__name__) @staticmethod - def _Button(frame: ttk.Frame, # pylint:disable=arguments-differ + def _Button(frame: ttk.Frame, # pylint:disable=arguments-differ,arguments-renamed text: str, image_file: str, toggle: bool, - command) -> Union[ttk.Button, ttk.Checkbutton]: + command) -> ttk.Button | ttk.Checkbutton: """ Override the default button method to use our icons and ttk widgets for consistent GUI layout. @@ -558,7 +569,10 @@ def _Button(frame: ttk.Frame, # pylint:disable=arguments-differ img = get_images().icons[icon] if not toggle: - btn = ttk.Button(frame, text=text, image=img, command=command) + btn: ttk.Button | ttk.Checkbutton = ttk.Button(frame, + text=text, + image=img, + command=command) else: var = tk.IntVar(master=frame) btn = ttk.Checkbutton(frame, text=text, image=img, command=command, variable=var) diff --git a/lib/gui/display_page.py b/lib/gui/display_page.py index 32ebcb1bb2..5602c22f52 100644 --- a/lib/gui/display_page.py +++ b/lib/gui/display_page.py @@ -9,23 +9,21 @@ from .custom_widgets import Tooltip from .utils import get_images -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) _ = _LANG.gettext -class DisplayPage(ttk.Frame): # pylint: disable=too-many-ancestors +class DisplayPage(ttk.Frame): # pylint:disable=too-many-ancestors """ Parent frame holder for each tab. Defines uniform structure for each tab to inherit from """ def __init__(self, parent, tab_name, helptext): - logger.debug("Initializing %s: (tab_name: '%s', helptext: %s)", - self.__class__.__name__, tab_name, helptext) - ttk.Frame.__init__(self, parent) + super().__init__(parent) self._parent = parent - self.runningtask = parent.runningtask + self.running_task = parent.running_task self.helptext = helptext self.tabname = tab_name @@ -42,8 +40,6 @@ def __init__(self, parent, tab_name, helptext): self.pack(fill=tk.BOTH, side=tk.TOP, anchor=tk.NW) parent.add(self, text=self.tabname.title()) - logger.debug("Initialized %s", self.__class__.__name__,) - @property def _tab_is_active(self): """ bool: ``True`` if the tab currently has focus otherwise ``False`` """ @@ -56,12 +52,11 @@ def add_optional_vars(self, varsdict): logger.debug("Adding: (%s: %s)", key, val) self.vars[key] = val - @staticmethod - def set_vars(): + def set_vars(self): """ Override to return a dict of page specific variables """ return {} - def on_tab_select(self): # pylint:disable=no-self-use + def on_tab_select(self): """ Override for specific actions when the current tab is selected """ logger.debug("Returning as 'on_tab_select' not implemented for %s", self.__class__.__name__) @@ -164,13 +159,11 @@ def subnotebook_page_from_id(self, tab_id): return self.subnotebook.children[tab_name] -class DisplayOptionalPage(DisplayPage): # pylint: disable=too-many-ancestors +class DisplayOptionalPage(DisplayPage): # pylint:disable=too-many-ancestors """ Parent Context Sensitive Display Tab """ def __init__(self, parent, tab_name, helptext, wait_time, command=None): - logger.debug("%s: OptionalPage args: (wait_time: %s, command: %s)", - self.__class__.__name__, wait_time, command) - DisplayPage.__init__(self, parent, tab_name, helptext) + super().__init__(parent, tab_name, helptext) self._waittime = wait_time self.command = command @@ -183,8 +176,7 @@ def __init__(self, parent, tab_name, helptext, wait_time, command=None): self.update_idletasks() self._update_page() - @staticmethod - def set_vars(): + def set_vars(self): """ Analysis specific vars """ enabled = tk.BooleanVar() enabled.set(True) @@ -192,12 +184,8 @@ def set_vars(): ready = tk.BooleanVar() ready.set(False) - modified = tk.DoubleVar() - modified.set(None) - tk_vars = {"enabled": enabled, - "ready": ready, - "modified": modified} + "ready": ready} logger.debug(tk_vars) return tk_vars @@ -265,7 +253,7 @@ def on_chkenable_change(self): def _update_page(self): """ Update the latest preview item """ - if not self.runningtask.get() or not self._tab_is_active: + if not self.running_task.get() or not self._tab_is_active: return if self.vars["enabled"].get(): logger.trace("Updating page: %s", self.__class__.__name__) diff --git a/lib/gui/menu.py b/lib/gui/menu.py index f5f6aa16f1..460e08fd58 100644 --- a/lib/gui/menu.py +++ b/lib/gui/menu.py @@ -1,19 +1,17 @@ #!/usr/bin python3 """ The Menu Bars for faceswap GUI """ - +from __future__ import annotations import gettext -import locale import logging import os -import sys import tkinter as tk +import typing as T from tkinter import ttk import webbrowser -from subprocess import Popen, PIPE, STDOUT - +from lib.git import git from lib.multithreading import MultiThread -from lib.serializer import get_serializer +from lib.serializer import get_serializer, Serializer from lib.utils import FaceswapError import update_deps @@ -21,23 +19,31 @@ from .custom_widgets import Tooltip from .utils import get_config, get_images -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from scripts.gui import FaceswapGui + +logger = logging.getLogger(__name__) # LOCALES -_LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) +_LANG = gettext.translation("gui.menu", localedir="locales", fallback=True) _ = _LANG.gettext -_WORKING_DIR = os.path.dirname(os.path.realpath(sys.argv[0])) - -_RESOURCES = [(_("faceswap.dev - Guides and Forum"), "https://www.faceswap.dev"), - (_("Patreon - Support this project"), "https://www.patreon.com/faceswap"), - (_("Discord - The FaceSwap Discord server"), "https://discord.gg/VasFUAy"), - (_("Github - Our Source Code"), "https://github.com/deepfakes/faceswap")] +_RESOURCES: list[tuple[str, str]] = [ + (_("faceswap.dev - Guides and Forum"), "https://www.faceswap.dev"), + (_("Patreon - Support this project"), "https://www.patreon.com/faceswap"), + (_("Discord - The FaceSwap Discord server"), "https://discord.gg/VasFUAy"), + (_("Github - Our Source Code"), "https://github.com/deepfakes/faceswap")] class MainMenuBar(tk.Menu): # pylint:disable=too-many-ancestors - """ GUI Main Menu Bar """ - def __init__(self, master=None): + """ GUI Main Menu Bar + + Parameters + ---------- + master: :class:`tkinter.Tk` + The root tkinter object + """ + def __init__(self, master: FaceswapGui) -> None: logger.debug("Initializing %s", self.__class__.__name__) super().__init__(master) self.root = master @@ -46,99 +52,130 @@ def __init__(self, master=None): self.settings_menu = SettingsMenu(self) self.help_menu = HelpMenu(self) - self.add_cascade(label="File", menu=self.file_menu, underline=0) - self.add_cascade(label="Settings", menu=self.settings_menu, underline=0) - self.add_cascade(label="Help", menu=self.help_menu, underline=0) + self.add_cascade(label=_("File"), menu=self.file_menu, underline=0) + self.add_cascade(label=_("Settings"), menu=self.settings_menu, underline=0) + self.add_cascade(label=_("Help"), menu=self.help_menu, underline=0) logger.debug("Initialized %s", self.__class__.__name__) class SettingsMenu(tk.Menu): # pylint:disable=too-many-ancestors - """ Settings menu items and functions """ - def __init__(self, parent): + """ Settings menu items and functions + + Parameters + ---------- + parent: :class:`tkinter.Menu` + The main menu bar to hold this menu item + """ + def __init__(self, parent: MainMenuBar) -> None: logger.debug("Initializing %s", self.__class__.__name__) super().__init__(parent, tearoff=0) self.root = parent.root - self.build() + self._build() logger.debug("Initialized %s", self.__class__.__name__) - def build(self): + def _build(self) -> None: """ Add the settings menu to the menu bar """ - # pylint: disable=cell-var-from-loop + # pylint:disable=cell-var-from-loop logger.debug("Building settings menu") - self.add_command(label="Configure Settings...", + self.add_command(label=_("Configure Settings..."), underline=0, command=open_popup) logger.debug("Built settings menu") class FileMenu(tk.Menu): # pylint:disable=too-many-ancestors - """ File menu items and functions """ - def __init__(self, parent): + """ File menu items and functions + + Parameters + ---------- + parent: :class:`tkinter.Menu` + The main menu bar to hold this menu item + """ + def __init__(self, parent: MainMenuBar) -> None: logger.debug("Initializing %s", self.__class__.__name__) super().__init__(parent, tearoff=0) self.root = parent.root self._config = get_config() - self.recent_menu = tk.Menu(self, tearoff=0, postcommand=self.refresh_recent_menu) - self.build() + self.recent_menu = tk.Menu(self, tearoff=0, postcommand=self._refresh_recent_menu) + self._build() logger.debug("Initialized %s", self.__class__.__name__) - def build(self): + def _refresh_recent_menu(self) -> None: + """ Refresh recent menu on save/load of files """ + self.recent_menu.delete(0, "end") + self._build_recent_menu() + + def _build(self) -> None: """ Add the file menu to the menu bar """ logger.debug("Building File menu") - self.add_command(label="New Project...", + self.add_command(label=_("New Project..."), underline=0, accelerator="Ctrl+N", command=self._config.project.new) self.root.bind_all("", self._config.project.new) - self.add_command(label="Open Project...", + self.add_command(label=_("Open Project..."), underline=0, accelerator="Ctrl+O", command=self._config.project.load) self.root.bind_all("", self._config.project.load) - self.add_command(label="Save Project", + self.add_command(label=_("Save Project"), underline=0, accelerator="Ctrl+S", command=lambda: self._config.project.save(save_as=False)) self.root.bind_all("", lambda e: self._config.project.save(e, save_as=False)) - self.add_command(label="Save Project as...", + self.add_command(label=_("Save Project as..."), underline=13, accelerator="Ctrl+Alt+S", command=lambda: self._config.project.save(save_as=True)) self.root.bind_all("", lambda e: self._config.project.save(e, save_as=True)) - self.add_command(label="Reload Project from Disk", + self.add_command(label=_("Reload Project from Disk"), underline=0, accelerator="F5", command=self._config.project.reload) self.root.bind_all("", self._config.project.reload) - self.add_command(label="Close Project", + self.add_command(label=_("Close Project"), underline=0, accelerator="Ctrl+W", command=self._config.project.close) self.root.bind_all("", self._config.project.close) self.add_separator() - self.add_command(label="Open Task...", + self.add_command(label=_("Open Task..."), underline=5, accelerator="Ctrl+Alt+T", command=lambda: self._config.tasks.load(current_tab=False)) self.root.bind_all("", lambda e: self._config.tasks.load(e, current_tab=False)) self.add_separator() - self.add_cascade(label="Open recent", underline=6, menu=self.recent_menu) + self.add_cascade(label=_("Open recent"), underline=6, menu=self.recent_menu) self.add_separator() - self.add_command(label="Quit", + self.add_command(label=_("Quit"), underline=0, accelerator="Alt+F4", command=self.root.close_app) self.root.bind_all("", self.root.close_app) logger.debug("Built File menu") - def build_recent_menu(self): + @classmethod + def _clear_recent_files(cls, serializer: Serializer, menu_file: str) -> None: + """ Creates or clears recent file list + + Parameters + ---------- + serializer: :class:`~lib.serializer.Serializer` + The serializer to use for storing files + menu_file: str + The file name holding the recent files + """ + logger.debug("clearing recent files list: '%s'", menu_file) + serializer.save(menu_file, []) + + def _build_recent_menu(self) -> None: """ Load recent files into menu bar """ logger.debug("Building Recent Files menu") serializer = get_serializer("json") menu_file = os.path.join(self._config.pathcache, ".recent.json") if not os.path.isfile(menu_file) or os.path.getsize(menu_file) == 0: - self.clear_recent_files(serializer, menu_file) + self._clear_recent_files(serializer, menu_file) try: recent_files = serializer.load(menu_file) except FaceswapError as err: @@ -146,7 +183,7 @@ def build_recent_menu(self): # Some reports of corruption breaking menus logger.warning("There was an error opening the recent files list so it has been " "reset.") - self.clear_recent_files(serializer, menu_file) + self._clear_recent_files(serializer, menu_file) recent_files = [] logger.debug("Loaded recent files: %s", recent_files) @@ -163,14 +200,14 @@ def build_recent_menu(self): if command.lower() == "project": load_func = self._config.project.load lbl = command - kwargs = dict(filename=filename) + kwargs = {"filename": filename} else: - load_func = self._config.tasks.load - lbl = "{} Task".format(command) - kwargs = dict(filename=filename, current_tab=False) + load_func = self._config.tasks.load # type:ignore + lbl = _("{} Task").format(command) + kwargs = {"filename": filename, "current_tab": False} self.recent_menu.add_command( - label="{} ({})".format(filename, lbl.title()), - command=lambda kw=kwargs, fn=load_func: fn(**kw)) + label=f"{filename} ({lbl.title()})", + command=lambda kw=kwargs, fn=load_func: fn(**kw)) # type:ignore if removed_files: for recent_item in removed_files: logger.debug("Removing from recent files: `%s`", recent_item[0]) @@ -178,57 +215,179 @@ def build_recent_menu(self): serializer.save(menu_file, recent_files) self.recent_menu.add_separator() self.recent_menu.add_command( - label="Clear recent files", + label=_("Clear recent files"), underline=0, - command=lambda srl=serializer, mnu=menu_file: self.clear_recent_files(srl, mnu)) + command=lambda srl=serializer, mnu=menu_file: self._clear_recent_files( # type:ignore + srl, mnu)) logger.debug("Built Recent Files menu") - @staticmethod - def clear_recent_files(serializer, menu_file): - """ Creates or clears recent file list """ - logger.debug("clearing recent files list: '%s'", menu_file) - serializer.save(menu_file, list()) - - def refresh_recent_menu(self): - """ Refresh recent menu on save/load of files """ - self.recent_menu.delete(0, "end") - self.build_recent_menu() - class HelpMenu(tk.Menu): # pylint:disable=too-many-ancestors - """ Help menu items and functions """ - def __init__(self, parent): + """ Help menu items and functions + + Parameters + ---------- + parent: :class:`tkinter.Menu` + The main menu bar to hold this menu item + """ + def __init__(self, parent: MainMenuBar) -> None: logger.debug("Initializing %s", self.__class__.__name__) super().__init__(parent, tearoff=0) self.root = parent.root self.recources_menu = tk.Menu(self, tearoff=0) self._branches_menu = tk.Menu(self, tearoff=0) - self.build() + self._build() logger.debug("Initialized %s", self.__class__.__name__) - def build(self): + def _in_thread(self, action: str): + """ Perform selected action inside a thread + + Parameters + ---------- + action: str + The action to be performed. The action corresponds to the function name to be called + """ + logger.debug("Performing help action: %s", action) + thread = MultiThread(getattr(self, action), thread_count=1) + thread.start() + logger.debug("Performed help action: %s", action) + + def _output_sysinfo(self): + """ Output system information to console """ + logger.debug("Obtaining system information") + self.root.config(cursor="watch") + self._clear_console() + try: + from lib.sysinfo import sysinfo # pylint:disable=import-outside-toplevel + info = sysinfo + except Exception as err: # pylint:disable=broad-except + info = f"Error obtaining system info: {str(err)}" + self._clear_console() + logger.debug("Obtained system information: %s", info) + print(info) + self.root.config(cursor="") + + @classmethod + def _process_status_output(cls, status: list[str]) -> bool: + """ Process the output of a git status call and output information + + Parameters + ---------- + status : list[str] + The lines returned from a git status call + + Returns + ------- + bool + ``True`` if the repo can be updated otherwise ``False`` + """ + for line in status: + if line.lower().startswith("your branch is ahead"): + logger.warning("Your branch is ahead of the remote repo. Not updating") + return False + if line.lower().startswith("your branch is up to date"): + logger.info("Faceswap is up to date.") + return False + if "have diverged" in line.lower(): + logger.warning("Your branch has diverged from the remote repo. Not updating") + return False + if line.lower().startswith("your branch is behind"): + return True + + logger.warning("Unable to retrieve status of branch") + return False + + def _check_for_updates(self, check: bool = False) -> bool: + """ Check whether an update is required + + Parameters + ---------- + check: bool + ``True`` if we are just checking for updates ``False`` if a check and update is to be + performed. Default: ``False`` + + Returns + ------- + bool + ``True`` if an update is required + """ + # Do the check + logger.info("Checking for updates...") + msg = ("Git is not installed or you are not running a cloned repo. " + "Unable to check for updates") + + sync = git.update_remote() + if not sync: + logger.warning(msg) + return False + + status = git.status + if not status: + logger.warning(msg) + return False + + retval = self._process_status_output(status) + if retval and check: + logger.info("There are updates available") + return retval + + def _check(self) -> None: + """ Check for updates and clone repository """ + logger.debug("Checking for updates...") + self.root.config(cursor="watch") + self._check_for_updates(check=True) + self.root.config(cursor="") + + def _do_update(self) -> bool: + """ Update Faceswap + + Returns + ------- + bool + ``True`` if update was successful + """ + logger.info("A new version is available. Updating...") + success = git.pull() + if not success: + logger.info("An error occurred during update") + return success + + def _update(self) -> None: + """ Check for updates and clone repository """ + logger.debug("Updating Faceswap...") + self.root.config(cursor="watch") + success = False + if self._check_for_updates(): + success = self._do_update() + update_deps.main(is_gui=True) + if success: + logger.info("Please restart Faceswap to complete the update.") + self.root.config(cursor="") + + def _build(self) -> None: """ Build the help menu """ logger.debug("Building Help menu") - self.add_command(label="Check for updates...", + self.add_command(label=_("Check for updates..."), underline=0, - command=lambda action="check": self.in_thread(action)) - self.add_command(label="Update Faceswap...", + command=lambda action="_check": self._in_thread(action)) # type:ignore + self.add_command(label=_("Update Faceswap..."), underline=0, - command=lambda action="update": self.in_thread(action)) + command=lambda action="_update": self._in_thread(action)) # type:ignore if self._build_branches_menu(): - self.add_cascade(label="Switch Branch", underline=7, menu=self._branches_menu) + self.add_cascade(label=_("Switch Branch"), underline=7, menu=self._branches_menu) self.add_separator() self._build_recources_menu() - self.add_cascade(label="Resources", underline=0, menu=self.recources_menu) + self.add_cascade(label=_("Resources"), underline=0, menu=self.recources_menu) self.add_separator() - self.add_command(label="Output System Information", - underline=0, - command=lambda action="output_sysinfo": self.in_thread(action)) + self.add_command( + label=_("Output System Information"), + underline=0, + command=lambda action="_output_sysinfo": self._in_thread(action)) # type:ignore logger.debug("Built help menu") - def _build_branches_menu(self): + def _build_branches_menu(self) -> bool: """ Build branch selection menu. Queries git for available branches and builds a menu based on output. @@ -238,75 +397,56 @@ def _build_branches_menu(self): bool ``True`` if menu was successfully built otherwise ``False`` """ - stdout = self._get_branches() - if stdout is None: + branches = git.branches + if not branches: return False - branches = self._filter_branches(stdout) + branches = self._filter_branches(branches) if not branches: return False for branch in branches: self._branches_menu.add_command( label=branch, - command=lambda b=branch: self._switch_branch(b)) + command=lambda b=branch: self._switch_branch(b)) # type:ignore return True - @staticmethod - def _get_branches(): - """ Get the available github branches - - Returns - ------- - str - The list of branches available. If no branches were found or there was an - error then `None` is returned - """ - gitcmd = "git branch -a" - cmd = Popen(gitcmd, shell=True, stdout=PIPE, stderr=STDOUT, cwd=_WORKING_DIR) - stdout, _ = cmd.communicate() - retcode = cmd.poll() - if retcode != 0: - logger.debug("Unable to list git branches. return code: %s, message: %s", - retcode, stdout.decode().strip().replace("\n", " - ")) - return None - return stdout.decode(locale.getpreferredencoding()) - - @staticmethod - def _filter_branches(stdout): - """ Filter the branches, remove duplicates and the current branch and return a sorted - list. + @classmethod + def _filter_branches(cls, branches: list[str]) -> list[str]: + """ Filter the branches, remove any non-local branches Parameters ---------- - stdout: str - The output from the git branch query converted to a string + branches: list[str] + list of available git branches Returns ------- - list + list[str] Unique list of available branches sorted in alphabetical order """ current = None - branches = set() - for line in stdout.splitlines(): - branch = line[line.rfind("/") + 1:] if "/" in line else line.strip() + unique = set() + for line in branches: + branch = line.strip() + if branch.startswith("remotes"): + continue if branch.startswith("*"): branch = branch.replace("*", "").strip() current = branch continue - branches.add(branch) - logger.debug("Found branches: %s", branches) - if current in branches: + unique.add(branch) + logger.debug("Found branches: %s", unique) + if current in unique: logger.debug("Removing current branch from output: %s", current) - branches.remove(current) + unique.remove(current) - branches = sorted(list(branches), key=str.casefold) - logger.debug("Final branches: %s", branches) - return branches + retval = sorted(list(unique), key=str.casefold) + logger.debug("Final branches: %s", retval) + return retval - @staticmethod - def _switch_branch(branch): + @classmethod + def _switch_branch(cls, branch: str) -> None: """ Change the currently checked out branch, and return a notification. Parameters @@ -315,139 +455,38 @@ def _switch_branch(branch): The branch to switch to """ logger.info("Switching branch to '%s'...", branch) - gitcmd = "git checkout {}".format(branch) - cmd = Popen(gitcmd, shell=True, stdout=PIPE, stderr=STDOUT, cwd=_WORKING_DIR) - stdout, _ = cmd.communicate() - retcode = cmd.poll() - if retcode != 0: - logger.error("Unable to switch branch. return code: %s, message: %s", - retcode, stdout.decode().strip().replace("\n", " - ")) + if not git.checkout(branch): + logger.error("Unable to switch branch to '%s'", branch) return logger.info("Succesfully switched to '%s'. You may want to check for updates to make sure " "that you have the latest code.", branch) logger.info("Please restart Faceswap to complete the switch.") - def _build_recources_menu(self): + def _build_recources_menu(self) -> None: """ Build resources menu """ - # pylint: disable=cell-var-from-loop + # pylint:disable=cell-var-from-loop logger.debug("Building Resources Files menu") for resource in _RESOURCES: self.recources_menu.add_command( label=resource[0], - command=lambda link=resource[1]: webbrowser.open_new(link)) + command=lambda link=resource[1]: webbrowser.open_new(link)) # type:ignore logger.debug("Built resources menu") - def in_thread(self, action): - """ Perform selected action inside a thread """ - logger.debug("Performing help action: %s", action) - thread = MultiThread(getattr(self, action), thread_count=1) - thread.start() - logger.debug("Performed help action: %s", action) - - @staticmethod - def clear_console(): + @classmethod + def _clear_console(cls) -> None: """ Clear the console window """ - get_config().tk_vars["console_clear"].set(True) + get_config().tk_vars.console_clear.set(True) - def output_sysinfo(self): - """ Output system information to console """ - logger.debug("Obtaining system information") - self.root.config(cursor="watch") - self.clear_console() - try: - from lib.sysinfo import sysinfo # pylint:disable=import-outside-toplevel - info = sysinfo - except Exception as err: # pylint:disable=broad-except - info = "Error obtaining system info: {}".format(str(err)) - self.clear_console() - logger.debug("Obtained system information: %s", info) - print(info) - self.root.config(cursor="") - def check(self): - """ Check for updates and clone repository """ - logger.debug("Checking for updates...") - self.root.config(cursor="watch") - encoding = locale.getpreferredencoding() - logger.debug("Encoding: %s", encoding) - self.check_for_updates(encoding, check=True) - self.root.config(cursor="") +class TaskBar(ttk.Frame): # pylint:disable=too-many-ancestors + """ Task bar buttons - def update(self): - """ Check for updates and clone repository """ - logger.debug("Updating Faceswap...") - self.root.config(cursor="watch") - encoding = locale.getpreferredencoding() - logger.debug("Encoding: %s", encoding) - success = False - if self.check_for_updates(encoding): - success = self.do_update(encoding) - update_deps.main(logger=logger) - if success: - logger.info("Please restart Faceswap to complete the update.") - self.root.config(cursor="") - - @staticmethod - def check_for_updates(encoding, check=False): - """ Check whether an update is required """ - # Do the check - logger.info("Checking for updates...") - update = False - msg = "" - gitcmd = "git remote update && git status -uno" - cmd = Popen(gitcmd, shell=True, stdout=PIPE, stderr=STDOUT, cwd=_WORKING_DIR) - stdout, _ = cmd.communicate() - retcode = cmd.poll() - if retcode != 0: - msg = ("Git is not installed or you are not running a cloned repo. " - "Unable to check for updates") - else: - chk = stdout.decode(encoding).splitlines() - for line in chk: - if line.lower().startswith("your branch is ahead"): - msg = "Your branch is ahead of the remote repo. Not updating" - break - if line.lower().startswith("your branch is up to date"): - msg = "Faceswap is up to date." - break - if line.lower().startswith("your branch is behind"): - msg = "There are updates available" - update = True - break - if "have diverged" in line.lower(): - msg = "Your branch has diverged from the remote repo. Not updating" - break - if not update or check: - logger.info(msg) - logger.debug("Checked for update. Update required: %s", update) - return update - - @staticmethod - def do_update(encoding): - """ Update Faceswap """ - logger.info("A new version is available. Updating...") - gitcmd = "git pull" - cmd = Popen(gitcmd, shell=True, stdout=PIPE, stderr=STDOUT, bufsize=1, cwd=_WORKING_DIR) - while True: - output = cmd.stdout.readline().decode(encoding) - if output == "" and cmd.poll() is not None: - break - if output: - logger.debug("'%s' output: '%s'", gitcmd, output.strip()) - print(output.strip()) - retcode = cmd.poll() - logger.debug("'%s' returncode: %s", gitcmd, retcode) - if retcode != 0: - logger.info("An error occurred during update. return code: %s", retcode) - retval = False - else: - retval = True - return retval - - -class TaskBar(ttk.Frame): # pylint: disable=too-many-ancestors - """ Task bar buttons """ - def __init__(self, parent): + Parameters + ---------- + parent: :class:`tkinter.ttk.Frame` + The frame that holds the task bar + """ + def __init__(self, parent: ttk.Frame) -> None: super().__init__(parent) self._config = get_config() self.pack(side=tk.TOP, anchor=tk.W, fill=tk.X, expand=False) @@ -461,7 +500,65 @@ def __init__(self, parent): self._settings_btns() self._section_separator() - def _project_btns(self): + @classmethod + def _loader_and_kwargs(cls, btntype: str) -> tuple[str, dict[str, bool]]: + """ Get the loader name and key word arguments for the given button type + + Parameters + ---------- + btntype: str + The button type to obtain the information for + + Returns + ------- + loader: str + The name of the loader to use for the given button type + kwargs: dict[str, bool] + The keyword arguments to use for the returned loader + """ + if btntype == "save": + loader = btntype + kwargs = {"save_as": False} + elif btntype == "save_as": + loader = "save" + kwargs = {"save_as": True} + else: + loader = btntype + kwargs = {} + logger.debug("btntype: %s, loader: %s, kwargs: %s", btntype, loader, kwargs) + return loader, kwargs + + @classmethod + def _set_help(cls, btntype: str) -> str: + """ Set the helptext for option buttons + + Parameters + ---------- + btntype: str + The button type to set the help text for + """ + logger.debug("Setting help") + hlp = "" + task = _("currently selected Task") if btntype[-1] == "2" else _("Project") + if btntype.startswith("reload"): + hlp = _("Reload {} from disk").format(task) + if btntype == "new": + hlp = _("Create a new {}...").format(task) + if btntype.startswith("clear"): + hlp = _("Reset {} to default").format(task) + elif btntype.startswith("save") and "_" not in btntype: + hlp = _("Save {}").format(task) + elif btntype.startswith("save_as"): + hlp = _("Save {} as...").format(task) + elif btntype.startswith("load"): + msg = task + if msg.endswith("Task"): + msg += _(" from a task or project file") + hlp = _("Load {}...").format(msg) + return hlp + + def _project_btns(self) -> None: + """ Place the project buttons """ frame = ttk.Frame(self._btn_frame) frame.pack(side=tk.LEFT, anchor=tk.W, expand=False, padx=2) @@ -472,17 +569,18 @@ def _project_btns(self): cmd = getattr(self._config.project, loader) btn = ttk.Button(frame, image=get_images().icons[btntype], - command=lambda fn=cmd, kw=kwargs: fn(**kw)) + command=lambda fn=cmd, kw=kwargs: fn(**kw)) # type:ignore btn.pack(side=tk.LEFT, anchor=tk.W) - hlp = self.set_help(btntype) + hlp = self._set_help(btntype) Tooltip(btn, text=hlp, wrap_length=200) - def _task_btns(self): + def _task_btns(self) -> None: + """ Place the task buttons """ frame = ttk.Frame(self._btn_frame) frame.pack(side=tk.LEFT, anchor=tk.W, expand=False, padx=2) for loadtype in ("load", "save", "save_as", "clear", "reload"): - btntype = "{}2".format(loadtype) + btntype = f"{loadtype}2" logger.debug("Adding button: '%s'", btntype) loader, kwargs = self._loader_and_kwargs(loadtype) @@ -492,69 +590,35 @@ def _task_btns(self): btn = ttk.Button( frame, image=get_images().icons[btntype], - command=lambda fn=cmd, kw=kwargs: fn(**kw)) + command=lambda fn=cmd, kw=kwargs: fn(**kw)) # type:ignore btn.pack(side=tk.LEFT, anchor=tk.W) - hlp = self.set_help(btntype) + hlp = self._set_help(btntype) Tooltip(btn, text=hlp, wrap_length=200) - @staticmethod - def _loader_and_kwargs(btntype): - if btntype == "save": - loader = btntype - kwargs = dict(save_as=False) - elif btntype == "save_as": - loader = "save" - kwargs = dict(save_as=True) - else: - loader = btntype - kwargs = dict() - logger.debug("btntype: %s, loader: %s, kwargs: %s", btntype, loader, kwargs) - return loader, kwargs - - def _settings_btns(self): - # pylint: disable=cell-var-from-loop + def _settings_btns(self) -> None: + """ Place the settings buttons """ + # pylint:disable=cell-var-from-loop frame = ttk.Frame(self._btn_frame) frame.pack(side=tk.LEFT, anchor=tk.W, expand=False, padx=2) for name in ("extract", "train", "convert"): - btntype = "settings_{}".format(name) + btntype = f"settings_{name}" btntype = btntype if btntype in get_images().icons else "settings" logger.debug("Adding button: '%s'", btntype) btn = ttk.Button( frame, image=get_images().icons[btntype], - command=lambda n=name: open_popup(name=n)) + command=lambda n=name: open_popup(name=n)) # type:ignore btn.pack(side=tk.LEFT, anchor=tk.W) hlp = _("Configure {} settings...").format(name.title()) Tooltip(btn, text=hlp, wrap_length=200) - @staticmethod - def set_help(btntype): - """ Set the helptext for option buttons """ - logger.debug("Setting help") - hlp = "" - task = _("currently selected Task") if btntype[-1] == "2" else _("Project") - if btntype.startswith("reload"): - hlp = _("Reload {} from disk").format(task) - if btntype == "new": - hlp = _("Create a new {}...").format(task) - if btntype.startswith("clear"): - hlp = _("Reset {} to default").format(task) - elif btntype.startswith("save") and "_" not in btntype: - hlp = _("Save {}").format(task) - elif btntype.startswith("save_as"): - hlp = _("Save {} as...").format(task) - elif btntype.startswith("load"): - msg = task - if msg.endswith("Task"): - msg += _(" from a task or project file") - hlp = _("Load {}...").format(msg) - return hlp - - def _group_separator(self): + def _group_separator(self) -> None: + """ Place a group separator """ separator = ttk.Separator(self._btn_frame, orient="vertical") separator.pack(padx=(2, 1), fill=tk.Y, side=tk.LEFT) - def _section_separator(self): + def _section_separator(self) -> None: + """ Place a section separator """ frame = ttk.Frame(self) frame.pack(side=tk.BOTTOM, fill=tk.X) separator = ttk.Separator(frame, orient="horizontal") diff --git a/lib/gui/options.py b/lib/gui/options.py index a207c43bcc..2579d60233 100644 --- a/lib/gui/options.py +++ b/lib/gui/options.py @@ -1,185 +1,392 @@ #!/usr/bin python3 """ Cli Options for the GUI """ +from __future__ import annotations + import inspect from argparse import SUPPRESS +from dataclasses import dataclass from importlib import import_module import logging import os import re import sys -from collections import OrderedDict +import typing as T -from lib.cli import actions, args as cli +from lib.cli import actions from .utils import get_images from .control_helper import ControlPanelOption -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from tkinter import Variable + from types import ModuleType + from lib.cli.args import FaceSwapArgs + +logger = logging.getLogger(__name__) + + +@dataclass +class CliOption: + """ A parsed command line option + + Parameters + ---------- + cpanel_option: :class:`~lib.gui.control_helper.ControlPanelOption`: + Object to hold information of a command line item for displaying in a GUI + :class:`~lib.gui.control_helper.ControlPanel` + opts: tuple[str, ...]: + The short switch and long name (if exists) of the command line option + nargs: Literal["+"] | None: + ``None`` for not used. "+" for at least 1 argument required with values to be contained + in a list + """ + cpanel_option: ControlPanelOption + """:class:`~lib.gui.control_helper.ControlPanelOption`: Object to hold information of a command + line item for displaying in a GUI :class:`~lib.gui.control_helper.ControlPanel`""" + opts: tuple[str, ...] + """tuple[str, ...]: The short switch and long name (if exists) of cli option """ + nargs: T.Literal["+"] | None + """Literal["+"] | None: ``None`` for not used. "+" for at least 1 argument required with + values to be contained in a list """ class CliOptions(): """ Class and methods for the command line options """ - def __init__(self): + def __init__(self) -> None: logger.debug("Initializing %s", self.__class__.__name__) - self.categories = ("faceswap", "tools") - self.commands = dict() - self.opts = dict() - self.build_options() + self._base_path = os.path.realpath(os.path.dirname(sys.argv[0])) + self._commands: dict[T.Literal["faceswap", "tools"], list[str]] = {"faceswap": [], + "tools": []} + self._opts: dict[str, dict[str, CliOption | str]] = {} + self._build_options() logger.debug("Initialized %s", self.__class__.__name__) - def build_options(self): - """ Get the commands that belong to each category """ - for category in self.categories: - logger.debug("Building '%s'", category) - if category == "tools": - mod_classes = self._get_tools_cli_classes() - self.commands[category] = self.sort_commands(category, mod_classes) - for tool in sorted(mod_classes): - self.opts.update(self.extract_options(mod_classes[tool], [tool])) - else: - mod_classes = self.get_cli_classes(cli) - self.commands[category] = self.sort_commands(category, mod_classes) - self.opts.update(self.extract_options(cli, mod_classes)) - logger.debug("Built '%s'", category) + @property + def categories(self) -> tuple[T.Literal["faceswap", "tools"], ...]: + """tuple[str, str] The categories for faceswap's GUI """ + return tuple(self._commands) - @staticmethod - def get_cli_classes(cli_source): - """ Parse the cli scripts for the argument classes """ - mod_classes = [] - for name, obj in inspect.getmembers(cli_source): - if inspect.isclass(obj) and name.lower().endswith("args") \ - and name.lower() not in (("faceswapargs", - "extractconvertargs", - "guiargs")): - mod_classes.append(name) - logger.debug(mod_classes) - return mod_classes - - @staticmethod - def _get_tools_cli_classes(): - """ Parse the tools cli scripts for the argument classes """ - base_path = os.path.realpath(os.path.dirname(sys.argv[0])) - tools_dir = os.path.join(base_path, "tools") - mod_classes = dict() + @property + def commands(self) -> dict[T.Literal["faceswap", "tools"], list[str]]: + """dict[str, ]""" + return self._commands + + @property + def opts(self) -> dict[str, dict[str, CliOption | str]]: + """dict[str, dict[str, CliOption | str]] The command line options collected from faceswap's + cli files """ + return self._opts + + def _get_modules_tools(self) -> list[ModuleType]: + """ Parse the tools cli python files for the modules that contain the command line + arguments + + Returns + ------- + list[`types.ModuleType`] + The modules for each faceswap tool that exists in the project + """ + tools_dir = os.path.join(self._base_path, "tools") + logger.debug("Scanning '%s' for cli files", tools_dir) + retval: list[ModuleType] = [] for tool_name in sorted(os.listdir(tools_dir)): cli_file = os.path.join(tools_dir, tool_name, "cli.py") - if os.path.exists(cli_file): - mod = ".".join(("tools", tool_name, "cli")) - mod_classes["{}Args".format(tool_name.title())] = import_module(mod) - return mod_classes - - def sort_commands(self, category, classes): - """ Format classes into command names and sort: - Specific workflow order for faceswap. - Alphabetical for all others """ - commands = sorted(self.format_command_name(command) - for command in classes) + if not os.path.exists(cli_file): + logger.debug("File does not exist. Skipping: '%s'", cli_file) + continue + + mod = ".".join(("tools", tool_name, "cli")) + retval.append(import_module(mod)) + logger.debug("Collected: %s", retval[-1]) + return retval + + def _get_modules_faceswap(self) -> list[ModuleType]: + """ Parse the faceswap cli python files for the modules that contain the command line + arguments + + Returns + ------- + list[`types.ModuleType`] + The modules for each faceswap command line argument file that exists in the project + """ + base_dir = ["lib", "cli"] + cli_dir = os.path.join(self._base_path, *base_dir) + logger.debug("Scanning '%s' for cli files", cli_dir) + retval: list[ModuleType] = [] + + for fname in os.listdir(cli_dir): + if not fname.startswith("args"): + logger.debug("Skipping file '%s'", fname) + continue + mod = ".".join((*base_dir, os.path.splitext(fname)[0])) + retval.append(import_module(mod)) + logger.debug("Collected: '%s", retval[-1]) + return retval + + def _get_modules(self, category: T.Literal["faceswap", "tools"]) -> list[ModuleType]: + """ Parse the cli files for faceswap and tools and return the imported module + + Parameters + ---------- + category: Literal["faceswap", "tools"] + The faceswap category to obtain the cli modules + + Returns + ------- + list[`types.ModuleType`] + The modules for each faceswap command/tool that exists in the project for the given + category + """ + logger.debug("Getting '%s' cli modules", category) + if category == "tools": + return self._get_modules_tools() + return self._get_modules_faceswap() + + @classmethod + def _get_classes(cls, module: ModuleType) -> list[T.Type[FaceSwapArgs]]: + """ Obtain the classes from the given module that contain the command line + arguments + + Parameters + ---------- + module: :class:`types.ModuleType` + The imported module to parse for command line argument classes + + Returns + ------- + list[:class:`~lib.cli.args.FaceswapArgs`] + The command line argument class objects that exist in the module + """ + retval = [] + for name, obj in inspect.getmembers(module): + if not inspect.isclass(obj) or not name.lower().endswith("args"): + logger.debug("Skipping non-cli class object '%s'", name) + continue + if name.lower() in (("faceswapargs", "extractconvertargs", "guiargs")): + logger.debug("Skipping uneeded object '%s'", name) + continue + logger.debug("Collecting %s", obj) + retval.append(obj) + logger.debug("Collected from '%s': %s", module.__name__, [c.__name__ for c in retval]) + return retval + + def _get_all_classes(self, modules: list[ModuleType]) -> list[T.Type[FaceSwapArgs]]: + """Obtain the the command line options classes for the given modules + + Parameters + ---------- + modules : list[:class:`types.ModuleType`] + The imported modules to extract the command line argument classes from + + Returns + ------- + list[:class:`~lib.cli.args.FaceSwapArgs`] + The valid command line class objects for the given modules + """ + retval = [] + for module in modules: + mod_classes = self._get_classes(module) + if not mod_classes: + logger.debug("module '%s' contains no cli classes. Skipping", module) + continue + retval.extend(mod_classes) + logger.debug("Obtained %s cli classes from %s modules", len(retval), len(modules)) + return retval + + @classmethod + def _class_name_to_command(cls, class_name: str) -> str: + """ Format a FaceSwapArgs class name to a standardized command name + + Parameters + ---------- + class_name: str + The name of the class to convert to a command name + + Returns + ------- + str + The formatted command name + """ + return class_name.lower()[:-4] + + def _store_commands(self, + category: T.Literal["faceswap", "tools"], + classes: list[T.Type[FaceSwapArgs]]) -> None: + """ Format classes into command names and sort. Store in :attr:`commands`. + Sorting is in specific workflow order for faceswap and alphabetical for all others + + Parameters + ---------- + category: Literal["faceswap", "tools"] + The category to store the command names for + classes: list[:class:`~lib.cli.args.FaceSwapArgs`] + The valid command line class objects for the category + """ + class_names = [c.__name__ for c in classes] + commands = sorted(self._class_name_to_command(n) for n in class_names) + if category == "faceswap": ordered = ["extract", "train", "convert"] commands = ordered + [command for command in commands if command not in ordered] - logger.debug(commands) - return commands - - @staticmethod - def format_command_name(classname): - """ Format args class name to command """ - return classname.lower()[:-4] - - def extract_options(self, cli_source, mod_classes): - """ Extract the existing ArgParse Options - into master options Dictionary """ - subopts = dict() - for classname in mod_classes: - logger.debug("Processing: (classname: '%s')", classname) - command = self.format_command_name(classname) - info, options = self.get_cli_arguments(cli_source, classname, command) - options = self.process_options(options, command) - options["helptext"] = info - logger.debug("Processed: (classname: '%s', command: '%s', options: %s)", - classname, command, options) - subopts[command] = options - return subopts - - @staticmethod - def get_cli_arguments(cli_source, classname, command): - """ Extract the options from the main and tools cli files """ - meth = getattr(cli_source, classname)(None, command) - return meth.info, meth.argument_list + meth.optional_arguments + meth.global_arguments - - def process_options(self, command_options, command): - """ Process the options for a single command """ - gui_options = OrderedDict() - for opt in command_options: - logger.trace("Processing: %s", opt) - if opt.get("help", "") == SUPPRESS: - logger.trace("Skipping suppressed option: %s", opt) - continue - title = self.set_control_title(opt["opts"]) - cpanel_option = ControlPanelOption( - title, - self.get_data_type(opt), - group=opt.get("group", None), - default=opt.get("default", None), - choices=opt.get("choices", None), - is_radio=opt.get("action", "") == actions.Radio, - is_multi_option=opt.get("action", "") == actions.MultiOption, - rounding=self.get_rounding(opt), - min_max=opt.get("min_max", None), - sysbrowser=self.get_sysbrowser(opt, command_options, command), - helptext=opt["help"], - track_modified=True, - command=command) - gui_options[title] = dict(cpanel_option=cpanel_option, - opts=opt["opts"], - nargs=opt.get("nargs", None)) - logger.trace("Processed: %s", gui_options) - return gui_options + self._commands[category].extend(commands) + logger.debug("Set '%s' commands: %s", category, self._commands[category]) + + @classmethod + def _get_cli_arguments(cls, + arg_class: T.Type[FaceSwapArgs], + command: str) -> tuple[str, list[dict[str, T.Any]]]: + """ Extract the command line options from the given cli class - @staticmethod - def set_control_title(opts): - """ Take the option switch and format it nicely """ + Parameters + ---------- + arg_class: :class:`~lib.cli.args.FaceSwapArgs` + The class to extract the options from + command: str + The command name to extract the options for + + Returns + ------- + info: str + The helptext information for given command + options: list[dict. str, Any] + The command line options for the given command + """ + args = arg_class(None, command) + arg_list = args.argument_list + args.optional_arguments + args.global_arguments + logger.debug("Obtain options for '%s'. Info: '%s', options: %s", + command, args.info, len(arg_list)) + return args.info, arg_list + + @classmethod + def _set_control_title(cls, opts: tuple[str, ...]) -> str: + """ Take the option switch and format it nicely + + Parameters + ---------- + opts: tuple[str, ...] + The option switch for a command line option + + Returns + ------- + str + The option switch formatted for display + """ ctltitle = opts[1] if len(opts) == 2 else opts[0] - ctltitle = ctltitle.replace("-", " ").replace("_", " ").strip().title() - return ctltitle - - @staticmethod - def get_data_type(opt): - """ Return a datatype for passing into control_helper.py to get the correct control """ - if opt.get("type", None) is not None and isinstance(opt["type"], type): - retval = opt["type"] + retval = ctltitle.replace("-", " ").replace("_", " ").strip().title() + logger.debug("Formatted '%s' to '%s'", ctltitle, retval) + return retval + + @classmethod + def _get_data_type(cls, opt: dict[str, T.Any]) -> type: + """ Return a data type for passing into control_helper.py to get the correct control + + Parameters + ---------- + option: dict[str, Any] + The option to extract the data type from + + Returns + ------- + :class:`type` + The Python type for the option + """ + type_ = opt.get("type") + if type_ is not None and isinstance(opt["type"], type): + retval = type_ elif opt.get("action", "") in ("store_true", "store_false"): retval = bool else: retval = str + logger.debug("Setting type to %s for %s", retval, type_) return retval - @staticmethod - def get_rounding(opt): - """ Return rounding if correct data type, else None """ - dtype = opt.get("type", None) + @classmethod + def _get_rounding(cls, opt: dict[str, T.Any]) -> int | None: + """ Return rounding for the given option + + Parameters + ---------- + option: dict[str, Any] + The option to extract the rounding from + + Returns + ------- + int | None + int if the data type supports rounding otherwise ``None`` + """ + dtype = opt.get("type") if dtype == float: retval = opt.get("rounding", 2) elif dtype == int: retval = opt.get("rounding", 1) else: retval = None + logger.debug("Setting rounding to %s for type %s", retval, dtype) return retval - def get_sysbrowser(self, option, options, command): - """ Return the system file browser and file types if required else None """ + @classmethod + def _expand_action_option(cls, + option: dict[str, T.Any], + options: list[dict[str, T.Any]]) -> None: + """ Expand the action option to the full command name + + Parameters + ---------- + option: dict[str, Any] + The option to expand the action for + options: list[dict[str, Any]] + The full list of options for the command + """ + opts = {opt["opts"][0]: opt["opts"][-1] + for opt in options} + old_val = option["action_option"] + new_val = opts[old_val] + logger.debug("Updating action option from '%s' to '%s'", old_val, new_val) + option["action_option"] = new_val + + def _get_sysbrowser(self, + option: dict[str, T.Any], + options: list[dict[str, T.Any]], + command: str) -> dict[T.Literal["filetypes", + "browser", + "command", + "destination", + "action_option"], str | list[str]] | None: + """ Return the system file browser and file types if required + + Parameters + ---------- + option: dict[str, Any] + The option to obtain the system browser for + options: list[dict[str, Any]] + The full list of options for the command + command: str + The command that the options belong to + + Returns + ------- + dict[Literal["filetypes", "browser", "command", + "destination", "action_option"], list[str]] | None + The browser information, if valid, or ``None`` if browser not required + """ action = option.get("action", None) if action not in (actions.DirFullPaths, actions.FileFullPaths, actions.FilesFullPaths, actions.DirOrFileFullPaths, + actions.DirOrFilesFullPaths, actions.SaveFileFullPaths, actions.ContextFullPaths): return None - retval = dict() + retval: dict[T.Literal["filetypes", + "browser", + "command", + "destination", + "action_option"], str | list[str]] = {} action_option = None if option.get("action_option", None) is not None: - self.expand_action_option(option, options) + self._expand_action_option(option, options) action_option = option["action_option"] retval["filetypes"] = option.get("filetypes", "default") if action == actions.FileFullPaths: @@ -190,6 +397,8 @@ def get_sysbrowser(self, option, options, command): retval["browser"] = ["save"] elif action == actions.DirOrFileFullPaths: retval["browser"] = ["folder", "load"] + elif action == actions.DirOrFilesFullPaths: + retval["browser"] = ["folder", "multi_load"] elif action == actions.ContextFullPaths and action_option: retval["browser"] = ["context"] retval["command"] = command @@ -200,50 +409,147 @@ def get_sysbrowser(self, option, options, command): logger.debug(retval) return retval - @staticmethod - def expand_action_option(option, options): - """ Expand the action option to the full command name """ - opts = {opt["opts"][0]: opt["opts"][-1] - for opt in options} - old_val = option["action_option"] - new_val = opts[old_val] - logger.debug("Updating action option from '%s' to '%s'", old_val, new_val) - option["action_option"] = new_val + def _process_options(self, command_options: list[dict[str, T.Any]], command: str + ) -> dict[str, CliOption]: + """ Process the options for a single command + + Parameters + ---------- + command_options: list[dict. str, Any] + The command line options for the given command + command: str + The command name to process + + Returns + ------- + dict[str, :class:`CliOption`] + The collected command line options for handling by the GUI + """ + retval: dict[str, CliOption] = {} + for opt in command_options: + logger.debug("Processing: cli option: %s", opt["opts"]) + if opt.get("help", "") == SUPPRESS: + logger.debug("Skipping suppressed option: %s", opt) + continue + title = self._set_control_title(opt["opts"]) + cpanel_option = ControlPanelOption( + title, + self._get_data_type(opt), + group=opt.get("group", None), + default=opt.get("default", None), + choices=opt.get("choices", None), + is_radio=opt.get("action", "") == actions.Radio, + is_multi_option=opt.get("action", "") == actions.MultiOption, + rounding=self._get_rounding(opt), + min_max=opt.get("min_max", None), + sysbrowser=self._get_sysbrowser(opt, command_options, command), + helptext=opt["help"], + track_modified=True, + command=command) + retval[title] = CliOption(cpanel_option=cpanel_option, + opts=opt["opts"], + nargs=opt.get("nargs")) + logger.debug("Processed: %s", retval) + return retval + + def _extract_options(self, arguments: list[T.Type[FaceSwapArgs]]): + """ Extract the collected command line FaceSwapArg options into master options + :attr:`opts` dictionary + + Parameters + ---------- + arguments: list[:class:`~lib.cli.args.FaceSwapArgs`] + The command line class objects to process + """ + retval = {} + for arg_class in arguments: + logger.debug("Processing: '%s'", arg_class.__name__) + command = self._class_name_to_command(arg_class.__name__) + info, options = self._get_cli_arguments(arg_class, command) + opts = T.cast(dict[str, CliOption | str], self._process_options(options, command)) + opts["helptext"] = info + retval[command] = opts + self._opts.update(retval) - def gen_command_options(self, command): - """ Yield each option for specified command """ - for key, val in self.opts[command].items(): - if not isinstance(val, dict): + def _build_options(self) -> None: + """ Parse the command line argument modules and populate :attr:`commands` and :attr:`opts` + for each category """ + for category in self.categories: + modules = self._get_modules(category) + classes = self._get_all_classes(modules) + self._store_commands(category, classes) + self._extract_options(classes) + logger.debug("Built '%s'", category) + + def _gen_command_options(self, command: str + ) -> T.Generator[tuple[str, CliOption], None, None]: + """ Yield each option for specified command + + Parameters + ---------- + command: str + The faceswap command to generate the options for + + Yields + ------ + str + The option name for display + :class:`CliOption`: + The option object + """ + for key, val in self._opts.get(command, {}).items(): + if not isinstance(val, CliOption): continue yield key, val - def options_to_process(self, command=None): + def _options_to_process(self, command: str | None = None) -> list[CliOption]: """ Return a consistent object for processing regardless of whether processing all commands - or just one command for reset and clear. Removes helptext from return value """ + or just one command for reset and clear. Removes helptext from return value + + Parameters + ---------- + command: str | None, optional + The command to return the options for. ``None`` for all commands. Default ``None`` + + Returns + ------- + list[:class:`CliOption`] + The options to be processed + """ if command is None: - options = [opt for opts in self.opts.values() - for opt in opts.values() if isinstance(opt, dict)] - else: - options = [opt for opt in self.opts[command].values() if isinstance(opt, dict)] - return options + return [opt for opts in self._opts.values() + for opt in opts if isinstance(opt, CliOption)] + return [opt for opt in self._opts[command] if isinstance(opt, CliOption)] - def reset(self, command=None): - """ Reset the options for all or passed command - back to default value """ + def reset(self, command: str | None = None) -> None: + """ Reset the options for all or passed command back to default value + + Parameters + ---------- + command: str | None, optional + The command to reset the options for. ``None`` to reset for all commands. + Default: ``None`` + """ logger.debug("Resetting options to default. (command: '%s'", command) - for option in self.options_to_process(command): - cp_opt = option["cpanel_option"] + for option in self._options_to_process(command): + cp_opt = option.cpanel_option default = "" if cp_opt.default is None else cp_opt.default - if (option.get("nargs", None) - and isinstance(default, (list, tuple))): + if option.nargs is not None and isinstance(default, (list, tuple)): default = ' '.join(str(val) for val in default) cp_opt.set(default) - def clear(self, command=None): - """ Clear the options values for all or passed commands """ + def clear(self, command: str | None = None) -> None: + """ Clear the options values for all or passed commands + + Parameters + ---------- + command: str | None, optional + The command to clear the options for. ``None`` to clear options for all commands. + Default: ``None`` + """ logger.debug("Clearing options. (command: '%s'", command) - for option in self.options_to_process(command): - cp_opt = option["cpanel_option"] + for option in self._options_to_process(command): + cp_opt = option.cpanel_option if isinstance(cp_opt.get(), bool): cp_opt.set(False) elif isinstance(cp_opt.get(), (int, float)): @@ -251,47 +557,93 @@ def clear(self, command=None): else: cp_opt.set("") - def get_option_values(self, command=None): - """ Return all or single command control titles with the associated tk_var value """ - ctl_dict = dict() - for cmd, opts in self.opts.items(): + def get_option_values(self, command: str | None = None + ) -> dict[str, dict[str, bool | int | float | str]]: + """ Return all or single command control titles with the associated tk_var value + + Parameters + ---------- + command: str | None, optional + The command to get the option values for. ``None`` to get all option values. + Default: ``None`` + + Returns + ------- + dict[str, dict[str, bool | int | float | str]] + option values in the format {command: {option_name: option_value}} + """ + ctl_dict: dict[str, dict[str, bool | int | float | str]] = {} + for cmd, opts in self._opts.items(): if command and command != cmd: continue - cmd_dict = dict() + cmd_dict: dict[str, bool | int | float | str] = {} for key, val in opts.items(): - if not isinstance(val, dict): + if not isinstance(val, CliOption): continue - cmd_dict[key] = val["cpanel_option"].get() + cmd_dict[key] = val.cpanel_option.get() ctl_dict[cmd] = cmd_dict logger.debug("command: '%s', ctl_dict: %s", command, ctl_dict) return ctl_dict - def get_one_option_variable(self, command, title): - """ Return a single tk_var for the specified - command and control_title """ - for opt_title, option in self.gen_command_options(command): + def get_one_option_variable(self, command: str, title: str) -> Variable | None: + """ Return a single :class:`tkinter.Variable` tk_var for the specified command and + control_title + + Parameters + ---------- + command: str + The command to return the variable from + title: str + The option title to return the variable for + + Returns + ------- + :class:`tkinter.Variable` | None + The requested tkinter variable, or ``None`` if it could not be found + """ + for opt_title, option in self._gen_command_options(command): if opt_title == title: - return option["cpanel_option"].tk_var + return option.cpanel_option.tk_var return None - def gen_cli_arguments(self, command): - """ Return the generated cli arguments for the selected command """ - for _, option in self.gen_command_options(command): - optval = str(option["cpanel_option"].get()) - opt = option["opts"][0] - if command in ("extract", "convert") and opt == "-o": - get_images().set_faceswap_output_path(optval) - if optval in ("False", ""): + def gen_cli_arguments(self, command: str) -> T.Generator[tuple[str, ...], None, None]: + """ Yield the generated cli arguments for the selected command + + Parameters + ---------- + command: str + The command to generate the command line arguments for + + Yields + ------ + tuple[str, ...] + The generated command line arguments + """ + output_dir = None + for _, option in self._gen_command_options(command): + str_val = str(option.cpanel_option.get()) + switch = option.opts[0] + batch_mode = command == "extract" and switch == "-b" # Check for batch mode + if command in ("extract", "convert") and switch == "-o": # Output location for preview + output_dir = str_val + + if str_val in ("False", ""): # skip no value opts continue - if optval == "True": - yield (opt, ) - else: - if option.get("nargs", None): - if "\"" in optval: - optval = [arg[1:-1] for arg in re.findall(r"\".+?\"", optval)] - else: - optval = optval.split(" ") - opt = [opt] + optval + + if str_val == "True": # store_true just output the switch + yield (switch, ) + continue + + if option.nargs is not None: + if "\"" in str_val: + val = [arg[1:-1] for arg in re.findall(r"\".+?\"", str_val)] else: - opt = (opt, optval) - yield opt + val = str_val.split(" ") + retval = (switch, *val) + else: + retval = (switch, str_val) + yield retval + + if command in ("extract", "convert") and output_dir is not None: + get_images().preview_extract.set_faceswap_output_path(output_dir, + batch_mode=batch_mode) diff --git a/lib/gui/popup_configure.py b/lib/gui/popup_configure.py index 1119990026..6bdb725db4 100644 --- a/lib/gui/popup_configure.py +++ b/lib/gui/popup_configure.py @@ -1,6 +1,6 @@ #!/usr/bin python3 """ The pop-up window of the Faceswap GUI for the setting of configuration options. """ - +from __future__ import annotations from collections import OrderedDict from configparser import ConfigParser import gettext @@ -9,6 +9,8 @@ import sys import tkinter as tk from tkinter import ttk +import typing as T + from importlib import import_module from lib.serializer import get_serializer @@ -17,7 +19,10 @@ from .custom_widgets import Tooltip from .utils import FileHandler, get_config, get_images, PATHCACHE -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from lib.config import FaceswapConfig + +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) @@ -120,7 +125,7 @@ def __init__(self, name, configurations): super().__init__() self._root = get_config().root self._set_geometry() - self._tk_vars = dict(header=tk.StringVar()) + self._tk_vars = {"header": tk.StringVar()} theme = {**get_config().user_theme["group_panel"], **get_config().user_theme["group_settings"]} @@ -292,7 +297,7 @@ def _fix_styles(cls, theme): style = ttk.Style() # Fix a bug in Tree-view that doesn't show alternate foreground on selection - fix_map = lambda o: [elm for elm in style.map("Treeview", query_opt=o) # noqa + fix_map = lambda o: [elm for elm in style.map("Treeview", query_opt=o) # noqa[E731] # pylint:disable=C3001 if elm[:2] != ("!disabled", "!selected")] # Remove the Borders @@ -398,7 +403,7 @@ class DisplayArea(ttk.Frame): # pylint:disable=too-many-ancestors """ def __init__(self, top_level, parent, configurations, tree, theme): super().__init__(parent) - self._configs = configurations + self._configs: dict[str, FaceswapConfig] = configurations self._theme = theme self._tree = tree self._vars = {} @@ -439,30 +444,28 @@ def _get_config(self): sect = section.split(".")[-1] # Elevate global to root key = plugin if sect == "global" else f"{plugin}|{category}|{sect}" - retval[key] = dict(helptext=None, options=OrderedDict()) + retval[key] = {"helptext": None, "options": OrderedDict()} - for option, params in conf.defaults[section].items(): - if option == "helptext": - retval[key]["helptext"] = params - continue + retval[key]["helptext"] = conf.defaults[section].helptext + for option, params in conf.defaults[section].items.items(): initial_value = conf.config_dict[option] initial_value = "none" if initial_value is None else initial_value - if params["type"] == list and isinstance(initial_value, list): + if params.datatype == list and isinstance(initial_value, list): # Split multi-select lists into space separated strings for tk variables initial_value = " ".join(initial_value) retval[key]["options"][option] = ControlPanelOption( title=option, - dtype=params["type"], - group=params["group"], - default=params["default"], + dtype=params.datatype, + group=params.group, + default=params.default, initial_value=initial_value, - choices=params["choices"], - is_radio=params["gui_radio"], - is_multi_option=params["type"] == list, - rounding=params["rounding"], - min_max=params["min_max"], - helptext=params["helptext"]) + choices=params.choices, + is_radio=params.gui_radio, + is_multi_option=params.datatype == list, + rounding=params.rounding, + min_max=params.min_max, + helptext=params.helptext) logger.debug("Formatted Config for GUI: %s", retval) return retval @@ -580,7 +583,7 @@ def _create_links_page(self, key): cursor="hand2") lbl.pack(side=tk.TOP, fill=tk.X, padx=10, pady=(0, 5)) bind = f"{key}|{link}" - lbl.bind("", lambda e, l=bind: self._link_callback(l)) + lbl.bind("", lambda e, x=bind: self._link_callback(x)) return frame @@ -628,6 +631,59 @@ def reset(self, page_only=False): item.set(item.default) logger.debug("Reset config") + def _get_new_config(self, + page_only: bool, + config: FaceswapConfig, + category: str, + lookup: str) -> ConfigParser: + """ Obtain a new configuration file for saving + + Parameters + ---------- + page_only: bool + ``True`` saves just the currently selected page's options, ``False`` saves all the + plugins options within the currently selected config. + config: :class:`~lib.config.FaceswapConfig` + The original config that is to be addressed + category: str + The configuration category to update + lookup: str + The section of the configuration to update + + Returns + ------- + :class:`configparse.ConfigParser` + The newly created configuration object for saving + """ + new_config = ConfigParser(allow_no_value=True) + for section_name, section in config.defaults.items(): + logger.debug("Adding section: '%s')", section_name) + config.insert_config_section(section_name, section.helptext, config=new_config) + for item, options in section.items.items(): + if item == "helptext": + continue + if page_only and section_name != lookup: + # Keep existing values for pages we are not updating + new_opt = config.get(section_name, item) + logger.debug("Retain existing value '%s' for %s", + new_opt, ".".join([section_name, item])) + else: + # Get currently selected value + key = category + if section_name != "global": + key += f"|{section_name.replace('.', '|')}" + new_opt = self._config_cpanel_dict[key]["options"][item].get() + logger.debug("Updating value to '%s' for %s", + new_opt, ".".join([section_name, item])) + helptext = config.format_help(options.helptext, is_section=False) + new_config.set(section_name, helptext) + if options.datatype == list: # Comma separated multi select options + assert isinstance(new_opt, (list, str)) + new_opt = ", ".join(new_opt if isinstance(new_opt, list) else new_opt.split()) + new_config.set(section_name, item, str(new_opt)) + + return new_config + def save(self, page_only=False): """ Save the configuration file to disk. @@ -642,7 +698,6 @@ def save(self, page_only=False): category = selection.split("|")[0] config = self._configs[category] # Create a new config to pull through any defaults change - new_config = ConfigParser(allow_no_value=True) if "|" in selection: lookup = ".".join(selection.split("|")[1:]) @@ -653,36 +708,12 @@ def save(self, page_only=False): logger.info("No settings to save for the current page") return - for section, items in config.defaults.items(): - logger.debug("Adding section: '%s')", section) - config.insert_config_section(section, items["helptext"], config=new_config) - for item, options in items.items(): - if item == "helptext": - continue - if page_only and section != lookup: - # Keep existing values for pages we are not updating - new_opt = config.get(section, item) - logger.debug("Retain existing value '%s' for %s", - new_opt, ".".join([section, item])) - else: - # Get currently selected value - key = category - if section != "global": - key += f"|{section.replace('.', '|')}" - new_opt = self._config_cpanel_dict[key]["options"][item].get() - logger.debug("Updating value to '%s' for %s", - new_opt, ".".join([section, item])) - helptext = config.format_help(options["helptext"], is_section=False) - new_config.set(section, helptext) - if options["type"] == list: # Comma separated multi select options - new_opt = ", ".join(new_opt if isinstance(new_opt, list) else new_opt.split()) - new_config.set(section, item, str(new_opt)) - config.config = new_config + config.config = self._get_new_config(page_only, config, category, lookup) config.save_config() logger.info("Saved config: '%s'", config.configfile) if category == "gui": - if not get_config().tk_vars["runningtask"].get(): + if not get_config().tk_vars.running_task.get(): get_config().root.rebuild() else: logger.info("Can't redraw GUI whilst a task is running. GUI Settings will be " @@ -782,8 +813,9 @@ def _get_filename(self, action): return None args = ("save_filename", "json") if action == "save" else ("filename", "json") - kwargs = dict(title=f"{action.title()} Preset...", - initial_folder=self._preset_path) + kwargs = {"title": f"{action.title()} Preset...", + "initial_folder": self._preset_path, + "parent": self._parent} if action == "save": kwargs["initial_file"] = self._get_initial_filename() diff --git a/lib/gui/popup_session.py b/lib/gui/popup_session.py index cbf0e2361a..6ba9e2b4a6 100644 --- a/lib/gui/popup_session.py +++ b/lib/gui/popup_session.py @@ -5,8 +5,9 @@ import gettext import logging import tkinter as tk + +from dataclasses import dataclass, field from tkinter import ttk -from typing import List from .control_helper import ControlBuilder, ControlPanelOption from .custom_widgets import Tooltip @@ -14,13 +15,59 @@ from .analysis import Calculations, Session from .utils import FileHandler, get_images, LongRunningTask -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("gui.tooltips", localedir="locales", fallback=True) _ = _LANG.gettext +@dataclass +class SessionTKVars: + """ Dataclass for holding the tk variables required for the session popup + + Parameters + ---------- + buildgraph: :class:`tkinter.BooleanVar` + Trigger variable to indicate the graph should be rebuilt + status: :class:`tkinter.StringVar` + The variable holding the current status of the popup window + display: :class:`tkinter.StringVar` + Variable indicating the type of information to be displayed + scale: :class:`tkinter.StringVar` + Variable indicating whether to display as log or linear data + raw: :class:`tkinter.BooleanVar` + Variable to indicate raw data should be displayed + trend: :class:`tkinter.BooleanVar` + Variable to indicate that trend data should be displayed + avg: :class:`tkinter.BooleanVar` + Variable to indicate that rolling average data should be displayed + smoothed: :class:`tkinter.BooleanVar` + Variable to indicate that smoothed data should be displayed + outliers: :class:`tkinter.BooleanVar` + Variable to indicate that outliers should be displayed + loss_keys: dict + Dictionary of names to :class:`tkinter.BooleanVar` indicating whether specific loss items + should be displayed + avgiterations: :class:`tkinter.IntVar` + The number of iterations to use for rolling average + smoothamount: :class:`tkinter.DoubleVar` + The amount of smoothing to apply for smoothed data + """ + buildgraph: tk.BooleanVar + status: tk.StringVar + display: tk.StringVar + scale: tk.StringVar + raw: tk.BooleanVar + trend: tk.BooleanVar + avg: tk.BooleanVar + smoothed: tk.BooleanVar + outliers: tk.BooleanVar + avgiterations: tk.IntVar + smoothamount: tk.DoubleVar + loss_keys: dict[str, tk.BooleanVar] = field(default_factory=dict) + + class SessionPopUp(tk.Toplevel): """ Pop up for detailed graph/stats for selected session. @@ -34,15 +81,16 @@ def __init__(self, session_id: int, data_points: int) -> None: logger.debug("Initializing: %s: (session_id: %s, data_points: %s)", self.__class__.__name__, session_id, data_points) super().__init__() - self._thread = None # Thread for loading data in a background task + self._thread: LongRunningTask | None = None # Thread for loading data in background self._default_view = "avg" if data_points > 1000 else "smoothed" self._session_id = None if session_id == "Total" else int(session_id) - self._graph_frame = None - self._graph = None - self._display_data = None + self._graph_frame = ttk.Frame(self) + self._graph: SessionGraph | None = None + self._display_data: Calculations | None = None self._vars = self._set_vars() + self._graph_initialised = False optsframe = self._layout_frames() @@ -56,24 +104,29 @@ def __init__(self, session_id: int, data_points: int) -> None: logger.debug("Initialized: %s", self.__class__.__name__) - def _set_vars(self) -> dict: + def _set_vars(self) -> SessionTKVars: """ Set status tkinter String variable and tkinter Boolean variable to callback when the graph is ready to build. Returns ------- - dict + :class:`SessionTKVars` The tkinter Variables for the pop up graph """ logger.debug("Setting tk graph build variable and internal variables") - - retval = dict(status=tk.StringVar()) - - var = tk.BooleanVar() - var.set(False) - var.trace("w", self._graph_build) - - retval["buildgraph"] = var + retval = SessionTKVars(buildgraph=tk.BooleanVar(), + status=tk.StringVar(), + display=tk.StringVar(), + scale=tk.StringVar(), + raw=tk.BooleanVar(), + trend=tk.BooleanVar(), + avg=tk.BooleanVar(), + smoothed=tk.BooleanVar(), + outliers=tk.BooleanVar(), + avgiterations=tk.IntVar(), + smoothamount=tk.DoubleVar()) + retval.buildgraph.set(False) + retval.buildgraph.trace("w", self._graph_build) return retval def _layout_frames(self) -> ttk.Frame: @@ -82,7 +135,6 @@ def _layout_frames(self) -> ttk.Frame: leftframe = ttk.Frame(self) sep = ttk.Frame(self, width=2, relief=tk.RIDGE) - self._graph_frame = ttk.Frame(self) self._graph_frame.pack(side=tk.RIGHT, fill=tk.BOTH, pady=5, expand=True) sep.pack(fill=tk.Y, side=tk.LEFT) @@ -119,10 +171,10 @@ def _opts_combobox(self, frame: ttk.Frame) -> None: The frame that the options reside in """ logger.debug("Building Combo boxes") - choices = dict(Display=("Loss", "Rate"), Scale=("Linear", "Log")) + choices = {"Display": ("Loss", "Rate"), "Scale": ("Linear", "Log")} for item in ["Display", "Scale"]: - var = tk.StringVar() + var: tk.StringVar = getattr(self._vars, item.lower()) cmbframe = ttk.Frame(frame) lblcmb = ttk.Label(cmbframe, text=f"{item}:", width=7, anchor=tk.W) @@ -132,8 +184,6 @@ def _opts_combobox(self, frame: ttk.Frame) -> None: cmd = self._option_button_reload if item == "Display" else self._graph_scale var.trace("w", cmd) - self._vars[item.lower().strip()] = var - hlp = self._set_help(item) Tooltip(cmbframe, text=hlp, wrap_length=200) @@ -160,10 +210,9 @@ def _opts_checkbuttons(self, frame: ttk.Frame) -> None: else: text = f"Show {item.title()}" - var = tk.BooleanVar() + var: tk.BooleanVar = getattr(self._vars, item) if item == self._default_view: var.set(True) - self._vars[item] = var ctl = ttk.Checkbutton(frame, variable=var, text=text) hlp = self._set_help(item) @@ -207,7 +256,7 @@ def _opts_loss_keys(self, frame: ttk.Frame) -> None: Tooltip(ctl, text=helptext, wrap_length=200) ctl.pack(side=tk.TOP, padx=5, pady=5, anchor=tk.W) - self._vars["loss_keys"] = lk_vars + self._vars.loss_keys = lk_vars logger.debug("Built Loss Key Check Buttons") def _opts_slider(self, frame: ttk.Frame) -> None: @@ -223,11 +272,11 @@ def _opts_slider(self, frame: ttk.Frame) -> None: logger.debug("Building Slider Controls") for item in ("avgiterations", "smoothamount"): if item == "avgiterations": - dtype = int + dtype: type[int] | type[float] = int text = "Iterations to Average:" - default = 500 + default: int | float = 500 rounding = 25 - min_max = (25, 2500) + min_max: tuple[int, int | float] = (25, 2500) elif item == "smoothamount": dtype = float text = "Smoothing Amount:" @@ -240,7 +289,7 @@ def _opts_slider(self, frame: ttk.Frame) -> None: rounding=rounding, min_max=min_max, helptext=self._set_help(item)) - self._vars[item] = slider.tk_var + setattr(self._vars, item, slider.tk_var) ControlBuilder(frame, slider, 1, 19, None, "Analysis.", True) logger.debug("Built Sliders") @@ -256,7 +305,7 @@ def _opts_buttons(self, frame: ttk.Frame) -> None: btnframe = ttk.Frame(frame) lblstatus = ttk.Label(btnframe, width=40, - textvariable=self._vars["status"], + textvariable=self._vars.status, anchor=tk.W) for btntype in ("reload", "save"): @@ -297,6 +346,7 @@ def _option_button_save(self) -> None: logger.debug("Save Cancelled") return logger.debug("Saving to: %s", savefile) + assert self._display_data is not None save_data = self._display_data.stats fieldnames = sorted(key for key in save_data.keys()) @@ -305,7 +355,7 @@ def _option_button_save(self) -> None: csvout.writerow(fieldnames) csvout.writerows(zip(*[save_data[key] for key in fieldnames])) - def _option_button_reload(self, *args) -> None: # pylint: disable=unused-argument + def _option_button_reload(self, *args) -> None: # pylint:disable=unused-argument """ Action for reset button press and checkbox changes. Parameters @@ -320,12 +370,13 @@ def _option_button_reload(self, *args) -> None: # pylint: disable=unused-argume if not valid: logger.debug("Invalid data") return + assert self._graph is not None self._graph.refresh(self._display_data, - self._vars["display"].get(), - self._vars["scale"].get()) + self._vars.display.get(), + self._vars.scale.get()) logger.debug("Refreshed Graph") - def _graph_scale(self, *args) -> None: # pylint: disable=unused-argument + def _graph_scale(self, *args) -> None: # pylint:disable=unused-argument """ Action for changing graph scale. Parameters @@ -333,9 +384,10 @@ def _graph_scale(self, *args) -> None: # pylint: disable=unused-argument args: tuple Required for TK Callback but unused """ + assert self._graph is not None if not self._graph_initialised: return - self._graph.set_yscale_type(self._vars["scale"].get()) + self._graph.set_yscale_type(self._vars.scale.get()) @classmethod def _set_help(cls, action: str) -> str: @@ -351,32 +403,21 @@ def _set_help(cls, action: str) -> str: str The help text for the given action """ - hlp = "" - action = action.lower() - if action == "reload": - hlp = _("Refresh graph") - elif action == "save": - hlp = _("Save display data to csv") - elif action == "avgiterations": - hlp = _("Number of data points to sample for rolling average") - elif action == "smoothamount": - hlp = _("Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing") - elif action == "outliers": - hlp = _("Flatten data points that fall more than 1 standard deviation from the mean " - "to the mean value.") - elif action == "avg": - hlp = _("Display rolling average of the data") - elif action == "smoothed": - hlp = _("Smooth the data") - elif action == "raw": - hlp = _("Display raw data") - elif action == "trend": - hlp = _("Display polynormal data trend") - elif action == "display": - hlp = _("Set the data to display") - elif action == "scale": - hlp = _("Change y-axis scale") - return hlp + lookup = { + "reload": _("Refresh graph"), + "save": _("Save display data to csv"), + "avgiterations": _("Number of data points to sample for rolling average"), + "smoothamount": _("Set the smoothing amount. 0 is no smoothing, 0.99 is maximum " + "smoothing"), + "outliers": _("Flatten data points that fall more than 1 standard deviation from the " + "mean to the mean value."), + "avg": _("Display rolling average of the data"), + "smoothed": _("Smooth the data"), + "raw": _("Display raw data"), + "trend": _("Display polynormal data trend"), + "display": _("Set the data to display"), + "scale": _("Change y-axis scale")} + return lookup.get(action.lower(), "") def _compile_display_data(self) -> bool: """ Compile the data to be displayed. @@ -388,7 +429,7 @@ def _compile_display_data(self) -> bool: """ if self._thread is None: logger.debug("Compiling Display Data in background thread") - loss_keys = [key for key, val in self._vars["loss_keys"].items() + loss_keys = [key for key, val in self._vars.loss_keys.items() if val.get()] logger.debug("Selected loss_keys: %s", loss_keys) @@ -397,20 +438,20 @@ def _compile_display_data(self) -> bool: if not self._check_valid_selection(loss_keys, selections): logger.warning("No data to display. Not refreshing") return False - self._vars["status"].set("Loading Data...") + self._vars.status.set("Loading Data...") if self._graph is not None: self._graph.pack_forget() self._lbl_loading.pack(fill=tk.BOTH, expand=True) self.update_idletasks() - kwargs = dict(session_id=self._session_id, - display=self._vars["display"].get(), - loss_keys=loss_keys, - selections=selections, - avg_samples=self._vars["avgiterations"].get(), - smooth_amount=self._vars["smoothamount"].get(), - flatten_outliers=self._vars["outliers"].get()) + kwargs = {"session_id": self._session_id, + "display": self._vars.display.get(), + "loss_keys": loss_keys, + "selections": selections, + "avg_samples": self._vars.avgiterations.get(), + "smooth_amount": self._vars.smoothamount.get(), + "flatten_outliers": self._vars.outliers.get()} self._thread = LongRunningTask(target=self._get_display_data, kwargs=kwargs, widget=self) @@ -427,14 +468,14 @@ def _compile_display_data(self) -> bool: self._thread = None if not self._check_valid_data(): logger.warning("No valid data to display. Not refreshing") - self._vars["status"].set("") + self._vars.status.set("") return False logger.debug("Compiled Display Data") - self._vars["buildgraph"].set(True) + self._vars.buildgraph.set(True) return True @classmethod - def _get_display_data(cls, **kwargs) -> None: + def _get_display_data(cls, **kwargs) -> Calculations: """ Get the display data in a LongRunningTask. Parameters @@ -449,7 +490,7 @@ def _get_display_data(cls, **kwargs) -> None: """ return Calculations(**kwargs) - def _check_valid_selection(self, loss_keys: List[str], selections: List[str]) -> bool: + def _check_valid_selection(self, loss_keys: list[str], selections: list[str]) -> bool: """ Check that there will be data to display. Parameters @@ -464,7 +505,7 @@ def _check_valid_selection(self, loss_keys: List[str], selections: List[str]) -> bool ``True` if there is data to be displayed, otherwise ``False`` """ - display = self._vars["display"].get().lower() + display = self._vars.display.get().lower() logger.debug("Validating selection. (loss_keys: %s, selections: %s, display: %s)", loss_keys, selections, display) if not selections or (display == "loss" and not loss_keys): @@ -480,6 +521,7 @@ def _check_valid_data(self) -> bool: bool ``True` if there is data to be displayed, otherwise ``False`` """ + assert self._display_data is not None logger.debug("Validating data. %s", {key: len(val) for key, val in self._display_data.stats.items()}) if any(len(val) == 0 # pylint:disable=len-as-condition @@ -487,7 +529,7 @@ def _check_valid_data(self) -> bool: return False return True - def _selections_to_list(self) -> List[str]: + def _selections_to_list(self) -> list[str]: """ Compile checkbox selections to a list. Returns @@ -497,11 +539,10 @@ def _selections_to_list(self) -> List[str]: """ logger.debug("Compiling selections to list") selections = [] - for key, val in self._vars.items(): - if (isinstance(val, tk.BooleanVar) - and key != "outliers" - and val.get()): - selections.append(key) + for item in ("raw", "trend", "avg", "smoothed"): + var: tk.BooleanVar = getattr(self._vars, item) + if var.get(): + selections.append(item) logger.debug("Compiling selections to list: %s", selections) return selections @@ -513,25 +554,26 @@ def _graph_build(self, *args) -> None: # pylint:disable=unused-argument args: tuple Required for TK Callback but unused """ - if not self._vars["buildgraph"].get(): + if not self._vars.buildgraph.get(): return - self._vars["status"].set("Loading Data...") + self._vars.status.set("Loading Data...") logger.debug("Building Graph") self._lbl_loading.pack_forget() self.update_idletasks() if self._graph is None: - self._graph = SessionGraph(self._graph_frame, - self._display_data, - self._vars["display"].get(), - self._vars["scale"].get()) - self._graph.pack(expand=True, fill=tk.BOTH) - self._graph.build() + graph = SessionGraph(self._graph_frame, + self._display_data, + self._vars.display.get(), + self._vars.scale.get()) + graph.pack(expand=True, fill=tk.BOTH) + graph.build() + self._graph = graph self._graph_initialised = True else: self._graph.refresh(self._display_data, - self._vars["display"].get(), - self._vars["scale"].get()) + self._vars.display.get(), + self._vars.scale.get()) self._graph.pack(fill=tk.BOTH, expand=True) - self._vars["status"].set("") - self._vars["buildgraph"].set(False) + self._vars.status.set("") + self._vars.buildgraph.set(False) logger.debug("Built Graph") diff --git a/lib/gui/project.py b/lib/gui/project.py index 7967a70129..1b5b2528a5 100644 --- a/lib/gui/project.py +++ b/lib/gui/project.py @@ -8,7 +8,7 @@ from lib.serializer import get_serializer -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class _GuiSession(): # pylint:disable=too-few-public-methods @@ -90,11 +90,12 @@ def _stored_tab_name(self): def _selected_to_choices(self): """ dict: The selected value and valid choices for multi-option, radio or combo options. """ - valid_choices = {cmd: {opt: dict(choices=val["cpanel_option"].choices, - is_multi=val["cpanel_option"].is_multi_option) + valid_choices = {cmd: {opt: {"choices": val.cpanel_option.choices, + "is_multi": val.cpanel_option.is_multi_option} for opt, val in data.items() - if isinstance(val, dict) and "cpanel_option" in val - and val["cpanel_option"].choices is not None} + if hasattr(val, "cpanel_option") # Filter out helptext + and val.cpanel_option.choices is not None + } for cmd, data in self._config.cli_opts.opts.items()} logger.trace("valid_choices: %s", valid_choices) retval = {command: {option: {"value": value, @@ -144,7 +145,7 @@ def _set_filename(self, filename=None, sess_type="project"): bool: `True` if filename has been successfully set otherwise ``False`` """ logger.debug("filename: '%s', sess_type: '%s'", filename, sess_type) - handler = "config_{}".format(sess_type) + handler = f"config_{sess_type}" if filename is None: logger.debug("Popping file handler") @@ -156,7 +157,7 @@ def _set_filename(self, filename=None, sess_type="project"): cfgfile.close() if not os.path.isfile(filename): - msg = "File does not exist: '{}'".format(filename) + msg = f"File does not exist: '{filename}'" logger.error(msg) return False ext = os.path.splitext(filename)[1] @@ -209,7 +210,7 @@ def _get_options_for_command(self, command): opts = self._options.get(command, None) retval = {command: opts} if not opts: - self._config.tk_vars["console_clear"].set(True) + self._config.tk_vars.console_clear.set(True) logger.info("No %s section found in file", command) retval = None logger.debug(retval) @@ -380,10 +381,9 @@ def _save_as_to_filename(self, session_type): True if :attr:`filename` successfully set otherwise ``False`` """ logger.debug("Popping save as file handler. session_type: '%s'", session_type) - title = "Save {}As...".format("{} ".format(session_type.title()) - if session_type != "all" else "") + title = f"Save {f'{session_type.title()} ' if session_type != 'all' else ''}As..." cfgfile = self._file_handler("save", - "config_{}".format(session_type), + f"config_{session_type}", title=title, initial_folder=self._dirname).return_file if not cfgfile: @@ -432,7 +432,7 @@ class Tasks(_GuiSession): """ def __init__(self, config, file_handler): super().__init__(config, file_handler) - self._tasks = dict() + self._tasks = {} @property def _is_project(self): @@ -539,7 +539,7 @@ def _update_legacy_task(self, filename): logger.debug("Not a .fsw file: '%s'", filename) return filename - new_filename = "{}.fst".format(fname) + new_filename = f"{fname}.fst" logger.debug("Renaming '%s' to '%s'", filename, new_filename) os.rename(filename, new_filename) self._del_from_recent(filename, save=True) @@ -601,9 +601,9 @@ def _add_task(self, command): The tab that pertains to the currently active task """ - self._tasks[command] = dict(filename=self._filename, - options=self._options, - is_project=self._is_project) + self._tasks[command] = {"filename": self._filename, + "options": self._options, + "is_project": self._is_project} def clear_tasks(self): """ Clears all of the stored tasks. @@ -612,7 +612,7 @@ def clear_tasks(self): called by :class:`Project` when a project has been loaded which is in fact a task. """ logger.debug("Clearing stored tasks") - self._tasks = dict() + self._tasks = {} def add_project_task(self, filename, command, options): """ Add an individual task from a loaded :class:`Project` to the internal :attr:`_tasks` @@ -630,7 +630,7 @@ def add_project_task(self, filename, command, options): options: dict The options for this task loaded from the project """ - self._tasks[command] = dict(filename=filename, options=options, is_project=True) + self._tasks[command] = {"filename": filename, "options": options, "is_project": True} def _set_active_task(self, command=None): """ Set the active :attr:`_filename` and :attr:`_options` to currently selected tab's @@ -684,7 +684,7 @@ def cli_options(self): @property def _project_modified(self): """bool: ``True`` if the project has been modified otherwise ``False``. """ - return any([var.get() for var in self._modified_vars.values()]) + return any(var.get() for var in self._modified_vars.values()) @property def _tasks(self): diff --git a/lib/gui/theme.py b/lib/gui/theme.py index 2e29abbe62..cdb42cbdba 100644 --- a/lib/gui/theme.py +++ b/lib/gui/theme.py @@ -11,10 +11,10 @@ from lib.utils import FaceswapError -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -class Style(): # pylint:disable=too-few-public-methods +class Style(): """ Set the overarching theme and customize widgets. Parameters @@ -196,7 +196,7 @@ def combobox(self, key, control_color, active_color, arrow_color, control_border The color of the input field's border """ # All the stock down arrow images are bad - images = dict() + images = {} for state in ("active", "normal"): images[f"arrow_{state}"] = self._images.get_image( (20, 20), @@ -343,7 +343,7 @@ def scrollbar(self, key, trough_color, border_color, control_backgrounds, contro "control_backgrounds: %s, control_foregrounds: %s, control_borders: %s)", key, trough_color, border_color, control_backgrounds, control_foregrounds, control_borders) - images = dict() + images = {} for idx, state in enumerate(("normal", "disabled", "active")): # Create arrow and slider widgets for each state img_args = ((16, 16), control_backgrounds[idx]) @@ -370,7 +370,7 @@ def scrollbar(self, key, trough_color, border_color, control_backgrounds, contro ("disabled", images[f"img_{lookup}_disabled"]), ("pressed !disabled", images[f"img_{lookup}_active"]), ("active !disabled", images[f"img_{lookup}_active"])) - kwargs = dict(border=1, sticky="ns") if element == "thumb" else dict() + kwargs = {"border": 1, "sticky": "ns"} if element == "thumb" else {} self._style.element_create(*args, **kwargs) # Get a configurable trough @@ -439,7 +439,7 @@ def slider(self, key, control_color, active_color, trough_color): troughcolor=trough_color) -class _TkImage(): # pylint:disable=too-few-public-methods +class _TkImage(): """ Create a tk image for a given pattern and shape. """ def __init__(self): @@ -487,7 +487,7 @@ def _get_arrow(cls, dimensions, thickness, direction): crop_size = (square_size // 16) * 16 draw_rows = int(6 * crop_size / 16) start_row = dimensions[1] // 2 - draw_rows // 2 - initial_indent = (2 * (crop_size // 16) + (dimensions[0] - crop_size) // 2) + initial_indent = 2 * (crop_size // 16) + (dimensions[0] - crop_size) // 2 retval = np.zeros((dimensions[1], dimensions[0]), dtype="uint8") for i in range(start_row, start_row + draw_rows): diff --git a/lib/gui/utils.py b/lib/gui/utils.py deleted file mode 100644 index 8189fb7c08..0000000000 --- a/lib/gui/utils.py +++ /dev/null @@ -1,1331 +0,0 @@ -#!/usr/bin/env python3 -""" Utility functions for the GUI """ -import logging -import os -import platform -import sys -import tkinter as tk - -from tkinter import filedialog -from threading import Event, Thread -from queue import Queue - -import numpy as np - -from PIL import Image, ImageDraw, ImageTk - -from ._config import Config as UserConfig -from .project import Project, Tasks -from .theme import Style - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name -_CONFIG = None -_IMAGES = None -_PREVIEW_TRIGGER = None -PATHCACHE = os.path.join(os.path.realpath(os.path.dirname(sys.argv[0])), "lib", "gui", ".cache") - - -def initialize_config(root, cli_opts, statusbar): - """ Initialize the GUI Master :class:`Config` and add to global constant. - - This should only be called once on first GUI startup. Future access to :class:`Config` - should only be executed through :func:`get_config`. - - Parameters - ---------- - root: :class:`tkinter.Tk` - The root Tkinter object - cli_opts: :class:`lib.gui.options.CliOpts` - The command line options object - statusbar: :class:`lib.gui.custom_widgets.StatusBar` - The GUI Status bar - """ - global _CONFIG # pylint: disable=global-statement - if _CONFIG is not None: - return None - logger.debug("Initializing config: (root: %s, cli_opts: %s, " - "statusbar: %s)", root, cli_opts, statusbar) - _CONFIG = Config(root, cli_opts, statusbar) - return _CONFIG - - -def get_config(): - """ Get the Master GUI configuration. - - Returns - ------- - :class:`Config` - The Master GUI Config - """ - return _CONFIG - - -def initialize_images(): - """ Initialize the :class:`Images` handler and add to global constant. - - This should only be called once on first GUI startup. Future access to :class:`Images` - handler should only be executed through :func:`get_images`. - """ - global _IMAGES # pylint: disable=global-statement - if _IMAGES is not None: - return - logger.debug("Initializing images") - _IMAGES = Images() - - -def get_images(): - """ Get the Master GUI Images handler. - - Returns - ------- - :class:`Images` - The Master GUI Images handler - """ - return _IMAGES - - -class FileHandler(): # pylint:disable=too-few-public-methods - """ Handles all GUI File Dialog actions and tasks. - - Parameters - ---------- - handle_type: ['open', 'save', 'filename', 'filename_multi', 'save_filename', 'context', `dir`] - The type of file dialog to return. `open` and `save` will perform the open and save actions - and return the file. `filename` returns the filename from an `open` dialog. - `filename_multi` allows for multi-selection of files and returns a list of files selected. - `save_filename` returns the filename from a `save as` dialog. `context` is a context - sensitive parameter that returns a certain dialog based on the current options. `dir` asks - for a folder location. - file_type: ['default', 'alignments', 'config_project', 'config_task', 'config_all', 'csv', \ - 'image', 'ini', 'state', 'log', 'video'] - The type of file that this dialog is for. `default` allows selection of any files. Other - options limit the file type selection - title: str, optional - The title to display on the file dialog. If `None` then the default title will be used. - Default: ``None`` - initial_folder: str, optional - The folder to initially open with the file dialog. If `None` then tkinter will decide. - Default: ``None`` - initial_file: str, optional - The filename to set with the file dialog. If `None` then tkinter no initial filename is. - specified. Default: ``None`` - command: str, optional - Required for context handling file dialog, otherwise unused. Default: ``None`` - action: str, optional - Required for context handling file dialog, otherwise unused. Default: ``None`` - variable: :class:`tkinter.StringVar`, optional - Required for context handling file dialog, otherwise unused. The variable to associate - with this file dialog. Default: ``None`` - - Attributes - ---------- - return_file: str or object - The return value from the file dialog - - Example - ------- - >>> handler = FileHandler('filename', 'video', title='Select a video...') - >>> video_file = handler.return_file - >>> print(video_file) - '/path/to/selected/video.mp4' - """ - - def __init__(self, handle_type, file_type, title=None, initial_folder=None, initial_file=None, - command=None, action=None, variable=None): - logger.debug("Initializing %s: (handle_type: '%s', file_type: '%s', title: '%s', " - "initial_folder: '%s', initial_file: '%s', command: '%s', action: '%s', " - "variable: %s)", self.__class__.__name__, handle_type, file_type, title, - initial_folder, initial_file, command, action, variable) - self._handletype = handle_type - self._dummy_master = self._set_dummy_master() - self._defaults = self._set_defaults() - self._kwargs = self._set_kwargs(title, - initial_folder, - initial_file, - file_type, - command, - action, - variable) - self.return_file = getattr(self, f"_{self._handletype.lower()}")() - self._remove_dummy_master() - - logger.debug("Initialized %s", self.__class__.__name__) - - @property - def _filetypes(self): - """ dict: The accepted extensions for each file type for opening/saving """ - all_files = ("All files", "*.*") - filetypes = dict( - default=(all_files,), - alignments=[("Faceswap Alignments", "*.fsa"), all_files], - config_project=[("Faceswap Project files", "*.fsw"), all_files], - config_task=[("Faceswap Task files", "*.fst"), all_files], - config_all=[("Faceswap Project and Task files", "*.fst *.fsw"), all_files], - csv=[("Comma separated values", "*.csv"), all_files], - image=[("Bitmap", "*.bmp"), - ("JPG", "*.jpeg *.jpg"), - ("PNG", "*.png"), - ("TIFF", "*.tif *.tiff"), - all_files], - ini=[("Faceswap config files", "*.ini"), all_files], - json=[("JSON file", "*.json"), all_files], - model=[("Keras model files", "*.h5"), all_files], - state=[("State files", "*.json"), all_files], - log=[("Log files", "*.log"), all_files], - video=[("Audio Video Interleave", "*.avi"), - ("Flash Video", "*.flv"), - ("Matroska", "*.mkv"), - ("MOV", "*.mov"), - ("MP4", "*.mp4"), - ("MPEG", "*.mpeg *.mpg *.ts *.vob"), - ("WebM", "*.webm"), - ("Windows Media Video", "*.wmv"), - all_files]) - - # Add in multi-select options and upper case extensions for Linux - for key in filetypes: - if platform.system() == "Linux": - filetypes[key] = [item - if item[0] == "All files" - else (item[0], f"{item[1]} {item[1].upper()}") - for item in filetypes[key]] - if len(filetypes[key]) > 2: - multi = [f"{key.title()} Files"] - multi.append(" ".join([ftype[1] - for ftype in filetypes[key] if ftype[0] != "All files"])) - filetypes[key].insert(0, tuple(multi)) - return filetypes - - @property - def _contexts(self): - """dict: Mapping of commands, actions and their corresponding file dialog for context - handle types. """ - return dict(effmpeg=dict(input={"extract": "filename", - "gen-vid": "dir", - "get-fps": "filename", - "get-info": "filename", - "mux-audio": "filename", - "rescale": "filename", - "rotate": "filename", - "slice": "filename"}, - output={"extract": "dir", - "gen-vid": "save_filename", - "get-fps": "nothing", - "get-info": "nothing", - "mux-audio": "save_filename", - "rescale": "save_filename", - "rotate": "save_filename", - "slice": "save_filename"})) - - @classmethod - def _set_dummy_master(cls): - """ Add an option to force black font on Linux file dialogs KDE issue that displays light - font on white background). - - This is a pretty hacky solution, but tkinter does not allow direct editing of file dialogs, - so we create a dummy frame and add the foreground option there, so that the file dialog can - inherit the foreground. - - Returns - ------- - tkinter.Frame or ``None`` - The dummy master frame for Linux systems, otherwise ``None`` - """ - if platform.system().lower() == "linux": - retval = tk.Frame() - retval.option_add("*foreground", "black") - else: - retval = None - return retval - - def _remove_dummy_master(self): - """ Destroy the dummy master widget on Linux systems. """ - if platform.system().lower() != "linux": - return - self._dummy_master.destroy() - del self._dummy_master - self._dummy_master = None - - def _set_defaults(self): - """ Set the default file type for the file dialog. Generally the first found file type - will be used, but this is overridden if it is not appropriate. - - Returns - ------- - dict: - The default file extension for each file type - """ - defaults = {key: next(ext for ext in val[0][1].split(" ")).replace("*", "") - for key, val in self._filetypes.items()} - defaults["default"] = None - defaults["video"] = ".mp4" - defaults["image"] = ".png" - logger.debug(defaults) - return defaults - - def _set_kwargs(self, title, initial_folder, initial_file, file_type, command, action, - variable=None): - """ Generate the required kwargs for the requested file dialog browser. - - Parameters - ---------- - title: str - The title to display on the file dialog. If `None` then the default title will be used. - initial_folder: str - The folder to initially open with the file dialog. If `None` then tkinter will decide. - initial_file: str - The filename to set with the file dialog. If `None` then tkinter no initial filename - is. - file_type: ['default', 'alignments', 'config_project', 'config_task', 'config_all', \ - 'csv', 'image', 'ini', 'state', 'log', 'video'] - The type of file that this dialog is for. `default` allows selection of any files. - Other options limit the file type selection - command: str - Required for context handling file dialog, otherwise unused. - action: str - Required for context handling file dialog, otherwise unused. - variable: :class:`tkinter.StringVar`, optional - Required for context handling file dialog, otherwise unused. The variable to associate - with this file dialog. Default: ``None`` - - Returns - ------- - dict: - The key word arguments for the file dialog to be launched - """ - logger.debug("Setting Kwargs: (title: %s, initial_folder: %s, initial_file: '%s', " - "file_type: '%s', command: '%s': action: '%s', variable: '%s')", - title, initial_folder, initial_file, file_type, command, action, variable) - - kwargs = dict(master=self._dummy_master) - - if self._handletype.lower() == "context": - self._set_context_handletype(command, action, variable) - - if title is not None: - kwargs["title"] = title - - if initial_folder is not None: - kwargs["initialdir"] = initial_folder - - if initial_file is not None: - kwargs["initialfile"] = initial_file - - if self._handletype.lower() in ( - "open", "save", "filename", "filename_multi", "save_filename"): - kwargs["filetypes"] = self._filetypes[file_type] - if self._defaults.get(file_type): - kwargs['defaultextension'] = self._defaults[file_type] - if self._handletype.lower() == "save": - kwargs["mode"] = "w" - if self._handletype.lower() == "open": - kwargs["mode"] = "r" - logger.debug("Set Kwargs: %s", kwargs) - return kwargs - - def _set_context_handletype(self, command, action, variable): - """ Sets the correct handle type based on context. - - Parameters - ---------- - command: str - The command that is being executed. Used to look up the context actions - action: str - The action that is being performed. Used to look up the correct file dialog - variable: :class:`tkinter.StringVar` - The variable associated with this file dialog - """ - if self._contexts[command].get(variable, None) is not None: - handletype = self._contexts[command][variable][action] - else: - handletype = self._contexts[command][action] - logger.debug(handletype) - self._handletype = handletype - - def _open(self): - """ Open a file. """ - logger.debug("Popping Open browser") - return filedialog.askopenfile(**self._kwargs) - - def _save(self): - """ Save a file. """ - logger.debug("Popping Save browser") - return filedialog.asksaveasfile(**self._kwargs) - - def _dir(self): - """ Get a directory location. """ - logger.debug("Popping Dir browser") - return filedialog.askdirectory(**self._kwargs) - - def _savedir(self): - """ Get a save directory location. """ - logger.debug("Popping SaveDir browser") - return filedialog.askdirectory(**self._kwargs) - - def _filename(self): - """ Get an existing file location. """ - logger.debug("Popping Filename browser") - return filedialog.askopenfilename(**self._kwargs) - - def _filename_multi(self): - """ Get multiple existing file locations. """ - logger.debug("Popping Filename browser") - return filedialog.askopenfilenames(**self._kwargs) - - def _save_filename(self): - """ Get a save file location. """ - logger.debug("Popping Save Filename browser") - return filedialog.asksaveasfilename(**self._kwargs) - - @staticmethod - def _nothing(): # pylint: disable=useless-return - """ Method that does nothing, used for disabling open/save pop up. """ - logger.debug("Popping Nothing browser") - return - - -class Images(): - """ The centralized image repository for holding all icons and images required by the GUI. - - This class should be initialized on GUI startup through :func:`initialize_images`. Any further - access to this class should be through :func:`get_images`. - """ - def __init__(self): - logger.debug("Initializing %s", self.__class__.__name__) - self._pathpreview = os.path.join(PATHCACHE, "preview") - self._pathoutput = None - self._previewoutput = None - self._previewtrain = {} - self._previewcache = dict(modified=None, # cache for extract and convert - images=None, - filenames=[], - placeholder=None) - self._errcount = 0 - self._icons = self._load_icons() - logger.debug("Initialized %s", self.__class__.__name__) - - @property - def previewoutput(self): - """ Tuple or ``None``: First item in the tuple is the extract or convert preview image - (:class:`PIL.Image`), the second item is the image in a format that tkinter can display - (:class:`PIL.ImageTK.PhotoImage`). - - The value of the property is ``None`` if no extract or convert task is running or there are - no files available in the output folder. """ - return self._previewoutput - - @property - def previewtrain(self): - """ dict or ``None``: The training preview images. Dictionary key is the image name - (`str`). Dictionary values are a `list` of the training image (:class:`PIL.Image`), the - image formatted for tkinter display (:class:`PIL.ImageTK.PhotoImage`), the last - modification time of the image (`float`). - - The value of this property is ``None`` if training is not running or there are no preview - images available. - """ - return self._previewtrain - - @property - def icons(self): - """ dict: The faceswap icons for all parts of the GUI. The dictionary key is the icon - name (`str`) the value is the icon sized and formatted for display - (:class:`PIL.ImageTK.PhotoImage`). - - Example - ------- - >>> icons = get_images().icons - >>> save = icons["save"] - >>> button = ttk.Button(parent, image=save) - >>> button.pack() - """ - return self._icons - - @staticmethod - def _load_icons(): - """ Scan the icons cache folder and load the icons into :attr:`icons` for retrieval - throughout the GUI. - - Returns - ------- - dict: - The icons formatted as described in :attr:`icons` - - """ - size = get_config().user_config_dict.get("icon_size", 16) - size = int(round(size * get_config().scaling_factor)) - icons = {} - pathicons = os.path.join(PATHCACHE, "icons") - for fname in os.listdir(pathicons): - name, ext = os.path.splitext(fname) - if ext != ".png": - continue - img = Image.open(os.path.join(pathicons, fname)) - img = ImageTk.PhotoImage(img.resize((size, size), resample=Image.HAMMING)) - icons[name] = img - logger.debug(icons) - return icons - - def set_faceswap_output_path(self, location): - """ Set the path that will contain the output from an Extract or Convert task. - - Required so that the GUI can fetch output images to display for return in - :attr:`previewoutput`. - - Parameters - ---------- - location: str - The output location that has been specified for an Extract or Convert task - """ - self._pathoutput = location - - def delete_preview(self): - """ Delete the preview files in the cache folder and reset the image cache. - - Should be called when terminating tasks, or when Faceswap starts up or shuts down. - """ - logger.debug("Deleting previews") - for item in os.listdir(self._pathpreview): - if item.startswith(".gui_training_preview") and item.endswith(".jpg"): - fullitem = os.path.join(self._pathpreview, item) - logger.debug("Deleting: '%s'", fullitem) - os.remove(fullitem) - for fname in self._previewcache["filenames"]: - if os.path.basename(fname) == ".gui_preview.jpg": - logger.debug("Deleting: '%s'", fname) - try: - os.remove(fname) - except FileNotFoundError: - logger.debug("File does not exist: %s", fname) - self._clear_image_cache() - - def _clear_image_cache(self): - """ Clear all cached images. """ - logger.debug("Clearing image cache") - self._pathoutput = None - self._previewoutput = None - self._previewtrain = {} - self._previewcache = dict(modified=None, # cache for extract and convert - images=None, - filenames=[], - placeholder=None) - - @staticmethod - def _get_images(image_path): - """ Get the images stored within the given directory. - - Parameters - ---------- - image_path: str - The folder containing images to be scanned - - Returns - ------- - list: - The image filenames stored within the given folder - - """ - logger.debug("Getting images: '%s'", image_path) - if not os.path.isdir(image_path): - logger.debug("Folder does not exist") - return None - files = [os.path.join(image_path, f) - for f in os.listdir(image_path) if f.lower().endswith((".png", ".jpg"))] - logger.debug("Image files: %s", files) - return files - - def load_latest_preview(self, thumbnail_size, frame_dims): - """ Load the latest preview image for extract and convert. - - Retrieves the latest preview images from the faceswap output folder, resizes to thumbnails - and lays out for display. Places the images into :attr:`previewoutput` for loading into - the display panel. - - Parameters - ---------- - thumbnail_size: int - The size of each thumbnail that should be created - frame_dims: tuple - The (width (`int`), height (`int`)) of the display panel that will display the preview - """ - logger.debug("Loading preview image: (thumbnail_size: %s, frame_dims: %s)", - thumbnail_size, frame_dims) - image_files = self._get_images(self._pathoutput) - gui_preview = os.path.join(self._pathoutput, ".gui_preview.jpg") - if not image_files or (len(image_files) == 1 and gui_preview not in image_files): - logger.debug("No preview to display") - return - # Filter to just the gui_preview if it exists in folder output - image_files = [gui_preview] if gui_preview in image_files else image_files - logger.debug("Image Files: %s", len(image_files)) - - image_files = self._get_newest_filenames(image_files) - if not image_files: - return - - if not self._load_images_to_cache(image_files, frame_dims, thumbnail_size): - logger.debug("Failed to load any preview images") - if gui_preview in image_files: - # Reset last modified for failed loading of a gui preview image so it is picked - # up next time - self._previewcache["modified"] = None - return - - if image_files == [gui_preview]: - # Delete the preview image so that the main scripts know to output another - logger.debug("Deleting preview image") - os.remove(image_files[0]) - show_image = self._place_previews(frame_dims) - if not show_image: - self._previewoutput = None - return - logger.debug("Displaying preview: %s", self._previewcache["filenames"]) - self._previewoutput = (show_image, ImageTk.PhotoImage(show_image)) - - def _get_newest_filenames(self, image_files): - """ Return image filenames that have been modified since the last check. - - Parameters - ---------- - image_files: list - The list of image files to check the modification date for - - Returns - ------- - list: - A list of images that have been modified since the last check - """ - if self._previewcache["modified"] is None: - retval = image_files - else: - retval = [fname for fname in image_files - if os.path.getmtime(fname) > self._previewcache["modified"]] - if not retval: - logger.debug("No new images in output folder") - else: - self._previewcache["modified"] = max([os.path.getmtime(img) for img in retval]) - logger.debug("Number new images: %s, Last Modified: %s", - len(retval), self._previewcache["modified"]) - return retval - - def _load_images_to_cache(self, image_files, frame_dims, thumbnail_size): - """ Load preview images to the image cache. - - Load new images and append to cache, filtering the cache the number of thumbnails that will - fit inside the display panel. - - Parameters - ---------- - image_files: list - A list of new image files that have been modified since the last check - frame_dims: tuple - The (width (`int`), height (`int`)) of the display panel that will display the preview - thumbnail_size: int - The size of each thumbnail that should be created - - Returns - ------- - bool - ``True`` if images were successfully loaded to cache otherwise ``False`` - """ - logger.debug("Number image_files: %s, frame_dims: %s, thumbnail_size: %s", - len(image_files), frame_dims, thumbnail_size) - num_images = (frame_dims[0] // thumbnail_size) * (frame_dims[1] // thumbnail_size) - logger.debug("num_images: %s", num_images) - if num_images == 0: - return False - samples = [] - start_idx = len(image_files) - num_images if len(image_files) > num_images else 0 - show_files = sorted(image_files, key=os.path.getctime)[start_idx:] - dropped_files = [] - for fname in show_files: - try: - img = Image.open(fname) - except PermissionError as err: - logger.debug("Permission error opening preview file: '%s'. Original error: %s", - fname, str(err)) - dropped_files.append(fname) - continue - except Exception as err: # pylint:disable=broad-except - # Swallow any issues with opening an image rather than spamming console - # Can happen when trying to read partially saved images - logger.debug("Error opening preview file: '%s'. Original error: %s", - fname, str(err)) - dropped_files.append(fname) - continue - - width, height = img.size - scaling = thumbnail_size / max(width, height) - logger.debug("image width: %s, height: %s, scaling: %s", width, height, scaling) - - try: - img = img.resize((int(width * scaling), int(height * scaling))) - except OSError as err: - # Image only gets loaded when we call a method, so may error on partial loads - logger.debug("OS Error resizing preview image: '%s'. Original error: %s", - fname, err) - dropped_files.append(fname) - continue - - if img.size[0] != img.size[1]: - # Pad to square - new_img = Image.new("RGB", (thumbnail_size, thumbnail_size)) - new_img.paste(img, ((thumbnail_size - img.size[0])//2, - (thumbnail_size - img.size[1])//2)) - img = new_img - draw = ImageDraw.Draw(img) - draw.rectangle(((0, 0), (thumbnail_size, thumbnail_size)), outline="#E5E5E5", width=1) - samples.append(np.array(img)) - - samples = np.array(samples) - if not np.any(samples): - logger.debug("No preview images collected.") - return False - - if dropped_files: - logger.debug("Removing dropped files: %s", dropped_files) - show_files = [fname for fname in show_files if fname not in dropped_files] - - self._previewcache["filenames"] = (self._previewcache["filenames"] + - show_files)[-num_images:] - cache = self._previewcache["images"] - if cache is None: - logger.debug("Creating new cache") - cache = samples[-num_images:] - else: - logger.debug("Appending to existing cache") - cache = np.concatenate((cache, samples))[-num_images:] - self._previewcache["images"] = cache - logger.debug("Cache shape: %s", self._previewcache["images"].shape) - return True - - def _place_previews(self, frame_dims): - """ Format the preview thumbnails stored in the cache into a grid fitting the display - panel. - - Parameters - ---------- - frame_dims: tuple - The (width (`int`), height (`int`)) of the display panel that will display the preview - - Returns - :class:`PIL.Image`: - The final preview display image - """ - if self._previewcache.get("images", None) is None: - logger.debug("No images in cache. Returning None") - return None - samples = self._previewcache["images"].copy() - num_images, thumbnail_size = samples.shape[:2] - if self._previewcache["placeholder"] is None: - self._create_placeholder(thumbnail_size) - - logger.debug("num_images: %s, thumbnail_size: %s", num_images, thumbnail_size) - cols, rows = frame_dims[0] // thumbnail_size, frame_dims[1] // thumbnail_size - logger.debug("cols: %s, rows: %s", cols, rows) - if cols == 0 or rows == 0: - logger.debug("Cols or Rows is zero. No items to display") - return None - remainder = (cols * rows) - num_images - if remainder != 0: - logger.debug("Padding sample display. Remainder: %s", remainder) - placeholder = np.concatenate([np.expand_dims(self._previewcache["placeholder"], - 0)] * remainder) - samples = np.concatenate((samples, placeholder)) - - display = np.vstack([np.hstack(samples[row * cols: (row + 1) * cols]) - for row in range(rows)]) - logger.debug("display shape: %s", display.shape) - return Image.fromarray(display) - - def _create_placeholder(self, thumbnail_size): - """ Create a placeholder image for when there are fewer thumbnails available - than columns to display them. - - Parameters - ---------- - thumbnail_size: int - The size of the thumbnail that the placeholder should replicate - """ - logger.debug("Creating placeholder. thumbnail_size: %s", thumbnail_size) - placeholder = Image.new("RGB", (thumbnail_size, thumbnail_size)) - draw = ImageDraw.Draw(placeholder) - draw.rectangle(((0, 0), (thumbnail_size, thumbnail_size)), outline="#E5E5E5", width=1) - placeholder = np.array(placeholder) - self._previewcache["placeholder"] = placeholder - logger.debug("Created placeholder. shape: %s", placeholder.shape) - - def load_training_preview(self): - """ Load the training preview images. - - Reads the training image currently stored in the cache folder and loads them to - :attr:`previewtrain` for retrieval in the GUI. - """ - logger.debug("Loading Training preview images") - image_files = self._get_images(self._pathpreview) - modified = None - if not image_files: - logger.debug("No preview to display") - self._previewtrain = {} - return - for img in image_files: - modified = os.path.getmtime(img) if modified is None else modified - name = os.path.basename(img) - name = os.path.splitext(name)[0] - name = name[name.rfind("_") + 1:].title() - try: - logger.debug("Displaying preview: '%s'", img) - size = self._get_current_size(name) - self._previewtrain[name] = [Image.open(img), None, modified] - self.resize_image(name, size) - self._errcount = 0 - except ValueError: - # This is probably an error reading the file whilst it's - # being saved so ignore it for now and only pick up if - # there have been multiple consecutive fails - logger.warning("Unable to display preview: (image: '%s', attempt: %s)", - img, self._errcount) - if self._errcount < 10: - self._errcount += 1 - else: - logger.error("Error reading the preview file for '%s'", img) - print(f"Error reading the preview file for {name}") - self._previewtrain[name] = None - - def _get_current_size(self, name): - """ Return the size of the currently displayed training preview image. - - Parameters - ---------- - name: str - The name of the training image to get the size for - - Returns - ------- - width: int - The width of the training image - height: int - The height of the training image - """ - logger.debug("Getting size: '%s'", name) - if not self._previewtrain.get(name, None): - return None - img = self._previewtrain[name][1] - if not img: - return None - logger.debug("Got size: (name: '%s', width: '%s', height: '%s')", - name, img.width(), img.height()) - return img.width(), img.height() - - def resize_image(self, name, frame_dims): - """ Resize the training preview image based on the passed in frame size. - - If the canvas that holds the preview image changes, update the image size - to fit the new canvas and refresh :attr:`previewtrain`. - - Parameters - ---------- - name: str - The name of the training image to be resized - frame_dims: tuple - The (width (`int`), height (`int`)) of the display panel that will display the preview - """ - logger.debug("Resizing image: (name: '%s', frame_dims: %s", name, frame_dims) - displayimg = self._previewtrain[name][0] - if frame_dims: - frameratio = float(frame_dims[0]) / float(frame_dims[1]) - imgratio = float(displayimg.size[0]) / float(displayimg.size[1]) - - if frameratio <= imgratio: - scale = frame_dims[0] / float(displayimg.size[0]) - size = (frame_dims[0], int(displayimg.size[1] * scale)) - else: - scale = frame_dims[1] / float(displayimg.size[1]) - size = (int(displayimg.size[0] * scale), frame_dims[1]) - logger.debug("Scaling: (scale: %s, size: %s", scale, size) - - # Hacky fix to force a reload if it happens to find corrupted - # data, probably due to reading the image whilst it is partially - # saved. If it continues to fail, then eventually raise. - for i in range(0, 1000): - try: - displayimg = displayimg.resize(size, Image.ANTIALIAS) - except OSError: - if i == 999: - raise - continue - break - self._previewtrain[name][1] = ImageTk.PhotoImage(displayimg) - - -class Config(): - """ The centralized configuration class for holding items that should be made available to all - parts of the GUI. - - This class should be initialized on GUI startup through :func:`initialize_config`. Any further - access to this class should be through :func:`get_config`. - - Parameters - ---------- - root: :class:`tkinter.Tk` - The root Tkinter object - cli_opts: :class:`lib.gui.options.CliOpts` - The command line options object - statusbar: :class:`lib.gui.custom_widgets.StatusBar` - The GUI Status bar - """ - def __init__(self, root, cli_opts, statusbar): - logger.debug("Initializing %s: (root %s, cli_opts: %s, statusbar: %s)", - self.__class__.__name__, root, cli_opts, statusbar) - self._default_font = tk.font.nametofont("TkDefaultFont").configure()["family"] - self._constants = dict( - root=root, - scaling_factor=self._get_scaling(root), - default_font=self._default_font) - self._gui_objects = dict( - cli_opts=cli_opts, - tk_vars=self._set_tk_vars(), - project=Project(self, FileHandler), - tasks=Tasks(self, FileHandler), - default_options=None, - status_bar=statusbar, - command_notebook=None) # set in command.py - self._user_config = UserConfig(None) - self._style = Style(self.default_font, root, PATHCACHE) - self._user_theme = self._style.user_theme - logger.debug("Initialized %s", self.__class__.__name__) - - # Constants - @property - def root(self): - """ :class:`tkinter.Tk`: The root tkinter window. """ - return self._constants["root"] - - @property - def scaling_factor(self): - """ float: The scaling factor for current display. """ - return self._constants["scaling_factor"] - - @property - def pathcache(self): - """ str: The path to the GUI cache folder """ - return PATHCACHE - - # GUI Objects - @property - def cli_opts(self): - """ :class:`lib.gui.options.CliOptions`: The command line options for this GUI Session. """ - return self._gui_objects["cli_opts"] - - @property - def tk_vars(self): - """ dict: The global tkinter variables. """ - return self._gui_objects["tk_vars"] - - @property - def project(self): - """ :class:`lib.gui.project.Project`: The project session handler. """ - return self._gui_objects["project"] - - @property - def tasks(self): - """ :class:`lib.gui.project.Tasks`: The session tasks handler. """ - return self._gui_objects["tasks"] - - @property - def default_options(self): - """ dict: The default options for all tabs """ - return self._gui_objects["default_options"] - - @property - def statusbar(self): - """ :class:`lib.gui.custom_widgets.StatusBar`: The GUI StatusBar - :class:`tkinter.ttk.Frame`. """ - return self._gui_objects["status_bar"] - - @property - def command_notebook(self): - """ :class:`lib.gui.command.CommandNoteboook`: The main Faceswap Command Notebook. """ - return self._gui_objects["command_notebook"] - - # Convenience GUI Objects - @property - def tools_notebook(self): - """ :class:`lib.gui.command.ToolsNotebook`: The Faceswap Tools sub-Notebook. """ - return self.command_notebook.tools_notebook - - @property - def modified_vars(self): - """ dict: The command notebook modified tkinter variables. """ - return self.command_notebook.modified_vars - - @property - def _command_tabs(self): - """ dict: Command tab titles with their IDs. """ - return self.command_notebook.tab_names - - @property - def _tools_tabs(self): - """ dict: Tools command tab titles with their IDs. """ - return self.command_notebook.tools_tab_names - - # Config - @property - def user_config(self): - """ dict: The GUI config in dict form. """ - return self._user_config - - @property - def user_config_dict(self): - """ dict: The GUI config in dict form. """ - return self._user_config.config_dict - - @property - def user_theme(self): - """ dict: The GUI theme selection options. """ - return self._user_theme - - @property - def default_font(self): - """ tuple: The selected font as configured in user settings. First item is the font (`str`) - second item the font size (`int`). """ - font = self.user_config_dict["font"] - font = self._default_font if font == "default" else font - return (font, self.user_config_dict["font_size"]) - - @staticmethod - def _get_scaling(root): - """ Get the display DPI. - - Returns - ------- - float: - The scaling factor - """ - dpi = root.winfo_fpixels("1i") - scaling = dpi / 72.0 - logger.debug("dpi: %s, scaling: %s'", dpi, scaling) - return scaling - - def set_default_options(self): - """ Set the default options for :mod:`lib.gui.projects` - - The Default GUI options are stored on Faceswap startup. - - Exposed as the :attr:`_default_opts` for a project cannot be set until after the main - Command Tabs have been loaded. - """ - default = self.cli_opts.get_option_values() - logger.debug(default) - self._gui_objects["default_options"] = default - self.project.set_default_options() - - def set_command_notebook(self, notebook): - """ Set the command notebook to the :attr:`command_notebook` attribute - and enable the modified callback for :attr:`project`. - - Parameters - ---------- - notebook: :class:`lib.gui.command.CommandNotebook` - The main command notebook for the Faceswap GUI - """ - logger.debug("Setting commane notebook: %s", notebook) - self._gui_objects["command_notebook"] = notebook - self.project.set_modified_callback() - - def set_active_tab_by_name(self, name): - """ Sets the :attr:`command_notebook` or :attr:`tools_notebook` to active based on given - name. - - Parameters - ---------- - name: str - The name of the tab to set active - """ - name = name.lower() - if name in self._command_tabs: - tab_id = self._command_tabs[name] - logger.debug("Setting active tab to: (name: %s, id: %s)", name, tab_id) - self.command_notebook.select(tab_id) - elif name in self._tools_tabs: - self.command_notebook.select(self._command_tabs["tools"]) - tab_id = self._tools_tabs[name] - logger.debug("Setting active Tools tab to: (name: %s, id: %s)", name, tab_id) - self.tools_notebook.select() - else: - logger.debug("Name couldn't be found. Setting to id 0: %s", name) - self.command_notebook.select(0) - - def set_modified_true(self, command): - """ Set the modified variable to ``True`` for the given command in :attr:`modified_vars`. - - Parameters - ---------- - command: str - The command to set the modified state to ``True`` - - """ - tkvar = self.modified_vars.get(command, None) - if tkvar is None: - logger.debug("No tkvar for command: '%s'", command) - return - tkvar.set(True) - logger.debug("Set modified var to True for: '%s'", command) - - def refresh_config(self): - """ Reload the user config from file. """ - self._user_config = UserConfig(None) - - def set_cursor_busy(self, widget=None): - """ Set the root or widget cursor to busy. - - Parameters - ---------- - widget: tkinter object, optional - The widget to set busy cursor for. If the provided value is ``None`` then sets the - cursor busy for the whole of the GUI. Default: ``None``. - """ - logger.debug("Setting cursor to busy. widget: %s", widget) - widget = self.root if widget is None else widget - widget.config(cursor="watch") - widget.update_idletasks() - - def set_cursor_default(self, widget=None): - """ Set the root or widget cursor to default. - - Parameters - ---------- - widget: tkinter object, optional - The widget to set default cursor for. If the provided value is ``None`` then sets the - cursor busy for the whole of the GUI. Default: ``None`` - """ - logger.debug("Setting cursor to default. widget: %s", widget) - widget = self.root if widget is None else widget - widget.config(cursor="") - widget.update_idletasks() - - @staticmethod - def _set_tk_vars(): - """ Set the global tkinter variables stored for easy access in :class:`Config`. - - The variables are available through :attr:`tk_vars`. - """ - display = tk.StringVar() - display.set(None) - - runningtask = tk.BooleanVar() - runningtask.set(False) - - istraining = tk.BooleanVar() - istraining.set(False) - - actioncommand = tk.StringVar() - actioncommand.set(None) - - generatecommand = tk.StringVar() - generatecommand.set(None) - - console_clear = tk.BooleanVar() - console_clear.set(False) - - refreshgraph = tk.BooleanVar() - refreshgraph.set(False) - - updatepreview = tk.BooleanVar() - updatepreview.set(False) - - analysis_folder = tk.StringVar() - analysis_folder.set(None) - - tk_vars = dict(display=display, - runningtask=runningtask, - istraining=istraining, - action=actioncommand, - generate=generatecommand, - console_clear=console_clear, - refreshgraph=refreshgraph, - updatepreview=updatepreview, - analysis_folder=analysis_folder) - logger.debug(tk_vars) - return tk_vars - - def set_root_title(self, text=None): - """ Set the main title text for Faceswap. - - The title will always begin with 'Faceswap.py'. Additional text can be appended. - - Parameters - ---------- - text: str, optional - Additional text to be appended to the GUI title bar. Default: ``None`` - """ - title = "Faceswap.py" - title += f" - {text}" if text is not None and text else "" - self.root.title(title) - - def set_geometry(self, width, height, fullscreen=False): - """ Set the geometry for the root tkinter object. - - Parameters - ---------- - width: int - The width to set the window to (prior to scaling) - height: int - The height to set the window to (prior to scaling) - fullscreen: bool, optional - Whether to set the window to full-screen mode. If ``True`` then :attr:`width` and - :attr:`height` are ignored. Default: ``False`` - """ - self.root.tk.call("tk", "scaling", self.scaling_factor) - if fullscreen: - initial_dimensions = (self.root.winfo_screenwidth(), self.root.winfo_screenheight()) - else: - initial_dimensions = (round(width * self.scaling_factor), - round(height * self.scaling_factor)) - - if fullscreen and sys.platform in ("win32", "darwin"): - self.root.state('zoomed') - elif fullscreen: - self.root.attributes('-zoomed', True) - else: - self.root.geometry(f"{str(initial_dimensions[0])}x{str(initial_dimensions[1])}+80+80") - logger.debug("Geometry: %sx%s", *initial_dimensions) - - -class LongRunningTask(Thread): - """ Runs long running tasks in a background thread to prevent the GUI from becoming - unresponsive. - - This is sub-classed from :class:`Threading.Thread` so check documentation there for base - parameters. Additional parameters listed below. - - Parameters - ---------- - widget: tkinter object, optional - The widget that this :class:`LongRunningTask` is associated with. Used for setting the busy - cursor in the correct location. Default: ``None``. - """ - def __init__(self, group=None, target=None, name=None, args=(), kwargs=None, *, daemon=True, - widget=None): - logger.debug("Initializing %s: (group: %s, target: %s, name: %s, args: %s, kwargs: %s, " - "daemon: %s)", self.__class__.__name__, group, target, name, args, kwargs, - daemon) - super().__init__(group=group, target=target, name=name, args=args, kwargs=kwargs, - daemon=daemon) - self.err = None - self._widget = widget - self._config = get_config() - self._config.set_cursor_busy(widget=self._widget) - self._complete = Event() - self._queue = Queue() - logger.debug("Initialized %s", self.__class__.__name__,) - - @property - def complete(self): - """ :class:`threading.Event`: Event is set if the thread has completed its task, - otherwise it is unset. - """ - return self._complete - - def run(self): - """ Commence the given task in a background thread. """ - try: - if self._target: - retval = self._target(*self._args, **self._kwargs) - self._queue.put(retval) - except Exception: # pylint: disable=broad-except - self.err = sys.exc_info() - logger.debug("Error in thread (%s): %s", self._name, - self.err[1].with_traceback(self.err[2])) - finally: - self._complete.set() - # Avoid a ref-cycle if the thread is running a function with - # an argument that has a member that points to the thread. - del self._target, self._args, self._kwargs - - def get_result(self): - """ Return the result from the given task. - - Returns - ------- - varies: - The result of the thread will depend on the given task. If a call is made to - :func:`get_result` prior to the thread completing its task then ``None`` will be - returned - """ - if not self._complete.is_set(): - logger.warning("Aborting attempt to retrieve result from a LongRunningTask that is " - "still running") - return None - if self.err: - logger.debug("Error caught in thread") - self._config.set_cursor_default(widget=self._widget) - raise self.err[1].with_traceback(self.err[2]) - - logger.debug("Getting result from thread") - retval = self._queue.get() - logger.debug("Got result from thread") - self._config.set_cursor_default(widget=self._widget) - return retval - - -class PreviewTrigger(): - """ Triggers to indicate to underlying Faceswap process that the preview image should - be updated. - - Writes a file to the cache folder that is picked up by the main process. - """ - def __init__(self): - logger.debug("Initializing: %s", self.__class__.__name__) - self._trigger_files = dict(update=os.path.join(PATHCACHE, ".preview_trigger"), - mask_toggle=os.path.join(PATHCACHE, ".preview_mask_toggle")) - logger.debug("Initialized: %s (trigger_files: %s)", - self.__class__.__name__, self._trigger_files) - - def set(self, trigger_type): - """ Place the trigger file into the cache folder - - Parameters - ---------- - trigger_type: ["update", "mask_toggle"] - The type of action to trigger. 'update': Full preview update. 'mask_toggle': toggle - mask on and off - """ - trigger = self._trigger_files[trigger_type] - if not os.path.isfile(trigger): - with open(trigger, "w", encoding="utf8"): - pass - logger.debug("Set preview trigger: %s", trigger) - - def clear(self, trigger_type=None): - """ Remove the trigger file from the cache folder. - - Parameters - ---------- - trigger_type: ["update", "mask_toggle", ``None``], optional - The trigger to clear. 'update': Full preview update. 'mask_toggle': toggle mask on - and off. ``None`` - clear all triggers. Default: ``None`` - """ - if trigger_type is None: - triggers = list(self._trigger_files.values()) - else: - triggers = [self._trigger_files[trigger_type]] - for trigger in triggers: - if os.path.isfile(trigger): - os.remove(trigger) - logger.debug("Removed preview trigger: %s", trigger) - - -def preview_trigger(): - """ Set the global preview trigger if it has not already been set and return. - - Returns - ------- - :class:`PreviewTrigger` - The trigger to indicate to the main faceswap process that it should perform a training - preview update - """ - global _PREVIEW_TRIGGER # pylint:disable=global-statement - if _PREVIEW_TRIGGER is None: - _PREVIEW_TRIGGER = PreviewTrigger() - return _PREVIEW_TRIGGER diff --git a/lib/gui/utils/__init__.py b/lib/gui/utils/__init__.py new file mode 100644 index 0000000000..24e46983d7 --- /dev/null +++ b/lib/gui/utils/__init__.py @@ -0,0 +1,7 @@ +#!/usr/bin python3 +""" Utilities for the Faceswap GUI """ + +from .config import get_config, initialize_config, PATHCACHE +from .file_handler import FileHandler +from .image import get_images, initialize_images, preview_trigger +from .misc import LongRunningTask diff --git a/lib/gui/utils/config.py b/lib/gui/utils/config.py new file mode 100644 index 0000000000..a0c2aa72f0 --- /dev/null +++ b/lib/gui/utils/config.py @@ -0,0 +1,457 @@ +#!/usr/bin python3 +""" Global configuration optiopns for the Faceswap GUI """ +from __future__ import annotations +import logging +import os +import sys +import tkinter as tk +import typing as T + +from dataclasses import dataclass, field + +from lib.gui._config import Config as UserConfig +from lib.gui.project import Project, Tasks +from lib.gui.theme import Style +from .file_handler import FileHandler + +if T.TYPE_CHECKING: + from lib.gui.options import CliOptions + from lib.gui.custom_widgets import StatusBar + from lib.gui.command import CommandNotebook + from lib.gui.command import ToolsNotebook + +logger = logging.getLogger(__name__) + +PATHCACHE = os.path.join(os.path.realpath(os.path.dirname(sys.argv[0])), "lib", "gui", ".cache") +_CONFIG: Config | None = None + + +def initialize_config(root: tk.Tk, + cli_opts: CliOptions | None, + statusbar: StatusBar | None) -> Config | None: + """ Initialize the GUI Master :class:`Config` and add to global constant. + + This should only be called once on first GUI startup. Future access to :class:`Config` + should only be executed through :func:`get_config`. + + Parameters + ---------- + root: :class:`tkinter.Tk` + The root Tkinter object + cli_opts: :class:`lib.gui.options.CliOptions` or ``None`` + The command line options object. Must be provided for main GUI. Must be ``None`` for tools + statusbar: :class:`lib.gui.custom_widgets.StatusBar` or ``None`` + The GUI Status bar. Must be provided for main GUI. Must be ``None`` for tools + + Returns + ------- + :class:`Config` or ``None`` + ``None`` if the config has already been initialized otherwise the global configuration + options + """ + global _CONFIG # pylint:disable=global-statement + if _CONFIG is not None: + return None + logger.debug("Initializing config: (root: %s, cli_opts: %s, " + "statusbar: %s)", root, cli_opts, statusbar) + _CONFIG = Config(root, cli_opts, statusbar) + return _CONFIG + + +def get_config() -> "Config": + """ Get the Master GUI configuration. + + Returns + ------- + :class:`Config` + The Master GUI Config + """ + assert _CONFIG is not None + return _CONFIG + + +class GlobalVariables(): + """ Global tkinter variables accessible from all parts of the GUI. Should only be accessed from + :attr:`get_config().tk_vars` """ + def __init__(self) -> None: + logger.debug("Initializing %s", self.__class__.__name__) + self._display = tk.StringVar() + self._running_task = tk.BooleanVar() + self._is_training = tk.BooleanVar() + self._action_command = tk.StringVar() + self._generate_command = tk.StringVar() + self._console_clear = tk.BooleanVar() + self._refresh_graph = tk.BooleanVar() + self._analysis_folder = tk.StringVar() + + self._initialize_variables() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def display(self) -> tk.StringVar: + """ :class:`tkinter.StringVar`: The current Faceswap command running """ + return self._display + + @property + def running_task(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: ``True`` if a Faceswap task is running otherwise + ``False`` """ + return self._running_task + + @property + def is_training(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: ``True`` if Faceswap is currently training otherwise + ``False`` """ + return self._is_training + + @property + def action_command(self) -> tk.StringVar: + """ :class:`tkinter.StringVar`: The command line action to perform """ + return self._action_command + + @property + def generate_command(self) -> tk.StringVar: + """ :class:`tkinter.StringVar`: The command line action to generate """ + return self._generate_command + + @property + def console_clear(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: ``True`` if the console should be cleared otherwise + ``False`` """ + return self._console_clear + + @property + def refresh_graph(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: ``True`` if the training graph should be refreshed + otherwise ``False`` """ + return self._refresh_graph + + @property + def analysis_folder(self) -> tk.StringVar: + """ :class:`tkinter.StringVar`: Full path the analysis folder""" + return self._analysis_folder + + def _initialize_variables(self) -> None: + """ Initialize the default variable values""" + self._display.set("") + self._running_task.set(False) + self._is_training.set(False) + self._action_command.set("") + self._generate_command.set("") + self._console_clear.set(False) + self._refresh_graph.set(False) + self._analysis_folder.set("") + + +@dataclass +class _GuiObjects: + """ Data class for commonly accessed GUI Objects """ + cli_opts: CliOptions | None + tk_vars: GlobalVariables + project: Project + tasks: Tasks + status_bar: StatusBar | None + default_options: dict[str, dict[str, T.Any]] = field(default_factory=dict) + command_notebook: CommandNotebook | None = None + + +class Config(): + """ The centralized configuration class for holding items that should be made available to all + parts of the GUI. + + This class should be initialized on GUI startup through :func:`initialize_config`. Any further + access to this class should be through :func:`get_config`. + + Parameters + ---------- + root: :class:`tkinter.Tk` + The root Tkinter object + cli_opts: :class:`lib.gui.options.CliOptions` or ``None`` + The command line options object. Must be provided for main GUI. Must be ``None`` for tools + statusbar: :class:`lib.gui.custom_widgets.StatusBar` or ``None`` + The GUI Status bar. Must be provided for main GUI. Must be ``None`` for tools + """ + def __init__(self, + root: tk.Tk, + cli_opts: CliOptions | None, + statusbar: StatusBar | None) -> None: + logger.debug("Initializing %s: (root %s, cli_opts: %s, statusbar: %s)", + self.__class__.__name__, root, cli_opts, statusbar) + self._default_font = T.cast(dict, + tk.font.nametofont("TkDefaultFont").configure())["family"] + self._constants = {"root": root, + "scaling_factor": self._get_scaling(root), + "default_font": self._default_font} + self._gui_objects = _GuiObjects( + cli_opts=cli_opts, + tk_vars=GlobalVariables(), + project=Project(self, FileHandler), + tasks=Tasks(self, FileHandler), + status_bar=statusbar) + + self._user_config = UserConfig(None) + self._style = Style(self.default_font, root, PATHCACHE) + self._user_theme = self._style.user_theme + logger.debug("Initialized %s", self.__class__.__name__) + + # Constants + @property + def root(self) -> tk.Tk: + """ :class:`tkinter.Tk`: The root tkinter window. """ + return self._constants["root"] + + @property + def scaling_factor(self) -> float: + """ float: The scaling factor for current display. """ + return self._constants["scaling_factor"] + + @property + def pathcache(self) -> str: + """ str: The path to the GUI cache folder """ + return PATHCACHE + + # GUI Objects + @property + def cli_opts(self) -> CliOptions: + """ :class:`lib.gui.options.CliOptions`: The command line options for this GUI Session. """ + # This should only be None when a separate tool (not main GUI) is used, at which point + # cli_opts do not exist + assert self._gui_objects.cli_opts is not None + return self._gui_objects.cli_opts + + @property + def tk_vars(self) -> GlobalVariables: + """ dict: The global tkinter variables. """ + return self._gui_objects.tk_vars + + @property + def project(self) -> Project: + """ :class:`lib.gui.project.Project`: The project session handler. """ + return self._gui_objects.project + + @property + def tasks(self) -> Tasks: + """ :class:`lib.gui.project.Tasks`: The session tasks handler. """ + return self._gui_objects.tasks + + @property + def default_options(self) -> dict[str, dict[str, T.Any]]: + """ dict: The default options for all tabs """ + return self._gui_objects.default_options + + @property + def statusbar(self) -> StatusBar: + """ :class:`lib.gui.custom_widgets.StatusBar`: The GUI StatusBar + :class:`tkinter.ttk.Frame`. """ + # This should only be None when a separate tool (not main GUI) is used, at which point + # this statusbar does not exist + assert self._gui_objects.status_bar is not None + return self._gui_objects.status_bar + + @property + def command_notebook(self) -> CommandNotebook | None: + """ :class:`lib.gui.command.CommandNotebook`: The main Faceswap Command Notebook. """ + return self._gui_objects.command_notebook + + # Convenience GUI Objects + @property + def tools_notebook(self) -> ToolsNotebook: + """ :class:`lib.gui.command.ToolsNotebook`: The Faceswap Tools sub-Notebook. """ + assert self.command_notebook is not None + return self.command_notebook.tools_notebook + + @property + def modified_vars(self) -> dict[str, tk.BooleanVar]: + """ dict: The command notebook modified tkinter variables. """ + assert self.command_notebook is not None + return self.command_notebook.modified_vars + + @property + def _command_tabs(self) -> dict[str, int]: + """ dict: Command tab titles with their IDs. """ + assert self.command_notebook is not None + return self.command_notebook.tab_names + + @property + def _tools_tabs(self) -> dict[str, int]: + """ dict: Tools command tab titles with their IDs. """ + assert self.command_notebook is not None + return self.command_notebook.tools_tab_names + + # Config + @property + def user_config(self) -> UserConfig: + """ dict: The GUI config in dict form. """ + return self._user_config + + @property + def user_config_dict(self) -> dict[str, T.Any]: # TODO Dataclass + """ dict: The GUI config in dict form. """ + return self._user_config.config_dict + + @property + def user_theme(self) -> dict[str, T.Any]: # TODO Dataclass + """ dict: The GUI theme selection options. """ + return self._user_theme + + @property + def default_font(self) -> tuple[str, int]: + """ tuple: The selected font as configured in user settings. First item is the font (`str`) + second item the font size (`int`). """ + font = self.user_config_dict["font"] + font = self._default_font if font == "default" else font + return (font, self.user_config_dict["font_size"]) + + @staticmethod + def _get_scaling(root) -> float: + """ Get the display DPI. + + Returns + ------- + float: + The scaling factor + """ + dpi = root.winfo_fpixels("1i") + scaling = dpi / 72.0 + logger.debug("dpi: %s, scaling: %s'", dpi, scaling) + return scaling + + def set_default_options(self) -> None: + """ Set the default options for :mod:`lib.gui.projects` + + The Default GUI options are stored on Faceswap startup. + + Exposed as the :attr:`_default_opts` for a project cannot be set until after the main + Command Tabs have been loaded. + """ + default = self.cli_opts.get_option_values() + logger.debug(default) + self._gui_objects.default_options = default + self.project.set_default_options() + + def set_command_notebook(self, notebook: CommandNotebook) -> None: + """ Set the command notebook to the :attr:`command_notebook` attribute + and enable the modified callback for :attr:`project`. + + Parameters + ---------- + notebook: :class:`lib.gui.command.CommandNotebook` + The main command notebook for the Faceswap GUI + """ + logger.debug("Setting commane notebook: %s", notebook) + self._gui_objects.command_notebook = notebook + self.project.set_modified_callback() + + def set_active_tab_by_name(self, name: str) -> None: + """ Sets the :attr:`command_notebook` or :attr:`tools_notebook` to active based on given + name. + + Parameters + ---------- + name: str + The name of the tab to set active + """ + assert self.command_notebook is not None + name = name.lower() + if name in self._command_tabs: + tab_id = self._command_tabs[name] + logger.debug("Setting active tab to: (name: %s, id: %s)", name, tab_id) + self.command_notebook.select(tab_id) + elif name in self._tools_tabs: + self.command_notebook.select(self._command_tabs["tools"]) + tab_id = self._tools_tabs[name] + logger.debug("Setting active Tools tab to: (name: %s, id: %s)", name, tab_id) + self.tools_notebook.select() + else: + logger.debug("Name couldn't be found. Setting to id 0: %s", name) + self.command_notebook.select(0) + + def set_modified_true(self, command: str) -> None: + """ Set the modified variable to ``True`` for the given command in :attr:`modified_vars`. + + Parameters + ---------- + command: str + The command to set the modified state to ``True`` + + """ + tkvar = self.modified_vars.get(command, None) + if tkvar is None: + logger.debug("No tkvar for command: '%s'", command) + return + tkvar.set(True) + logger.debug("Set modified var to True for: '%s'", command) + + def refresh_config(self) -> None: + """ Reload the user config from file. """ + self._user_config = UserConfig(None) + + def set_cursor_busy(self, widget: tk.Widget | None = None) -> None: + """ Set the root or widget cursor to busy. + + Parameters + ---------- + widget: tkinter object, optional + The widget to set busy cursor for. If the provided value is ``None`` then sets the + cursor busy for the whole of the GUI. Default: ``None``. + """ + logger.debug("Setting cursor to busy. widget: %s", widget) + component = self.root if widget is None else widget + component.config(cursor="watch") # type: ignore + component.update_idletasks() + + def set_cursor_default(self, widget: tk.Widget | None = None) -> None: + """ Set the root or widget cursor to default. + + Parameters + ---------- + widget: tkinter object, optional + The widget to set default cursor for. If the provided value is ``None`` then sets the + cursor busy for the whole of the GUI. Default: ``None`` + """ + logger.debug("Setting cursor to default. widget: %s", widget) + component = self.root if widget is None else widget + component.config(cursor="") # type: ignore + component.update_idletasks() + + def set_root_title(self, text: str | None = None) -> None: + """ Set the main title text for Faceswap. + + The title will always begin with 'Faceswap.py'. Additional text can be appended. + + Parameters + ---------- + text: str, optional + Additional text to be appended to the GUI title bar. Default: ``None`` + """ + title = "Faceswap.py" + title += f" - {text}" if text is not None and text else "" + self.root.title(title) + + def set_geometry(self, width: int, height: int, fullscreen: bool = False) -> None: + """ Set the geometry for the root tkinter object. + + Parameters + ---------- + width: int + The width to set the window to (prior to scaling) + height: int + The height to set the window to (prior to scaling) + fullscreen: bool, optional + Whether to set the window to full-screen mode. If ``True`` then :attr:`width` and + :attr:`height` are ignored. Default: ``False`` + """ + self.root.tk.call("tk", "scaling", self.scaling_factor) + if fullscreen: + initial_dimensions = (self.root.winfo_screenwidth(), self.root.winfo_screenheight()) + else: + initial_dimensions = (round(width * self.scaling_factor), + round(height * self.scaling_factor)) + + if fullscreen and sys.platform in ("win32", "darwin"): + self.root.state('zoomed') + elif fullscreen: + self.root.attributes('-zoomed', True) + else: + self.root.geometry(f"{str(initial_dimensions[0])}x{str(initial_dimensions[1])}+80+80") + logger.debug("Geometry: %sx%s", *initial_dimensions) diff --git a/lib/gui/utils/file_handler.py b/lib/gui/utils/file_handler.py new file mode 100644 index 0000000000..4364eece3e --- /dev/null +++ b/lib/gui/utils/file_handler.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +""" File browser utility functions for the Faceswap GUI. """ +import logging +import platform +import tkinter as tk +from tkinter import filedialog +import typing as T + +logger = logging.getLogger(__name__) +_FILETYPE = T.Literal["default", "alignments", "config_project", "config_task", + "config_all", "csv", "image", "ini", "state", "log", "video"] +_HANDLETYPE = T.Literal["open", "save", "filename", "filename_multi", "save_filename", + "context", "dir"] + + +class FileHandler(): # pylint:disable=too-few-public-methods + """ Handles all GUI File Dialog actions and tasks. + + Parameters + ---------- + handle_type: ['open', 'save', 'filename', 'filename_multi', 'save_filename', 'context', 'dir'] + The type of file dialog to return. `open` and `save` will perform the open and save actions + and return the file. `filename` returns the filename from an `open` dialog. + `filename_multi` allows for multi-selection of files and returns a list of files selected. + `save_filename` returns the filename from a `save as` dialog. `context` is a context + sensitive parameter that returns a certain dialog based on the current options. `dir` asks + for a folder location. + file_type: ['default', 'alignments', 'config_project', 'config_task', 'config_all', 'csv', \ + 'image', 'ini', 'state', 'log', 'video'] or ``None`` + The type of file that this dialog is for. `default` allows selection of any files. Other + options limit the file type selection + title: str, optional + The title to display on the file dialog. If `None` then the default title will be used. + Default: ``None`` + initial_folder: str, optional + The folder to initially open with the file dialog. If `None` then tkinter will decide. + Default: ``None`` + initial_file: str, optional + The filename to set with the file dialog. If `None` then tkinter no initial filename is. + specified. Default: ``None`` + command: str, optional + Required for context handling file dialog, otherwise unused. Default: ``None`` + action: str, optional + Required for context handling file dialog, otherwise unused. Default: ``None`` + variable: str, optional + Required for context handling file dialog, otherwise unused. The variable to associate + with this file dialog. Default: ``None`` + parent: :class:`tkinter.Frame`, optional + The parent that is launching the file dialog. ``None`` sets this to root. Default: ``None`` + + Attributes + ---------- + return_file: str or object + The return value from the file dialog + + Example + ------- + >>> handler = FileHandler('filename', 'video', title='Select a video...') + >>> video_file = handler.return_file + >>> print(video_file) + '/path/to/selected/video.mp4' + """ + + def __init__(self, + handle_type: _HANDLETYPE, + file_type: _FILETYPE | None, + title: str | None = None, + initial_folder: str | None = None, + initial_file: str | None = None, + command: str | None = None, + action: str | None = None, + variable: str | None = None, + parent: tk.Frame | None = None) -> None: + logger.debug("Initializing %s: (handle_type: '%s', file_type: '%s', title: '%s', " + "initial_folder: '%s', initial_file: '%s', command: '%s', action: '%s', " + "variable: %s, parent: %s)", self.__class__.__name__, handle_type, file_type, + title, initial_folder, initial_file, command, action, variable, parent) + self._handletype = handle_type + self._dummy_master = self._set_dummy_master() + self._defaults = self._set_defaults() + self._kwargs = self._set_kwargs(title, + initial_folder, + initial_file, + file_type, + command, + action, + variable, + parent) + self.return_file = getattr(self, f"_{self._handletype.lower()}")() + self._remove_dummy_master() + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def _filetypes(self) -> dict[str, list[tuple[str, str]]]: + """ dict: The accepted extensions for each file type for opening/saving """ + all_files = ("All files", "*.*") + filetypes = { + "default": [all_files], + "alignments": [("Faceswap Alignments", "*.fsa"), all_files], + "config_project": [("Faceswap Project files", "*.fsw"), all_files], + "config_task": [("Faceswap Task files", "*.fst"), all_files], + "config_all": [("Faceswap Project and Task files", "*.fst *.fsw"), all_files], + "csv": [("Comma separated values", "*.csv"), all_files], + "image": [("Bitmap", "*.bmp"), + ("JPG", "*.jpeg *.jpg"), + ("PNG", "*.png"), + ("TIFF", "*.tif *.tiff"), + all_files], + "ini": [("Faceswap config files", "*.ini"), all_files], + "json": [("JSON file", "*.json"), all_files], + "model": [("Keras model files", "*.h5"), all_files], + "state": [("State files", "*.json"), all_files], + "log": [("Log files", "*.log"), all_files], + "video": [("Audio Video Interleave", "*.avi"), + ("Flash Video", "*.flv"), + ("Matroska", "*.mkv"), + ("MOV", "*.mov"), + ("MP4", "*.mp4"), + ("MPEG", "*.mpeg *.mpg *.ts *.vob"), + ("WebM", "*.webm"), + ("Windows Media Video", "*.wmv"), + all_files]} + + # Add in multi-select options and upper case extensions for Linux + for key in filetypes: + if platform.system() == "Linux": + filetypes[key] = [item + if item[0] == "All files" + else (item[0], f"{item[1]} {item[1].upper()}") + for item in filetypes[key]] + if len(filetypes[key]) > 2: + multi = [f"{key.title()} Files"] + multi.append(" ".join([ftype[1] + for ftype in filetypes[key] if ftype[0] != "All files"])) + filetypes[key].insert(0, T.cast(tuple[str, str], tuple(multi))) + return filetypes + + @property + def _contexts(self) -> dict[str, dict[str, str | dict[str, str]]]: + """dict: Mapping of commands, actions and their corresponding file dialog for context + handle types. """ + return {"effmpeg": {"input": {"extract": "filename", + "gen-vid": "dir", + "get-fps": "filename", + "get-info": "filename", + "mux-audio": "filename", + "rescale": "filename", + "rotate": "filename", + "slice": "filename"}, + "output": {"extract": "dir", + "gen-vid": "save_filename", + "get-fps": "nothing", + "get-info": "nothing", + "mux-audio": "save_filename", + "rescale": "save_filename", + "rotate": "save_filename", + "slice": "save_filename"}}} + + @classmethod + def _set_dummy_master(cls) -> tk.Frame | None: + """ Add an option to force black font on Linux file dialogs KDE issue that displays light + font on white background). + + This is a pretty hacky solution, but tkinter does not allow direct editing of file dialogs, + so we create a dummy frame and add the foreground option there, so that the file dialog can + inherit the foreground. + + Returns + ------- + tkinter.Frame or ``None`` + The dummy master frame for Linux systems, otherwise ``None`` + """ + if platform.system().lower() == "linux": + frame = tk.Frame() + frame.option_add("*foreground", "black") + retval: tk.Frame | None = frame + else: + retval = None + return retval + + def _remove_dummy_master(self) -> None: + """ Destroy the dummy master widget on Linux systems. """ + if platform.system().lower() != "linux" or self._dummy_master is None: + return + self._dummy_master.destroy() + del self._dummy_master + self._dummy_master = None + + def _set_defaults(self) -> dict[str, str | None]: + """ Set the default file type for the file dialog. Generally the first found file type + will be used, but this is overridden if it is not appropriate. + + Returns + ------- + dict: + The default file extension for each file type + """ + defaults: dict[str, str | None] = { + key: next(ext for ext in val[0][1].split(" ")).replace("*", "") + for key, val in self._filetypes.items()} + defaults["default"] = None + defaults["video"] = ".mp4" + defaults["image"] = ".png" + logger.debug(defaults) + return defaults + + def _set_kwargs(self, + title: str | None, + initial_folder: str | None, + initial_file: str | None, + file_type: _FILETYPE | None, + command: str | None, + action: str | None, + variable: str | None, + parent: tk.Frame | None + ) -> dict[str, None | tk.Frame | str | list[tuple[str, str]]]: + """ Generate the required kwargs for the requested file dialog browser. + + Parameters + ---------- + title: str + The title to display on the file dialog. If `None` then the default title will be used. + initial_folder: str + The folder to initially open with the file dialog. If `None` then tkinter will decide. + initial_file: str + The filename to set with the file dialog. If `None` then tkinter no initial filename + is. + file_type: ['default', 'alignments', 'config_project', 'config_task', 'config_all', \ + 'csv', 'image', 'ini', 'state', 'log', 'video'] or ``None`` + The type of file that this dialog is for. `default` allows selection of any files. + Other options limit the file type selection + command: str + Required for context handling file dialog, otherwise unused. + action: str + Required for context handling file dialog, otherwise unused. + variable: str, optional + Required for context handling file dialog, otherwise unused. The variable to associate + with this file dialog. Default: ``None`` + parent: :class:`tkinter.Frame` + The parent that is launching the file dialog. ``None`` sets this to root + + Returns + ------- + dict: + The key word arguments for the file dialog to be launched + """ + logger.debug("Setting Kwargs: (title: %s, initial_folder: %s, initial_file: '%s', " + "file_type: '%s', command: '%s': action: '%s', variable: '%s', parent: %s)", + title, initial_folder, initial_file, file_type, command, action, variable, + parent) + + kwargs: dict[str, None | tk.Frame | str | list[tuple[str, str]]] = { + "master": self._dummy_master} + + if self._handletype.lower() == "context": + assert command is not None and action is not None and variable is not None + self._set_context_handletype(command, action, variable) + + if title is not None: + kwargs["title"] = title + + if initial_folder is not None: + kwargs["initialdir"] = initial_folder + + if initial_file is not None: + kwargs["initialfile"] = initial_file + + if parent is not None: + kwargs["parent"] = parent + + if self._handletype.lower() in ( + "open", "save", "filename", "filename_multi", "save_filename"): + assert file_type is not None + kwargs["filetypes"] = self._filetypes[file_type] + if self._defaults.get(file_type): + kwargs['defaultextension'] = self._defaults[file_type] + if self._handletype.lower() == "save": + kwargs["mode"] = "w" + if self._handletype.lower() == "open": + kwargs["mode"] = "r" + logger.debug("Set Kwargs: %s", kwargs) + return kwargs + + def _set_context_handletype(self, command: str, action: str, variable: str) -> None: + """ Sets the correct handle type based on context. + + Parameters + ---------- + command: str + The command that is being executed. Used to look up the context actions + action: str + The action that is being performed. Used to look up the correct file dialog + variable: str + The variable associated with this file dialog + """ + if self._contexts[command].get(variable, None) is not None: + handletype = T.cast(dict[str, dict[str, dict[str, str]]], + self._contexts)[command][variable][action] + else: + handletype = T.cast(dict[str, dict[str, str]], + self._contexts)[command][action] + logger.debug(handletype) + self._handletype = T.cast(_HANDLETYPE, handletype) + + def _open(self) -> T.IO | None: + """ Open a file. """ + logger.debug("Popping Open browser") + return filedialog.askopenfile(**self._kwargs) # type: ignore + + def _save(self) -> T.IO | None: + """ Save a file. """ + logger.debug("Popping Save browser") + return filedialog.asksaveasfile(**self._kwargs) # type: ignore + + def _dir(self) -> str: + """ Get a directory location. """ + logger.debug("Popping Dir browser") + return filedialog.askdirectory(**self._kwargs) # type: ignore + + def _savedir(self) -> str: + """ Get a save directory location. """ + logger.debug("Popping SaveDir browser") + return filedialog.askdirectory(**self._kwargs) # type: ignore + + def _filename(self) -> str: + """ Get an existing file location. """ + logger.debug("Popping Filename browser") + return filedialog.askopenfilename(**self._kwargs) # type: ignore + + def _filename_multi(self) -> tuple[str, ...]: + """ Get multiple existing file locations. """ + logger.debug("Popping Filename browser") + return filedialog.askopenfilenames(**self._kwargs) # type: ignore + + def _save_filename(self) -> str: + """ Get a save file location. """ + logger.debug("Popping Save Filename browser") + return filedialog.asksaveasfilename(**self._kwargs) # type: ignore + + @staticmethod + def _nothing() -> None: # pylint:disable=useless-return + """ Method that does nothing, used for disabling open/save pop up. """ + logger.debug("Popping Nothing browser") + return diff --git a/lib/gui/utils/image.py b/lib/gui/utils/image.py new file mode 100644 index 0000000000..816ba19257 --- /dev/null +++ b/lib/gui/utils/image.py @@ -0,0 +1,659 @@ +#!/usr/bin python3 +""" Utilities for handling images in the Faceswap GUI """ +from __future__ import annotations +import logging +import os +import typing as T + +import cv2 +import numpy as np +from PIL import Image, ImageDraw, ImageTk + +from lib.training.preview_cv import PreviewBuffer + +from .config import get_config, PATHCACHE + +if T.TYPE_CHECKING: + from collections.abc import Sequence + +logger = logging.getLogger(__name__) +_IMAGES: "Images" | None = None +_PREVIEW_TRIGGER: "PreviewTrigger" | None = None +TRAININGPREVIEW = ".gui_training_preview.png" + + +def initialize_images() -> None: + """ Initialize the :class:`Images` handler and add to global constant. + + This should only be called once on first GUI startup. Future access to :class:`Images` + handler should only be executed through :func:`get_images`. + """ + global _IMAGES # pylint:disable=global-statement + if _IMAGES is not None: + return + logger.debug("Initializing images") + _IMAGES = Images() + + +def get_images() -> "Images": + """ Get the Master GUI Images handler. + + Returns + ------- + :class:`Images` + The Master GUI Images handler + """ + assert _IMAGES is not None + return _IMAGES + + +def _get_previews(image_path: str) -> list[str]: + """ Get the images stored within the given directory. + + Parameters + ---------- + image_path: str + The folder containing images to be scanned + + Returns + ------- + list: + The image filenames stored within the given folder + + """ + logger.debug("Getting images: '%s'", image_path) + if not os.path.isdir(image_path): + logger.debug("Folder does not exist") + return [] + files = [os.path.join(image_path, f) + for f in os.listdir(image_path) if f.lower().endswith((".png", ".jpg"))] + logger.debug("Image files: %s", files) + return files + + +class PreviewTrain(): + """ Handles the loading of the training preview image(s) and adding to the display buffer + + Parameters + ---------- + cache_path: str + Full path to the cache folder that contains the preview images + """ + def __init__(self, cache_path: str) -> None: + logger.debug("Initializing %s: (cache_path: '%s')", self.__class__.__name__, cache_path) + self._buffer = PreviewBuffer() + self._cache_path = cache_path + self._modified: float = 0.0 + self._error_count: int = 0 + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def buffer(self) -> PreviewBuffer: + """ :class:`~lib.training.PreviewBuffer` The preview buffer for the training preview + image. """ + return self._buffer + + def load(self) -> bool: + """ Load the latest training preview image(s) from disk and add to :attr:`buffer` """ + logger.trace("Loading Training preview images") # type:ignore + image_files = _get_previews(self._cache_path) + filename = next((fname for fname in image_files + if os.path.basename(fname) == TRAININGPREVIEW), "") + if not filename: + logger.trace("No preview to display") # type:ignore + return False + try: + modified = os.path.getmtime(filename) + if modified <= self._modified: + logger.trace("preview '%s' not updated. Current timestamp: %s, " # type:ignore + "existing timestamp: %s", filename, modified, self._modified) + return False + + logger.debug("Loading preview: '%s'", filename) + img = cv2.imread(filename, cv2.IMREAD_UNCHANGED) + assert img is not None + self._modified = modified + self._buffer.add_image(os.path.basename(filename), img) + self._error_count = 0 + except (ValueError, AssertionError): + # This is probably an error reading the file whilst it's being saved so ignore it + # for now and only pick up if there have been multiple consecutive fails + logger.debug("Unable to display preview: (image: '%s', attempt: %s)", + img, self._error_count) + if self._error_count < 10: + self._error_count += 1 + else: + logger.error("Error reading the preview file for '%s'", filename) + return False + + logger.debug("Loaded preview: '%s' (%s)", filename, img.shape) + return True + + def reset(self) -> None: + """ Reset the preview buffer when the display page has been disabled. + + Notes + ----- + The buffer requires resetting, otherwise the re-enabled preview window hangs waiting for a + training image that has already been marked as processed + """ + logger.debug("Resetting training preview") + del self._buffer + self._buffer = PreviewBuffer() + self._modified = 0.0 + self._error_count = 0 + + +class PreviewExtract(): + """ Handles the loading of preview images for extract and convert + + Parameters + ---------- + cache_path: str + Full path to the cache folder that contains the preview images + """ + def __init__(self, cache_path: str) -> None: + logger.debug("Initializing %s: (cache_path: '%s')", self.__class__.__name__, cache_path) + self._cache_path = cache_path + + self._batch_mode = False + self._output_path = "" + + self._modified: float = 0.0 + self._filenames: list[str] = [] + self._images: np.ndarray | None = None + self._placeholder: np.ndarray | None = None + + self._preview_image: Image.Image | None = None + self._preview_image_tk: ImageTk.PhotoImage | None = None + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def image(self) -> ImageTk.PhotoImage: + """:class:`PIL.ImageTk.PhotoImage` The preview image for displaying in a tkinter canvas """ + assert self._preview_image_tk is not None + return self._preview_image_tk + + def save(self, filename: str) -> None: + """ Save the currently displaying preview image to the given location + + Parameters + ---------- + filename: str + The full path to the filename to save the preview image to + """ + logger.debug("Saving preview to %s", filename) + assert self._preview_image is not None + self._preview_image.save(filename) + + def set_faceswap_output_path(self, location: str, batch_mode: bool = False) -> None: + """ Set the path that will contain the output from an Extract or Convert task. + + Required so that the GUI can fetch output images to display for return in + :attr:`preview_image`. + + Parameters + ---------- + location: str + The output location that has been specified for an Extract or Convert task + batch_mode: bool + ``True`` if extracting in batch mode otherwise False + """ + self._output_path = location + self._batch_mode = batch_mode + + def _get_newest_folder(self) -> str: + """ Obtain the most recent folder created in the extraction output folder when processing + in batch mode. + + Returns + ------- + str + The most recently modified folder within the parent output folder. If no folders have + been created, returns the parent output folder + + """ + folders = [] if not os.path.exists(self._output_path) else [ + os.path.join(self._output_path, folder) + for folder in os.listdir(self._output_path) + if os.path.isdir(os.path.join(self._output_path, folder))] + + folders.sort(key=os.path.getmtime) + retval = folders[-1] if folders else self._output_path + logger.debug("sorted folders: %s, return value: %s", folders, retval) + return retval + + def _get_newest_filenames(self, image_files: list[str]) -> list[str]: + """ Return image filenames that have been modified since the last check. + + Parameters + ---------- + image_files: list + The list of image files to check the modification date for + + Returns + ------- + list: + A list of images that have been modified since the last check + """ + if not self._modified: + retval = image_files + else: + retval = [fname for fname in image_files + if os.path.getmtime(fname) > self._modified] + if not retval: + logger.debug("No new images in output folder") + else: + self._modified = max(os.path.getmtime(img) for img in retval) + logger.debug("Number new images: %s, Last Modified: %s", + len(retval), self._modified) + return retval + + def _pad_and_border(self, image: Image.Image, size: int) -> np.ndarray: + """ Pad rectangle images to a square and draw borders + + Parameters + ---------- + image: :class:`PIL.Image` + The image to process + size: int + The size of the image as it should be displayed + + Returns + ------- + :class:`numpy.ndarray`: + The processed image + """ + if image.size[0] != image.size[1]: + # Pad to square + new_img = Image.new("RGB", (size, size)) + new_img.paste(image, ((size - image.size[0]) // 2, (size - image.size[1]) // 2)) + image = new_img + draw = ImageDraw.Draw(image) + draw.rectangle(((0, 0), (size, size)), outline="#E5E5E5", width=1) + retval = np.array(image) + logger.trace("image shape: %s", retval.shape) # type: ignore + return retval + + def _process_samples(self, + samples: list[np.ndarray], + filenames: list[str], + num_images: int) -> bool: + """ Process the latest sample images into a displayable image. + + Parameters + ---------- + samples: list + The list of extract/convert preview images to display + filenames: list + The full path to the filenames corresponding to the images + num_images: int + The number of images that should be displayed + + Returns + ------- + bool + ``True`` if samples succesfully compiled otherwise ``False`` + """ + asamples = np.array(samples) + if not np.any(asamples): + logger.debug("No preview images collected.") + return False + + self._filenames = (self._filenames + filenames)[-num_images:] + cache = self._images + + if cache is None: + logger.debug("Creating new cache") + cache = asamples[-num_images:] + else: + logger.debug("Appending to existing cache") + cache = np.concatenate((cache, asamples))[-num_images:] + + self._images = cache + assert self._images is not None + logger.debug("Cache shape: %s", self._images.shape) + return True + + def _load_images_to_cache(self, + image_files: list[str], + frame_dims: tuple[int, int], + thumbnail_size: int) -> bool: + """ Load preview images to the image cache. + + Load new images and append to cache, filtering the cache to the number of thumbnails that + will fit inside the display panel. + + Parameters + ---------- + image_files: list + A list of new image files that have been modified since the last check + frame_dims: tuple + The (width (`int`), height (`int`)) of the display panel that will display the preview + thumbnail_size: int + The size of each thumbnail that should be created + + Returns + ------- + bool + ``True`` if images were successfully loaded to cache otherwise ``False`` + """ + logger.debug("Number image_files: %s, frame_dims: %s, thumbnail_size: %s", + len(image_files), frame_dims, thumbnail_size) + num_images = (frame_dims[0] // thumbnail_size) * (frame_dims[1] // thumbnail_size) + logger.debug("num_images: %s", num_images) + if num_images == 0: + return False + samples: list[np.ndarray] = [] + start_idx = len(image_files) - num_images if len(image_files) > num_images else 0 + show_files = sorted(image_files, key=os.path.getctime)[start_idx:] + dropped_files = [] + for fname in show_files: + try: + img = Image.open(fname) + except PermissionError as err: + logger.debug("Permission error opening preview file: '%s'. Original error: %s", + fname, str(err)) + dropped_files.append(fname) + continue + except Exception as err: # pylint:disable=broad-except + # Swallow any issues with opening an image rather than spamming console + # Can happen when trying to read partially saved images + logger.debug("Error opening preview file: '%s'. Original error: %s", + fname, str(err)) + dropped_files.append(fname) + continue + + width, height = img.size + scaling = thumbnail_size / max(width, height) + logger.debug("image width: %s, height: %s, scaling: %s", width, height, scaling) + + try: + img = img.resize((int(width * scaling), int(height * scaling))) + except OSError as err: + # Image only gets loaded when we call a method, so may error on partial loads + logger.debug("OS Error resizing preview image: '%s'. Original error: %s", + fname, err) + dropped_files.append(fname) + continue + + samples.append(self._pad_and_border(img, thumbnail_size)) + + return self._process_samples(samples, + [fname for fname in show_files if fname not in dropped_files], + num_images) + + def _create_placeholder(self, thumbnail_size: int) -> None: + """ Create a placeholder image for when there are fewer thumbnails available + than columns to display them. + + Parameters + ---------- + thumbnail_size: int + The size of the thumbnail that the placeholder should replicate + """ + logger.debug("Creating placeholder. thumbnail_size: %s", thumbnail_size) + placeholder = Image.new("RGB", (thumbnail_size, thumbnail_size)) + draw = ImageDraw.Draw(placeholder) + draw.rectangle(((0, 0), (thumbnail_size, thumbnail_size)), outline="#E5E5E5", width=1) + placeholder = np.array(placeholder) + self._placeholder = placeholder + logger.debug("Created placeholder. shape: %s", placeholder.shape) + + def _place_previews(self, frame_dims: tuple[int, int]) -> Image.Image: + """ Format the preview thumbnails stored in the cache into a grid fitting the display + panel. + + Parameters + ---------- + frame_dims: tuple + The (width (`int`), height (`int`)) of the display panel that will display the preview + + Returns + ------- + :class:`PIL.Image`: + The final preview display image + """ + if self._images is None: + logger.debug("No images in cache. Returning None") + return None + samples = self._images.copy() + num_images, thumbnail_size = samples.shape[:2] + if self._placeholder is None: + self._create_placeholder(thumbnail_size) + + logger.debug("num_images: %s, thumbnail_size: %s", num_images, thumbnail_size) + cols, rows = frame_dims[0] // thumbnail_size, frame_dims[1] // thumbnail_size + logger.debug("cols: %s, rows: %s", cols, rows) + if cols == 0 or rows == 0: + logger.debug("Cols or Rows is zero. No items to display") + return None + + remainder = (cols * rows) - num_images + if remainder != 0: + logger.debug("Padding sample display. Remainder: %s", remainder) + assert self._placeholder is not None + placeholder = np.concatenate([np.expand_dims(self._placeholder, 0)] * remainder) + samples = np.concatenate((samples, placeholder)) + + display = np.vstack([np.hstack(T.cast("Sequence", samples[row * cols: (row + 1) * cols])) + for row in range(rows)]) + logger.debug("display shape: %s", display.shape) + return Image.fromarray(display) + + def load_latest_preview(self, thumbnail_size: int, frame_dims: tuple[int, int]) -> bool: + """ Load the latest preview image for extract and convert. + + Retrieves the latest preview images from the faceswap output folder, resizes to thumbnails + and lays out for display. Places the images into :attr:`preview_image` for loading into + the display panel. + + Parameters + ---------- + thumbnail_size: int + The size of each thumbnail that should be created + frame_dims: tuple + The (width (`int`), height (`int`)) of the display panel that will display the preview + + Returns + ------- + bool + ``True`` if a preview was succesfully loaded otherwise ``False`` + """ + logger.debug("Loading preview image: (thumbnail_size: %s, frame_dims: %s)", + thumbnail_size, frame_dims) + image_path = self._get_newest_folder() if self._batch_mode else self._output_path + image_files = _get_previews(image_path) + gui_preview = os.path.join(self._output_path, ".gui_preview.jpg") + if not image_files or (len(image_files) == 1 and gui_preview not in image_files): + logger.debug("No preview to display") + return False + # Filter to just the gui_preview if it exists in folder output + image_files = [gui_preview] if gui_preview in image_files else image_files + logger.debug("Image Files: %s", len(image_files)) + + image_files = self._get_newest_filenames(image_files) + if not image_files: + return False + + if not self._load_images_to_cache(image_files, frame_dims, thumbnail_size): + logger.debug("Failed to load any preview images") + if gui_preview in image_files: + # Reset last modified for failed loading of a gui preview image so it is picked + # up next time + self._modified = 0.0 + return False + + if image_files == [gui_preview]: + # Delete the preview image so that the main scripts know to output another + logger.debug("Deleting preview image") + os.remove(image_files[0]) + show_image = self._place_previews(frame_dims) + if not show_image: + self._preview_image = None + self._preview_image_tk = None + return False + + logger.debug("Displaying preview: %s", self._filenames) + self._preview_image = show_image + self._preview_image_tk = ImageTk.PhotoImage(show_image) + return True + + def delete_previews(self) -> None: + """ Remove any image preview files """ + for fname in self._filenames: + if os.path.basename(fname) == ".gui_preview.jpg": + logger.debug("Deleting: '%s'", fname) + try: + os.remove(fname) + except FileNotFoundError: + logger.debug("File does not exist: %s", fname) + + +class Images(): + """ The centralized image repository for holding all icons and images required by the GUI. + + This class should be initialized on GUI startup through :func:`initialize_images`. Any further + access to this class should be through :func:`get_images`. + """ + def __init__(self) -> None: + logger.debug("Initializing %s", self.__class__.__name__) + self._pathpreview = os.path.join(PATHCACHE, "preview") + self._pathoutput: str | None = None + self._batch_mode = False + self._preview_train = PreviewTrain(self._pathpreview) + self._preview_extract = PreviewExtract(self._pathpreview) + self._icons = self._load_icons() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def preview_train(self) -> PreviewTrain: + """ :class:`PreviewTrain` The object handling the training preview images """ + return self._preview_train + + @property + def preview_extract(self) -> PreviewExtract: + """ :class:`PreviewTrain` The object handling the training preview images """ + return self._preview_extract + + @property + def icons(self) -> dict[str, ImageTk.PhotoImage]: + """ dict: The faceswap icons for all parts of the GUI. The dictionary key is the icon + name (`str`) the value is the icon sized and formatted for display + (:class:`PIL.ImageTK.PhotoImage`). + + Example + ------- + >>> icons = get_images().icons + >>> save = icons["save"] + >>> button = ttk.Button(parent, image=save) + >>> button.pack() + """ + return self._icons + + @staticmethod + def _load_icons() -> dict[str, ImageTk.PhotoImage]: + """ Scan the icons cache folder and load the icons into :attr:`icons` for retrieval + throughout the GUI. + + Returns + ------- + dict: + The icons formatted as described in :attr:`icons` + + """ + size = get_config().user_config_dict.get("icon_size", 16) + size = int(round(size * get_config().scaling_factor)) + icons: dict[str, ImageTk.PhotoImage] = {} + pathicons = os.path.join(PATHCACHE, "icons") + for fname in os.listdir(pathicons): + name, ext = os.path.splitext(fname) + if ext != ".png": + continue + img = Image.open(os.path.join(pathicons, fname)) + img = ImageTk.PhotoImage(img.resize((size, size), resample=Image.HAMMING)) + icons[name] = img + logger.debug(icons) + return icons + + def delete_preview(self) -> None: + """ Delete the preview files in the cache folder and reset the image cache. + + Should be called when terminating tasks, or when Faceswap starts up or shuts down. + """ + logger.debug("Deleting previews") + for item in os.listdir(self._pathpreview): + if item.startswith(os.path.splitext(TRAININGPREVIEW)[0]) and item.endswith((".jpg", + ".png")): + fullitem = os.path.join(self._pathpreview, item) + logger.debug("Deleting: '%s'", fullitem) + os.remove(fullitem) + + self._preview_extract.delete_previews() + del self._preview_train + del self._preview_extract + self._preview_train = PreviewTrain(self._pathpreview) + self._preview_extract = PreviewExtract(self._pathpreview) + + +class PreviewTrigger(): + """ Triggers to indicate to underlying Faceswap process that the preview image should + be updated. + + Writes a file to the cache folder that is picked up by the main process. + """ + def __init__(self) -> None: + logger.debug("Initializing: %s", self.__class__.__name__) + self._trigger_files = {"update": os.path.join(PATHCACHE, ".preview_trigger"), + "mask_toggle": os.path.join(PATHCACHE, ".preview_mask_toggle")} + logger.debug("Initialized: %s (trigger_files: %s)", + self.__class__.__name__, self._trigger_files) + + def set(self, trigger_type: T.Literal["update", "mask_toggle"]): + """ Place the trigger file into the cache folder + + Parameters + ---------- + trigger_type: ["update", "mask_toggle"] + The type of action to trigger. 'update': Full preview update. 'mask_toggle': toggle + mask on and off + """ + trigger = self._trigger_files[trigger_type] + if not os.path.isfile(trigger): + with open(trigger, "w", encoding="utf8"): + pass + logger.debug("Set preview trigger: %s", trigger) + + def clear(self, trigger_type: T.Literal["update", "mask_toggle"] | None = None) -> None: + """ Remove the trigger file from the cache folder. + + Parameters + ---------- + trigger_type: ["update", "mask_toggle", ``None``], optional + The trigger to clear. 'update': Full preview update. 'mask_toggle': toggle mask on + and off. ``None`` - clear all triggers. Default: ``None`` + """ + if trigger_type is None: + triggers = list(self._trigger_files.values()) + else: + triggers = [self._trigger_files[trigger_type]] + for trigger in triggers: + if os.path.isfile(trigger): + os.remove(trigger) + logger.debug("Removed preview trigger: %s", trigger) + + +def preview_trigger() -> PreviewTrigger: + """ Set the global preview trigger if it has not already been set and return. + + Returns + ------- + :class:`PreviewTrigger` + The trigger to indicate to the main faceswap process that it should perform a training + preview update + """ + global _PREVIEW_TRIGGER # pylint:disable=global-statement + if _PREVIEW_TRIGGER is None: + _PREVIEW_TRIGGER = PreviewTrigger() + return _PREVIEW_TRIGGER diff --git a/lib/gui/utils/misc.py b/lib/gui/utils/misc.py new file mode 100644 index 0000000000..4d60c80214 --- /dev/null +++ b/lib/gui/utils/misc.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +""" Miscellaneous Utility functions for the GUI. Includes LongRunningTask object """ +from __future__ import annotations +import logging +import sys +import typing as T + +from threading import Event, Thread +from queue import Queue + +from .config import get_config + +if T.TYPE_CHECKING: + from collections.abc import Callable + from types import TracebackType + from lib.multithreading import _ErrorType + + +logger = logging.getLogger(__name__) + + +class LongRunningTask(Thread): + """ Runs long running tasks in a background thread to prevent the GUI from becoming + unresponsive. + + This is sub-classed from :class:`Threading.Thread` so check documentation there for base + parameters. Additional parameters listed below. + + Parameters + ---------- + widget: tkinter object, optional + The widget that this :class:`LongRunningTask` is associated with. Used for setting the busy + cursor in the correct location. Default: ``None``. + """ + _target: Callable + _args: tuple + _kwargs: dict[str, T.Any] + _name: str + + def __init__(self, + target: Callable | None = None, + name: str | None = None, + args: tuple = (), + kwargs: dict[str, T.Any] | None = None, + *, + daemon: bool = True, + widget=None): + logger.debug("Initializing %s: (target: %s, name: %s, args: %s, kwargs: %s, " + "daemon: %s)", self.__class__.__name__, target, name, args, kwargs, + daemon) + super().__init__(target=target, name=name, args=args, kwargs=kwargs, + daemon=daemon) + self.err: _ErrorType = None + self._widget = widget + self._config = get_config() + self._config.set_cursor_busy(widget=self._widget) + self._complete = Event() + self._queue: Queue = Queue() + logger.debug("Initialized %s", self.__class__.__name__,) + + @property + def complete(self) -> Event: + """ :class:`threading.Event`: Event is set if the thread has completed its task, + otherwise it is unset. + """ + return self._complete + + def run(self) -> None: + """ Commence the given task in a background thread. """ + try: + if self._target is not None: + retval = self._target(*self._args, **self._kwargs) + self._queue.put(retval) + except Exception: # pylint:disable=broad-except + self.err = T.cast(tuple[type[BaseException], BaseException, "TracebackType"], + sys.exc_info()) + assert self.err is not None + logger.debug("Error in thread (%s): %s", self._name, + self.err[1].with_traceback(self.err[2])) + finally: + self._complete.set() + # Avoid a ref-cycle if the thread is running a function with + # an argument that has a member that points to the thread. + del self._target, self._args, self._kwargs + + def get_result(self) -> T.Any: + """ Return the result from the given task. + + Returns + ------- + varies: + The result of the thread will depend on the given task. If a call is made to + :func:`get_result` prior to the thread completing its task then ``None`` will be + returned + """ + if not self._complete.is_set(): + logger.warning("Aborting attempt to retrieve result from a LongRunningTask that is " + "still running") + return None + if self.err: + logger.debug("Error caught in thread") + self._config.set_cursor_default(widget=self._widget) + raise self.err[1].with_traceback(self.err[2]) + + logger.debug("Getting result from thread") + retval = self._queue.get() + logger.debug("Got result from thread") + self._config.set_cursor_default(widget=self._widget) + return retval diff --git a/lib/gui/wrapper.py b/lib/gui/wrapper.py index 276e34b773..88ee11f646 100644 --- a/lib/gui/wrapper.py +++ b/lib/gui/wrapper.py @@ -1,11 +1,14 @@ #!/usr/bin python3 """ Process wrapper for underlying faceswap commands for the GUI """ +from __future__ import annotations import os import logging import re import signal -from subprocess import PIPE, Popen import sys +import typing as T + +from subprocess import PIPE, Popen from threading import Thread from time import time @@ -15,87 +18,144 @@ from .utils import get_config, get_images, LongRunningTask, preview_trigger if os.name == "nt": - import win32console # pylint: disable=import-error - + import win32console # pylint:disable=import-error -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class ProcessWrapper(): """ Builds command, launches and terminates the underlying faceswap process. Updates GUI display depending on state """ - def __init__(self): + def __init__(self) -> None: logger.debug("Initializing %s", self.__class__.__name__) - self.tk_vars = get_config().tk_vars - self.set_callbacks() - self.pathscript = os.path.realpath(os.path.dirname(sys.argv[0])) - self.command = None - self.statusbar = get_config().statusbar - self._training_session_location = dict() - self.task = FaceswapControl(self) + self._tk_vars = get_config().tk_vars + self._set_callbacks() + self._command: str | None = None + """ str | None: The currently executing command, when process running or ``None`` """ + + self._statusbar = get_config().statusbar + self._training_session_location: dict[T.Literal["model_name", "model_folder"], str] = {} + self._task = FaceswapControl(self) logger.debug("Initialized %s", self.__class__.__name__) - def set_callbacks(self): - """ Set the tkinter variable callbacks """ + @property + def task(self) -> FaceswapControl: + """ :class:`FaceswapControl`: The object that controls the underlying faceswap process """ + return self._task + + def _set_callbacks(self) -> None: + """ Set the tkinter variable callbacks for performing an action or generating a command """ logger.debug("Setting tk variable traces") - self.tk_vars["action"].trace("w", self.action_command) - self.tk_vars["generate"].trace("w", self.generate_command) + self._tk_vars.action_command.trace("w", self._action_command) + self._tk_vars.generate_command.trace("w", self._generate_command) + + def _action_command(self, *args: tuple[str, str, str]): # pylint:disable=unused-argument + """ Callback for when the Action button is pressed. Process command line options and + launches the action - def action_command(self, *args): - """ The action to perform when the action button is pressed """ - if not self.tk_vars["action"].get(): + Parameters + ---------- + args: + tuple[str, str, str] + Tkinter variable callback args. Required but unused + """ + if not self._tk_vars.action_command.get(): return - category, command = self.tk_vars["action"].get().split(",") + category, command = self._tk_vars.action_command.get().split(",") - if self.tk_vars["runningtask"].get(): - self.task.terminate() + if self._tk_vars.running_task.get(): + self._task.terminate() else: - self.command = command - args = self.prepare(category) - self.task.execute_script(command, args) - self.tk_vars["action"].set(None) - - def generate_command(self, *args): - """ Generate the command line arguments and output """ - if not self.tk_vars["generate"].get(): + self._command = command + fs_args = self._prepare(T.cast(T.Literal["faceswap", "tools"], category)) + self._task.execute_script(command, fs_args) + self._tk_vars.action_command.set("") + + def _generate_command(self, # pylint:disable=unused-argument + *args: tuple[str, str, str]) -> None: + """ Callback for when the Generate button is pressed. Process command line options and + output the cli command + + Parameters + ---------- + args: + tuple[str, str, str] + Tkinter variable callback args. Required but unused + """ + if not self._tk_vars.generate_command.get(): return - category, command = self.tk_vars["generate"].get().split(",") - args = self.build_args(category, command=command, generate=True) - self.tk_vars["console_clear"].set(True) - logger.debug(" ".join(args)) - print(" ".join(args)) - self.tk_vars["generate"].set(None) - - def prepare(self, category): - """ Prepare the environment for execution """ + category, command = self._tk_vars.generate_command.get().split(",") + fs_args = self._build_args(category, command=command, generate=True) + self._tk_vars.console_clear.set(True) + logger.debug(" ".join(fs_args)) + print(" ".join(fs_args)) + self._tk_vars.generate_command.set("") + + def _prepare(self, category: T.Literal["faceswap", "tools"]) -> list[str]: + """ Prepare the environment for execution, Sets the 'running task' and 'console clear' + global tkinter variables. If training, sets the 'is training' variable + + Parameters + ---------- + category: str, ["faceswap", "tools"] + The script that is executing the command + + Returns + ------- + list[str] + The command line arguments to execute for the faceswap job + """ logger.debug("Preparing for execution") - self.tk_vars["runningtask"].set(True) - self.tk_vars["console_clear"].set(True) - if self.command == "train": - self.tk_vars["istraining"].set(True) + assert self._command is not None + self._tk_vars.running_task.set(True) + self._tk_vars.console_clear.set(True) + if self._command == "train": + self._tk_vars.is_training.set(True) print("Loading...") - self.statusbar.message.set("Executing - {}.py".format(self.command)) - mode = "indeterminate" if self.command in ("effmpeg", "train") else "determinate" - self.statusbar.start(mode) + self._statusbar.message.set(f"Executing - {self._command}.py") + mode: T.Literal["indeterminate", + "determinate"] = ("indeterminate" if self._command in ("effmpeg", "train") + else "determinate") + self._statusbar.start(mode) - args = self.build_args(category) - self.tk_vars["display"].set(self.command) + args = self._build_args(category) + self._tk_vars.display.set(self._command) logger.debug("Prepared for execution") return args - def build_args(self, category, command=None, generate=False): + def _build_args(self, + category: str, + command: str | None = None, + generate: bool = False) -> list[str]: """ Build the faceswap command and arguments list. If training, pass the model folder and name to the training :class:`lib.gui.analysis.Session` for the GUI. + + Parameters + ---------- + category: str, ["faceswap", "tools"] + The script that is executing the command + command: str, optional + The main faceswap command to execute, if provided. The currently running task if + ``None``. Default: ``None`` + generate: bool, optional + ``True`` if the command is just to be generated for display. ``False`` if the command + is to be executed + + Returns + ------- + list[str] + The full faceswap command to be executed or displayed """ logger.debug("Build cli arguments: (category: %s, command: %s, generate: %s)", category, command, generate) - command = self.command if not command else command - script = "{}.{}".format(category, "py") - pathexecscript = os.path.join(self.pathscript, script) + command = self._command if not command else command + assert command is not None + script = f"{category}.py" + pathexecscript = os.path.join(os.path.realpath(os.path.dirname(sys.argv[0])), script) args = [sys.executable] if generate else [sys.executable, "-u"] args.extend([pathexecscript, command]) @@ -110,19 +170,19 @@ def build_args(self, category, command=None, generate=False): args.append("-gui") # Indicate to Faceswap that we are running the GUI if generate: # Delimit args with spaces - args = ['"{}"'.format(arg) if " " in arg and not arg.startswith(("[", "(")) + args = [f'"{arg}"' if " " in arg and not arg.startswith(("[", "(")) and not arg.endswith(("]", ")")) else arg for arg in args] logger.debug("Built cli arguments: (%s)", args) return args - def _get_training_session_info(self, cli_option): + def _get_training_session_info(self, cli_option: list[str]) -> None: """ Set the model folder and model name to :`attr:_training_session_location` so the global session picks them up for logging to the graph and analysis tab. Parameters ---------- - cli_option: list + cli_option: list[str] The command line option to be checked for model folder or name """ if cli_option[0] == "-t": @@ -132,118 +192,211 @@ def _get_training_session_info(self, cli_option): self._training_session_location["model_folder"] = cli_option[1] logger.debug("model_folder: '%s'", self._training_session_location["model_folder"]) - def terminate(self, message): - """ Finalize wrapper when process has exited """ + def terminate(self, message: str) -> None: + """ Finalize wrapper when process has exited. Stops the progress bar, sets the status + message. If the terminating task is 'train', then triggers the training close down actions + + Parameters + ---------- + message: str + The message to display in the status bar + """ logger.debug("Terminating Faceswap processes") - self.tk_vars["runningtask"].set(False) - if self.task.command == "train": - self.tk_vars["istraining"].set(False) + self._tk_vars.running_task.set(False) + if self._task.command == "train": + self._tk_vars.is_training.set(False) Session.stop_training() - self.statusbar.stop() - self.statusbar.message.set(message) - self.tk_vars["display"].set(None) + self._statusbar.stop() + self._statusbar.message.set(message) + self._tk_vars.display.set("") get_images().delete_preview() preview_trigger().clear(trigger_type=None) - self.command = None + self._command = None logger.debug("Terminated Faceswap processes") print("Process exited.") class FaceswapControl(): - """ Control the underlying Faceswap tasks """ - def __init__(self, wrapper): - logger.debug("Initializing %s", self.__class__.__name__) - self.wrapper = wrapper + """ Control the underlying Faceswap tasks. + + wrapper: :class:`ProcessWrapper` + The object responsible for managing this faceswap task + """ + def __init__(self, wrapper: ProcessWrapper) -> None: + logger.debug("Initializing %s (wrapper: %s)", self.__class__.__name__, wrapper) + self._wrapper = wrapper self._session_info = wrapper._training_session_location - self.config = get_config() - self.statusbar = self.config.statusbar - self.command = None - self.args = None - self.process = None - self.thread = None # Thread for LongRunningTask termination - self.train_stats = {"iterations": 0, "timestamp": None} - self.consoleregex = { + self._config = get_config() + self._statusbar = self._config.statusbar + self._command: str | None = None + self._process: Popen | None = None + self._thread: LongRunningTask | None = None + self._train_stats: dict[T.Literal["iterations", "timestamp"], + int | float | None] = {"iterations": 0, "timestamp": None} + self._consoleregex: dict[T.Literal["loss", "tqdm", "ffmpeg"], re.Pattern] = { "loss": re.compile(r"[\W]+(\d+)?[\W]+([a-zA-Z\s]*)[\W]+?(\d+\.\d+)"), "tqdm": re.compile(r"(?P.*?)(?P\d+%).*?(?P\S+/\S+)\W\[" r"(?P[\d+:]+<.*),\W(?P.*)[a-zA-Z/]*\]"), "ffmpeg": re.compile(r"([a-zA-Z]+)=\s*(-?[\d|N/A]\S+)")} + self._first_loss_seen = False logger.debug("Initialized %s", self.__class__.__name__) - def execute_script(self, command, args): - """ Execute the requested Faceswap Script """ + @property + def command(self) -> str | None: + """ str | None: The currently executing command, when process running or ``None`` """ + return self._command + + def execute_script(self, command: str, args: list[str]) -> None: + """ Execute the requested Faceswap Script + + Parameters + ---------- + command: str + The faceswap command that is to be run + args: list[str] + The full command line arguments to be executed + """ logger.debug("Executing Faceswap: (command: '%s', args: %s)", command, args) - self.thread = None - self.command = command - kwargs = {"stdout": PIPE, - "stderr": PIPE, - "bufsize": 1, - "universal_newlines": True} - - self.process = Popen(args, **kwargs, stdin=PIPE) - self.thread_stdout() - self.thread_stderr() + self._thread = None + self._command = command + + proc = Popen(args, # pylint:disable=consider-using-with + stdout=PIPE, + stderr=PIPE, + bufsize=1, + text=True, + stdin=PIPE, + errors="backslashreplace") + self._process = proc + self._thread_stdout() + self._thread_stderr() logger.debug("Executed Faceswap") - def read_stdout(self): - """ Read stdout from the subprocess. If training, pass the loss - values to Queue """ + def _process_training_determinate_function(self, output: str) -> bool: + """ Process an stdout/stderr message to check for determinate TQDM output when training + + Parameters + ---------- + output: str + The stdout/stderr string to test + + Returns + ------- + bool + ``True`` if a determinate TQDM line was parsed when training otherwise ``False`` + """ + if self._command == "train" and not self._first_loss_seen and self._capture_tqdm(output): + self._statusbar.set_mode("determinate") + return True + return False + + def _process_progress_stdout(self, output: str) -> bool: + """ Process stdout for any faceswap processes that update the status/progress bar(s) + + Parameters + ---------- + output: str + The output line read from stdout + + Returns + ------- + bool + ``True`` if all actions have been completed on the output line otherwise ``False`` + """ + if self._process_training_determinate_function(output): + return True + + if self._command == "train" and self._capture_loss(output): + return True + + if self._command == "effmpeg" and self._capture_ffmpeg(output): + return True + + if self._command not in ("train", "effmpeg") and self._capture_tqdm(output): + return True + + return False + + def _process_training_stdout(self, output: str) -> None: + """ Process any triggers that are required to update the GUI when Faceswap is running a + training session. + + Parameters + ---------- + output: str + The output line read from stdout + """ + tk_vars = get_config().tk_vars + if self._command != "train" or not tk_vars.is_training.get(): + return + + t_output = output.strip().lower() + if "[saved model]" not in t_output or t_output.endswith("[saved model]"): + # Not a saved model line or saving the model for a reason other than standard saving + return + + logger.debug("Trigger GUI Training update") + logger.trace("tk_vars: %s", {itm: var.get() # type:ignore[attr-defined] + for itm, var in tk_vars.__dict__.items()}) + if not Session.is_training: + # Don't initialize session until after the first save as state file must exist first + logger.debug("Initializing curret training session") + Session.initialize_session(self._session_info["model_folder"], + self._session_info["model_name"], + is_training=True) + tk_vars.refresh_graph.set(True) + + def _read_stdout(self) -> None: + """ Read stdout from the subprocess. """ logger.debug("Opening stdout reader") + assert self._process is not None while True: try: - output = self.process.stdout.readline() + buff = self._process.stdout + assert buff is not None + output: str = buff.readline() except ValueError as err: if str(err).lower().startswith("i/o operation on closed file"): break raise - if output == "" and self.process.poll() is not None: + + if output == "" and self._process.poll() is not None: break + + if output and self._process_progress_stdout(output): + continue + if output: - if ((self.command == "train" and self.capture_loss(output)) or - (self.command == "effmpeg" and self.capture_ffmpeg(output)) or - (self.command not in ("train", "effmpeg") and self.capture_tqdm(output))): - continue - if self.command == "train" and self.wrapper.tk_vars["istraining"].get(): - if "[saved models]" in output.strip().lower(): - logger.debug("Trigger GUI Training update") - logger.trace("tk_vars: %s", {itm: var.get() - for itm, var in self.wrapper.tk_vars.items()}) - if not Session.is_training: - # Don't initialize session until after the first save as state - # file must exist first - logger.debug("Initializing curret training session") - Session.initialize_session( - self._session_info["model_folder"], - self._session_info["model_name"], - is_training=True) - self.wrapper.tk_vars["updatepreview"].set(True) - self.wrapper.tk_vars["refreshgraph"].set(True) - if "[preview updated]" in output.strip().lower(): - self.wrapper.tk_vars["updatepreview"].set(True) - continue + self._process_training_stdout(output) print(output.rstrip()) - returncode = self.process.poll() - message = self.set_final_status(returncode) - self.wrapper.terminate(message) + + returncode = self._process.poll() + assert returncode is not None + self._first_loss_seen = False + message = self._set_final_status(returncode) + self._wrapper.terminate(message) logger.debug("Terminated stdout reader. returncode: %s", returncode) - def read_stderr(self): + def _read_stderr(self) -> None: """ Read stdout from the subprocess. If training, pass the loss values to Queue """ logger.debug("Opening stderr reader") + assert self._process is not None while True: try: - output = self.process.stderr.readline() + buff = self._process.stderr + assert buff is not None + output: str = buff.readline() except ValueError as err: if str(err).lower().startswith("i/o operation on closed file"): break raise - if output == "" and self.process.poll() is not None: + if output == "" and self._process.poll() is not None: break if output: - if self.command != "train" and self.capture_tqdm(output): + if self._command != "train" and self._capture_tqdm(output): continue - if self.command == "train" and output.startswith("Reading training images"): - print(output.strip(), file=sys.stdout) + if self._process_training_determinate_function(output): continue if os.name == "nt" and "Call to CreateProcess failed. Error code: 2" in output: # Suppress ptxas errors on Tensorflow for Windows @@ -252,145 +405,199 @@ def read_stderr(self): print(output.strip(), file=sys.stderr) logger.debug("Terminated stderr reader") - def thread_stdout(self): - """ Put the subprocess stdout so that it can be read without - blocking """ + def _thread_stdout(self) -> None: + """ Put the subprocess stdout so that it can be read without blocking """ logger.debug("Threading stdout") - thread = Thread(target=self.read_stdout) + thread = Thread(target=self._read_stdout) thread.daemon = True thread.start() logger.debug("Threaded stdout") - def thread_stderr(self): - """ Put the subprocess stderr so that it can be read without - blocking """ + def _thread_stderr(self) -> None: + """ Put the subprocess stderr so that it can be read without blocking """ logger.debug("Threading stderr") - thread = Thread(target=self.read_stderr) + thread = Thread(target=self._read_stderr) thread.daemon = True thread.start() logger.debug("Threaded stderr") - def capture_loss(self, string): - """ Capture loss values from stdout """ - logger.trace("Capturing loss") + def _capture_loss(self, string: str) -> bool: + """ Capture loss values from stdout + + Parameters + ---------- + string: str + An output line read from stdout + + Returns + ------- + bool + ``True`` if a loss line was captured from stdout, otherwise ``False`` + """ + logger.trace("Capturing loss") # type:ignore[attr-defined] if not str.startswith(string, "["): - logger.trace("Not loss message. Returning False") + logger.trace("Not loss message. Returning False") # type:ignore[attr-defined] return False - loss = self.consoleregex["loss"].findall(string) + loss = self._consoleregex["loss"].findall(string) if len(loss) != 2 or not all(len(itm) == 3 for itm in loss): - logger.trace("Not loss message. Returning False") + logger.trace("Not loss message. Returning False") # type:ignore[attr-defined] return False - message = "Total Iterations: {} | ".format(int(loss[0][0])) - message += " ".join(["{}: {}".format(itm[1], itm[2]) for itm in loss]) + message = f"Total Iterations: {int(loss[0][0])} | " + message += " ".join([f"{itm[1]}: {itm[2]}" for itm in loss]) if not message: - logger.trace("Error creating loss message. Returning False") + logger.trace( # type:ignore[attr-defined] + "Error creating loss message. Returning False") return False - iterations = self.train_stats["iterations"] + iterations = self._train_stats["iterations"] + assert isinstance(iterations, int) if iterations == 0: # Set initial timestamp - self.train_stats["timestamp"] = time() + self._train_stats["timestamp"] = time() iterations += 1 - self.train_stats["iterations"] = iterations - - elapsed = self.calc_elapsed() - message = "Elapsed: {} | Session Iterations: {} {}".format( - elapsed, - self.train_stats["iterations"], message) - self.statusbar.progress_update(message, 0, False) - logger.trace("Succesfully captured loss: %s", message) + self._train_stats["iterations"] = iterations + + elapsed = self._calculate_elapsed() + message = (f"Elapsed: {elapsed} | " + f"Session Iterations: {self._train_stats['iterations']} {message}") + + if not self._first_loss_seen: + self._statusbar.set_mode("indeterminate") + self._first_loss_seen = True + + self._statusbar.progress_update(message, 0, False) + logger.trace("Succesfully captured loss: %s", message) # type:ignore[attr-defined] return True - def calc_elapsed(self): - """ Calculate and format time since training started """ + def _calculate_elapsed(self) -> str: + """ Calculate and format time since training started + + Returns + ------- + str + The amount of time elapsed since training started in HH:mm:ss format + """ now = time() - elapsed_time = now - self.train_stats["timestamp"] + timestamp = self._train_stats["timestamp"] + assert isinstance(timestamp, float) + elapsed_time = now - timestamp try: - hrs = int(elapsed_time // 3600) - if hrs < 10: - hrs = "{0:02d}".format(hrs) - mins = "{0:02d}".format((int(elapsed_time % 3600) // 60)) - secs = "{0:02d}".format((int(elapsed_time % 3600) % 60)) + i_hrs = int(elapsed_time // 3600) + hrs = f"{i_hrs:02d}" if i_hrs < 10 else str(i_hrs) + mins = f"{(int(elapsed_time % 3600) // 60):02d}" + secs = f"{(int(elapsed_time % 3600) % 60):02d}" except ZeroDivisionError: - hrs = "00" - mins = "00" - secs = "00" - return "{}:{}:{}".format(hrs, mins, secs) - - def capture_tqdm(self, string): - """ Capture tqdm output for progress bar """ - logger.trace("Capturing tqdm") - tqdm = self.consoleregex["tqdm"].match(string) - if not tqdm: + hrs = mins = secs = "00" + return f"{hrs}:{mins}:{secs}" + + def _capture_tqdm(self, string: str) -> bool: + """ Capture tqdm output for progress bar + + Parameters + ---------- + string: str + An output line read from stdout + + Returns + ------- + bool + ``True`` if a tqdm line was captured from stdout, otherwise ``False`` + """ + logger.trace("Capturing tqdm") # type:ignore[attr-defined] + mtqdm = self._consoleregex["tqdm"].match(string) + if not mtqdm: return False - tqdm = tqdm.groupdict() + tqdm = mtqdm.groupdict() if any("?" in val for val in tqdm.values()): - logger.trace("tqdm initializing. Skipping") + logger.trace("tqdm initializing. Skipping") # type:ignore[attr-defined] return True description = tqdm["dsc"].strip() - description = description if description == "" else "{} | ".format(description[:-1]) - processtime = "Elapsed: {} Remaining: {}".format(tqdm["tme"].split("<")[0], - tqdm["tme"].split("<")[1]) - message = "{}{} | {} | {} | {}".format(description, - processtime, - tqdm["rte"], - tqdm["itm"], - tqdm["pct"]) + description = description if description == "" else f"{description[:-1]} | " + processtime = (f"Elapsed: {tqdm['tme'].split('<')[0]} " + f"Remaining: {tqdm['tme'].split('<')[1]}") + msg = f"{description}{processtime} | {tqdm['rte']} | {tqdm['itm']} | {tqdm['pct']}" position = tqdm["pct"].replace("%", "") position = int(position) if position.isdigit() else 0 - self.statusbar.progress_update(message, position, True) - logger.trace("Succesfully captured tqdm message: %s", message) + self._statusbar.progress_update(msg, position, True) + logger.trace("Succesfully captured tqdm message: %s", msg) # type:ignore[attr-defined] return True - def capture_ffmpeg(self, string): - """ Capture tqdm output for progress bar """ - logger.trace("Capturing ffmpeg") - ffmpeg = self.consoleregex["ffmpeg"].findall(string) + def _capture_ffmpeg(self, string: str) -> bool: + """ Capture ffmpeg output for progress bar + + Parameters + ---------- + string: str + An output line read from stdout + + Returns + ------- + bool + ``True`` if an ffmpeg line was captured from stdout, otherwise ``False`` + """ + logger.trace("Capturing ffmpeg") # type:ignore[attr-defined] + ffmpeg = self._consoleregex["ffmpeg"].findall(string) if len(ffmpeg) < 7: - logger.trace("Not ffmpeg message. Returning False") + logger.trace("Not ffmpeg message. Returning False") # type:ignore[attr-defined] return False message = "" for item in ffmpeg: - message += "{}: {} ".format(item[0], item[1]) + message += f"{item[0]}: {item[1]} " if not message: - logger.trace("Error creating ffmpeg message. Returning False") + logger.trace( # type:ignore[attr-defined] + "Error creating ffmpeg message. Returning False") return False - self.statusbar.progress_update(message, 0, False) - logger.trace("Succesfully captured ffmpeg message: %s", message) + self._statusbar.progress_update(message, 0, False) + logger.trace("Succesfully captured ffmpeg message: %s", # type:ignore[attr-defined] + message) return True - def terminate(self): - """ Terminate the running process in a LongRunningTask so we can still - output to console """ - if self.thread is None: + def terminate(self) -> None: + """ Terminate the running process in a LongRunningTask so console can still be updated + console """ + if self._thread is None: logger.debug("Terminating wrapper in LongRunningTask") - self.thread = LongRunningTask(target=self.terminate_in_thread, - args=(self.command, self.process)) - if self.command == "train": - self.wrapper.tk_vars["istraining"].set(False) - self.thread.start() - self.config.root.after(1000, self.terminate) - elif not self.thread.complete.is_set(): + self._thread = LongRunningTask(target=self._terminate_in_thread, + args=(self._command, self._process)) + if self._command == "train": + get_config().tk_vars.is_training.set(False) + self._thread.start() + self._config.root.after(1000, self.terminate) + elif not self._thread.complete.is_set(): logger.debug("Not finished terminating") - self.config.root.after(1000, self.terminate) + self._config.root.after(1000, self.terminate) else: logger.debug("Termination Complete. Cleaning up") - _ = self.thread.get_result() # Terminate the LongRunningTask object - self.thread = None + _ = self._thread.get_result() # Terminate the LongRunningTask object + self._thread = None + + def _terminate_in_thread(self, command: str, process: Popen) -> bool: + """ Terminate the subprocess + + Parameters + ---------- + command: str + The command that is running + + process: :class:`subprocess.Popen` + The running process - def terminate_in_thread(self, command, process): - """ Terminate the subprocess """ + Returns + ------- + bool + ``True`` when this function exits + """ logger.debug("Terminating wrapper") if command == "train": - timeout = self.config.user_config_dict.get("timeout", 120) + timeout = self._config.user_config_dict.get("timeout", 120) logger.debug("Sending Exit Signal") print("Sending Exit Signal", flush=True) now = time() @@ -398,7 +605,7 @@ def terminate_in_thread(self, command, process): logger.debug("Sending carriage return to process") con_in = win32console.GetStdHandle( # pylint:disable=c-extension-no-member win32console.STD_INPUT_HANDLE) # pylint:disable=c-extension-no-member - keypress = self.generate_windows_keypress("\n") + keypress = self._generate_windows_keypress("\n") con_in.WriteConsoleInput([keypress]) else: logger.debug("Sending SIGINT to process") @@ -409,14 +616,25 @@ def terminate_in_thread(self, command, process): break if timeelapsed > timeout: logger.error("Timeout reached sending Exit Signal") - self.terminate_all_children() + self._terminate_all_children() else: - self.terminate_all_children() + self._terminate_all_children() return True - @staticmethod - def generate_windows_keypress(character): - """ Generate an 'Enter' key press to terminate Windows training """ + @classmethod + def _generate_windows_keypress(cls, character: str) -> bytes: + """ Generate a Windows keypress + + Parameters + ---------- + character: str + The caracter to generate the keypress for + + Returns + ------- + bytes + The generated Windows keypress + """ buf = win32console.PyINPUT_RECORDType( # pylint:disable=c-extension-no-member win32console.KEY_EVENT) # pylint:disable=c-extension-no-member buf.KeyDown = 1 @@ -424,8 +642,8 @@ def generate_windows_keypress(character): buf.Char = character return buf - @staticmethod - def terminate_all_children(): + @classmethod + def _terminate_all_children(cls) -> None: """ Terminates all children """ logger.debug("Terminating Process...") print("Terminating Process...", flush=True) @@ -448,24 +666,34 @@ def terminate_all_children(): print("Killed") else: for child in alive: - msg = "Process {} survived SIGKILL. Giving up".format(child) + msg = f"Process {child} survived SIGKILL. Giving up" logger.debug(msg) print(msg) - def set_final_status(self, returncode): - """ Set the status bar output based on subprocess return code - and reset training stats """ + def _set_final_status(self, returncode: int) -> str: + """ Set the status bar output based on subprocess return code and reset training stats + + Parameters + ---------- + returncode: int + The returncode from the terminated process + + Returns + ------- + str + The final statusbar text + """ logger.debug("Setting final status. returncode: %s", returncode) - self.train_stats = {"iterations": 0, "timestamp": None} + self._train_stats = {"iterations": 0, "timestamp": None} if returncode in (0, 3221225786): status = "Ready" elif returncode == -15: - status = "Terminated - {}.py".format(self.command) + status = f"Terminated - {self._command}.py" elif returncode == -9: - status = "Killed - {}.py".format(self.command) + status = f"Killed - {self._command}.py" elif returncode == -6: - status = "Aborted - {}.py".format(self.command) + status = f"Aborted - {self._command}.py" else: - status = "Failed - {}.py. Return Code: {}".format(self.command, returncode) + status = f"Failed - {self._command}.py. Return Code: {returncode}" logger.debug("Set final status: %s", status) return status diff --git a/lib/image.py b/lib/image.py index c7dc31b244..897d8e2daa 100644 --- a/lib/image.py +++ b/lib/image.py @@ -1,12 +1,14 @@ #!/usr/bin python3 """ Utilities for working with images and videos """ - +from __future__ import annotations +import json import logging import re import subprocess import os import struct import sys +import typing as T from ast import literal_eval from bisect import bisect @@ -21,9 +23,12 @@ from lib.multithreading import MultiThread from lib.queue_manager import queue_manager, QueueEmpty -from lib.utils import convert_to_secs, FaceswapError, _video_extensions, get_image_paths +from lib.utils import convert_to_secs, FaceswapError, VIDEO_EXTENSIONS, get_image_paths + +if T.TYPE_CHECKING: + from lib.align.alignments import PNGHeaderDict -logger = logging.getLogger(__name__) # pylint:disable=invalid-name +logger = logging.getLogger(__name__) # ################### # # <<< IMAGE UTILS >>> # @@ -32,7 +37,7 @@ # <<< IMAGE IO >>> # -class FfmpegReader(imageio.plugins.ffmpeg.FfmpegFormat.Reader): +class FfmpegReader(imageio.plugins.ffmpeg.FfmpegFormat.Reader): # type:ignore """ Monkey patch imageio ffmpeg to use keyframes whilst seeking """ def __init__(self, format, request): super().__init__(format, request) @@ -143,7 +148,7 @@ def _previous_keyframe_info(self, index=0): logger.trace("keyframe pts_time: %s, keyframe: %s", prev_pts_time, prev_keyframe) return prev_pts_time, prev_keyframe - def _initialize(self, index=0): + def _initialize(self, index=0): # noqa:C901 """ Replace ImageIO _initialize with a version that explictly uses keyframes. Notes @@ -154,7 +159,7 @@ def _initialize(self, index=0): correct frame for all videos. Navigating to the previous keyframe then discarding frames until the correct frame is reached appears to work well. """ - # pylint: disable-all + # pylint:disable-all if self._read_gen is not None: self._read_gen.close() @@ -250,7 +255,7 @@ def _initialize(self, index=0): self._read_gen.__next__() # we already have meta data -imageio.plugins.ffmpeg.FfmpegFormat.Reader = FfmpegReader +imageio.plugins.ffmpeg.FfmpegFormat.Reader = FfmpegReader # type: ignore def read_image(filename, raise_error=False, with_metadata=False): @@ -291,16 +296,16 @@ def read_image(filename, raise_error=False, with_metadata=False): success = True image = None try: - if not with_metadata: - retval = cv2.imread(filename) - if retval is None: + with open(filename, "rb") as infile: + raw_file = infile.read() + image = cv2.imdecode(np.frombuffer(raw_file, dtype="uint8"), cv2.IMREAD_COLOR) + if image is None: raise ValueError("Image is None") - else: - with open(filename, "rb") as infile: - raw_file = infile.read() + if with_metadata: metadata = png_read_meta(raw_file) - image = cv2.imdecode(np.frombuffer(raw_file, dtype="uint8"), cv2.IMREAD_UNCHANGED) - retval = (image, metadata) + retval = (image, metadata) + else: + retval = image except TypeError as err: success = False msg = "Error while reading image (TypeError): '{}'".format(filename) @@ -356,26 +361,20 @@ def read_image_batch(filenames, with_metadata=False): >>> images = read_image_batch(image_filenames) """ logger.trace("Requested batch: '%s'", filenames) - executor = futures.ThreadPoolExecutor() - with executor: + batch = [None for _ in range(len(filenames))] + if with_metadata: + meta = [None for _ in range(len(filenames))] + + with futures.ThreadPoolExecutor() as executor: images = {executor.submit(read_image, filename, - raise_error=True, with_metadata=with_metadata): filename - for filename in filenames} - batch = [None for _ in range(len(filenames))] - if with_metadata: - meta = [None for _ in range(len(filenames))] - # There is no guarantee that the same filename will not be passed through multiple times - # (and when shuffle is true this can definitely happen), so we can't just call - # filenames.index(). - return_indices = {filename: [idx for idx, fname in enumerate(filenames) - if fname == filename] - for filename in set(filenames)} + raise_error=True, with_metadata=with_metadata): idx + for idx, filename in enumerate(filenames)} for future in futures.as_completed(images): - return_idx = return_indices[images[future]].pop() + ret_idx = images[future] if with_metadata: - batch[return_idx], meta[return_idx] = future.result() + batch[ret_idx], meta[ret_idx] = future.result() else: - batch[return_idx] = future.result() + batch[ret_idx] = future.result() batch = np.array(batch) retval = (batch, meta) if with_metadata else batch @@ -435,10 +434,11 @@ def read_image_meta(filename): elif field == b"iTXt": keyword, value = infile.read(length).split(b"\0", 1) if keyword == b"faceswap": - retval["itxt"] = literal_eval(value[4:].decode("utf-8")) + retval["itxt"] = literal_eval(value[4:].decode("utf-8", errors="replace")) break else: - logger.trace("Skipping iTXt chunk: '%s'", keyword.decode("latin-1", "ignore")) + logger.trace("Skipping iTXt chunk: '%s'", keyword.decode("latin-1", + errors="ignore")) length = 0 # Reset marker for next chunk infile.seek(length + 4, 1) logger.trace("filename: %s, metadata: %s", filename, retval) @@ -557,7 +557,10 @@ def update_existing_metadata(filename, metadata): os.replace(tmp_filename, filename) -def encode_image(image, extension, metadata=None): +def encode_image(image: np.ndarray, + extension: str, + encoding_args: tuple[int, ...] | None = None, + metadata: PNGHeaderDict | dict[str, T.Any] | bytes | None = None) -> bytes: """ Encode an image. Parameters @@ -566,14 +569,17 @@ def encode_image(image, extension, metadata=None): The image to be encoded in `BGR` channel order. extension: str A compatible `cv2` image file extension that the final image is to be saved to. - metadata: dict, optional - Metadata for the image. If provided, and the extension is png, this information will be - written to the PNG itxt header. Default:``None`` + encoding_args: tuple[int, ...], optional + Any encoding arguments to pass to cv2's imencode function + metadata: dict or bytes, optional + Metadata for the image. If provided, and the extension is png or tiff, this information + will be written to the PNG itxt header. Default:``None`` Can be provided as a python dict + or pre-encoded Returns ------- encoded_image: bytes - The image encoded into the correct file format + The image encoded into the correct file format as bytes Example ------- @@ -581,20 +587,23 @@ def encode_image(image, extension, metadata=None): >>> image = read_image(image_file) >>> encoded_image = encode_image(image, ".jpg") """ - if metadata and extension.lower() != ".png": - raise ValueError("Metadata is only supported for .png images") - retval = cv2.imencode(extension, image)[1] + if metadata and extension.lower() not in (".png", ".tif"): + raise ValueError("Metadata is only supported for .png and .tif images") + args = tuple() if encoding_args is None else encoding_args + + retval = cv2.imencode(extension, image, args)[1].tobytes() if metadata: - retval = np.frombuffer(png_write_meta(retval.tobytes(), metadata), dtype="uint8") + func = {".png": png_write_meta, ".tif": tiff_write_meta}[extension] + retval = func(retval, metadata) return retval -def png_write_meta(png, data): +def png_write_meta(image: bytes, data: PNGHeaderDict | dict[str, T.Any] | bytes) -> bytes: """ Write Faceswap information to a png's iTXt field. Parameters ---------- - png: bytes + image: bytes The bytes encoded png file to write header data to data: dict or bytes The dictionary to write to the header. Can be pre-encoded as utf-8. @@ -610,17 +619,115 @@ def png_write_meta(png, data): PNG Specification: https://www.w3.org/TR/2003/REC-PNG-20031110/ """ - split = png.find(b"IDAT") - 4 - retval = png[:split] + pack_to_itxt(data) + png[split:] + split = image.find(b"IDAT") - 4 + retval = image[:split] + pack_to_itxt(data) + image[split:] return retval -def png_read_meta(png): - """ Read the Faceswap information stored in a png's iTXt field. +def tiff_write_meta(image: bytes, data: PNGHeaderDict | dict[str, T.Any] | bytes) -> bytes: + """ Write Faceswap information to a tiff's image_description field. Parameters ---------- png: bytes + The bytes encoded tiff file to write header data to + data: dict or bytes + The data to write to the image-description field. If provided as a dict, then it should be + a json serializable object, otherwise it should be data encoded as ascii bytes + + Notes + ----- + This handles a very specific task of adding, and populating, an ImageDescription field in a + Tiff file generated by OpenCV. For any other usecases it will likely fail + """ + if not isinstance(data, bytes): + data = json.dumps(data, ensure_ascii=True).encode("ascii") + + assert image[:2] == b"II", "Not a supported TIFF file" + assert struct.unpack(" 270: + insert_idx = i # Log insert location of image description + + if size <= 4: # value in offset column + ifd += tag + continue + + ifd += tag[:8] + tag_offset = struct.unpack(" dict[str, T.Any]: + """ Read information stored in a Tiff's Image Description field """ + assert image[:2] == b"II", "Not a supported TIFF file" + assert struct.unpack("I", png[pointer:pointer + 4])[0] + length = struct.unpack(">I", image[pointer:pointer + 4])[0] pointer += 8 - keyword, value = png[pointer:pointer + length].split(b"\0", 1) + keyword, value = image[pointer:pointer + length].split(b"\0", 1) if keyword == b"faceswap": - retval = literal_eval(value[4:].decode("utf-8")) + retval = literal_eval(value[4:].decode("utf-8", errors="ignore")) break - logger.trace("Skipping iTXt chunk: '%s'", keyword.decode("latin-1", "ignore")) + logger.trace("Skipping iTXt chunk: '%s'", keyword.decode("latin-1", errors="ignore")) pointer += length + 4 return retval @@ -799,7 +906,7 @@ def count_frames(filename, fast=False): process = subprocess.Popen(cmd, stderr=subprocess.STDOUT, stdout=subprocess.PIPE, - universal_newlines=True) + universal_newlines=True, encoding="utf8") pbar = None duration = None init_tqdm = False @@ -894,9 +1001,10 @@ def _check_location_exists(self): def _set_thread(self): """ Set the background thread for the load and save iterators and launch it. """ - logger.debug("Setting thread") + logger.trace("Setting thread") # type:ignore[attr-defined] if self._thread is not None and self._thread.is_alive(): - logger.debug("Thread pre-exists and is alive: %s", self._thread) + logger.trace("Thread pre-exists and is alive: %s", # type:ignore[attr-defined] + self._thread) return self._thread = MultiThread(self._process, self._queue, @@ -920,6 +1028,7 @@ def close(self): logger.debug("Received Close") if self._thread is not None: self._thread.join() + del self._thread self._thread = None logger.debug("Closed") @@ -1037,9 +1146,9 @@ def _check_for_video(self): If the given location is a file and does not have a valid video extension. """ - if os.path.isdir(self.location): + if not isinstance(self.location, str) or os.path.isdir(self.location): retval = False - elif os.path.splitext(self.location)[1].lower() in _video_extensions: + elif os.path.splitext(self.location)[1].lower() in VIDEO_EXTENSIONS: retval = True else: raise FaceswapError("The input file '{}' is not a valid video".format(self.location)) @@ -1144,7 +1253,8 @@ def _from_video(self): reader.close() def _dummy_video_framename(self, index): - """ Return a dummy filename for video files + """ Return a dummy filename for video files. The file name is made up of: + _. Parameters ---------- @@ -1159,8 +1269,8 @@ def _dummy_video_framename(self, index): Returns ------- str: A dummied filename for a video frame """ - vidname = os.path.splitext(os.path.basename(self.location))[0] - return "{}_{:06d}.png".format(vidname, index + 1) + vidname, ext = os.path.splitext(os.path.basename(self.location)) + return f"{vidname}_{index + 1:06d}{ext}" def _from_folder(self): """ Generator for loading images from a folder @@ -1233,6 +1343,29 @@ def __init__(self, path, skip_list=None, count=None): path, count) super().__init__(path, queue_size=8, skip_list=skip_list, count=count) + def _get_count_and_filelist(self, fast_count, count): + """ Override default implementation to only return png files from the source folder + + Parameters + ---------- + fast_count: bool + Not used for faces loader + count: int + The number of images that the loader will encounter if already known, otherwise + ``None`` + """ + if isinstance(self.location, (list, tuple)): + file_list = self.location + else: + file_list = get_image_paths(self.location) + + self._file_list = [fname for fname in file_list + if os.path.splitext(fname)[-1].lower() == ".png"] + self._count = len(self.file_list) if count is None else count + + logger.debug("count: %s", self.count) + logger.trace("filelist: %s", self.file_list) + def _from_folder(self): """ Generator for loading images from a folder Faces will only ever be loaded from a folder, so this is the only function requiring @@ -1333,7 +1466,10 @@ def image_from_index(self, index): image = self._reader.get_data(index)[..., ::-1] filename = self._dummy_video_framename(index) else: - filename = self.file_list[index] + file_list = [f for idx, f in enumerate(self._file_list) + if idx not in self._skip_list] if self._skip_list else self._file_list + + filename = file_list[index] image = read_image(filename, raise_error=True) filename = os.path.basename(filename) logger.trace("index: %s, filename: %s image shape: %s", index, filename, image.shape) @@ -1405,7 +1541,10 @@ def _process(self, queue): executor.submit(self._save, *item) executor.shutdown() - def _save(self, filename, image): + def _save(self, + filename: str, + image: bytes | np.ndarray, + sub_folder: str | None) -> None: """ Save a single image inside a ThreadPoolExecutor Parameters @@ -1413,21 +1552,35 @@ def _save(self, filename, image): filename: str The filename of the image to be saved. NB: Any folders passed in with the filename will be stripped and replaced with :attr:`location`. - image: numpy.ndarray - The image to be saved + image: bytes or :class:`numpy.ndarray` + The encoded image or numpy array to be saved + subfolder: str or ``None`` + If the file should be saved in a subfolder in the output location, the subfolder should + be provided here. ``None`` for no subfolder. """ - filename = os.path.join(self.location, os.path.basename(filename)) + location = os.path.join(self.location, sub_folder) if sub_folder else self._location + if sub_folder and not os.path.exists(location): + os.makedirs(location) + + filename = os.path.join(location, os.path.basename(filename)) try: if self._as_bytes: + assert isinstance(image, bytes) with open(filename, "wb") as out_file: out_file.write(image) else: + assert isinstance(image, np.ndarray) cv2.imwrite(filename, image) - logger.trace("Saved image: '%s'", filename) - except Exception as err: # pylint: disable=broad-except - logger.error("Failed to save image '%s'. Original Error: %s", filename, err) - - def save(self, filename, image): + logger.trace("Saved image: '%s'", filename) # type:ignore + except Exception as err: # pylint:disable=broad-except + logger.error("Failed to save image '%s'. Original Error: %s", filename, str(err)) + del image + del filename + + def save(self, + filename: str, + image: bytes | np.ndarray, + sub_folder: str | None = None) -> None: """ Save the given image in the background thread Ensure that :func:`close` is called once all save operations are complete. @@ -1435,13 +1588,17 @@ def save(self, filename, image): Parameters ---------- filename: str - The filename of the image to be saved - image: numpy.ndarray - The image to be saved + The filename of the image to be saved. NB: Any folders passed in with the filename + will be stripped and replaced with :attr:`location`. + image: bytes + The encoded image to be saved + subfolder: str, optional + If the file should be saved in a subfolder in the output location, the subfolder should + be provided here. ``None`` for no subfolder. Default: ``None`` """ self._set_thread() - logger.trace("Putting to save queue: '%s'", filename) - self._queue.put((filename, image)) + logger.trace("Putting to save queue: '%s'", filename) # type:ignore + self._queue.put((filename, image, sub_folder)) def close(self): """ Signal to the Save Threads that they should be closed and cleanly shutdown diff --git a/lib/keras_utils.py b/lib/keras_utils.py new file mode 100644 index 0000000000..af47a3e466 --- /dev/null +++ b/lib/keras_utils.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python3 +""" Common multi-backend Keras utilities """ +from __future__ import annotations +import typing as T + +import numpy as np + +import tensorflow.keras.backend as K # pylint:disable=import-error + +if T.TYPE_CHECKING: + from tensorflow import Tensor + + +def frobenius_norm(matrix: Tensor, + axis: int = -1, + keep_dims: bool = True, + epsilon: float = 1e-15) -> Tensor: + """ Frobenius normalization for Keras Tensor + + Parameters + ---------- + matrix: Tensor + The matrix to normalize + axis: int, optional + The axis to normalize. Default: `-1` + keep_dims: bool, Optional + Whether to retain the original matrix shape or not. Default:``True`` + epsilon: flot, optional + Epsilon to apply to the normalization to preven NaN errors on zero values + + Returns + ------- + Tensor + The normalized output + """ + return K.sqrt(K.sum(K.pow(matrix, 2), axis=axis, keepdims=keep_dims) + epsilon) + + +def replicate_pad(image: Tensor, padding: int) -> Tensor: + """ Apply replication padding to an input batch of images. Expects 4D tensor in BHWC format. + + Notes + ----- + At the time of writing Keras/Tensorflow does not have a native replication padding method. + The implementation here is probably not the most efficient, but it is a pure keras method + which should work on TF. + + Parameters + ---------- + image: Tensor + Image tensor to pad + pad: int + The amount of padding to apply to each side of the input image + + Returns + ------- + Tensor + The input image with replication padding applied + """ + top_pad = K.tile(image[:, :1, ...], (1, padding, 1, 1)) + bottom_pad = K.tile(image[:, -1:, ...], (1, padding, 1, 1)) + pad_top_bottom = K.concatenate([top_pad, image, bottom_pad], axis=1) + left_pad = K.tile(pad_top_bottom[..., :1, :], (1, 1, padding, 1)) + right_pad = K.tile(pad_top_bottom[..., -1:, :], (1, 1, padding, 1)) + padded = K.concatenate([left_pad, pad_top_bottom, right_pad], axis=2) + return padded + + +class ColorSpaceConvert(): + """ Transforms inputs between different color spaces on the GPU + + Notes + ----- + The following color space transformations are implemented: + - rgb to lab + - rgb to xyz + - srgb to _rgb + - srgb to ycxcz + - xyz to ycxcz + - xyz to lab + - xyz to rgb + - ycxcz to rgb + - ycxcz to xyz + + Parameters + ---------- + from_space: str + One of `"srgb"`, `"rgb"`, `"xyz"` + to_space: str + One of `"lab"`, `"rgb"`, `"ycxcz"`, `"xyz"` + + Raises + ------ + ValueError + If the requested color space conversion is not defined + """ + def __init__(self, from_space: str, to_space: str) -> None: + functions = {"rgb_lab": self._rgb_to_lab, + "rgb_xyz": self._rgb_to_xyz, + "srgb_rgb": self._srgb_to_rgb, + "srgb_ycxcz": self._srgb_to_ycxcz, + "xyz_ycxcz": self._xyz_to_ycxcz, + "xyz_lab": self._xyz_to_lab, + "xyz_to_rgb": self._xyz_to_rgb, + "ycxcz_rgb": self._ycxcz_to_rgb, + "ycxcz_xyz": self._ycxcz_to_xyz} + func_name = f"{from_space.lower()}_{to_space.lower()}" + if func_name not in functions: + raise ValueError(f"The color transform {from_space} to {to_space} is not defined.") + + self._func = functions[func_name] + self._ref_illuminant = K.constant(np.array([[[0.950428545, 1.000000000, 1.088900371]]]), + dtype="float32") + self._inv_ref_illuminant = 1. / self._ref_illuminant + + self._rgb_xyz_map = self._get_rgb_xyz_map() + self._xyz_multipliers = K.constant([116, 500, 200], dtype="float32") + + @classmethod + def _get_rgb_xyz_map(cls) -> tuple[Tensor, Tensor]: + """ Obtain the mapping and inverse mapping for rgb to xyz color space conversion. + + Returns + ------- + tuple + The mapping and inverse Tensors for rgb to xyz color space conversion + """ + mapping = np.array([[10135552 / 24577794, 8788810 / 24577794, 4435075 / 24577794], + [2613072 / 12288897, 8788810 / 12288897, 887015 / 12288897], + [1425312 / 73733382, 8788810 / 73733382, 70074185 / 73733382]]) + inverse = np.linalg.inv(mapping) + return (K.constant(mapping, dtype="float32"), K.constant(inverse, dtype="float32")) + + def __call__(self, image: Tensor) -> Tensor: + """ Call the colorspace conversion function. + + Parameters + ---------- + image: Tensor + The image tensor in the colorspace defined by :param:`from_space` + + Returns + ------- + Tensor + The image tensor in the colorspace defined by :param:`to_space` + """ + return self._func(image) + + def _rgb_to_lab(self, image: Tensor) -> Tensor: + """ RGB to LAB conversion. + + Parameters + ---------- + image: Tensor + The image tensor in RGB format + + Returns + ------- + Tensor + The image tensor in LAB format + """ + converted = self._rgb_to_xyz(image) + return self._xyz_to_lab(converted) + + def _rgb_xyz_rgb(self, image: Tensor, mapping: Tensor) -> Tensor: + """ RGB to XYZ or XYZ to RGB conversion. + + Notes + ----- + The conversion in both directions is the same, but the mappping matrix for XYZ to RGB is + the inverse of RGB to XYZ. + + References + ---------- + https://www.image-engineering.de/library/technotes/958-how-to-convert-between-srgb-and-ciexyz + + Parameters + ---------- + mapping: Tensor + The mapping matrix to perform either the XYZ to RGB or RGB to XYZ color space + conversion + + image: Tensor + The image tensor in RGB format + + Returns + ------- + Tensor + The image tensor in XYZ format + """ + dim = K.int_shape(image) + image = K.permute_dimensions(image, (0, 3, 1, 2)) + image = K.reshape(image, (dim[0], dim[3], dim[1] * dim[2])) + converted = K.permute_dimensions(K.dot(mapping, image), (1, 2, 0)) + return K.reshape(converted, dim) + + def _rgb_to_xyz(self, image: Tensor) -> Tensor: + """ RGB to XYZ conversion. + + Parameters + ---------- + image: Tensor + The image tensor in RGB format + + Returns + ------- + Tensor + The image tensor in XYZ format + """ + return self._rgb_xyz_rgb(image, self._rgb_xyz_map[0]) + + @classmethod + def _srgb_to_rgb(cls, image: Tensor) -> Tensor: + """ SRGB to RGB conversion. + + Notes + ----- + RGB Image is clipped to a small epsilon to stabalize training + + Parameters + ---------- + image: Tensor + The image tensor in SRGB format + + Returns + ------- + Tensor + The image tensor in RGB format + """ + limit = 0.04045 + return K.switch(image > limit, + K.pow((K.clip(image, limit, None) + 0.055) / 1.055, 2.4), + image / 12.92) + + def _srgb_to_ycxcz(self, image: Tensor) -> Tensor: + """ SRGB to YcXcZ conversion. + + Parameters + ---------- + image: Tensor + The image tensor in SRGB format + + Returns + ------- + Tensor + The image tensor in YcXcZ format + """ + converted = self._srgb_to_rgb(image) + converted = self._rgb_to_xyz(converted) + return self._xyz_to_ycxcz(converted) + + def _xyz_to_lab(self, image: Tensor) -> Tensor: + """ XYZ to LAB conversion. + + Parameters + ---------- + image: Tensor + The image tensor in XYZ format + + Returns + ------- + Tensor + The image tensor in LAB format + """ + image = image * self._inv_ref_illuminant + delta = 6 / 29 + delta_cube = delta ** 3 + factor = 1 / (3 * (delta ** 2)) + + clamped_term = K.pow(K.clip(image, delta_cube, None), 1.0 / 3.0) + div = factor * image + (4 / 29) + + image = K.switch(image > delta_cube, clamped_term, div) + return K.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16., + self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])], + axis=-1) + + def _xyz_to_rgb(self, image: Tensor) -> Tensor: + """ XYZ to YcXcZ conversion. + + Parameters + ---------- + image: Tensor + The image tensor in XYZ format + + Returns + ------- + Tensor + The image tensor in RGB format + """ + return self._rgb_xyz_rgb(image, self._rgb_xyz_map[1]) + + def _xyz_to_ycxcz(self, image: Tensor) -> Tensor: + """ XYZ to YcXcZ conversion. + + Parameters + ---------- + image: Tensor + The image tensor in XYZ format + + Returns + ------- + Tensor + The image tensor in YcXcZ format + """ + image = image * self._inv_ref_illuminant + return K.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16., + self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])], + axis=-1) + + def _ycxcz_to_rgb(self, image: Tensor) -> Tensor: + """ YcXcZ to RGB conversion. + + Parameters + ---------- + image: Tensor + The image tensor in YcXcZ format + + Returns + ------- + Tensor + The image tensor in RGB format + """ + converted = self._ycxcz_to_xyz(image) + return self._xyz_to_rgb(converted) + + def _ycxcz_to_xyz(self, image: Tensor) -> Tensor: + """ YcXcZ to XYZ conversion. + + Parameters + ---------- + image: Tensor + The image tensor in YcXcZ format + + Returns + ------- + Tensor + The image tensor in XYZ format + """ + ch_y = (image[..., 0:1] + 16.) / self._xyz_multipliers[0] + return K.concatenate([ch_y + (image[..., 1:2] / self._xyz_multipliers[1]), + ch_y, + ch_y - (image[..., 2:3] / self._xyz_multipliers[2])], + axis=-1) * self._ref_illuminant diff --git a/lib/keypress.py b/lib/keypress.py index 55c4450a5b..c5c2030216 100644 --- a/lib/keypress.py +++ b/lib/keypress.py @@ -21,7 +21,7 @@ # Windows if os.name == "nt": - import msvcrt # pylint: disable=import-error + import msvcrt # pylint:disable=import-error # Posix (Linux, OS X) else: @@ -43,7 +43,7 @@ def __init__(self, is_gui=False): self.old_term = termios.tcgetattr(self.file_desc) # New terminal setting unbuffered - self.new_term[3] = (self.new_term[3] & ~termios.ICANON & ~termios.ECHO) + self.new_term[3] = self.new_term[3] & ~termios.ICANON & ~termios.ECHO termios.tcsetattr(self.file_desc, termios.TCSAFLUSH, self.new_term) # Support normal-terminal reset at exit @@ -62,7 +62,7 @@ def getch(self): if (self.is_gui or not sys.stdout.isatty()) and os.name != "nt": return None if os.name == "nt": - return msvcrt.getch().decode("utf-8") + return msvcrt.getch().decode("utf-8", errors="replace") return sys.stdin.read(1) def getarrow(self): @@ -83,7 +83,7 @@ def getarrow(self): char = sys.stdin.read(3)[2] vals = [65, 67, 66, 68] - return vals.index(ord(char.decode("utf-8"))) + return vals.index(ord(char.decode("utf-8", errors="replace"))) def kbhit(self): """ Returns True if keyboard character was hit, False otherwise. """ diff --git a/lib/logger.py b/lib/logger.py index 452080c3b1..89fb036627 100644 --- a/lib/logger.py +++ b/lib/logger.py @@ -1,19 +1,34 @@ #!/usr/bin/python """ Logging Functions for Faceswap. """ +# NOTE: Don't import non stdlib packages. This module is accessed by setup.py import collections import logging from logging.handlers import RotatingFileHandler import os +import platform +import re import sys +import typing as T +import time import traceback from datetime import datetime -from tqdm import tqdm + + +# TODO - Remove this monkey patch when TF autograph fixed to handle newer logging lib +def _patched_format(self, record): + """ Autograph tf-2.10 has a bug with the 3.10 version of logging.PercentStyle._format(). It is + non-critical but spits out warnings. This is the Python 3.9 version of the function and should + be removed once fixed """ + return self._fmt % record.__dict__ # pylint:disable=protected-access + + +setattr(logging.PercentStyle, "_format", _patched_format) class FaceswapLogger(logging.Logger): """ A standard :class:`logging.logger` with additional "verbose" and "trace" levels added. """ - def __init__(self, name): + def __init__(self, name: str) -> None: for new_level in (("VERBOSE", 15), ("TRACE", 5)): level_name, level_num = new_level if hasattr(logging, level_name): @@ -22,7 +37,7 @@ def __init__(self, name): setattr(logging, level_name, level_num) super().__init__(name) - def verbose(self, msg, *args, **kwargs): + def verbose(self, msg: str, *args, **kwargs) -> None: # pylint:disable=wrong-spelling-in-docstring """ Create a log message at severity level 15. @@ -38,7 +53,7 @@ def verbose(self, msg, *args, **kwargs): if self.isEnabledFor(15): self._log(15, msg, args, **kwargs) - def trace(self, msg, *args, **kwargs): + def trace(self, msg: str, *args, **kwargs) -> None: # pylint:disable=wrong-spelling-in-docstring """ Create a log message at severity level 5. @@ -55,6 +70,122 @@ def trace(self, msg, *args, **kwargs): self._log(5, msg, args, **kwargs) +class ColoredFormatter(logging.Formatter): + """ Overrides the stand :class:`logging.Formatter` to enable colored labels for message level + labels on supported platforms + + Parameters + ---------- + fmt: str + The format string for the message as a whole + pad_newlines: bool, Optional + If ``True`` new lines will be padded to appear in line with the log message, if ``False`` + they will be left aligned + + kwargs: dict + Standard :class:`logging.Formatter` keyword arguments + """ + def __init__(self, fmt: str, pad_newlines: bool = False, **kwargs) -> None: + super().__init__(fmt, **kwargs) + self._use_color = self._get_color_compatibility() + self._level_colors = {"CRITICAL": "\033[31m", # red + "ERROR": "\033[31m", # red + "WARNING": "\033[33m", # yellow + "INFO": "\033[32m", # green + "VERBOSE": "\033[34m"} # blue + self._default_color = "\033[0m" + self._newline_padding = self._get_newline_padding(pad_newlines, fmt) + + @classmethod + def _get_color_compatibility(cls) -> bool: + """ Return whether the system supports color ansi codes. Most OSes do other than Windows + below Windows 10 version 1511. + + Returns + ------- + bool + ``True`` if the system supports color ansi codes otherwise ``False`` + """ + if platform.system().lower() != "windows": + return True + try: + win = sys.getwindowsversion() # type:ignore # pylint:disable=no-member + if win.major >= 10 and win.build >= 10586: + return True + except Exception: # pylint:disable=broad-except + return False + return False + + def _get_newline_padding(self, pad_newlines: bool, fmt: str) -> int: + """ Parses the format string to obtain padding for newlines if requested + + Parameters + ---------- + fmt: str + The format string for the message as a whole + pad_newlines: bool, Optional + If ``True`` new lines will be padded to appear in line with the log message, if + ``False`` they will be left aligned + + Returns + ------- + int + The amount of padding to apply to the front of newlines + """ + if not pad_newlines: + return 0 + msg_idx = fmt.find("%(message)") + 1 + filtered = fmt[:msg_idx - 1] + spaces = filtered.count(" ") + pads = [int(pad.replace("s", "")) for pad in re.findall(r"\ds", filtered)] + if "asctime" in filtered: + pads.append(self._get_sample_time_string()) + return sum(pads) + spaces + + def _get_sample_time_string(self) -> int: + """ Obtain a sample time string and calculate correct padding. + + This may be inaccurate when ticking over an integer from single to double digits, but that + shouldn't be a huge issue. + + Returns + ------- + int + The length of the formatted date-time string + """ + sample_time = time.time() + date_format = self.datefmt if self.datefmt else self.default_time_format + datestring = time.strftime(date_format, logging.Formatter.converter(sample_time)) + if not self.datefmt and self.default_msec_format: + msecs = (sample_time - int(sample_time)) * 1000 + datestring = self.default_msec_format % (datestring, msecs) + return len(datestring) + + def format(self, record: logging.LogRecord) -> str: + """ Color the log message level if supported otherwise return the standard log message. + + Parameters + ---------- + record: :class:`logging.LogRecord` + The incoming log record to be formatted for entry into the logger. + + Returns + ------- + str + The formatted log message + """ + formatted = super().format(record) + levelname = record.levelname + if self._use_color and levelname in self._level_colors: + formatted = re.sub(levelname, + f"{self._level_colors[levelname]}{levelname}{self._default_color}", + formatted, + 1) + if self._newline_padding: + formatted = formatted.replace("\n", f"\n{' ' * self._newline_padding}") + return formatted + + class FaceswapFormatter(logging.Formatter): """ Overrides the standard :class:`logging.Formatter`. @@ -63,7 +194,7 @@ class FaceswapFormatter(logging.Formatter): Rewrites some upstream warning messages to debug level to avoid spamming the console. """ - def format(self, record): + def format(self, record: logging.LogRecord) -> str: """ Strip new lines from log records and rewrite certain warning messages to debug level. Parameters @@ -102,7 +233,7 @@ def format(self, record): return msg @classmethod - def _rewrite_warnings(cls, record): + def _rewrite_warnings(cls, record: logging.LogRecord) -> logging.LogRecord: """ Change certain warning messages from WARNING to DEBUG to avoid passing non-important information to output. @@ -133,7 +264,7 @@ def _rewrite_warnings(cls, record): return record @classmethod - def _lower_external(cls, record): + def _lower_external(cls, record: logging.LogRecord) -> logging.LogRecord: """ Some external libs log at a higher level than we would really like, so lower their log level. @@ -162,7 +293,7 @@ class RollingBuffer(collections.deque): """File-like that keeps a certain number of lines of text in memory for writing out to the crash log. """ - def write(self, buffer): + def write(self, buffer: str) -> None: """ Splits lines from the incoming buffer and writes them out to the rolling buffer. Parameters @@ -178,7 +309,7 @@ class TqdmHandler(logging.StreamHandler): """ Overrides :class:`logging.StreamHandler` to use :func:`tqdm.tqdm.write` rather than writing to :func:`sys.stderr` so that log messages do not mess up tqdm progress bars. """ - def emit(self, record): + def emit(self, record: logging.LogRecord) -> None: """ Format the incoming message and pass to :func:`tqdm.tqdm.write`. Parameters @@ -186,11 +317,13 @@ def emit(self, record): record : :class:`logging.LogRecord` The incoming log record to be formatted for entry into the logger. """ + # tqdm is imported here as it won't be installed when setup.py is running + from tqdm import tqdm # pylint:disable=import-outside-toplevel msg = self.format(record) tqdm.write(msg) -def _set_root_logger(loglevel=logging.INFO): +def _set_root_logger(loglevel: int = logging.INFO) -> logging.Logger: """ Setup the root logger. Parameters @@ -208,7 +341,7 @@ def _set_root_logger(loglevel=logging.INFO): return rootlogger -def log_setup(loglevel, log_file, command, is_gui=False): +def log_setup(loglevel, log_file: str, command: str, is_gui: bool = False) -> None: """ Set up logging for Faceswap. Sets up the root logger, the formatting for the crash logger and the file logger, and sets up @@ -230,19 +363,32 @@ def log_setup(loglevel, log_file, command, is_gui=False): numeric_loglevel = get_loglevel(loglevel) root_loglevel = min(logging.DEBUG, numeric_loglevel) rootlogger = _set_root_logger(loglevel=root_loglevel) - log_format = FaceswapFormatter("%(asctime)s %(processName)-15s %(threadName)-30s " - "%(module)-15s %(funcName)-30s %(levelname)-8s %(message)s", - datefmt="%m/%d/%Y %H:%M:%S") - f_handler = _file_handler(numeric_loglevel, log_file, log_format, command) - s_handler = _stream_handler(numeric_loglevel, is_gui) - c_handler = _crash_handler(log_format) + + if command == "setup": + log_format = FaceswapFormatter("%(asctime)s %(module)-16s %(funcName)-30s %(levelname)-8s " + "%(message)s", datefmt="%m/%d/%Y %H:%M:%S") + s_handler = _stream_setup_handler(numeric_loglevel) + f_handler = _file_handler(root_loglevel, log_file, log_format, command) + else: + log_format = FaceswapFormatter("%(asctime)s %(processName)-15s %(threadName)-30s " + "%(module)-15s %(funcName)-30s %(levelname)-8s %(message)s", + datefmt="%m/%d/%Y %H:%M:%S") + s_handler = _stream_handler(numeric_loglevel, is_gui) + f_handler = _file_handler(numeric_loglevel, log_file, log_format, command) + rootlogger.addHandler(f_handler) rootlogger.addHandler(s_handler) - rootlogger.addHandler(c_handler) - logging.info("Log level set to: %s", loglevel.upper()) + + if command != "setup": + c_handler = _crash_handler(log_format) + rootlogger.addHandler(c_handler) + logging.info("Log level set to: %s", loglevel.upper()) -def _file_handler(loglevel, log_file, log_format, command): +def _file_handler(loglevel, + log_file: str, + log_format: FaceswapFormatter, + command: str) -> RotatingFileHandler: """ Add a rotating file handler for the current Faceswap session. 1 backup is always kept. Parameters @@ -262,7 +408,7 @@ def _file_handler(loglevel, log_file, log_format, command): :class:`logging.RotatingFileHandler` The logging file handler """ - if log_file is not None: + if log_file: filename = log_file else: filename = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "faceswap") @@ -270,21 +416,21 @@ def _file_handler(loglevel, log_file, log_format, command): filename += "_gui.log" if command == "gui" else ".log" should_rotate = os.path.isfile(filename) - log_file = RotatingFileHandler(filename, backupCount=1, encoding="utf-8") + handler = RotatingFileHandler(filename, backupCount=1, encoding="utf-8") if should_rotate: - log_file.doRollover() - log_file.setFormatter(log_format) - log_file.setLevel(loglevel) - return log_file + handler.doRollover() + handler.setFormatter(log_format) + handler.setLevel(loglevel) + return handler -def _stream_handler(loglevel, is_gui): +def _stream_handler(loglevel: int, is_gui: bool) -> logging.StreamHandler | TqdmHandler: """ Add a stream handler for the current Faceswap session. The stream handler will only ever output at a maximum of VERBOSE level to avoid spamming the console. Parameters ---------- - loglevel: str + loglevel: int The requested log level that messages should be logged at. is_gui: bool, optional Whether Faceswap is running in the GUI or not. Dictates where the stream handler should @@ -311,7 +457,30 @@ def _stream_handler(loglevel, is_gui): return log_console -def _crash_handler(log_format): +def _stream_setup_handler(loglevel: int) -> logging.StreamHandler: + """ Add a stream handler for faceswap's setup.py script + This stream handler outputs a limited set of easy to use information using colored labels + if available. It will only ever output at a minimum of INFO level + + Parameters + ---------- + loglevel: int + The requested log level that messages should be logged at. + + Returns + ------- + :class:`logging.StreamHandler` + The stream handler to use + """ + loglevel = max(loglevel, 15) + log_format = ColoredFormatter("%(levelname)-8s %(message)s", pad_newlines=True) + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(log_format) + handler.setLevel(loglevel) + return handler + + +def _crash_handler(log_format: FaceswapFormatter) -> logging.StreamHandler: """ Add a handler that stores the last 100 debug lines to :attr:'_DEBUG_BUFFER' for use in crash reports. @@ -331,7 +500,7 @@ def _crash_handler(log_format): return log_crash -def get_loglevel(loglevel): +def get_loglevel(loglevel: str) -> int: """ Check whether a valid log level has been supplied, and return the numeric log level that corresponds to the given string level. @@ -351,7 +520,7 @@ def get_loglevel(loglevel): return numeric_level -def crash_log(): +def crash_log() -> str: """ On a crash, write out the contents of :func:`_DEBUG_BUFFER` containing the last 100 lines of debug messages to a crash report in the root Faceswap folder. @@ -376,14 +545,60 @@ def crash_log(): return filename +def _process_value(value: T.Any) -> T.Any: + """ Process the values from a local dict and return in a loggable format + + Parameters + ---------- + value: Any + The dictionary value + + Returns + ------- + Any + The original or ammended value + """ + if isinstance(value, str): + return f'"{value}"' + if isinstance(value, (list, tuple, set)) and len(value) > 10: + return f'[type: "{type(value).__name__}" len: {len(value)}' + + try: + import numpy as np # pylint:disable=import-outside-toplevel + except ImportError: + return value + + if isinstance(value, np.ndarray) and np.prod(value.shape) > 10: + return f'[type: "{type(value).__name__}" shape: {value.shape}, dtype: "{value.dtype}"]' + return value + + +def parse_class_init(locals_dict: dict[str, T.Any]) -> str: + """ Parse a locals dict from a class and return in a format suitable for logging + Parameters + ---------- + locals_dict: dict[str, T.Any] + A locals() dictionary from a newly initialized class + Returns + ------- + str + The locals information suitable for logging + """ + delimit = {k: _process_value(v) + for k, v in locals_dict.items() if k != "self"} + dsp = ", ".join(f"{k}: {v}" for k, v in delimit.items()) + dsp = f" ({dsp})" if dsp else "" + return f"Initializing {locals_dict['self'].__class__.__name__}{dsp}" + + _OLD_FACTORY = logging.getLogRecordFactory() -def _faceswap_logrecord(*args, **kwargs): +def _faceswap_logrecord(*args, **kwargs) -> logging.LogRecord: """ Add a flag to :class:`logging.LogRecord` to not strip formatting from particular records. """ record = _OLD_FACTORY(*args, **kwargs) - record.strip_spaces = True + record.strip_spaces = True # type:ignore return record diff --git a/lib/model/__init__.py b/lib/model/__init__.py index 9ac90e31cb..e69de29bb2 100644 --- a/lib/model/__init__.py +++ b/lib/model/__init__.py @@ -1,12 +0,0 @@ -#!/usr/bin/env python3 -""" Conditional imports depending on whether the AMD version is installed or not """ - -from lib.utils import get_backend - -from .normalization import * -if get_backend() == "amd": - from . import losses_plaid as losses - from . import optimizers_plaid as optimizers -else: - from . import losses_tf as losses - from . import optimizers_tf as optimizers diff --git a/lib/model/autoclip.py b/lib/model/autoclip.py new file mode 100644 index 0000000000..826959d66b --- /dev/null +++ b/lib/model/autoclip.py @@ -0,0 +1,111 @@ +""" Auto clipper for clipping gradients. """ +import numpy as np +import tensorflow as tf + + +class AutoClipper(): + """ AutoClip: Adaptive Gradient Clipping for Source Separation Networks + + Parameters + ---------- + clip_percentile: int + The percentile to clip the gradients at + history_size: int, optional + The number of iterations of data to use to calculate the norm + Default: ``10000`` + + References + ---------- + tf implementation: https://github.com/pseeth/autoclip + original paper: https://arxiv.org/abs/2007.14469 + """ + def __init__(self, clip_percentile: int, history_size: int = 10000): + self._clip_percentile = tf.cast(clip_percentile, tf.float64) + self._grad_history = tf.Variable(tf.zeros(history_size), trainable=False) + self._index = tf.Variable(0, trainable=False) + self._history_size = history_size + + def _percentile(self, grad_history: tf.Tensor) -> tf.Tensor: + """ Compute the clip percentile of the gradient history + + Parameters + ---------- + grad_history: :class:`tensorflow.Tensor` + Tge gradient history to calculate the clip percentile for + + Returns + ------- + :class:`tensorflow.Tensor` + A rank(:attr:`clip_percentile`) `Tensor` + + Notes + ----- + Adapted from + https://github.com/tensorflow/probability/blob/r0.14/tensorflow_probability/python/stats/quantiles.py + to remove reliance on full tensorflow_probability libraray + """ + with tf.name_scope("percentile"): + frac_at_q_or_below = self._clip_percentile / 100. + sorted_hist = tf.sort(grad_history, axis=-1, direction="ASCENDING") + + num = tf.cast(tf.shape(grad_history)[-1], tf.float64) + + # get indices + indices = tf.round((num - 1) * frac_at_q_or_below) + indices = tf.clip_by_value(tf.cast(indices, tf.int32), + 0, + tf.shape(grad_history)[-1] - 1) + gathered_hist = tf.gather(sorted_hist, indices, axis=-1) + + # Propagate NaNs. Apparently tf.is_nan doesn't like other dtypes + nan_batch_members = tf.reduce_any(tf.math.is_nan(grad_history), axis=None) + right_rank_matched_shape = tf.pad(tf.shape(nan_batch_members), + paddings=[[0, tf.rank(self._clip_percentile)]], + constant_values=1) + nan_batch_members = tf.reshape(nan_batch_members, shape=right_rank_matched_shape) + + nan = np.array(np.nan, gathered_hist.dtype.as_numpy_dtype) + gathered_hist = tf.where(nan_batch_members, nan, gathered_hist) + + return gathered_hist + + def __call__(self, grads_and_vars: list[tf.Tensor]) -> list[tf.Tensor]: + """ Call the AutoClip function. + + Parameters + ---------- + grads_and_vars: list + The list of gradient tensors and variables for the optimizer + """ + grad_norms = [self._get_grad_norm(g) for g, _ in grads_and_vars] + total_norm = tf.norm(grad_norms) + assign_idx = tf.math.mod(self._index, self._history_size) + self._grad_history = self._grad_history[assign_idx].assign(total_norm) + self._index = self._index.assign_add(1) + clip_value = self._percentile(self._grad_history[: self._index]) + return [(tf.clip_by_norm(g, clip_value), v) for g, v in grads_and_vars] + + @classmethod + def _get_grad_norm(cls, gradients: tf.Tensor) -> tf.Tensor: + """ Obtain the L2 Norm for the gradients + + Parameters + ---------- + gradients: :class:`tensorflow.Tensor` + The gradients to calculate the L2 norm for + + Returns + ------- + :class:`tensorflow.Tensor` + The L2 Norm of the given gradients + """ + values = tf.convert_to_tensor(gradients.values + if isinstance(gradients, tf.IndexedSlices) + else gradients, name="t") + + # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm + l2sum = tf.math.reduce_sum(values * values, axis=None, keepdims=True) + pred = l2sum > 0 + # Two-tap tf.where trick to bypass NaN gradients + l2sum_safe = tf.where(pred, l2sum, tf.ones_like(l2sum)) + return tf.squeeze(tf.where(pred, tf.math.sqrt(l2sum_safe), l2sum)) diff --git a/lib/model/backup_restore.py b/lib/model/backup_restore.py index 0408040b85..e143266c9e 100644 --- a/lib/model/backup_restore.py +++ b/lib/model/backup_restore.py @@ -10,7 +10,7 @@ from lib.serializer import get_serializer from lib.utils import get_folder -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Backup(): diff --git a/lib/model/initializers.py b/lib/model/initializers.py index c1d0204710..41f7682371 100644 --- a/lib/model/initializers.py +++ b/lib/model/initializers.py @@ -8,18 +8,12 @@ import numpy as np import tensorflow as tf -from lib.utils import get_backend +# Fix intellisense/linting for tf.keras' thoroughly broken import system +keras = tf.keras +K = keras.backend -if get_backend() == "amd": - from keras.utils import get_custom_objects # pylint:disable=no-name-in-module - from keras import backend as K - from keras import initializers -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.utils import get_custom_objects # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.keras import initializers, backend as K # noqa pylint:disable=no-name-in-module,import-error -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) def compute_fans(shape, data_format='channels_last'): @@ -70,7 +64,7 @@ def compute_fans(shape, data_format='channels_last'): return fan_in, fan_out -class ICNR(initializers.Initializer): # pylint: disable=invalid-name,no-member +class ICNR(keras.initializers.Initializer): # type:ignore[name-defined] """ ICNR initializer for checkerboard artifact free sub pixel convolution Parameters @@ -100,7 +94,7 @@ def __init__(self, initializer, scale=2): self.scale = scale self.initializer = initializer - def __call__(self, shape, dtype="float32"): + def __call__(self, shape, dtype="float32", **kwargs): """ Call function for the ICNR initializer. Parameters @@ -120,7 +114,7 @@ def __call__(self, shape, dtype="float32"): return self.initializer(shape) new_shape = shape[:3] + [shape[3] // (self.scale ** 2)] if isinstance(self.initializer, dict): - self.initializer = initializers.deserialize(self.initializer) + self.initializer = keras.initializers.deserialize(self.initializer) var_x = self.initializer(new_shape, dtype) var_x = K.permute_dimensions(var_x, [2, 0, 1, 3]) var_x = K.resize_images(var_x, @@ -136,9 +130,6 @@ def __call__(self, shape, dtype="float32"): def _space_to_depth(self, input_tensor): """ Space to depth implementation. - PlaidML does not have a space to depth operation, so calculate if backend is amd - otherwise returns the :func:`tensorflow.space_to_depth` operation. - Parameters ---------- input_tensor: tensor @@ -149,16 +140,7 @@ def _space_to_depth(self, input_tensor): tensor The manipulated input tensor """ - if get_backend() == "amd": - batch, height, width, depth = input_tensor.shape.dims - new_height = height // self.scale - new_width = width // self.scale - reshaped = K.reshape(input_tensor, - (batch, new_height, self.scale, new_width, self.scale, depth)) - retval = K.reshape(K.permute_dimensions(reshaped, [0, 1, 3, 2, 4, 5]), - (batch, new_height, new_width, -1)) - else: - retval = tf.nn.space_to_depth(input_tensor, block_size=self.scale, data_format="NHWC") + retval = tf.nn.space_to_depth(input_tensor, block_size=self.scale, data_format="NHWC") logger.debug("Input shape: %s, Output shape: %s", input_tensor.shape, retval.shape) return retval @@ -177,7 +159,7 @@ def get_config(self): return dict(list(base_config.items()) + list(config.items())) -class ConvolutionAware(initializers.Initializer): # pylint: disable=no-member +class ConvolutionAware(keras.initializers.Initializer): # type:ignore[name-defined] """ Initializer that generates orthogonal convolution filters in the Fourier space. If this initializer is passed a shape that is not 3D or 4D, orthogonal initialization will be used. @@ -210,11 +192,11 @@ class ConvolutionAware(initializers.Initializer): # pylint: disable=no-member def __init__(self, eps_std=0.05, seed=None, initialized=False): self.eps_std = eps_std self.seed = seed - self.orthogonal = initializers.Orthogonal() # pylint:disable=no-member - self.he_uniform = initializers.he_uniform() # pylint:disable=no-member + self.orthogonal = keras.initializers.Orthogonal() + self.he_uniform = keras.initializers.he_uniform() self.initialized = initialized - def __call__(self, shape, dtype=None): + def __call__(self, shape, dtype=None, **kwargs): """ Call function for the ICNR initializer. Parameters @@ -248,7 +230,7 @@ def __call__(self, shape, dtype=None): transpose_dimensions = (2, 1, 0) kernel_shape = (row,) - correct_ifft = lambda shape, s=[None]: np.fft.irfft(shape, s[0]) # noqa + correct_ifft = lambda shape, s=[None]: np.fft.irfft(shape, s[0]) # noqa:E501,E731 # pylint:disable=unnecessary-lambda-assignment correct_fft = np.fft.rfft elif rank == 4: @@ -317,12 +299,12 @@ def get_config(self): dict The configuration for ICNR Initialization """ - return dict(eps_std=self.eps_std, - seed=self.seed, - initialized=self.initialized) + return {"eps_std": self.eps_std, + "seed": self.seed, + "initialized": self.initialized} # Update initializers into Keras custom objects for name, obj in inspect.getmembers(sys.modules[__name__]): if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) + keras.utils.get_custom_objects().update({name: obj}) diff --git a/lib/model/layers.py b/lib/model/layers.py index 9fccc66aa5..4f2c9824c6 100644 --- a/lib/model/layers.py +++ b/lib/model/layers.py @@ -1,201 +1,106 @@ #!/usr/bin/env python3 """ Custom Layers for faceswap.py. """ - -from __future__ import absolute_import +from __future__ import annotations import sys import inspect +import typing as T import tensorflow as tf -from lib.utils import get_backend +# Fix intellisense/linting for tf.keras' thoroughly broken import system +from tensorflow.python.keras.utils import conv_utils # pylint:disable=no-name-in-module +keras = tf.keras +layers = keras.layers +K = keras.backend -if get_backend() == "amd": - from lib.plaidml_utils import pad - from keras.utils import get_custom_objects, conv_utils # pylint:disable=no-name-in-module - import keras.backend as K - from keras.layers import InputSpec, Layer -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.utils import get_custom_objects # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.keras import backend as K # pylint:disable=import-error - from tensorflow.keras.layers import InputSpec, Layer # noqa pylint:disable=no-name-in-module,import-error - from tensorflow import pad - from tensorflow.python.keras.utils import conv_utils # pylint:disable=no-name-in-module +class _GlobalPooling2D(tf.keras.layers.Layer): + """Abstract class for different global pooling 2D layers. -class PixelShuffler(Layer): - """ PixelShuffler layer for Keras. + From keras as access to pooling is trickier in tensorflow.keras + """ + def __init__(self, data_format: str | None = None, **kwargs) -> None: + super().__init__(**kwargs) + self.data_format = conv_utils.normalize_data_format(data_format) + self.input_spec = keras.layers.InputSpec(ndim=4) - This layer requires a Convolution2D prior to it, having output filters computed according to - the formula :math:`filters = k * (scale_factor * scale_factor)` where `k` is a user defined - number of filters (generally larger than 32) and `scale_factor` is the up-scaling factor - (generally 2). + def compute_output_shape(self, input_shape): + """ Compute the output shape based on the input shape. - This layer performs the depth to space operation on the convolution filters, and returns a - tensor with the size as defined below. + Parameters + ---------- + input_shape: tuple + The input shape to the layer + """ + if self.data_format == 'channels_last': + return (input_shape[0], input_shape[3]) + return (input_shape[0], input_shape[1]) - Notes - ----- - In practice, it is useful to have a second convolution layer after the - :class:`PixelShuffler` layer to speed up the learning process. However, if you are stacking - multiple :class:`PixelShuffler` blocks, it may increase the number of parameters greatly, - so the Convolution layer after :class:`PixelShuffler` layer can be removed. + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """ Override to call the layer. - Example - ------- - >>> # A standard sub-pixel up-scaling block - >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(...) - >>> u = PixelShuffler(size=(2, 2))(x) - [Optional] - >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(u) + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input to the layer + """ + raise NotImplementedError - Parameters - ---------- - size: tuple, optional - The (`h`, `w`) scaling factor for up-scaling. Default: `(2, 2)` - data_format: ["channels_first", "channels_last", ``None``], optional - The data format for the input. Default: ``None`` - kwargs: dict - The standard Keras Layer keyword arguments (if any) + def get_config(self) -> dict[str, T.Any]: + """ Set the Keras config """ + config = {'data_format': self.data_format} + base_config = super().get_config() + return dict(list(base_config.items()) + list(config.items())) - References - ---------- - https://gist.github.com/t-ae/6e1016cc188104d123676ccef3264981 - """ - def __init__(self, size=(2, 2), data_format=None, **kwargs): - super().__init__(**kwargs) - if get_backend() == "amd": - self.data_format = K.normalize_data_format(data_format) # pylint:disable=no-member - else: - self.data_format = conv_utils.normalize_data_format(data_format) - self.size = conv_utils.normalize_tuple(size, 2, 'size') - def call(self, inputs, *args, **kwargs): +class GlobalMinPooling2D(_GlobalPooling2D): + """Global minimum pooling operation for spatial data. """ + + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: """This is where the layer's logic lives. Parameters ---------- - inputs: tensor + inputs: :class:`tf.Tensor` Input tensor, or list/tuple of input tensors - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments Returns ------- tensor A tensor or list/tuple of tensors """ - input_shape = K.int_shape(inputs) - if len(input_shape) != 4: - raise ValueError('Inputs should have rank ' + - str(4) + - '; Received input shape:', str(input_shape)) - - if self.data_format == 'channels_first': - batch_size, channels, height, width = input_shape - if batch_size is None: - batch_size = -1 - r_height, r_width = self.size - o_height, o_width = height * r_height, width * r_width - o_channels = channels // (r_height * r_width) - - out = K.reshape(inputs, (batch_size, r_height, r_width, o_channels, height, width)) - out = K.permute_dimensions(out, (0, 3, 4, 1, 5, 2)) - out = K.reshape(out, (batch_size, o_channels, o_height, o_width)) - elif self.data_format == 'channels_last': - batch_size, height, width, channels = input_shape - if batch_size is None: - batch_size = -1 - r_height, r_width = self.size - o_height, o_width = height * r_height, width * r_width - o_channels = channels // (r_height * r_width) + if self.data_format == 'channels_last': + pooled = K.min(inputs, axis=[1, 2]) + else: + pooled = K.min(inputs, axis=[2, 3]) + return pooled - out = K.reshape(inputs, (batch_size, height, width, r_height, r_width, o_channels)) - out = K.permute_dimensions(out, (0, 1, 3, 2, 4, 5)) - out = K.reshape(out, (batch_size, o_height, o_width, o_channels)) - return out - def compute_output_shape(self, input_shape): - """Computes the output shape of the layer. +class GlobalStdDevPooling2D(_GlobalPooling2D): + """Global standard deviation pooling operation for spatial data. """ - Assumes that the layer will be built to match that input shape provided. + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """This is where the layer's logic lives. Parameters ---------- - input_shape: tuple or list of tuples - Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the - layer). Shape tuples can include None for free dimensions, instead of an integer. + inputs: tensor + Input tensor, or list/tuple of input tensors Returns ------- - tuple - An input shape tuple - """ - if len(input_shape) != 4: - raise ValueError('Inputs should have rank ' + - str(4) + - '; Received input shape:', str(input_shape)) - - if self.data_format == 'channels_first': - height = None - width = None - if input_shape[2] is not None: - height = input_shape[2] * self.size[0] - if input_shape[3] is not None: - width = input_shape[3] * self.size[1] - channels = input_shape[1] // self.size[0] // self.size[1] - - if channels * self.size[0] * self.size[1] != input_shape[1]: - raise ValueError('channels of input and size are incompatible') - - retval = (input_shape[0], - channels, - height, - width) - elif self.data_format == 'channels_last': - height = None - width = None - if input_shape[1] is not None: - height = input_shape[1] * self.size[0] - if input_shape[2] is not None: - width = input_shape[2] * self.size[1] - channels = input_shape[3] // self.size[0] // self.size[1] - - if channels * self.size[0] * self.size[1] != input_shape[3]: - raise ValueError('channels of input and size are incompatible') - - retval = (input_shape[0], - height, - width, - channels) - return retval - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) containing the configuration of a - layer. The same layer can be reinstated later (without its trained weights) from this - configuration. - - The configuration of a layer does not include connectivity information, nor the layer - class name. These are handled by `Network` (one layer of abstraction above). - - Returns - -------- - dict - A python dictionary containing the layer configuration + tensor + A tensor or list/tuple of tensors """ - config = {'size': self.size, - 'data_format': self.data_format} - base_config = super().get_config() - - return dict(list(base_config.items()) + list(config.items())) + if self.data_format == 'channels_last': + pooled = K.std(inputs, axis=[1, 2]) + else: + pooled = K.std(inputs, axis=[2, 3]) + return pooled -class KResizeImages(Layer): +class KResizeImages(tf.keras.layers.Layer): """ A custom upscale function that uses :class:`keras.backend.resize_images` to upsample. Parameters @@ -207,26 +112,25 @@ class KResizeImages(Layer): kwargs: dict The standard Keras Layer keyword arguments (if any) """ - def __init__(self, size=2, interpolation="nearest", **kwargs): + def __init__(self, + size: int = 2, + interpolation: T.Literal["nearest", "bilinear"] = "nearest", + **kwargs) -> None: super().__init__(**kwargs) self.size = size self.interpolation = interpolation - def call(self, inputs, *args, **kwargs): + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: """ Call the upsample layer Parameters ---------- - inputs: tensor + inputs: :class:`tf.Tensor` Input tensor, or list/tuple of input tensors - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments Returns ------- - tensor + :class:`tf.Tensor` A tensor or list/tuple of tensors """ if isinstance(self.size, int): @@ -238,13 +142,10 @@ def call(self, inputs, *args, **kwargs): else: # Arbitrary resizing size = int(round(K.int_shape(inputs)[1] * self.size)) - if get_backend() != "amd": - retval = tf.image.resize(inputs, (size, size), method=self.interpolation) - else: - raise NotImplementedError + retval = tf.image.resize(inputs, (size, size), method=self.interpolation) return retval - def compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the output shape of the layer. This is the input shape with size dimensions multiplied by :attr:`size` @@ -263,7 +164,7 @@ def compute_output_shape(self, input_shape): batch, height, width, channels = input_shape return (batch, height * self.size, width * self.size, channels) - def get_config(self): + def get_config(self) -> dict[str, T.Any]: """Returns the config of the layer. Returns @@ -271,13 +172,62 @@ def get_config(self): dict A python dictionary containing the layer configuration """ - config = dict(size=self.size, interpolation=self.interpolation) + config = {"size": self.size, "interpolation": self.interpolation} base_config = super().get_config() return dict(list(base_config.items()) + list(config.items())) -class SubPixelUpscaling(Layer): - """ Sub-pixel convolutional up-scaling layer. +class L2_normalize(tf.keras.layers.Layer): # pylint:disable=invalid-name + """ Normalizes a tensor w.r.t. the L2 norm alongside the specified axis. + + Parameters + ---------- + axis: int + The axis to perform normalization across + kwargs: dict + The standard Keras Layer keyword arguments (if any) + """ + def __init__(self, axis: int, **kwargs) -> None: + self.axis = axis + super().__init__(**kwargs) + + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """This is where the layer's logic lives. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + Input tensor, or list/tuple of input tensors + + Returns + ------- + :class:`tf.Tensor` + A tensor or list/tuple of tensors + """ + return K.l2_normalize(inputs, self.axis) + + def get_config(self) -> dict[str, T.Any]: + """Returns the config of the layer. + + A layer config is a Python dictionary (serializable) containing the configuration of a + layer. The same layer can be reinstated later (without its trained weights) from this + configuration. + + The configuration of a layer does not include connectivity information, nor the layer + class name. These are handled by `Network` (one layer of abstraction above). + + Returns + -------- + dict + A python dictionary containing the layer configuration + """ + config = super().get_config() + config["axis"] = self.axis + return config + + +class PixelShuffler(tf.keras.layers.Layer): + """ PixelShuffler layer for Keras. This layer requires a Convolution2D prior to it, having output filters computed according to the formula :math:`filters = k * (scale_factor * scale_factor)` where `k` is a user defined @@ -289,27 +239,23 @@ class SubPixelUpscaling(Layer): Notes ----- - This method is deprecated as it just performs the same as :class:`PixelShuffler` - using explicit Tensorflow ops. The method is kept in the repository to support legacy - models that have been created with this layer. - In practice, it is useful to have a second convolution layer after the - :class:`SubPixelUpscaling` layer to speed up the learning process. However, if you are stacking - multiple :class:`SubPixelUpscaling` blocks, it may increase the number of parameters greatly, - so the Convolution layer after :class:`SubPixelUpscaling` layer can be removed. + :class:`PixelShuffler` layer to speed up the learning process. However, if you are stacking + multiple :class:`PixelShuffler` blocks, it may increase the number of parameters greatly, + so the Convolution layer after :class:`PixelShuffler` layer can be removed. Example ------- >>> # A standard sub-pixel up-scaling block >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(...) - >>> u = SubPixelUpscaling(scale_factor=2)(x) + >>> u = PixelShuffler(size=(2, 2))(x) [Optional] >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(u) Parameters ---------- - size: int, optional - The up-scaling factor. Default: `2` + size: tuple, optional + The (`h`, `w`) scaling factor for up-scaling. Default: `(2, 2)` data_format: ["channels_first", "channels_last", ``None``], optional The data format for the input. Default: ``None`` kwargs: dict @@ -317,53 +263,60 @@ class SubPixelUpscaling(Layer): References ---------- - based on the paper "Real-Time Single Image and Video Super-Resolution Using an Efficient - Sub-Pixel Convolutional Neural Network" (https://arxiv.org/abs/1609.05158). + https://gist.github.com/t-ae/6e1016cc188104d123676ccef3264981 """ - - def __init__(self, scale_factor=2, data_format=None, **kwargs): + def __init__(self, + size: int | tuple[int, int] = (2, 2), + data_format: str | None = None, + **kwargs) -> None: super().__init__(**kwargs) + self.data_format = conv_utils.normalize_data_format(data_format) + self.size = conv_utils.normalize_tuple(size, 2, 'size') - self.scale_factor = scale_factor - if get_backend() == "amd": - self.data_format = K.normalize_data_format(data_format) # pylint:disable=no-member - else: - self.data_format = conv_utils.normalize_data_format(data_format) - - def build(self, input_shape): - """Creates the layer weights. - - Must be implemented on all layers that have weights. - - Parameters - ---------- - input_shape: tensor - Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to - reference for weight shape computations. - """ - pass # pylint: disable=unnecessary-pass - - def call(self, inputs, *args, **kwargs): + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: """This is where the layer's logic lives. Parameters ---------- - inputs: tensor + inputs: :class:`tf.Tensor` Input tensor, or list/tuple of input tensors - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments Returns ------- - tensor + :class:`tf.Tensor` A tensor or list/tuple of tensors """ - retval = self._depth_to_space(inputs, self.scale_factor, self.data_format) - return retval + input_shape = K.int_shape(inputs) + if len(input_shape) != 4: + raise ValueError('Inputs should have rank ' + + str(4) + + '; Received input shape:', str(input_shape)) - def compute_output_shape(self, input_shape): + if self.data_format == 'channels_first': + batch_size, channels, height, width = input_shape + if batch_size is None: + batch_size = -1 + r_height, r_width = self.size + o_height, o_width = height * r_height, width * r_width + o_channels = channels // (r_height * r_width) + + out = K.reshape(inputs, (batch_size, r_height, r_width, o_channels, height, width)) + out = K.permute_dimensions(out, (0, 3, 4, 1, 5, 2)) + out = K.reshape(out, (batch_size, o_channels, o_height, o_width)) + elif self.data_format == 'channels_last': + batch_size, height, width, channels = input_shape + if batch_size is None: + batch_size = -1 + r_height, r_width = self.size + o_height, o_width = height * r_height, width * r_width + o_channels = channels // (r_height * r_width) + + out = K.reshape(inputs, (batch_size, height, width, r_height, r_width, o_channels)) + out = K.permute_dimensions(out, (0, 1, 3, 2, 4, 5)) + out = K.reshape(out, (batch_size, o_height, o_width, o_channels)) + return out + + def compute_output_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the output shape of the layer. Assumes that the layer will be built to match that input shape provided. @@ -379,79 +332,46 @@ def compute_output_shape(self, input_shape): tuple An input shape tuple """ - if self.data_format == "channels_first": - batch, channels, rows, columns = input_shape - return (batch, - channels // (self.scale_factor ** 2), - rows * self.scale_factor, - columns * self.scale_factor) - batch, rows, columns, channels = input_shape - return (batch, - rows * self.scale_factor, - columns * self.scale_factor, - channels // (self.scale_factor ** 2)) - - @classmethod - def _depth_to_space(cls, ipt, scale, data_format=None): - """ Uses phase shift algorithm to convert channels/depth for spatial resolution """ - if data_format is None: - data_format = K.image_data_format() - data_format = data_format.lower() - ipt = cls._preprocess_conv2d_input(ipt, data_format) - out = tf.nn.depth_to_space(ipt, scale) - out = cls._postprocess_conv2d_output(out, data_format) - return out - - @staticmethod - def _postprocess_conv2d_output(inputs, data_format): - """Transpose and cast the output from conv2d if needed. - - Parameters - ---------- - inputs: tensor - The input that requires transposing and casting - data_format: str - `"channels_last"` or `"channels_first"` - - Returns - ------- - tensor - The transposed and cast input tensor - """ + if len(input_shape) != 4: + raise ValueError('Inputs should have rank ' + + str(4) + + '; Received input shape:', str(input_shape)) - if data_format == "channels_first": - inputs = tf.transpose(inputs, (0, 3, 1, 2)) + if self.data_format == 'channels_first': + height = None + width = None + if input_shape[2] is not None: + height = input_shape[2] * self.size[0] + if input_shape[3] is not None: + width = input_shape[3] * self.size[1] + channels = input_shape[1] // self.size[0] // self.size[1] - if K.floatx() == "float64": - inputs = tf.cast(inputs, "float64") - return inputs + if channels * self.size[0] * self.size[1] != input_shape[1]: + raise ValueError('channels of input and size are incompatible') - @staticmethod - def _preprocess_conv2d_input(inputs, data_format): - """Transpose and cast the input before the conv2d. + retval = (input_shape[0], + channels, + height, + width) + elif self.data_format == 'channels_last': + height = None + width = None + if input_shape[1] is not None: + height = input_shape[1] * self.size[0] + if input_shape[2] is not None: + width = input_shape[2] * self.size[1] + channels = input_shape[3] // self.size[0] // self.size[1] - Parameters - ---------- - inputs: tensor - The input that requires transposing and casting - data_format: str - `"channels_last"` or `"channels_first"` + if channels * self.size[0] * self.size[1] != input_shape[3]: + raise ValueError('channels of input and size are incompatible') - Returns - ------- - tensor - The transposed and cast input tensor - """ - if K.dtype(inputs) == "float64": - inputs = tf.cast(inputs, "float32") - if data_format == "channels_first": - # Tensorflow uses the last dimension as channel dimension, instead of the 2nd one. - # Theano input shape: (samples, input_depth, rows, cols) - # Tensorflow input shape: (samples, rows, cols, input_depth) - inputs = tf.transpose(inputs, (0, 2, 3, 1)) - return inputs + retval = (input_shape[0], + height, + width, + channels) + return retval - def get_config(self): + def get_config(self) -> dict[str, T.Any]: """Returns the config of the layer. A layer config is a Python dictionary (serializable) containing the configuration of a @@ -466,13 +386,44 @@ class name. These are handled by `Network` (one layer of abstraction above). dict A python dictionary containing the layer configuration """ - config = {"scale_factor": self.scale_factor, - "data_format": self.data_format} + config = {'size': self.size, + 'data_format': self.data_format} base_config = super().get_config() + return dict(list(base_config.items()) + list(config.items())) -class ReflectionPadding2D(Layer): +class QuickGELU(tf.keras.layers.Layer): + """ Applies GELU approximation that is fast but somewhat inaccurate. + + Parameters + ---------- + name: str, optional + The name for the layer. Default: "QuickGELU" + kwargs: dict + The standard Keras Layer keyword arguments (if any) + """ + + def __init__(self, name: str = "QuickGELU", **kwargs) -> None: + super().__init__(name=name, **kwargs) + + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """ Call the QuickGELU layerr + + Parameters + ---------- + inputs : :class:`tf.Tensor` + The input Tensor + + Returns + ------- + :class:`tf.Tensor` + The output Tensor + """ + return inputs * K.sigmoid(1.702 * inputs) + + +class ReflectionPadding2D(tf.keras.layers.Layer): """Reflection-padding layer for 2D input (e.g. picture). This layer can add rows and columns at the top, bottom, left and right side of an image tensor. @@ -486,30 +437,30 @@ class ReflectionPadding2D(Layer): kwargs: dict The standard Keras Layer keyword arguments (if any) """ - def __init__(self, stride=2, kernel_size=5, **kwargs): + def __init__(self, stride: int = 2, kernel_size: int = 5, **kwargs) -> None: if isinstance(stride, (tuple, list)): assert len(stride) == 2 and stride[0] == stride[1] stride = stride[0] self.stride = stride self.kernel_size = kernel_size - self.input_spec = None + self.input_spec: list[tf.Tensor] | None = None super().__init__(**kwargs) - def build(self, input_shape): + def build(self, input_shape: tf.Tensor) -> None: """Creates the layer weights. Must be implemented on all layers that have weights. Parameters ---------- - input_shape: tensor + input_shape: :class:`tf.Tensor` Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to reference for weight shape computations. """ - self.input_spec = [InputSpec(shape=input_shape)] + self.input_spec = [keras.layers.InputSpec(shape=input_shape)] super().build(input_shape) - def compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the output shape of the layer. Assumes that the layer will be built to match that input shape provided. @@ -525,6 +476,7 @@ def compute_output_shape(self, input_shape): tuple An input shape tuple """ + assert self.input_spec is not None input_shape = self.input_spec[0].shape in_width, in_height = input_shape[2], input_shape[1] kernel_width, kernel_height = self.kernel_size, self.kernel_size @@ -543,21 +495,20 @@ def compute_output_shape(self, input_shape): input_shape[2] + padding_width, input_shape[3]) - def call(self, var_x, mask=None): # pylint:disable=unused-argument,arguments-differ + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: """This is where the layer's logic lives. Parameters ---------- - inputs: tensor + inputs: :class:`tf.Tensor` Input tensor, or list/tuple of input tensors - kwargs: dict - Additional keyword arguments Returns ------- - tensor + :class:`tf.Tensor` A tensor or list/tuple of tensors """ + assert self.input_spec is not None input_shape = self.input_spec[0].shape in_width, in_height = input_shape[2], input_shape[1] kernel_width, kernel_height = self.kernel_size, self.kernel_size @@ -576,14 +527,14 @@ def call(self, var_x, mask=None): # pylint:disable=unused-argument,arguments-di padding_left = padding_width // 2 padding_right = padding_width - padding_left - return pad(var_x, - [[0, 0], - [padding_top, padding_bot], - [padding_left, padding_right], - [0, 0]], - 'REFLECT') + return tf.pad(inputs, + [[0, 0], + [padding_top, padding_bot], + [padding_left, padding_right], + [0, 0]], + 'REFLECT') - def get_config(self): + def get_config(self) -> dict[str, T.Any]: """Returns the config of the layer. A layer config is a Python dictionary (serializable) containing the configuration of a @@ -604,138 +555,193 @@ class name. These are handled by `Network` (one layer of abstraction above). return dict(list(base_config.items()) + list(config.items())) -class _GlobalPooling2D(Layer): - """Abstract class for different global pooling 2D layers. +class SubPixelUpscaling(tf.keras.layers.Layer): + """ Sub-pixel convolutional up-scaling layer. - From keras as access to pooling is trickier in tensorflow.keras + This layer requires a Convolution2D prior to it, having output filters computed according to + the formula :math:`filters = k * (scale_factor * scale_factor)` where `k` is a user defined + number of filters (generally larger than 32) and `scale_factor` is the up-scaling factor + (generally 2). + + This layer performs the depth to space operation on the convolution filters, and returns a + tensor with the size as defined below. + + Notes + ----- + This method is deprecated as it just performs the same as :class:`PixelShuffler` + using explicit Tensorflow ops. The method is kept in the repository to support legacy + models that have been created with this layer. + + In practice, it is useful to have a second convolution layer after the + :class:`SubPixelUpscaling` layer to speed up the learning process. However, if you are stacking + multiple :class:`SubPixelUpscaling` blocks, it may increase the number of parameters greatly, + so the Convolution layer after :class:`SubPixelUpscaling` layer can be removed. + + Example + ------- + >>> # A standard sub-pixel up-scaling block + >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(...) + >>> u = SubPixelUpscaling(scale_factor=2)(x) + [Optional] + >>> x = Convolution2D(256, 3, 3, padding="same", activation="relu")(u) + + Parameters + ---------- + size: int, optional + The up-scaling factor. Default: `2` + data_format: ["channels_first", "channels_last", ``None``], optional + The data format for the input. Default: ``None`` + kwargs: dict + The standard Keras Layer keyword arguments (if any) + + References + ---------- + based on the paper "Real-Time Single Image and Video Super-Resolution Using an Efficient + Sub-Pixel Convolutional Neural Network" (https://arxiv.org/abs/1609.05158). """ - def __init__(self, data_format=None, **kwargs): + + def __init__(self, scale_factor: int = 2, data_format: str | None = None, **kwargs) -> None: super().__init__(**kwargs) - if get_backend() == "amd": - self.data_format = K.normalize_data_format(data_format) # pylint:disable=no-member - else: - self.data_format = conv_utils.normalize_data_format(data_format) - self.input_spec = InputSpec(ndim=4) - def compute_output_shape(self, input_shape): - """ Compute the output shape based on the input shape. + self.scale_factor = scale_factor + self.data_format = conv_utils.normalize_data_format(data_format) - Parameters - ---------- - input_shape: tuple - The input shape to the layer - """ - if self.data_format == 'channels_last': - return (input_shape[0], input_shape[3]) - return (input_shape[0], input_shape[1]) + def build(self, input_shape: tuple[int, ...]) -> None: + """Creates the layer weights. - def call(self, inputs, *args, **kwargs): - """ Override to call the layer. + Must be implemented on all layers that have weights. Parameters ---------- - inputs: Tensor - The input to the layer - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments + input_shape: tensor + Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to + reference for weight shape computations. """ - raise NotImplementedError - - def get_config(self): - """ Set the Keras config """ - config = {'data_format': self.data_format} - base_config = super().get_config() - return dict(list(base_config.items()) + list(config.items())) - - -class GlobalMinPooling2D(_GlobalPooling2D): - """Global minimum pooling operation for spatial data. """ + pass # pylint:disable=unnecessary-pass - def call(self, inputs, *args, **kwargs): + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: """This is where the layer's logic lives. Parameters ---------- - inputs: tensor + inputs: :class:`tf.Tensor` Input tensor, or list/tuple of input tensors - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments Returns ------- - tensor + :class:`tf.Tensor` A tensor or list/tuple of tensors """ - if self.data_format == 'channels_last': - pooled = K.min(inputs, axis=[1, 2]) - else: - pooled = K.min(inputs, axis=[2, 3]) - return pooled + retval = self._depth_to_space(inputs, self.scale_factor, self.data_format) + return retval + def compute_output_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]: + """Computes the output shape of the layer. -class GlobalStdDevPooling2D(_GlobalPooling2D): - """Global standard deviation pooling operation for spatial data. """ + Assumes that the layer will be built to match that input shape provided. - def call(self, inputs, *args, **kwargs): - """This is where the layer's logic lives. + Parameters + ---------- + input_shape: tuple or list of tuples + Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the + layer). Shape tuples can include None for free dimensions, instead of an integer. + + Returns + ------- + tuple + An input shape tuple + """ + if self.data_format == "channels_first": + batch, channels, rows, columns = input_shape + return (batch, + channels // (self.scale_factor ** 2), + rows * self.scale_factor, + columns * self.scale_factor) + batch, rows, columns, channels = input_shape + return (batch, + rows * self.scale_factor, + columns * self.scale_factor, + channels // (self.scale_factor ** 2)) + + @classmethod + def _depth_to_space(cls, + inputs: tf.Tensor, + scale: int, + data_format: str | None = None) -> tf.Tensor: + """ Uses phase shift algorithm to convert channels/depth for spatial resolution Parameters ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - args: tuple - Additional standard keras Layer arguments - kwargs: dict - Additional standard keras Layer keyword arguments + inputs : :class:`tf.Tensor` + The input Tensor + scale : int + Scale factor + data_format : str | None, optional + "channels_first" or "channels_last" Returns ------- - tensor - A tensor or list/tuple of tensors + :class:`tf.Tensor` + The output Tensor """ - if self.data_format == 'channels_last': - pooled = K.std(inputs, axis=[1, 2]) - else: - pooled = K.std(inputs, axis=[2, 3]) - return pooled + if data_format is None: + data_format = K.image_data_format() + data_format = data_format.lower() + inputs = cls._preprocess_conv2d_input(inputs, data_format) + out = tf.nn.depth_to_space(inputs, scale) + out = cls._postprocess_conv2d_output(out, data_format) + return out + @staticmethod + def _postprocess_conv2d_output(inputs: tf.Tensor, data_format: str | None) -> tf.Tensor: + """Transpose and cast the output from conv2d if needed. -class L2_normalize(Layer): # pylint:disable=invalid-name - """ Normalizes a tensor w.r.t. the L2 norm alongside the specified axis. + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input that requires transposing and casting + data_format: str + `"channels_last"` or `"channels_first"` - Parameters - ---------- - axis: int - The axis to perform normalization across - kwargs: dict - The standard Keras Layer keyword arguments (if any) - """ - def __init__(self, axis, **kwargs): - self.axis = axis - super().__init__(**kwargs) + Returns + ------- + :class:`tf.Tensor` + The transposed and cast input tensor + """ - def call(self, inputs): # pylint:disable=arguments-differ - """This is where the layer's logic lives. + if data_format == "channels_first": + inputs = tf.transpose(inputs, (0, 3, 1, 2)) + + if K.floatx() == "float64": + inputs = tf.cast(inputs, "float64") + return inputs + + @staticmethod + def _preprocess_conv2d_input(inputs: tf.Tensor, data_format: str | None) -> tf.Tensor: + """Transpose and cast the input before the conv2d. Parameters ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - kwargs: dict - Additional keyword arguments + inputs: :class:`tf.Tensor` + The input that requires transposing and casting + data_format: str + `"channels_last"` or `"channels_first"` Returns ------- - tensor - A tensor or list/tuple of tensors + :class:`tf.Tensor` + The transposed and cast input tensor """ - return K.l2_normalize(inputs, self.axis) + if K.dtype(inputs) == "float64": + inputs = tf.cast(inputs, "float32") + if data_format == "channels_first": + # Tensorflow uses the last dimension as channel dimension, instead of the 2nd one. + # Theano input shape: (samples, input_depth, rows, cols) + # Tensorflow input shape: (samples, rows, cols, input_depth) + inputs = tf.transpose(inputs, (0, 2, 3, 1)) + return inputs - def get_config(self): + def get_config(self) -> dict[str, T.Any]: """Returns the config of the layer. A layer config is a Python dictionary (serializable) containing the configuration of a @@ -750,12 +756,13 @@ class name. These are handled by `Network` (one layer of abstraction above). dict A python dictionary containing the layer configuration """ - config = super().get_config() - config["axis"] = self.axis - return config + config = {"scale_factor": self.scale_factor, + "data_format": self.data_format} + base_config = super().get_config() + return dict(list(base_config.items()) + list(config.items())) -class Swish(Layer): +class Swish(tf.keras.layers.Layer): """ Swish Activation Layer implementation for Keras. Parameters @@ -769,21 +776,23 @@ class Swish(Layer): ----------- Swish: a Self-Gated Activation Function: https://arxiv.org/abs/1710.05941v1 """ - def __init__(self, beta=1.0, **kwargs): + def __init__(self, beta: float = 1.0, **kwargs) -> None: super().__init__(**kwargs) self.beta = beta - def call(self, inputs): # pylint:disable=arguments-differ + def call(self, inputs, *args, **kwargs): """ Call the Swish Activation function. Parameters ---------- inputs: tensor Input tensor, or list/tuple of input tensors + + Returns + ------- + :class:`tf.Tensor` + A tensor or list/tuple of tensors """ - if get_backend() == "amd": - return inputs * K.sigmoid(inputs * self.beta) - # Native TF Implementation has more memory-efficient gradients return tf.nn.swish(inputs * self.beta) def get_config(self): @@ -802,6 +811,6 @@ def get_config(self): # Update layers into Keras custom objects -for name, obj in inspect.getmembers(sys.modules[__name__]): +for name_, obj in inspect.getmembers(sys.modules[__name__]): if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) + keras.utils.get_custom_objects().update({name_: obj}) diff --git a/lib/model/losses/__init__.py b/lib/model/losses/__init__.py new file mode 100644 index 0000000000..751e791011 --- /dev/null +++ b/lib/model/losses/__init__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 +""" Custom Loss Functions for Faceswap """ + +from .feature_loss import LPIPSLoss +from .loss import (FocalFrequencyLoss, GeneralizedLoss, GradientLoss, + LaplacianPyramidLoss, LInfNorm, LossWrapper) +from .perceptual_loss import DSSIMObjective, GMSDLoss, LDRFLIPLoss, MSSIMLoss diff --git a/lib/model/losses/feature_loss.py b/lib/model/losses/feature_loss.py new file mode 100644 index 0000000000..a23060f3e5 --- /dev/null +++ b/lib/model/losses/feature_loss.py @@ -0,0 +1,403 @@ +#!/usr/bin/env python3 +""" Custom Feature Map Loss Functions for faceswap.py """ +from __future__ import annotations +from dataclasses import dataclass, field +import logging +import typing as T + +# Ignore linting errors from Tensorflow's thoroughly broken import system +import tensorflow as tf +from tensorflow.keras import applications as kapp # pylint:disable=import-error +from tensorflow.keras.layers import Dropout, Conv2D, Input, Layer, Resizing # noqa,pylint:disable=no-name-in-module,import-error +from tensorflow.keras.models import Model # pylint:disable=no-name-in-module,import-error +import tensorflow.keras.backend as K # pylint:disable=no-name-in-module,import-error + +import numpy as np + +from lib.model.networks import AlexNet, SqueezeNet +from lib.utils import GetModel + +if T.TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) + + +@dataclass +class NetInfo: + """ Data class for holding information about Trunk and Linear Layer nets. + + Parameters + ---------- + model_id: int + The model ID for the model stored in the deepfakes Model repo + model_name: str + The filename of the decompressed model/weights file + net: callable, Optional + The net definition to load, if any. Default:``None`` + init_kwargs: dict, optional + Keyword arguments to initialize any :attr:`net`. Default: empty ``dict`` + needs_init: bool, optional + True if the net needs initializing otherwise False. Default: ``True`` + """ + model_id: int = 0 + model_name: str = "" + net: Callable | None = None + init_kwargs: dict[str, T.Any] = field(default_factory=dict) + needs_init: bool = True + outputs: list[Layer] = field(default_factory=list) + + +class _LPIPSTrunkNet(): + """ Trunk neural network loader for LPIPS Loss function. + + Parameters + ---------- + net_name: str + The name of the trunk network to load. One of "alex", "squeeze" or "vgg16" + eval_mode: bool + ``True`` for evaluation mode, ``False`` for training mode + load_weights: bool + ``True`` if pretrained trunk network weights should be loaded, otherwise ``False`` + """ + def __init__(self, net_name: str, eval_mode: bool, load_weights: bool) -> None: + logger.debug("Initializing: %s (net_name '%s', eval_mode: %s, load_weights: %s)", + self.__class__.__name__, net_name, eval_mode, load_weights) + self._eval_mode = eval_mode + self._load_weights = load_weights + self._net_name = net_name + self._net = self._nets[net_name] + logger.debug("Initialized: %s ", self.__class__.__name__) + + @property + def _nets(self) -> dict[str, NetInfo]: + """ :class:`NetInfo`: The Information about the requested net.""" + return { + "alex": NetInfo(model_id=15, + model_name="alexnet_imagenet_no_top_v1.h5", + net=AlexNet, + outputs=[f"features.{idx}" for idx in (0, 3, 6, 8, 10)]), + "squeeze": NetInfo(model_id=16, + model_name="squeezenet_imagenet_no_top_v1.h5", + net=SqueezeNet, + outputs=[f"features.{idx}" for idx in (0, 4, 7, 9, 10, 11, 12)]), + "vgg16": NetInfo(model_id=17, + model_name="vgg16_imagenet_no_top_v1.h5", + net=kapp.vgg16.VGG16, + init_kwargs={"include_top": False, "weights": None}, + outputs=[f"block{i + 1}_conv{2 if i < 2 else 3}" for i in range(5)])} + + @classmethod + def _normalize_output(cls, inputs: tf.Tensor, epsilon: float = 1e-10) -> tf.Tensor: + """ Normalize the output tensors from the trunk network. + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + An output tensor from the trunk model + epsilon: float, optional + Epsilon to apply to the normalization operation. Default: `1e-10` + """ + norm_factor = K.sqrt(K.sum(K.square(inputs), axis=-1, keepdims=True)) + return inputs / (norm_factor + epsilon) + + def _process_weights(self, model: Model) -> Model: + """ Save and lock weights if requested. + + Parameters + ---------- + model :class:`keras.models.Model` + The loaded trunk or linear network + + Returns + ------- + :class:`keras.models.Model` + The network with weights loaded/not loaded and layers locked/unlocked + """ + if self._load_weights: + weights = GetModel(self._net.model_name, self._net.model_id).model_path + model.load_weights(weights) + + if self._eval_mode: + model.trainable = False + for layer in model.layers: + layer.trainable = False + return model + + def __call__(self) -> Model: + """ Load the Trunk net, add normalization to feature outputs, load weights and set + trainable state. + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The trunk net with normalized feature output layers + """ + if self._net.net is None: + raise ValueError("No net loaded") + + model = self._net.net(**self._net.init_kwargs) + model = model if self._net_name == "vgg16" else model() + out_layers = [self._normalize_output(model.get_layer(name).output) + for name in self._net.outputs] + model = Model(inputs=model.input, outputs=out_layers) + model = self._process_weights(model) + return model + + +class _LPIPSLinearNet(_LPIPSTrunkNet): + """ The Linear Network to be applied to the difference between the true and predicted outputs + of the trunk network. + + Parameters + ---------- + net_name: str + The name of the trunk network in use. One of "alex", "squeeze" or "vgg16" + eval_mode: bool + ``True`` for evaluation mode, ``False`` for training mode + load_weights: bool + ``True`` if pretrained linear network weights should be loaded, otherwise ``False`` + trunk_net: :class:`keras.models.Model` + The trunk net to place the linear layer on. + use_dropout: bool + ``True`` if a dropout layer should be used in the Linear network otherwise ``False`` + """ + def __init__(self, + net_name: str, + eval_mode: bool, + load_weights: bool, + trunk_net: Model, + use_dropout: bool) -> None: + logger.debug( + "Initializing: %s (trunk_net: %s, use_dropout: %s)", self.__class__.__name__, + trunk_net, use_dropout) + super().__init__(net_name=net_name, eval_mode=eval_mode, load_weights=load_weights) + + self._trunk = trunk_net + self._use_dropout = use_dropout + + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def _nets(self) -> dict[str, NetInfo]: + """ :class:`NetInfo`: The Information about the requested net.""" + return { + "alex": NetInfo(model_id=18, + model_name="alexnet_lpips_v1.h5",), + "squeeze": NetInfo(model_id=19, + model_name="squeezenet_lpips_v1.h5"), + "vgg16": NetInfo(model_id=20, + model_name="vgg16_lpips_v1.h5")} + + def _linear_block(self, net_output_layer: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: + """ Build a linear block for a trunk network output. + + Parameters + ---------- + net_output_layer: :class:`tensorflow.Tensor` + An output from the selected trunk network + + Returns + ------- + :class:`tensorflow.Tensor` + The input to the linear block + :class:`tensorflow.Tensor` + The output from the linear block + """ + in_shape = K.int_shape(net_output_layer)[1:] + input_ = Input(in_shape) + var_x = Dropout(rate=0.5)(input_) if self._use_dropout else input_ + var_x = Conv2D(1, 1, strides=1, padding="valid", use_bias=False)(var_x) + return input_, var_x + + def __call__(self) -> Model: + """ Build the linear network for the given trunk network's outputs. Load in trained weights + and set the model's trainable parameters. + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The compiled Linear Net model + """ + inputs = [] + outputs = [] + + for input_ in self._trunk.outputs: + in_, out = self._linear_block(input_) + inputs.append(in_) + outputs.append(out) + + model = Model(inputs=inputs, outputs=outputs) + model = self._process_weights(model) + return model + + +class LPIPSLoss(): + """ LPIPS Loss Function. + + A perceptual loss function that uses linear outputs from pretrained CNNs feature layers. + + Notes + ----- + Channels Last implementation. All trunks implemented from the original paper. + + References + ---------- + https://richzhang.github.io/PerceptualSimilarity/ + + Parameters + ---------- + trunk_network: str + The name of the trunk network to use. One of "alex", "squeeze" or "vgg16" + trunk_pretrained: bool, optional + ``True`` Load the imagenet pretrained weights for the trunk network. ``False`` randomly + initialize the trunk network. Default: ``True`` + trunk_eval_mode: bool, optional + ``True`` for running inference on the trunk network (standard mode), ``False`` for training + the trunk network. Default: ``True`` + linear_pretrained: bool, optional + ``True`` loads the pretrained weights for the linear network layers. ``False`` randomly + initializes the layers. Default: ``True`` + linear_eval_mode: bool, optional + ``True`` for running inference on the linear network (standard mode), ``False`` for + training the linear network. Default: ``True`` + linear_use_dropout: bool, optional + ``True`` if a dropout layer should be used in the Linear network otherwise ``False``. + Default: ``True`` + lpips: bool, optional + ``True`` to use linear network on top of the trunk network. ``False`` to just average the + output from the trunk network. Default ``True`` + spatial: bool, optional + ``True`` output the loss in the spatial domain (i.e. as a grayscale tensor of height and + width of the input image). ``Bool`` reduce the spatial dimensions for loss calculation. + Default: ``False`` + normalize: bool, optional + ``True`` if the input Tensor needs to be normalized from the 0. to 1. range to the -1. to + 1. range. Default: ``True`` + ret_per_layer: bool, optional + ``True`` to return the loss value per feature output layer otherwise ``False``. + Default: ``False`` + """ + def __init__(self, # pylint:disable=too-many-arguments + trunk_network: str, + trunk_pretrained: bool = True, + trunk_eval_mode: bool = True, + linear_pretrained: bool = True, + linear_eval_mode: bool = True, + linear_use_dropout: bool = True, + lpips: bool = True, + spatial: bool = False, + normalize: bool = True, + ret_per_layer: bool = False) -> None: + logger.debug( + "Initializing: %s (trunk_network '%s', trunk_pretrained: %s, trunk_eval_mode: %s, " + "linear_pretrained: %s, linear_eval_mode: %s, linear_use_dropout: %s, lpips: %s, " + "spatial: %s, normalize: %s, ret_per_layer: %s)", self.__class__.__name__, + trunk_network, trunk_pretrained, trunk_eval_mode, linear_pretrained, linear_eval_mode, + linear_use_dropout, lpips, spatial, normalize, ret_per_layer) + + self._spatial = spatial + self._use_lpips = lpips + self._normalize = normalize + self._ret_per_layer = ret_per_layer + self._shift = K.constant(np.array([-.030, -.088, -.188], + dtype="float32")[None, None, None, :]) + self._scale = K.constant(np.array([.458, .448, .450], + dtype="float32")[None, None, None, :]) + + # Loss needs to be done as fp32. We could cast at output, but better to update the model + switch_mixed_precision = tf.keras.mixed_precision.global_policy().name == "mixed_float16" + if switch_mixed_precision: + logger.debug("Temporarily disabling mixed precision") + tf.keras.mixed_precision.set_global_policy("float32") + + self._trunk_net = _LPIPSTrunkNet(trunk_network, trunk_eval_mode, trunk_pretrained)() + self._linear_net = _LPIPSLinearNet(trunk_network, + linear_eval_mode, + linear_pretrained, + self._trunk_net, + linear_use_dropout)() + if switch_mixed_precision: + logger.debug("Re-enabling mixed precision") + tf.keras.mixed_precision.set_global_policy("mixed_float16") + logger.debug("Initialized: %s", self.__class__.__name__) + + def _process_diffs(self, inputs: list[tf.Tensor]) -> list[tf.Tensor]: + """ Perform processing on the Trunk Network outputs. + + If :attr:`use_ldip` is enabled, process the diff values through the linear network, + otherwise return the diff values summed on the channels axis. + + Parameters + ---------- + inputs: list + List of the squared difference of the true and predicted outputs from the trunk network + + Returns + ------- + list + List of either the linear network outputs (when using lpips) or summed network outputs + """ + if self._use_lpips: + return self._linear_net(inputs) + return [K.sum(x, axis=-1) for x in inputs] + + def _process_output(self, inputs: tf.Tensor, output_dims: tuple) -> tf.Tensor: + """ Process an individual output based on whether :attr:`is_spatial` has been selected. + + When spatial output is selected, all outputs are sized to the shape of the original True + input Tensor. When not selected, the mean across the spatial axes (h, w) are returned + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + An individual diff output tensor from the linear network or summed output + output_dims: tuple + The (height, width) of the original true image + + Returns + ------- + :class:`tensorflow.Tensor` + Either the original tensor resized to the true image dimensions, or the mean + value across the height, width axes. + """ + if self._spatial: + return Resizing(*output_dims, interpolation="bilinear")(inputs) + return K.mean(inputs, axis=(1, 2), keepdims=True) + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Perform the LPIPS Loss Function. + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images + y_pred: :class:`tensorflow.Tensor` + The predicted batch of images + + Returns + ------- + :class:`tensorflow.Tensor` + The final loss value + """ + if self._normalize: + y_true = (y_true * 2.0) - 1.0 + y_pred = (y_pred * 2.0) - 1.0 + + y_true = (y_true - self._shift) / self._scale + y_pred = (y_pred - self._shift) / self._scale + + net_true = self._trunk_net(y_true) + net_pred = self._trunk_net(y_pred) + + diffs = [(out_true - out_pred) ** 2 + for out_true, out_pred in zip(net_true, net_pred)] + + dims = K.int_shape(y_true)[1:3] + res = [self._process_output(diff, dims) for diff in self._process_diffs(diffs)] + + axis = 0 if self._spatial else None + val = K.sum(res, axis=axis) + + retval = (val, res) if self._ret_per_layer else val + return retval / 10.0 # Reduce by factor of 10 'cos this loss is STRONG diff --git a/lib/model/losses/loss.py b/lib/model/losses/loss.py new file mode 100644 index 0000000000..ab03ff53fd --- /dev/null +++ b/lib/model/losses/loss.py @@ -0,0 +1,668 @@ +#!/usr/bin/env python3 +""" Custom Loss Functions for faceswap.py """ + +from __future__ import annotations +import logging +import typing as T + +import numpy as np +import tensorflow as tf + +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.python.keras.engine import compile_utils # pylint:disable=no-name-in-module +from tensorflow.keras import backend as K # pylint:disable=import-error + +if T.TYPE_CHECKING: + from collections.abc import Callable + +logger = logging.getLogger(__name__) + + +class FocalFrequencyLoss(): # pylint:disable=too-few-public-methods + """ Focal Frequencey Loss Function. + + A channels last implementation. + + Notes + ----- + There is a bug in this implementation that will do an incorrect FFT if + :attr:`patch_factor` > ``1``, which means incorrect loss will be returned, so keep + patch factor at 1. + + Parameters + ---------- + alpha: float, Optional + Scaling factor of the spectrum weight matrix for flexibility. Default: ``1.0`` + patch_factor: int, Optional + Factor to crop image patches for patch-based focal frequency loss. + Default: ``1`` + ave_spectrum: bool, Optional + ``True`` to use minibatch average spectrum otherwise ``False``. Default: ``False`` + log_matrix: bool, Optional + ``True`` to adjust the spectrum weight matrix by logarithm otherwise ``False``. + Default: ``False`` + batch_matrix: bool, Optional + ``True`` to calculate the spectrum weight matrix using batch-based statistics otherwise + ``False``. Default: ``False`` + + References + ---------- + https://arxiv.org/pdf/2012.12821.pdf + https://github.com/EndlessSora/focal-frequency-loss + """ + + def __init__(self, + alpha: float = 1.0, + patch_factor: int = 1, + ave_spectrum: bool = False, + log_matrix: bool = False, + batch_matrix: bool = False) -> None: + self._alpha = alpha + # TODO Fix bug where FFT will be incorrect if patch_factor > 1 + self._patch_factor = patch_factor + self._ave_spectrum = ave_spectrum + self._log_matrix = log_matrix + self._batch_matrix = batch_matrix + self._dims: tuple[int, int] = (0, 0) + + def _get_patches(self, inputs: tf.Tensor) -> tf.Tensor: + """ Crop the incoming batch of images into patches as defined by :attr:`_patch_factor. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + A batch of images to be converted into patches + + Returns + ------- + :class`tf.Tensor`` + The incoming batch converted into patches + """ + rows, cols = self._dims + patch_list = [] + patch_rows = cols // self._patch_factor + patch_cols = rows // self._patch_factor + for i in range(self._patch_factor): + for j in range(self._patch_factor): + row_from = i * patch_rows + row_to = (i + 1) * patch_rows + col_from = j * patch_cols + col_to = (j + 1) * patch_cols + patch_list.append(inputs[:, row_from: row_to, col_from: col_to, :]) + + retval = K.stack(patch_list, axis=1) + return retval + + def _tensor_to_frequency_spectrum(self, patch: tf.Tensor) -> tf.Tensor: + """ Perform FFT to create the orthonomalized DFT frequencies. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The incoming batch of patches to convert to the frequency spectrum + + Returns + ------- + :class:`tf.Tensor` + The DFT frequencies split into real and imaginary numbers as float32 + """ + # TODO fix this for when self._patch_factor != 1. + rows, cols = self._dims + patch = K.permute_dimensions(patch, (0, 1, 4, 2, 3)) # move channels to first + + patch = patch / np.sqrt(rows * cols) # Orthonormalization + + patch = K.cast(patch, "complex64") + freq = tf.signal.fft2d(patch)[..., None] + + freq = K.concatenate([tf.math.real(freq), tf.math.imag(freq)], axis=-1) + freq = K.cast(freq, "float32") + + freq = K.permute_dimensions(freq, (0, 1, 3, 4, 2, 5)) # channels to last + + return freq + + def _get_weight_matrix(self, freq_true: tf.Tensor, freq_pred: tf.Tensor) -> tf.Tensor: + """ Calculate a continuous, dynamic weight matrix based on current Euclidean distance. + + Parameters + ---------- + freq_true: :class:`tf.Tensor` + The real and imaginary DFT frequencies for the true batch of images + freq_pred: :class:`tf.Tensor` + The real and imaginary DFT frequencies for the predicted batch of images + + Returns + ------- + :class:`tf.Tensor` + The weights matrix for prioritizing hard frequencies + """ + weights = K.square(freq_pred - freq_true) + weights = K.sqrt(weights[..., 0] + weights[..., 1]) + weights = K.pow(weights, self._alpha) + + if self._log_matrix: # adjust the spectrum weight matrix by logarithm + weights = K.log(weights + 1.0) + + if self._batch_matrix: # calculate the spectrum weight matrix using batch-based statistics + weights = weights / K.max(weights) + else: + weights = weights / K.max(K.max(weights, axis=-2), axis=-2)[..., None, None, :] + + weights = K.switch(tf.math.is_nan(weights), K.zeros_like(weights), weights) + weights = K.clip(weights, min_value=0.0, max_value=1.0) + + return weights + + @classmethod + def _calculate_loss(cls, + freq_true: tf.Tensor, + freq_pred: tf.Tensor, + weight_matrix: tf.Tensor) -> tf.Tensor: + """ Perform the loss calculation on the DFT spectrum applying the weights matrix. + + Parameters + ---------- + freq_true: :class:`tf.Tensor` + The real and imaginary DFT frequencies for the true batch of images + freq_pred: :class:`tf.Tensor` + The real and imaginary DFT frequencies for the predicted batch of images + + Returns + :class:`tf.Tensor` + The final loss matrix + """ + + tmp = K.square(freq_pred - freq_true) # freq distance using squared Euclidean distance + + freq_distance = tmp[..., 0] + tmp[..., 1] + loss = weight_matrix * freq_distance # dynamic spectrum weighting (Hadamard product) + + return loss + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the Focal Frequency Loss Function. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth batch of images + y_pred: :class:`tf.Tensor` + The predicted batch of images + + Returns + ------- + :class:`tf.Tensor` + The loss for this batch of images + """ + if not all(self._dims): + rows, cols = K.int_shape(y_true)[1:3] + assert cols % self._patch_factor == 0 and rows % self._patch_factor == 0, ( + "Patch factor must be a divisor of the image height and width") + self._dims = (rows, cols) + + patches_true = self._get_patches(y_true) + patches_pred = self._get_patches(y_pred) + + freq_true = self._tensor_to_frequency_spectrum(patches_true) + freq_pred = self._tensor_to_frequency_spectrum(patches_pred) + + if self._ave_spectrum: # whether to use minibatch average spectrum + freq_true = K.mean(freq_true, axis=0, keepdims=True) + freq_pred = K.mean(freq_pred, axis=0, keepdims=True) + + weight_matrix = self._get_weight_matrix(freq_true, freq_pred) + return self._calculate_loss(freq_true, freq_pred, weight_matrix) + + +class GeneralizedLoss(): # pylint:disable=too-few-public-methods + """ Generalized function used to return a large variety of mathematical loss functions. + + The primary benefit is a smooth, differentiable version of L1 loss. + + References + ---------- + Barron, J. A General and Adaptive Robust Loss Function - https://arxiv.org/pdf/1701.03077.pdf + + Example + ------- + >>> a=1.0, x>>c , c=1.0/255.0 # will give a smoothly differentiable version of L1 / MAE loss + >>> a=1.999999 (limit as a->2), beta=1.0/255.0 # will give L2 / RMSE loss + + Parameters + ---------- + alpha: float, optional + Penalty factor. Larger number give larger weight to large deviations. Default: `1.0` + beta: float, optional + Scale factor used to adjust to the input scale (i.e. inputs of mean `1e-4` or `256`). + Default: `1.0/255.0` + """ + def __init__(self, alpha: float = 1.0, beta: float = 1.0/255.0) -> None: + self._alpha = alpha + self._beta = beta + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the Generalized Loss Function + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class:`tf.Tensor` + The loss value from the results of function(y_pred - y_true) + """ + diff = y_pred - y_true + second = (K.pow(K.pow(diff/self._beta, 2.) / K.abs(2. - self._alpha) + 1., + (self._alpha / 2.)) - 1.) + loss = (K.abs(2. - self._alpha)/self._alpha) * second + loss = K.mean(loss, axis=-1) * self._beta + return loss + + +class GradientLoss(): # pylint:disable=too-few-public-methods + """ Gradient Loss Function. + + Calculates the first and second order gradient difference between pixels of an image in the x + and y dimensions. These gradients are then compared between the ground truth and the predicted + image and the difference is taken. When used as a loss, its minimization will result in + predicted images approaching the same level of sharpness / blurriness as the ground truth. + + References + ---------- + TV+TV2 Regularization with Non-Convex Sparseness-Inducing Penalty for Image Restoration, + Chengwu Lu & Hua Huang, 2014 - http://downloads.hindawi.com/journals/mpe/2014/790547.pdf + """ + def __init__(self) -> None: + self.generalized_loss = GeneralizedLoss(alpha=1.9999) + self._tv_weight = 1.0 + self._tv2_weight = 1.0 + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the gradient loss function. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class:`tf.Tensor` + The loss value + """ + loss = 0.0 + loss += self._tv_weight * (self.generalized_loss(self._diff_x(y_true), + self._diff_x(y_pred)) + + self.generalized_loss(self._diff_y(y_true), + self._diff_y(y_pred))) + loss += self._tv2_weight * (self.generalized_loss(self._diff_xx(y_true), + self._diff_xx(y_pred)) + + self.generalized_loss(self._diff_yy(y_true), + self._diff_yy(y_pred)) + + self.generalized_loss(self._diff_xy(y_true), + self._diff_xy(y_pred)) * 2.) + loss = loss / (self._tv_weight + self._tv2_weight) + # TODO simplify to use MSE instead + return loss + + @classmethod + def _diff_x(cls, img: tf.Tensor) -> tf.Tensor: + """ X Difference """ + x_left = img[:, :, 1:2, :] - img[:, :, 0:1, :] + x_inner = img[:, :, 2:, :] - img[:, :, :-2, :] + x_right = img[:, :, -1:, :] - img[:, :, -2:-1, :] + x_out = K.concatenate([x_left, x_inner, x_right], axis=2) + return x_out * 0.5 + + @classmethod + def _diff_y(cls, img: tf.Tensor) -> tf.Tensor: + """ Y Difference """ + y_top = img[:, 1:2, :, :] - img[:, 0:1, :, :] + y_inner = img[:, 2:, :, :] - img[:, :-2, :, :] + y_bot = img[:, -1:, :, :] - img[:, -2:-1, :, :] + y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) + return y_out * 0.5 + + @classmethod + def _diff_xx(cls, img: tf.Tensor) -> tf.Tensor: + """ X-X Difference """ + x_left = img[:, :, 1:2, :] + img[:, :, 0:1, :] + x_inner = img[:, :, 2:, :] + img[:, :, :-2, :] + x_right = img[:, :, -1:, :] + img[:, :, -2:-1, :] + x_out = K.concatenate([x_left, x_inner, x_right], axis=2) + return x_out - 2.0 * img + + @classmethod + def _diff_yy(cls, img: tf.Tensor) -> tf.Tensor: + """ Y-Y Difference """ + y_top = img[:, 1:2, :, :] + img[:, 0:1, :, :] + y_inner = img[:, 2:, :, :] + img[:, :-2, :, :] + y_bot = img[:, -1:, :, :] + img[:, -2:-1, :, :] + y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) + return y_out - 2.0 * img + + @classmethod + def _diff_xy(cls, img: tf.Tensor) -> tf.Tensor: + """ X-Y Difference """ + # xout1 + # Left + top = img[:, 1:2, 1:2, :] + img[:, 0:1, 0:1, :] + inner = img[:, 2:, 1:2, :] + img[:, :-2, 0:1, :] + bottom = img[:, -1:, 1:2, :] + img[:, -2:-1, 0:1, :] + xy_left = K.concatenate([top, inner, bottom], axis=1) + # Mid + top = img[:, 1:2, 2:, :] + img[:, 0:1, :-2, :] + mid = img[:, 2:, 2:, :] + img[:, :-2, :-2, :] + bottom = img[:, -1:, 2:, :] + img[:, -2:-1, :-2, :] + xy_mid = K.concatenate([top, mid, bottom], axis=1) + # Right + top = img[:, 1:2, -1:, :] + img[:, 0:1, -2:-1, :] + inner = img[:, 2:, -1:, :] + img[:, :-2, -2:-1, :] + bottom = img[:, -1:, -1:, :] + img[:, -2:-1, -2:-1, :] + xy_right = K.concatenate([top, inner, bottom], axis=1) + + # Xout2 + # Left + top = img[:, 0:1, 1:2, :] + img[:, 1:2, 0:1, :] + inner = img[:, :-2, 1:2, :] + img[:, 2:, 0:1, :] + bottom = img[:, -2:-1, 1:2, :] + img[:, -1:, 0:1, :] + xy_left = K.concatenate([top, inner, bottom], axis=1) + # Mid + top = img[:, 0:1, 2:, :] + img[:, 1:2, :-2, :] + mid = img[:, :-2, 2:, :] + img[:, 2:, :-2, :] + bottom = img[:, -2:-1, 2:, :] + img[:, -1:, :-2, :] + xy_mid = K.concatenate([top, mid, bottom], axis=1) + # Right + top = img[:, 0:1, -1:, :] + img[:, 1:2, -2:-1, :] + inner = img[:, :-2, -1:, :] + img[:, 2:, -2:-1, :] + bottom = img[:, -2:-1, -1:, :] + img[:, -1:, -2:-1, :] + xy_right = K.concatenate([top, inner, bottom], axis=1) + + xy_out1 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) + xy_out2 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) + return (xy_out1 - xy_out2) * 0.25 + + +class LaplacianPyramidLoss(): # pylint:disable=too-few-public-methods + """ Laplacian Pyramid Loss Function + + Notes + ----- + Channels last implementation on square images only. + + Parameters + ---------- + max_levels: int, Optional + The max number of laplacian pyramid levels to use. Default: `5` + gaussian_size: int, Optional + The size of the gaussian kernel. Default: `5` + gaussian_sigma: float, optional + The gaussian sigma. Default: 2.0 + + References + ---------- + https://arxiv.org/abs/1707.05776 + https://github.com/nathanaelbosch/generative-latent-optimization/blob/master/utils.py + """ + def __init__(self, + max_levels: int = 5, + gaussian_size: int = 5, + gaussian_sigma: float = 1.0) -> None: + self._max_levels = max_levels + self._weights = K.constant([np.power(2., -2 * idx) for idx in range(max_levels + 1)]) + self._gaussian_kernel = self._get_gaussian_kernel(gaussian_size, gaussian_sigma) + + @classmethod + def _get_gaussian_kernel(cls, size: int, sigma: float) -> tf.Tensor: + """ Obtain the base gaussian kernel for the Laplacian Pyramid. + + Parameters + ---------- + size: int, Optional + The size of the gaussian kernel + sigma: float + The gaussian sigma + + Returns + ------- + :class:`tf.Tensor` + The base single channel Gaussian kernel + """ + assert size % 2 == 1, ("kernel size must be uneven") + x_1 = np.linspace(- (size // 2), size // 2, size, dtype="float32") + x_1 /= np.sqrt(2)*sigma + x_2 = x_1 ** 2 + kernel = np.exp(- x_2[:, None] - x_2[None, :]) + kernel /= kernel.sum() + kernel = np.reshape(kernel, (size, size, 1, 1)) + return K.constant(kernel) + + def _conv_gaussian(self, inputs: tf.Tensor) -> tf.Tensor: + """ Perform Gaussian convolution on a batch of images. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input batch of images to perform Gaussian convolution on. + + Returns + ------- + :class:`tf.Tensor` + The convolved images + """ + channels = K.int_shape(inputs)[-1] + gauss = K.tile(self._gaussian_kernel, (1, 1, 1, channels)) + + # TF doesn't implement replication padding like pytorch. This is an inefficient way to + # implement it for a square guassian kernel + size = self._gaussian_kernel.shape[1] // 2 + padded_inputs = inputs + for _ in range(size): + padded_inputs = tf.pad(padded_inputs, # noqa,pylint:disable=no-value-for-parameter,unexpected-keyword-arg + ([0, 0], [1, 1], [1, 1], [0, 0]), + mode="SYMMETRIC") + + retval = K.conv2d(padded_inputs, gauss, strides=1, padding="valid") + return retval + + def _get_laplacian_pyramid(self, inputs: tf.Tensor) -> list[tf.Tensor]: + """ Obtain the Laplacian Pyramid. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input batch of images to run through the Laplacian Pyramid + + Returns + ------- + list + The tensors produced from the Laplacian Pyramid + """ + pyramid = [] + current = inputs + for _ in range(self._max_levels): + gauss = self._conv_gaussian(current) + diff = current - gauss + pyramid.append(diff) + current = K.pool2d(gauss, (2, 2), strides=(2, 2), padding="valid", pool_mode="avg") + pyramid.append(current) + return pyramid + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Calculate the Laplacian Pyramid Loss. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class: `tf.Tensor` + The loss value + """ + pyramid_true = self._get_laplacian_pyramid(y_true) + pyramid_pred = self._get_laplacian_pyramid(y_pred) + + losses = K.stack([K.sum(K.abs(ppred - ptrue)) / K.cast(K.prod(K.shape(ptrue)), "float32") + for ptrue, ppred in zip(pyramid_true, pyramid_pred)]) + loss = K.sum(losses * self._weights) + + return loss + + +class LInfNorm(): # pylint:disable=too-few-public-methods + """ Calculate the L-inf norm as a loss function. """ + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the L-inf norm loss function. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class:`tf.Tensor` + The loss value + """ + diff = K.abs(y_true - y_pred) + max_loss = K.max(diff, axis=(1, 2), keepdims=True) + loss = K.mean(max_loss, axis=-1) + return loss + + +class LossWrapper(tf.keras.losses.Loss): + """ A wrapper class for multiple keras losses to enable multiple masked weighted loss + functions on a single output. + + Notes + ----- + Whilst Keras does allow for applying multiple weighted loss functions, it does not allow + for an easy mechanism to add additional data (in our case masks) that are batch specific + but are not fed in to the model. + + This wrapper receives this additional mask data for the batch stacked onto the end of the + color channels of the received :attr:`y_true` batch of images. These masks are then split + off the batch of images and applied to both the :attr:`y_true` and :attr:`y_pred` tensors + prior to feeding into the loss functions. + + For example, for an image of shape (4, 128, 128, 3) 3 additional masks may be stacked onto + the end of y_true, meaning we receive an input of shape (4, 128, 128, 6). This wrapper then + splits off (4, 128, 128, 3:6) from the end of the tensor, leaving the original y_true of + shape (4, 128, 128, 3) ready for masking and feeding through the loss functions. + """ + def __init__(self) -> None: + logger.debug("Initializing: %s", self.__class__.__name__) + super().__init__(name="LossWrapper") + self._loss_functions: list[compile_utils.LossesContainer] = [] + self._loss_weights: list[float] = [] + self._mask_channels: list[int] = [] + logger.debug("Initialized: %s", self.__class__.__name__) + + def add_loss(self, + function: Callable, + weight: float = 1.0, + mask_channel: int = -1) -> None: + """ Add the given loss function with the given weight to the loss function chain. + + Parameters + ---------- + function: :class:`tf.keras.losses.Loss` + The loss function to add to the loss chain + weight: float, optional + The weighting to apply to the loss function. Default: `1.0` + mask_channel: int, optional + The channel in the `y_true` image that the mask exists in. Set to `-1` if there is no + mask for the given loss function. Default: `-1` + """ + logger.debug("Adding loss: (function: %s, weight: %s, mask_channel: %s)", + function, weight, mask_channel) + # Loss must be compiled inside LossContainer for keras to handle distibuted strategies + self._loss_functions.append(compile_utils.LossesContainer(function)) + self._loss_weights.append(weight) + self._mask_channels.append(mask_channel) + + def call(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the sub loss functions for the loss wrapper. + + Loss is returned as the weighted sum of the chosen losses. + + If masks are being applied to the loss function inputs, then they should be included as + additional channels at the end of :attr:`y_true`, so that they can be split off and + applied to the actual inputs to the selected loss function(s). + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images, with any required masks stacked on the end + y_pred: :class:`tensorflow.Tensor` + The batch of model predictions + + Returns + ------- + :class:`tensorflow.Tensor` + The final weighted loss + """ + loss = 0.0 + for func, weight, mask_channel in zip(self._loss_functions, + self._loss_weights, + self._mask_channels): + logger.debug("Processing loss function: (func: %s, weight: %s, mask_channel: %s)", + func, weight, mask_channel) + n_true, n_pred = self._apply_mask(y_true, y_pred, mask_channel) + loss += (func(n_true, n_pred) * weight) + return loss + + @classmethod + def _apply_mask(cls, + y_true: tf.Tensor, + y_pred: tf.Tensor, + mask_channel: int, + mask_prop: float = 1.0) -> tuple[tf.Tensor, tf.Tensor]: + """ Apply the mask to the input y_true and y_pred. If a mask is not required then + return the unmasked inputs. + + Parameters + ---------- + y_true: tensor or variable + The ground truth value + y_pred: tensor or variable + The predicted value + mask_channel: int + The channel within y_true that the required mask resides in + mask_prop: float, optional + The amount of mask propagation. Default: `1.0` + + Returns + ------- + tf.Tensor + The ground truth batch of images, with the required mask applied + tf.Tensor + The predicted batch of images with the required mask applied + """ + if mask_channel == -1: + logger.debug("No mask to apply") + return y_true[..., :3], y_pred[..., :3] + + logger.debug("Applying mask from channel %s", mask_channel) + + mask = K.tile(K.expand_dims(y_true[..., mask_channel], axis=-1), (1, 1, 1, 3)) + mask_as_k_inv_prop = 1 - mask_prop + mask = (mask * mask_prop) + mask_as_k_inv_prop + + m_true = y_true[..., :3] * mask + m_pred = y_pred[..., :3] * mask + + return m_true, m_pred diff --git a/lib/model/losses/perceptual_loss.py b/lib/model/losses/perceptual_loss.py new file mode 100644 index 0000000000..0fc09b81d7 --- /dev/null +++ b/lib/model/losses/perceptual_loss.py @@ -0,0 +1,750 @@ +#!/usr/bin/env python3 +""" TF Keras implementation of Perceptual Loss Functions for faceswap.py """ + +import logging +import typing as T + +import numpy as np +import tensorflow as tf + +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import backend as K # pylint:disable=import-error + +from lib.keras_utils import ColorSpaceConvert, frobenius_norm, replicate_pad + +logger = logging.getLogger(__name__) + + +class DSSIMObjective(): # pylint:disable=too-few-public-methods + """ DSSIM Loss Functions + + Difference of Structural Similarity (DSSIM loss function). + + Adapted from :func:`tensorflow.image.ssim` for a pure keras implentation. + + Notes + ----- + Channels last only. Assumes all input images are the same size and square + + Parameters + ---------- + k_1: float, optional + Parameter of the SSIM. Default: `0.01` + k_2: float, optional + Parameter of the SSIM. Default: `0.03` + filter_size: int, optional + size of gaussian filter Default: `11` + filter_sigma: float, optional + Width of gaussian filter Default: `1.5` + max_value: float, optional + Max value of the output. Default: `1.0` + + Notes + ------ + You should add a regularization term like a l2 loss in addition to this one. + """ + def __init__(self, + k_1: float = 0.01, + k_2: float = 0.03, + filter_size: int = 11, + filter_sigma: float = 1.5, + max_value: float = 1.0) -> None: + self._filter_size = filter_size + self._filter_sigma = filter_sigma + self._kernel = self._get_kernel() + + compensation = 1.0 + self._c1 = (k_1 * max_value) ** 2 + self._c2 = ((k_2 * max_value) ** 2) * compensation + + def _get_kernel(self) -> tf.Tensor: + """ Obtain the base kernel for performing depthwise convolution. + + Returns + ------- + :class:`tf.Tensor` + The gaussian kernel based on selected size and sigma + """ + coords = np.arange(self._filter_size, dtype="float32") + coords -= (self._filter_size - 1) / 2. + + kernel = np.square(coords) + kernel *= -0.5 / np.square(self._filter_sigma) + kernel = np.reshape(kernel, (1, -1)) + np.reshape(kernel, (-1, 1)) + kernel = K.constant(np.reshape(kernel, (1, -1))) + kernel = K.softmax(kernel) + kernel = K.reshape(kernel, (self._filter_size, self._filter_size, 1, 1)) + return kernel + + @classmethod + def _depthwise_conv2d(cls, image: tf.Tensor, kernel: tf.Tensor) -> tf.Tensor: + """ Perform a standardized depthwise convolution. + + Parameters + ---------- + image: :class:`tf.Tensor` + Batch of images, channels last, to perform depthwise convolution + kernel: :class:`tf.Tensor` + convolution kernel + + Returns + ------- + :class:`tf.Tensor` + The output from the convolution + """ + return K.depthwise_conv2d(image, kernel, strides=(1, 1), padding="valid") + + def _get_ssim(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]: + """ Obtain the structural similarity between a batch of true and predicted images. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The input batch of ground truth images + y_pred: :class:`tf.Tensor` + The input batch of predicted images + + Returns + ------- + :class:`tf.Tensor` + The SSIM for the given images + :class:`tf.Tensor` + The Contrast for the given images + """ + channels = K.int_shape(y_true)[-1] + kernel = K.tile(self._kernel, (1, 1, channels, 1)) + + # SSIM luminance measure is (2 * mu_x * mu_y + c1) / (mu_x ** 2 + mu_y ** 2 + c1) + mean_true = self._depthwise_conv2d(y_true, kernel) + mean_pred = self._depthwise_conv2d(y_pred, kernel) + num_lum = mean_true * mean_pred * 2.0 + den_lum = K.square(mean_true) + K.square(mean_pred) + luminance = (num_lum + self._c1) / (den_lum + self._c1) + + # SSIM contrast-structure measure is (2 * cov_{xy} + c2) / (cov_{xx} + cov_{yy} + c2) + num_con = self._depthwise_conv2d(y_true * y_pred, kernel) * 2.0 + den_con = self._depthwise_conv2d(K.square(y_true) + K.square(y_pred), kernel) + + contrast = (num_con - num_lum + self._c2) / (den_con - den_lum + self._c2) + + # Average over the height x width dimensions + axes = (-3, -2) + ssim = K.mean(luminance * contrast, axis=axes) + contrast = K.mean(contrast, axis=axes) + + return ssim, contrast + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the DSSIM or MS-DSSIM Loss Function. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The input batch of ground truth images + y_pred: :class:`tf.Tensor` + The input batch of predicted images + + Returns + ------- + :class:`tf.Tensor` + The DSSIM or MS-DSSIM for the given images + """ + ssim = self._get_ssim(y_true, y_pred)[0] + retval = (1. - ssim) / 2.0 + return K.mean(retval) + + +class GMSDLoss(): # pylint:disable=too-few-public-methods + """ Gradient Magnitude Similarity Deviation Loss. + + Improved image quality metric over MS-SSIM with easier calculations + + References + ---------- + http://www4.comp.polyu.edu.hk/~cslzhang/IQA/GMSD/GMSD.htm + https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf + """ + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Return the Gradient Magnitude Similarity Deviation Loss. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class:`tf.Tensor` + The loss value + """ + true_edge = self._scharr_edges(y_true, True) + pred_edge = self._scharr_edges(y_pred, True) + ephsilon = 0.0025 + upper = 2.0 * true_edge * pred_edge + lower = K.square(true_edge) + K.square(pred_edge) + gms = (upper + ephsilon) / (lower + ephsilon) + gmsd = K.std(gms, axis=(1, 2, 3), keepdims=True) + gmsd = K.squeeze(gmsd, axis=-1) + return gmsd + + @classmethod + def _scharr_edges(cls, image: tf.Tensor, magnitude: bool) -> tf.Tensor: + """ Returns a tensor holding modified Scharr edge maps. + + Parameters + ---------- + image: :class:`tf.Tensor` + Image tensor with shape [batch_size, h, w, d] and type float32. The image(s) must be + 2x2 or larger. + magnitude: bool + Boolean to determine if the edge magnitude or edge direction is returned + + Returns + ------- + :class:`tf.Tensor` + Tensor holding edge maps for each channel. Returns a tensor with shape `[batch_size, h, + w, d, 2]` where the last two dimensions hold `[[dy[0], dx[0]], [dy[1], dx[1]], ..., + [dy[d-1], dx[d-1]]]` calculated using the Scharr filter. + """ + + # Define vertical and horizontal Scharr filters. + static_image_shape = image.get_shape() + image_shape = K.shape(image) + + # 5x5 modified Scharr kernel ( reshape to (5,5,1,2) ) + matrix = np.array([[[[0.00070, 0.00070]], + [[0.00520, 0.00370]], + [[0.03700, 0.00000]], + [[0.00520, -0.0037]], + [[0.00070, -0.0007]]], + [[[0.00370, 0.00520]], + [[0.11870, 0.11870]], + [[0.25890, 0.00000]], + [[0.11870, -0.1187]], + [[0.00370, -0.0052]]], + [[[0.00000, 0.03700]], + [[0.00000, 0.25890]], + [[0.00000, 0.00000]], + [[0.00000, -0.2589]], + [[0.00000, -0.0370]]], + [[[-0.0037, 0.00520]], + [[-0.1187, 0.11870]], + [[-0.2589, 0.00000]], + [[-0.1187, -0.1187]], + [[-0.0037, -0.0052]]], + [[[-0.0007, 0.00070]], + [[-0.0052, 0.00370]], + [[-0.0370, 0.00000]], + [[-0.0052, -0.0037]], + [[-0.0007, -0.0007]]]]) + num_kernels = [2] + kernels = K.constant(matrix, dtype='float32') + kernels = K.tile(kernels, [1, 1, image_shape[-1], 1]) + + # Use depth-wise convolution to calculate edge maps per channel. + # Output tensor has shape [batch_size, h, w, d * num_kernels]. + pad_sizes = [[0, 0], [2, 2], [2, 2], [0, 0]] + padded = tf.pad(image, # pylint:disable=unexpected-keyword-arg,no-value-for-parameter + pad_sizes, + mode='REFLECT') + output = K.depthwise_conv2d(padded, kernels) + + if not magnitude: # direction of edges + # Reshape to [batch_size, h, w, d, num_kernels]. + shape = K.concatenate([image_shape, num_kernels], axis=0) + output = K.reshape(output, shape=shape) + output.set_shape(static_image_shape.concatenate(num_kernels)) + output = tf.atan(K.squeeze(output[:, :, :, :, 0] / output[:, :, :, :, 1], axis=None)) + # magnitude of edges -- unified x & y edges don't work well with Neural Networks + return output + + +class LDRFLIPLoss(): # pylint:disable=too-few-public-methods + """ Computes the LDR-FLIP error map between two LDR images, assuming the images are observed + at a certain number of pixels per degree of visual angle. + + References + ---------- + https://research.nvidia.com/sites/default/files/node/3260/FLIP_Paper.pdf + https://github.com/NVlabs/flip + + License + ------- + BSD 3-Clause License + Copyright (c) 2020-2022, NVIDIA Corporation & AFFILIATES. All rights reserved. + Redistribution and use in source and binary forms, with or without modification, are permitted + provided that the following conditions are met: + Redistributions of source code must retain the above copyright notice, this list of conditions + and the following disclaimer. + Redistributions in binary form must reproduce the above copyright notice, this list of + conditions and the following disclaimer in the documentation and/or other materials provided + with the distribution. + Neither the name of the copyright holder nor the names of its contributors may be used to + endorse or promote products derived from this software without specific prior written + permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY + AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR + CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR + OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + POSSIBILITY OF SUCH DAMAGE. + + Parameters + ---------- + computed_distance_exponent: float, Optional + The computed distance exponent to apply to Hunt adjusted, filtered colors. + (`qc` in original paper). Default: `0.7` + feature_exponent: float, Optional + The feature exponent to apply for increasing the impact of feature difference on the + final loss value. (`qf` in original paper). Default: `0.5` + lower_threshold_exponent: float, Optional + The `pc` exponent for the color pipeline as described in the original paper: Default: `0.4` + upper_threshold_exponent: float, Optional + The `pt` exponent for the color pipeline as described in the original paper. + Default: `0.95` + epsilon: float + A small value to improve training stability. Default: `1e-15` + pixels_per_degree: float, Optional + The estimated number of pixels per degree of visual angle of the observer. This effectively + impacts the tolerance when calculating loss. The default corresponds to viewing images on a + 0.7m wide 4K monitor at 0.7m from the display. Default: ``None`` + color_order: str + The `"BGR"` or `"RGB"` color order of the incoming images + """ + def __init__(self, + computed_distance_exponent: float = 0.7, + feature_exponent: float = 0.5, + lower_threshold_exponent: float = 0.4, + upper_threshold_exponent: float = 0.95, + epsilon: float = 1e-15, + pixels_per_degree: float | None = None, + color_order: T.Literal["bgr", "rgb"] = "bgr") -> None: + logger.debug("Initializing: %s (computed_distance_exponent '%s', feature_exponent: %s, " + "lower_threshold_exponent: %s, upper_threshold_exponent: %s, epsilon: %s, " + "pixels_per_degree: %s, color_order: %s)", self.__class__.__name__, + computed_distance_exponent, feature_exponent, lower_threshold_exponent, + upper_threshold_exponent, epsilon, pixels_per_degree, color_order) + + self._computed_distance_exponent = computed_distance_exponent + self._feature_exponent = feature_exponent + self._pc = lower_threshold_exponent + self._pt = upper_threshold_exponent + self._epsilon = epsilon + self._color_order = color_order.lower() + + if pixels_per_degree is None: + pixels_per_degree = (0.7 * 3840 / 0.7) * np.pi / 180 + self._pixels_per_degree = pixels_per_degree + self._spatial_filters = _SpatialFilters(pixels_per_degree) + self._feature_detector = _FeatureDetection(pixels_per_degree) + logger.debug("Initialized: %s ", self.__class__.__name__) + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the LDR Flip Loss Function + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images + y_pred: :class:`tensorflow.Tensor` + The predicted batch of images + + Returns + ------- + :class::class:`tensorflow.Tensor` + The calculated Flip loss value + """ + if self._color_order == "bgr": # Switch models training in bgr order to rgb + y_true = y_true[..., 2::-1] + y_pred = y_pred[..., 2::-1] + + y_true = K.clip(y_true, 0, 1.) + y_pred = K.clip(y_pred, 0, 1.) + + rgb2ycxcz = ColorSpaceConvert("srgb", "ycxcz") + true_ycxcz = rgb2ycxcz(y_true) + pred_ycxcz = rgb2ycxcz(y_pred) + + delta_e_color = self._color_pipeline(true_ycxcz, pred_ycxcz) + delta_e_features = self._process_features(true_ycxcz, pred_ycxcz) + + loss = K.pow(delta_e_color, 1 - delta_e_features) + return loss + + def _color_pipeline(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Perform the color processing part of the FLIP loss function + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images in YCxCz color space + y_pred: :class:`tensorflow.Tensor` + The predicted batch of images in YCxCz color space + + Returns + ------- + :class:`tensorflow.Tensor` + The exponentiated, maximum HyAB difference between two colors in Hunt-adjusted + L*A*B* space + """ + filtered_true = self._spatial_filters(y_true) + filtered_pred = self._spatial_filters(y_pred) + + rgb2lab = ColorSpaceConvert(from_space="rgb", to_space="lab") + preprocessed_true = self._hunt_adjustment(rgb2lab(filtered_true)) + preprocessed_pred = self._hunt_adjustment(rgb2lab(filtered_pred)) + hunt_adjusted_green = self._hunt_adjustment(rgb2lab(K.constant([[[[0.0, 1.0, 0.0]]]], + dtype="float32"))) + hunt_adjusted_blue = self._hunt_adjustment(rgb2lab(K.constant([[[[0.0, 0.0, 1.0]]]], + dtype="float32"))) + + delta = self._hyab(preprocessed_true, preprocessed_pred) + power_delta = K.pow(delta, self._computed_distance_exponent) + cmax = K.pow(self._hyab(hunt_adjusted_green, hunt_adjusted_blue), + self._computed_distance_exponent) + return self._redistribute_errors(power_delta, cmax) + + def _process_features(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Perform the color processing part of the FLIP loss function + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images in YCxCz color space + y_pred: :class:`tensorflow.Tensor` + The predicted batch of images in YCxCz color space + + Returns + ------- + :class:`tensorflow.Tensor` + The exponentiated features delta + """ + col_y_true = (y_true[..., 0:1] + 16) / 116. + col_y_pred = (y_pred[..., 0:1] + 16) / 116. + + edges_true = self._feature_detector(col_y_true, "edge") + points_true = self._feature_detector(col_y_true, "point") + edges_pred = self._feature_detector(col_y_pred, "edge") + points_pred = self._feature_detector(col_y_pred, "point") + + delta = K.maximum(K.abs(frobenius_norm(edges_true) - frobenius_norm(edges_pred)), + K.abs(frobenius_norm(points_pred) - frobenius_norm(points_true))) + + delta = K.clip(delta, min_value=self._epsilon, max_value=None) + return K.pow(((1 / np.sqrt(2)) * delta), self._feature_exponent) + + @classmethod + def _hunt_adjustment(cls, image: tf.Tensor) -> tf.Tensor: + """ Apply Hunt-adjustment to an image in L*a*b* color space + + Parameters + ---------- + image: :class:`tensorflow.Tensor` + The batch of images in L*a*b* to adjust + + Returns + ------- + :class:`tensorflow.Tensor` + The hunt adjusted batch of images in L*a*b color space + """ + ch_l = image[..., 0:1] + adjusted = K.concatenate([ch_l, image[..., 1:] * (ch_l * 0.01)], axis=-1) + return adjusted + + def _hyab(self, y_true, y_pred): + """ Compute the HyAB distance between true and predicted images. + + Parameters + ---------- + y_true: :class:`tensorflow.Tensor` + The ground truth batch of images in standard or Hunt-adjusted L*A*B* color space + y_pred: :class:`tensorflow.Tensor` + The predicted batch of images in in standard or Hunt-adjusted L*A*B* color space + + Returns + ------- + :class:`tensorflow.Tensor` + image tensor containing the per-pixel HyAB distances between true and predicted images + """ + delta = y_true - y_pred + root = K.sqrt(K.clip(K.pow(delta[..., 0:1], 2), min_value=self._epsilon, max_value=None)) + delta_norm = frobenius_norm(delta[..., 1:3]) + return root + delta_norm + + def _redistribute_errors(self, power_delta_e_hyab, cmax): + """ Redistribute exponentiated HyAB errors to the [0,1] range + + Parameters + ---------- + power_delta_e_hyab: :class:`tensorflow.Tensor` + The exponentiated HyAb distance + cmax: :class:`tensorflow.Tensor` + The exponentiated, maximum HyAB difference between two colors in Hunt-adjusted + L*A*B* space + + Returns + ------- + :class:`tensorflow.Tensor` + The redistributed per-pixel HyAB distances (in range [0,1]) + """ + pccmax = self._pc * cmax + delta_e_c = K.switch( + power_delta_e_hyab < pccmax, + (self._pt / pccmax) * power_delta_e_hyab, + self._pt + ((power_delta_e_hyab - pccmax) / (cmax - pccmax)) * (1.0 - self._pt)) + return delta_e_c + + +class _SpatialFilters(): # pylint:disable=too-few-public-methods + """ Filters an image with channel specific spatial contrast sensitivity functions and clips + result to the unit cube in linear RGB. + + For use with LDRFlipLoss. + + Parameters + ---------- + pixels_per_degree: float + The estimated number of pixels per degree of visual angle of the observer. This effectively + impacts the tolerance when calculating loss. + """ + def __init__(self, pixels_per_degree: float) -> None: + self._pixels_per_degree = pixels_per_degree + self._spatial_filters, self._radius = self._generate_spatial_filters() + self._ycxcz2rgb = ColorSpaceConvert(from_space="ycxcz", to_space="rgb") + + def _generate_spatial_filters(self) -> tuple[tf.Tensor, int]: + """ Generates spatial contrast sensitivity filters with width depending on the number of + pixels per degree of visual angle of the observer for channels "A", "RG" and "BY" + + Returns + ------- + dict + the channels ("A" (Achromatic CSF), "RG" (Red-Green CSF) or "BY" (Blue-Yellow CSF)) as + key with the Filter kernel corresponding to the spatial contrast sensitivity function + of channel and kernel's radius + """ + mapping = {"A": {"a1": 1, "b1": 0.0047, "a2": 0, "b2": 1e-5}, + "RG": {"a1": 1, "b1": 0.0053, "a2": 0, "b2": 1e-5}, + "BY": {"a1": 34.1, "b1": 0.04, "a2": 13.5, "b2": 0.025}} + + domain, radius = self._get_evaluation_domain(mapping["A"]["b1"], + mapping["A"]["b2"], + mapping["RG"]["b1"], + mapping["RG"]["b2"], + mapping["BY"]["b1"], + mapping["BY"]["b2"]) + + weights = np.array([self._generate_weights(mapping[channel], domain) + for channel in ("A", "RG", "BY")]) + weights = K.constant(np.moveaxis(weights, 0, -1), dtype="float32") + + return weights, radius + + def _get_evaluation_domain(self, + b1_a: float, + b2_a: float, + b1_rg: float, + b2_rg: float, + b1_by: float, + b2_by: float) -> tuple[np.ndarray, int]: + """ TODO docstring """ + max_scale_parameter = max([b1_a, b2_a, b1_rg, b2_rg, b1_by, b2_by]) + delta_x = 1.0 / self._pixels_per_degree + radius = int(np.ceil(3 * np.sqrt(max_scale_parameter / (2 * np.pi**2)) + * self._pixels_per_degree)) + ax_x, ax_y = np.meshgrid(range(-radius, radius + 1), range(-radius, radius + 1)) + domain = (ax_x * delta_x) ** 2 + (ax_y * delta_x) ** 2 + return domain, radius + + @classmethod + def _generate_weights(cls, channel: dict[str, float], domain: np.ndarray) -> tf.Tensor: + """ TODO docstring """ + a_1, b_1, a_2, b_2 = channel["a1"], channel["b1"], channel["a2"], channel["b2"] + grad = (a_1 * np.sqrt(np.pi / b_1) * np.exp(-np.pi ** 2 * domain / b_1) + + a_2 * np.sqrt(np.pi / b_2) * np.exp(-np.pi ** 2 * domain / b_2)) + grad = grad / np.sum(grad) + grad = np.reshape(grad, (*grad.shape, 1)) + return grad + + def __call__(self, image: tf.Tensor) -> tf.Tensor: + """ Call the spacial filtering. + + Parameters + ---------- + image: Tensor + Image tensor to filter in YCxCz color space + + Returns + ------- + Tensor + The input image transformed to linear RGB after filtering with spatial contrast + sensitivity functions + """ + padded_image = replicate_pad(image, self._radius) + image_tilde_opponent = K.conv2d(padded_image, + self._spatial_filters, + strides=1, + padding="valid") + rgb = K.clip(self._ycxcz2rgb(image_tilde_opponent), 0., 1.) + return rgb + + +class _FeatureDetection(): # pylint:disable=too-few-public-methods + """ Detect features (i.e. edges and points) in an achromatic YCxCz image. + + For use with LDRFlipLoss. + + Parameters + ---------- + pixels_per_degree: float + The number of pixels per degree of visual angle of the observer + """ + def __init__(self, pixels_per_degree: float) -> None: + width = 0.082 + self._std = 0.5 * width * pixels_per_degree + self._radius = int(np.ceil(3 * self._std)) + self._grid = np.meshgrid(range(-self._radius, self._radius + 1), + range(-self._radius, self._radius + 1)) + self._gradient = np.exp(-(self._grid[0] ** 2 + self._grid[1] ** 2) + / (2 * (self._std ** 2))) + + def __call__(self, image: tf.Tensor, feature_type: str) -> tf.Tensor: + """ Run the feature detection + + Parameters + ---------- + image: Tensor + Batch of images in YCxCz color space with normalized Y values + feature_type: str + Type of features to detect (`"edge"` or `"point"`) + + Returns + ------- + Tensor + Detected features in the 0-1 range + """ + feature_type = feature_type.lower() + + if feature_type == 'edge': + grad_x = np.multiply(-self._grid[0], self._gradient) + else: + grad_x = np.multiply(self._grid[0] ** 2 / (self._std ** 2) - 1, self._gradient) + + negative_weights_sum = -np.sum(grad_x[grad_x < 0]) + positive_weights_sum = np.sum(grad_x[grad_x > 0]) + + grad_x = K.constant(grad_x) + grad_x = K.switch(grad_x < 0, grad_x / negative_weights_sum, grad_x / positive_weights_sum) + kernel = K.expand_dims(K.expand_dims(grad_x, axis=-1), axis=-1) + + features_x = K.conv2d(replicate_pad(image, self._radius), + kernel, + strides=1, + padding="valid") + kernel = K.permute_dimensions(kernel, (1, 0, 2, 3)) + features_y = K.conv2d(replicate_pad(image, self._radius), + kernel, + strides=1, + padding="valid") + features = K.concatenate([features_x, features_y], axis=-1) + return features + + +class MSSIMLoss(): # pylint:disable=too-few-public-methods + """ Multiscale Structural Similarity Loss Function + + Parameters + ---------- + k_1: float, optional + Parameter of the SSIM. Default: `0.01` + k_2: float, optional + Parameter of the SSIM. Default: `0.03` + filter_size: int, optional + size of gaussian filter Default: `11` + filter_sigma: float, optional + Width of gaussian filter Default: `1.5` + max_value: float, optional + Max value of the output. Default: `1.0` + power_factors: tuple, optional + Iterable of weights for each of the scales. The number of scales used is the length of the + list. Index 0 is the unscaled resolution's weight and each increasing scale corresponds to + the image being downsampled by 2. Defaults to the values obtained in the original paper. + Default: (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + + Notes + ------ + You should add a regularization term like a l2 loss in addition to this one. + """ + def __init__(self, + k_1: float = 0.01, + k_2: float = 0.03, + filter_size: int = 11, + filter_sigma: float = 1.5, + max_value: float = 1.0, + power_factors: tuple[float, ...] = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) + ) -> None: + self.filter_size = filter_size + self.filter_sigma = filter_sigma + self.k_1 = k_1 + self.k_2 = k_2 + self.max_value = max_value + self.power_factors = power_factors + + def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: + """ Call the MS-SSIM Loss Function. + + Parameters + ---------- + y_true: :class:`tf.Tensor` + The ground truth value + y_pred: :class:`tf.Tensor` + The predicted value + + Returns + ------- + :class:`tf.Tensor` + The MS-SSIM Loss value + """ + im_size = K.int_shape(y_true)[1] + # filter size cannot be larger than the smallest scale + smallest_scale = self._get_smallest_size(im_size, len(self.power_factors) - 1) + filter_size = min(self.filter_size, smallest_scale) + + ms_ssim = tf.image.ssim_multiscale(y_true, + y_pred, + self.max_value, + power_factors=self.power_factors, + filter_size=filter_size, + filter_sigma=self.filter_sigma, + k1=self.k_1, + k2=self.k_2) + ms_ssim_loss = 1. - ms_ssim + return K.mean(ms_ssim_loss) + + def _get_smallest_size(self, size: int, idx: int) -> int: + """ Recursive function to obtain the smallest size that the image will be scaled to. + + Parameters + ---------- + size: int + The current scaled size to iterate through + idx: int + The current iteration to be performed. When iteration hits zero the value will + be returned + + Returns + ------- + int + The smallest size the image will be scaled to based on the original image size and + the amount of scaling factors that will occur + """ + logger.debug("scale id: %s, size: %s", idx, size) + if idx > 0: + size = self._get_smallest_size(size // 2, idx - 1) + return size diff --git a/lib/model/losses_plaid.py b/lib/model/losses_plaid.py deleted file mode 100644 index 1d5cf5442b..0000000000 --- a/lib/model/losses_plaid.py +++ /dev/null @@ -1,658 +0,0 @@ -#!/usr/bin/env python3 -""" Custom Loss Functions for faceswap.py """ - -from __future__ import absolute_import - -import logging - -import numpy as np -import tensorflow as tf - -from keras import backend as K -from plaidml.op import extract_image_patches -from lib.plaidml_utils import pad -from lib.utils import FaceswapError - -logger = logging.getLogger(__name__) # pylint:disable=invalid-name - - -class DSSIMObjective(): - """ DSSIM Loss Function - - Difference of Structural Similarity (DSSIM loss function). Clipped between 0 and 0.5 - - Parameters - ---------- - k_1: float, optional - Parameter of the SSIM. Default: `0.01` - k_2: float, optional - Parameter of the SSIM. Default: `0.03` - kernel_size: int, optional - Size of the sliding window Default: `3` - max_value: float, optional - Max value of the output. Default: `1.0` - - Notes - ------ - You should add a regularization term like a l2 loss in addition to this one. - - References - ---------- - https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py - - MIT License - - Copyright (c) 2017 Fariz Rahman - - Permission is hereby granted, free of charge, to any person obtaining a copy - of this software and associated documentation files (the "Software"), to deal - in the Software without restriction, including without limitation the rights - to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - copies of the Software, and to permit persons to whom the Software is - furnished to do so, subject to the following conditions: - - The above copyright notice and this permission notice shall be included in all - copies or substantial portions of the Software. - - THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - SOFTWARE. - """ - def __init__(self, k_1=0.01, k_2=0.03, kernel_size=3, max_value=1.0): - self.__name__ = 'DSSIMObjective' - self.kernel_size = kernel_size - self.k_1 = k_1 - self.k_2 = k_2 - self.max_value = max_value - self.c_1 = (self.k_1 * self.max_value) ** 2 - self.c_2 = (self.k_2 * self.max_value) ** 2 - self.dim_ordering = K.image_data_format() - - @staticmethod - def _int_shape(input_tensor): - """ Returns the shape of tensor or variable as a tuple of int or None entries. - - Parameters - ---------- - input_tensor: tensor or variable - The input to return the shape for - - Returns - ------- - tuple - A tuple of integers (or None entries) - """ - return K.int_shape(input_tensor) - - def __call__(self, y_true, y_pred): - """ Call the DSSIM Loss Function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The DSSIM Loss value - - Notes - ----- - There are additional parameters for this function. some of the 'modes' for edge behavior - do not yet have a gradient definition in the Theano tree and cannot be used for learning - """ - - kernel = [self.kernel_size, self.kernel_size] - y_true = K.reshape(y_true, [-1] + list(self._int_shape(y_pred)[1:])) - y_pred = K.reshape(y_pred, [-1] + list(self._int_shape(y_pred)[1:])) - patches_pred = self.extract_image_patches(y_pred, - kernel, - kernel, - 'valid', - self.dim_ordering) - patches_true = self.extract_image_patches(y_true, - kernel, - kernel, - 'valid', - self.dim_ordering) - - # Get mean - u_true = K.mean(patches_true, axis=-1) - u_pred = K.mean(patches_pred, axis=-1) - # Get variance - var_true = K.var(patches_true, axis=-1) - var_pred = K.var(patches_pred, axis=-1) - # Get standard deviation - covar_true_pred = K.mean( - patches_true * patches_pred, axis=-1) - u_true * u_pred - - ssim = (2 * u_true * u_pred + self.c_1) * ( - 2 * covar_true_pred + self.c_2) - denom = (K.square(u_true) + K.square(u_pred) + self.c_1) * ( - var_pred + var_true + self.c_2) - ssim /= denom # no need for clipping, c_1 + c_2 make the denorm non-zero - return (1.0 - ssim) / 2.0 - - @staticmethod - def _preprocess_padding(padding): - """Convert keras padding to tensorflow padding. - - Parameters - ---------- - padding: string, - `"same"` or `"valid"`. - - Returns - ------- - str - `"SAME"` or `"VALID"`. - - Raises - ------ - ValueError - If `padding` is invalid. - """ - if padding == 'same': - padding = 'SAME' - elif padding == 'valid': - padding = 'VALID' - else: - raise ValueError('Invalid padding:', padding) - return padding - - def extract_image_patches(self, input_tensor, k_sizes, s_sizes, - padding='same', data_format='channels_last'): - """ Extract the patches from an image. - - Parameters - ---------- - input_tensor: tensor - The input image - k_sizes: tuple - 2-d tuple with the kernel size - s_sizes: tuple - 2-d tuple with the strides size - padding: str, optional - `"same"` or `"valid"`. Default: `"same"` - data_format: str, optional. - `"channels_last"` or `"channels_first"`. Default: `"channels_last"` - - Returns - ------- - The (k_w, k_h) patches extracted - Tensorflow ==> (batch_size, w, h, k_w, k_h, c) - Theano ==> (batch_size, w, h, c, k_w, k_h) - """ - kernel = [1, k_sizes[0], k_sizes[1], 1] - strides = [1, s_sizes[0], s_sizes[1], 1] - padding = self._preprocess_padding(padding) - if data_format == 'channels_first': - input_tensor = K.permute_dimensions(input_tensor, (0, 2, 3, 1)) - patches = extract_image_patches(input_tensor, kernel, strides, [1, 1, 1, 1], padding) - return patches - - -class MSSSIMLoss(): # pylint:disable=too-few-public-methods - """ Multiscale Structural Similarity Loss Function - - Parameters - ---------- - k_1: float, optional - Parameter of the SSIM. Default: `0.01` - k_2: float, optional - Parameter of the SSIM. Default: `0.03` - filter_size: int, optional - size of gaussian filter Default: `11` - filter_sigma: float, optional - Width of gaussian filter Default: `1.5` - max_value: float, optional - Max value of the output. Default: `1.0` - power_factors: tuple, optional - Iterable of weights for each of the scales. The number of scales used is the length of the - list. Index 0 is the unscaled resolution's weight and each increasing scale corresponds to - the image being downsampled by 2. Defaults to the values obtained in the original paper. - Default: (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) - - Notes - ------ - You should add a regularization term like a l2 loss in addition to this one. - """ - def __init__(self, - k_1=0.01, - k_2=0.03, - filter_size=4, - filter_sigma=1.5, - max_value=1.0, - power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): - self.filter_size = filter_size - self.filter_sigma = filter_sigma - self.k_1 = k_1 - self.k_2 = k_2 - self.max_value = max_value - self.power_factors = power_factors - - def __call__(self, y_true, y_pred): - """ Call the MS-SSIM Loss Function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The MS-SSIM Loss value - """ - raise FaceswapError("MS-SSIM Loss is not currently compatible with PlaidML. Please select " - "a different Loss method.") - - -class GeneralizedLoss(): # pylint:disable=too-few-public-methods - """ Generalized function used to return a large variety of mathematical loss functions. - - The primary benefit is a smooth, differentiable version of L1 loss. - - References - ---------- - Barron, J. A More General Robust Loss Function - https://arxiv.org/pdf/1701.03077.pdf - - Example - ------- - >>> a=1.0, x>>c , c=1.0/255.0 # will give a smoothly differentiable version of L1 / MAE loss - >>> a=1.999999 (limit as a->2), beta=1.0/255.0 # will give L2 / RMSE loss - - Parameters - ---------- - alpha: float, optional - Penalty factor. Larger number give larger weight to large deviations. Default: `1.0` - beta: float, optional - Scale factor used to adjust to the input scale (i.e. inputs of mean `1e-4` or `256`). - Default: `1.0/255.0` - """ - def __init__(self, alpha=1.0, beta=1.0/255.0): - self.alpha = alpha - self.beta = beta - - def __call__(self, y_true, y_pred): - """ Call the Generalized Loss Function - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value from the results of function(y_pred - y_true) - """ - diff = y_pred - y_true - second = (K.pow(K.pow(diff/self.beta, 2.) / K.abs(2. - self.alpha) + 1., - (self.alpha / 2.)) - 1.) - loss = (K.abs(2. - self.alpha)/self.alpha) * second - loss = K.mean(loss, axis=-1) * self.beta - return loss - - -class LInfNorm(): # pylint:disable=too-few-public-methods - """ Calculate the L-inf norm as a loss function. """ - - def __call__(self, y_true, y_pred): - """ Call the L-inf norm loss function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - diff = K.abs(y_true - y_pred) - max_loss = K.max(diff, axis=(1, 2), keepdims=True) - loss = K.mean(max_loss, axis=-1) - return loss - - -class GradientLoss(): # pylint:disable=too-few-public-methods - """ Gradient Loss Function. - - Calculates the first and second order gradient difference between pixels of an image in the x - and y dimensions. These gradients are then compared between the ground truth and the predicted - image and the difference is taken. When used as a loss, its minimization will result in - predicted images approaching the same level of sharpness / blurriness as the ground truth. - - References - ---------- - TV+TV2 Regularization with Non-Convex Sparseness-Inducing Penalty for Image Restoration, - Chengwu Lu & Hua Huang, 2014 - http://downloads.hindawi.com/journals/mpe/2014/790547.pdf - """ - def __init__(self): - self.generalized_loss = GeneralizedLoss(alpha=1.9999) - - def __call__(self, y_true, y_pred): - """ Call the gradient loss function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - tv_weight = 1.0 - tv2_weight = 1.0 - loss = 0.0 - loss += tv_weight * (self.generalized_loss(self._diff_x(y_true), self._diff_x(y_pred)) + - self.generalized_loss(self._diff_y(y_true), self._diff_y(y_pred))) - loss += tv2_weight * (self.generalized_loss(self._diff_xx(y_true), self._diff_xx(y_pred)) + - self.generalized_loss(self._diff_yy(y_true), self._diff_yy(y_pred)) + - self.generalized_loss(self._diff_xy(y_true), self._diff_xy(y_pred)) - * 2.) - loss = loss / (tv_weight + tv2_weight) - # TODO simplify to use MSE instead - return loss - - @classmethod - def _diff_x(cls, img): - """ X Difference """ - x_left = img[:, :, 1:2, :] - img[:, :, 0:1, :] - x_inner = img[:, :, 2:, :] - img[:, :, :-2, :] - x_right = img[:, :, -1:, :] - img[:, :, -2:-1, :] - x_out = K.concatenate([x_left, x_inner, x_right], axis=2) - return x_out * 0.5 - - @classmethod - def _diff_y(cls, img): - """ Y Difference """ - y_top = img[:, 1:2, :, :] - img[:, 0:1, :, :] - y_inner = img[:, 2:, :, :] - img[:, :-2, :, :] - y_bot = img[:, -1:, :, :] - img[:, -2:-1, :, :] - y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) - return y_out * 0.5 - - @classmethod - def _diff_xx(cls, img): - """ X-X Difference """ - x_left = img[:, :, 1:2, :] + img[:, :, 0:1, :] - x_inner = img[:, :, 2:, :] + img[:, :, :-2, :] - x_right = img[:, :, -1:, :] + img[:, :, -2:-1, :] - x_out = K.concatenate([x_left, x_inner, x_right], axis=2) - return x_out - 2.0 * img - - @classmethod - def _diff_yy(cls, img): - """ Y-Y Difference """ - y_top = img[:, 1:2, :, :] + img[:, 0:1, :, :] - y_inner = img[:, 2:, :, :] + img[:, :-2, :, :] - y_bot = img[:, -1:, :, :] + img[:, -2:-1, :, :] - y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) - return y_out - 2.0 * img - - @classmethod - def _diff_xy(cls, img): - """ X-Y Difference """ - # xout1 - top_left = img[:, 1:2, 1:2, :] + img[:, 0:1, 0:1, :] - inner_left = img[:, 2:, 1:2, :] + img[:, :-2, 0:1, :] - bot_left = img[:, -1:, 1:2, :] + img[:, -2:-1, 0:1, :] - xy_left = K.concatenate([top_left, inner_left, bot_left], axis=1) - - top_mid = img[:, 1:2, 2:, :] + img[:, 0:1, :-2, :] - mid_mid = img[:, 2:, 2:, :] + img[:, :-2, :-2, :] - bot_mid = img[:, -1:, 2:, :] + img[:, -2:-1, :-2, :] - xy_mid = K.concatenate([top_mid, mid_mid, bot_mid], axis=1) - - top_right = img[:, 1:2, -1:, :] + img[:, 0:1, -2:-1, :] - inner_right = img[:, 2:, -1:, :] + img[:, :-2, -2:-1, :] - bot_right = img[:, -1:, -1:, :] + img[:, -2:-1, -2:-1, :] - xy_right = K.concatenate([top_right, inner_right, bot_right], axis=1) - - # Xout2 - top_left = img[:, 0:1, 1:2, :] + img[:, 1:2, 0:1, :] - inner_left = img[:, :-2, 1:2, :] + img[:, 2:, 0:1, :] - bot_left = img[:, -2:-1, 1:2, :] + img[:, -1:, 0:1, :] - xy_left = K.concatenate([top_left, inner_left, bot_left], axis=1) - - top_mid = img[:, 0:1, 2:, :] + img[:, 1:2, :-2, :] - mid_mid = img[:, :-2, 2:, :] + img[:, 2:, :-2, :] - bot_mid = img[:, -2:-1, 2:, :] + img[:, -1:, :-2, :] - xy_mid = K.concatenate([top_mid, mid_mid, bot_mid], axis=1) - - top_right = img[:, 0:1, -1:, :] + img[:, 1:2, -2:-1, :] - inner_right = img[:, :-2, -1:, :] + img[:, 2:, -2:-1, :] - bot_right = img[:, -2:-1, -1:, :] + img[:, -1:, -2:-1, :] - xy_right = K.concatenate([top_right, inner_right, bot_right], axis=1) - - xy_out1 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) - xy_out2 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) - return (xy_out1 - xy_out2) * 0.25 - - -class GMSDLoss(): # pylint:disable=too-few-public-methods - """ Gradient Magnitude Similarity Deviation Loss. - - Improved image quality metric over MS-SSIM with easier calculations - - References - ---------- - http://www4.comp.polyu.edu.hk/~cslzhang/IQA/GMSD/GMSD.htm - https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf - """ - - def __call__(self, y_true, y_pred): - """ Return the Gradient Magnitude Similarity Deviation Loss. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - raise FaceswapError("GMSD Loss is not currently compatible with PlaidML. Please select a " - "different Loss method.") - - true_edge = self._scharr_edges(y_true, True) - pred_edge = self._scharr_edges(y_pred, True) - ephsilon = 0.0025 - upper = 2.0 * true_edge * pred_edge - lower = K.square(true_edge) + K.square(pred_edge) - gms = (upper + ephsilon) / (lower + ephsilon) - gmsd = K.std(gms, axis=(1, 2, 3), keepdims=True) - gmsd = K.squeeze(gmsd, axis=-1) - return gmsd - - @classmethod - def _scharr_edges(cls, image, magnitude): - """ Returns a tensor holding modified Scharr edge maps. - - Parameters - ---------- - image: tensor - Image tensor with shape [batch_size, h, w, d] and type float32. The image(s) must be - 2x2 or larger. - magnitude: bool - Boolean to determine if the edge magnitude or edge direction is returned - - Returns - ------- - tensor - Tensor holding edge maps for each channel. Returns a tensor with shape `[batch_size, h, - w, d, 2]` where the last two dimensions hold `[[dy[0], dx[0]], [dy[1], dx[1]], ..., - [dy[d-1], dx[d-1]]]` calculated using the Scharr filter. - """ - - # Define vertical and horizontal Scharr filters. - # TODO PlaidML: AttributeError: 'Value' object has no attribute 'get_shape' - static_image_shape = image.get_shape() - image_shape = K.shape(image) - - # 5x5 modified Scharr kernel ( reshape to (5,5,1,2) ) - matrix = np.array([[[[0.00070, 0.00070]], - [[0.00520, 0.00370]], - [[0.03700, 0.00000]], - [[0.00520, -0.0037]], - [[0.00070, -0.0007]]], - [[[0.00370, 0.00520]], - [[0.11870, 0.11870]], - [[0.25890, 0.00000]], - [[0.11870, -0.1187]], - [[0.00370, -0.0052]]], - [[[0.00000, 0.03700]], - [[0.00000, 0.25890]], - [[0.00000, 0.00000]], - [[0.00000, -0.2589]], - [[0.00000, -0.0370]]], - [[[-0.0037, 0.00520]], - [[-0.1187, 0.11870]], - [[-0.2589, 0.00000]], - [[-0.1187, -0.1187]], - [[-0.0037, -0.0052]]], - [[[-0.0007, 0.00070]], - [[-0.0052, 0.00370]], - [[-0.0370, 0.00000]], - [[-0.0052, -0.0037]], - [[-0.0007, -0.0007]]]]) - num_kernels = [2] - kernels = K.constant(matrix, dtype='float32') - kernels = K.tile(kernels, [1, 1, image_shape[-1], 1]) - - # Use depth-wise convolution to calculate edge maps per channel. - # Output tensor has shape [batch_size, h, w, d * num_kernels]. - pad_sizes = [[0, 0], [2, 2], [2, 2], [0, 0]] - padded = pad(image, pad_sizes, mode='REFLECT') - output = K.depthwise_conv2d(padded, kernels) - - if not magnitude: # direction of edges - # Reshape to [batch_size, h, w, d, num_kernels]. - shape = K.concatenate([image_shape, num_kernels], axis=0) - output = K.reshape(output, shape=shape) - output.set_shape(static_image_shape.concatenate(num_kernels)) - output = tf.atan(K.squeeze(output[:, :, :, :, 0] / output[:, :, :, :, 1], axis=None)) - # magnitude of edges -- unified x & y edges don't work well with Neural Networks - return output - - -class LossWrapper(): # pylint:disable=too-few-public-methods - """ A wrapper class for multiple keras losses to enable multiple weighted loss functions on a - single output and masking. - """ - def __init__(self): - logger.debug("Initializing: %s", self.__class__.__name__) - self._loss_functions = [] - self._loss_weights = [] - self._mask_channels = [] - logger.debug("Initialized: %s", self.__class__.__name__) - - def add_loss(self, function, weight=1.0, mask_channel=-1): - """ Add the given loss function with the given weight to the loss function chain. - - Parameters - ---------- - function: :class:`keras.losses.Loss` - The loss function to add to the loss chain - weight: float, optional - The weighting to apply to the loss function. Default: `1.0` - mask_channel: int, optional - The channel in the `y_true` image that the mask exists in. Set to `-1` if there is no - mask for the given loss function. Default: `-1` - """ - logger.debug("Adding loss: (function: %s, weight: %s, mask_channel: %s)", - function, weight, mask_channel) - self._loss_functions.append(function) - self._loss_weights.append(weight) - self._mask_channels.append(mask_channel) - - def __call__(self, y_true, y_pred): - """ Call the sub loss functions for the loss wrapper. - - Weights are returned as the weighted sum of the chosen losses. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The final loss value - """ - loss = 0.0 - for func, weight, mask_channel in zip(self._loss_functions, - self._loss_weights, - self._mask_channels): - logger.debug("Processing loss function: (func: %s, weight: %s, mask_channel: %s)", - func, weight, mask_channel) - n_true, n_pred = self._apply_mask(y_true, y_pred, mask_channel) - if isinstance(func, DSSIMObjective): - # Extract Image Patches in SSIM requires that y_pred be of a known shape, so - # specifically reshape the tensor. - n_pred = K.reshape(n_pred, K.int_shape(y_pred)) - this_loss = func(n_true, n_pred) - loss_dims = K.ndim(this_loss) - loss += (K.mean(this_loss, axis=list(range(1, loss_dims))) * weight) - return loss - - @classmethod - def _apply_mask(cls, y_true, y_pred, mask_channel, mask_prop=1.0): - """ Apply the mask to the input y_true and y_pred. If a mask is not required then - return the unmasked inputs. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - mask_channel: int - The channel within y_true that the required mask resides in - mask_prop: float, optional - The amount of mask propagation. Default: `1.0` - - Returns - ------- - tuple - (n_true, n_pred): The ground truth and predicted value tensors with the mask applied - """ - if mask_channel == -1: - logger.debug("No mask to apply") - return y_true[..., :3], y_pred[..., :3] - - logger.debug("Applying mask from channel %s", mask_channel) - mask = K.expand_dims(y_true[..., mask_channel], axis=-1) - mask_as_k_inv_prop = 1 - mask_prop - mask = (mask * mask_prop) + mask_as_k_inv_prop - - n_true = y_true[..., :3] * mask - n_pred = y_pred * mask - return n_true, n_pred diff --git a/lib/model/losses_tf.py b/lib/model/losses_tf.py deleted file mode 100644 index cd5aa132d9..0000000000 --- a/lib/model/losses_tf.py +++ /dev/null @@ -1,598 +0,0 @@ -#!/usr/bin/env python3 -""" Custom Loss Functions for faceswap.py """ - -from __future__ import absolute_import - -import logging -from typing import Tuple - -import numpy as np -import tensorflow as tf - -# Ignore linting errors from Tensorflow's thoroughly broken import system -from tensorflow.python.keras.engine import compile_utils # noqa pylint:disable=no-name-in-module,import-error -from tensorflow.keras import backend as K # pylint:disable=import-error - -logger = logging.getLogger(__name__) - - -class DSSIMObjective(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ DSSIM Loss Function - - Difference of Structural Similarity (DSSIM loss function). - - Parameters - ---------- - k_1: float, optional - Parameter of the SSIM. Default: `0.01` - k_2: float, optional - Parameter of the SSIM. Default: `0.03` - filter_size: int, optional - size of gaussian filter Default: `11` - filter_sigma: float, optional - Width of gaussian filter Default: `1.5` - max_value: float, optional - Max value of the output. Default: `1.0` - - Notes - ------ - You should add a regularization term like a l2 loss in addition to this one. - """ - def __init__(self, k_1=0.01, k_2=0.03, filter_size=11, filter_sigma=1.5, max_value=1.0): - super().__init__(name="DSSIMObjective", reduction=tf.keras.losses.Reduction.NONE) - self._filter_size = filter_size - self._filter_sigma = filter_sigma - self._k_1 = k_1 - self._k_2 = k_2 - self._max_value = max_value - - def call(self, y_true, y_pred): - """ Call the DSSIM Loss Function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The DSSIM Loss value - """ - ssim = tf.image.ssim(y_true, - y_pred, - self._max_value, - filter_size=self._filter_size, - filter_sigma=self._filter_sigma, - k1=self._k_1, - k2=self._k_2) - dssim_loss = (1. - ssim) / 2.0 - return dssim_loss - - -class MSSSIMLoss(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ Multiscale Structural Similarity Loss Function - - Parameters - ---------- - k_1: float, optional - Parameter of the SSIM. Default: `0.01` - k_2: float, optional - Parameter of the SSIM. Default: `0.03` - filter_size: int, optional - size of gaussian filter Default: `11` - filter_sigma: float, optional - Width of gaussian filter Default: `1.5` - max_value: float, optional - Max value of the output. Default: `1.0` - power_factors: tuple, optional - Iterable of weights for each of the scales. The number of scales used is the length of the - list. Index 0 is the unscaled resolution's weight and each increasing scale corresponds to - the image being downsampled by 2. Defaults to the values obtained in the original paper. - Default: (0.0448, 0.2856, 0.3001, 0.2363, 0.1333) - - Notes - ------ - You should add a regularization term like a l2 loss in addition to this one. - """ - def __init__(self, - k_1=0.01, - k_2=0.03, - filter_size=4, - filter_sigma=1.5, - max_value=1.0, - power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)): - super().__init__(name="SSIM_Multiscale_Loss", reduction=tf.keras.losses.Reduction.NONE) - self.filter_size = filter_size - self.filter_sigma = filter_sigma - self.k_1 = k_1 - self.k_2 = k_2 - self.max_value = max_value - self.power_factors = power_factors - - def call(self, y_true, y_pred): - """ Call the MS-SSIM Loss Function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The MS-SSIM Loss value - """ - im_size = K.int_shape(y_true)[1] - # filter size cannot be larger than the smallest scale - smallest_scale = self._get_smallest_size(im_size, len(self.power_factors) - 1) - filter_size = min(self.filter_size, smallest_scale) - - ms_ssim = tf.image.ssim_multiscale(y_true, - y_pred, - self.max_value, - power_factors=self.power_factors, - filter_size=filter_size, - filter_sigma=self.filter_sigma, - k1=self.k_1, - k2=self.k_2) - ms_ssim_loss = 1. - ms_ssim - return ms_ssim_loss - - def _get_smallest_size(self, size, idx): - """ Recursive function to obtain the smallest size that the image will be scaled to. - - Parameters - ---------- - size: int - The current scaled size to iterate through - idx: int - The current iteration to be performed. When iteration hits zero the value will - be returned - - Returns - ------- - int - The smallest size the image will be scaled to based on the original image size and - the amount of scaling factors that will occur - """ - logger.debug("scale id: %s, size: %s", idx, size) - if idx > 0: - size = self._get_smallest_size(size // 2, idx - 1) - return size - - -class GeneralizedLoss(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ Generalized function used to return a large variety of mathematical loss functions. - - The primary benefit is a smooth, differentiable version of L1 loss. - - References - ---------- - Barron, J. A More General Robust Loss Function - https://arxiv.org/pdf/1701.03077.pdf - - Example - ------- - >>> a=1.0, x>>c , c=1.0/255.0 # will give a smoothly differentiable version of L1 / MAE loss - >>> a=1.999999 (limit as a->2), beta=1.0/255.0 # will give L2 / RMSE loss - - Parameters - ---------- - alpha: float, optional - Penalty factor. Larger number give larger weight to large deviations. Default: `1.0` - beta: float, optional - Scale factor used to adjust to the input scale (i.e. inputs of mean `1e-4` or `256`). - Default: `1.0/255.0` - """ - def __init__(self, alpha=1.0, beta=1.0/255.0): - super().__init__(name="generalized_loss") - self.alpha = alpha - self.beta = beta - - def call(self, y_true, y_pred): - """ Call the Generalized Loss Function - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value from the results of function(y_pred - y_true) - """ - diff = y_pred - y_true - second = (K.pow(K.pow(diff/self.beta, 2.) / K.abs(2. - self.alpha) + 1., - (self.alpha / 2.)) - 1.) - loss = (K.abs(2. - self.alpha)/self.alpha) * second - loss = K.mean(loss, axis=-1) * self.beta - return loss - - -class LInfNorm(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ Calculate the L-inf norm as a loss function. """ - def __init__(self): - super().__init__(name="l_inf_norm_loss") - - @classmethod - def call(cls, y_true, y_pred): - """ Call the L-inf norm loss function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - diff = K.abs(y_true - y_pred) - max_loss = K.max(diff, axis=(1, 2), keepdims=True) - loss = K.mean(max_loss, axis=-1) - return loss - - -class GradientLoss(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ Gradient Loss Function. - - Calculates the first and second order gradient difference between pixels of an image in the x - and y dimensions. These gradients are then compared between the ground truth and the predicted - image and the difference is taken. When used as a loss, its minimization will result in - predicted images approaching the same level of sharpness / blurriness as the ground truth. - - References - ---------- - TV+TV2 Regularization with Non-Convex Sparseness-Inducing Penalty for Image Restoration, - Chengwu Lu & Hua Huang, 2014 - http://downloads.hindawi.com/journals/mpe/2014/790547.pdf - """ - def __init__(self): - super().__init__(name="generalized_loss") - self.generalized_loss = GeneralizedLoss(alpha=1.9999) - - def call(self, y_true, y_pred): - """ Call the gradient loss function. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - tv_weight = 1.0 - tv2_weight = 1.0 - loss = 0.0 - loss += tv_weight * (self.generalized_loss(self._diff_x(y_true), self._diff_x(y_pred)) + - self.generalized_loss(self._diff_y(y_true), self._diff_y(y_pred))) - loss += tv2_weight * (self.generalized_loss(self._diff_xx(y_true), self._diff_xx(y_pred)) + - self.generalized_loss(self._diff_yy(y_true), self._diff_yy(y_pred)) + - self.generalized_loss(self._diff_xy(y_true), self._diff_xy(y_pred)) - * 2.) - loss = loss / (tv_weight + tv2_weight) - # TODO simplify to use MSE instead - return loss - - @classmethod - def _diff_x(cls, img): - """ X Difference """ - x_left = img[:, :, 1:2, :] - img[:, :, 0:1, :] - x_inner = img[:, :, 2:, :] - img[:, :, :-2, :] - x_right = img[:, :, -1:, :] - img[:, :, -2:-1, :] - x_out = K.concatenate([x_left, x_inner, x_right], axis=2) - return x_out * 0.5 - - @classmethod - def _diff_y(cls, img): - """ Y Difference """ - y_top = img[:, 1:2, :, :] - img[:, 0:1, :, :] - y_inner = img[:, 2:, :, :] - img[:, :-2, :, :] - y_bot = img[:, -1:, :, :] - img[:, -2:-1, :, :] - y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) - return y_out * 0.5 - - @classmethod - def _diff_xx(cls, img): - """ X-X Difference """ - x_left = img[:, :, 1:2, :] + img[:, :, 0:1, :] - x_inner = img[:, :, 2:, :] + img[:, :, :-2, :] - x_right = img[:, :, -1:, :] + img[:, :, -2:-1, :] - x_out = K.concatenate([x_left, x_inner, x_right], axis=2) - return x_out - 2.0 * img - - @classmethod - def _diff_yy(cls, img): - """ Y-Y Difference """ - y_top = img[:, 1:2, :, :] + img[:, 0:1, :, :] - y_inner = img[:, 2:, :, :] + img[:, :-2, :, :] - y_bot = img[:, -1:, :, :] + img[:, -2:-1, :, :] - y_out = K.concatenate([y_top, y_inner, y_bot], axis=1) - return y_out - 2.0 * img - - @classmethod - def _diff_xy(cls, img): - """ X-Y Difference """ - # xout1 - top_left = img[:, 1:2, 1:2, :] + img[:, 0:1, 0:1, :] - inner_left = img[:, 2:, 1:2, :] + img[:, :-2, 0:1, :] - bot_left = img[:, -1:, 1:2, :] + img[:, -2:-1, 0:1, :] - xy_left = K.concatenate([top_left, inner_left, bot_left], axis=1) - - top_mid = img[:, 1:2, 2:, :] + img[:, 0:1, :-2, :] - mid_mid = img[:, 2:, 2:, :] + img[:, :-2, :-2, :] - bot_mid = img[:, -1:, 2:, :] + img[:, -2:-1, :-2, :] - xy_mid = K.concatenate([top_mid, mid_mid, bot_mid], axis=1) - - top_right = img[:, 1:2, -1:, :] + img[:, 0:1, -2:-1, :] - inner_right = img[:, 2:, -1:, :] + img[:, :-2, -2:-1, :] - bot_right = img[:, -1:, -1:, :] + img[:, -2:-1, -2:-1, :] - xy_right = K.concatenate([top_right, inner_right, bot_right], axis=1) - - # Xout2 - top_left = img[:, 0:1, 1:2, :] + img[:, 1:2, 0:1, :] - inner_left = img[:, :-2, 1:2, :] + img[:, 2:, 0:1, :] - bot_left = img[:, -2:-1, 1:2, :] + img[:, -1:, 0:1, :] - xy_left = K.concatenate([top_left, inner_left, bot_left], axis=1) - - top_mid = img[:, 0:1, 2:, :] + img[:, 1:2, :-2, :] - mid_mid = img[:, :-2, 2:, :] + img[:, 2:, :-2, :] - bot_mid = img[:, -2:-1, 2:, :] + img[:, -1:, :-2, :] - xy_mid = K.concatenate([top_mid, mid_mid, bot_mid], axis=1) - - top_right = img[:, 0:1, -1:, :] + img[:, 1:2, -2:-1, :] - inner_right = img[:, :-2, -1:, :] + img[:, 2:, -2:-1, :] - bot_right = img[:, -2:-1, -1:, :] + img[:, -1:, -2:-1, :] - xy_right = K.concatenate([top_right, inner_right, bot_right], axis=1) - - xy_out1 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) - xy_out2 = K.concatenate([xy_left, xy_mid, xy_right], axis=2) - return (xy_out1 - xy_out2) * 0.25 - - -class GMSDLoss(tf.keras.losses.Loss): # pylint:disable=too-few-public-methods - """ Gradient Magnitude Similarity Deviation Loss. - - Improved image quality metric over MS-SSIM with easier calculations - - References - ---------- - http://www4.comp.polyu.edu.hk/~cslzhang/IQA/GMSD/GMSD.htm - https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf - """ - def __init__(self): - super().__init__(name="gmsd_loss", reduction=tf.keras.losses.Reduction.NONE) - - def call(self, y_true, y_pred): - """ Return the Gradient Magnitude Similarity Deviation Loss. - - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - - Returns - ------- - tensor - The loss value - """ - true_edge = self._scharr_edges(y_true, True) - pred_edge = self._scharr_edges(y_pred, True) - ephsilon = 0.0025 - upper = 2.0 * true_edge * pred_edge - lower = K.square(true_edge) + K.square(pred_edge) - gms = (upper + ephsilon) / (lower + ephsilon) - gmsd = K.std(gms, axis=(1, 2, 3), keepdims=True) - gmsd = K.squeeze(gmsd, axis=-1) - return gmsd - - @classmethod - def _scharr_edges(cls, image, magnitude): - """ Returns a tensor holding modified Scharr edge maps. - - Parameters - ---------- - image: tensor - Image tensor with shape [batch_size, h, w, d] and type float32. The image(s) must be - 2x2 or larger. - magnitude: bool - Boolean to determine if the edge magnitude or edge direction is returned - - Returns - ------- - tensor - Tensor holding edge maps for each channel. Returns a tensor with shape `[batch_size, h, - w, d, 2]` where the last two dimensions hold `[[dy[0], dx[0]], [dy[1], dx[1]], ..., - [dy[d-1], dx[d-1]]]` calculated using the Scharr filter. - """ - - # Define vertical and horizontal Scharr filters. - static_image_shape = image.get_shape() - image_shape = K.shape(image) - - # 5x5 modified Scharr kernel ( reshape to (5,5,1,2) ) - matrix = np.array([[[[0.00070, 0.00070]], - [[0.00520, 0.00370]], - [[0.03700, 0.00000]], - [[0.00520, -0.0037]], - [[0.00070, -0.0007]]], - [[[0.00370, 0.00520]], - [[0.11870, 0.11870]], - [[0.25890, 0.00000]], - [[0.11870, -0.1187]], - [[0.00370, -0.0052]]], - [[[0.00000, 0.03700]], - [[0.00000, 0.25890]], - [[0.00000, 0.00000]], - [[0.00000, -0.2589]], - [[0.00000, -0.0370]]], - [[[-0.0037, 0.00520]], - [[-0.1187, 0.11870]], - [[-0.2589, 0.00000]], - [[-0.1187, -0.1187]], - [[-0.0037, -0.0052]]], - [[[-0.0007, 0.00070]], - [[-0.0052, 0.00370]], - [[-0.0370, 0.00000]], - [[-0.0052, -0.0037]], - [[-0.0007, -0.0007]]]]) - num_kernels = [2] - kernels = K.constant(matrix, dtype='float32') - kernels = K.tile(kernels, [1, 1, image_shape[-1], 1]) - - # Use depth-wise convolution to calculate edge maps per channel. - # Output tensor has shape [batch_size, h, w, d * num_kernels]. - pad_sizes = [[0, 0], [2, 2], [2, 2], [0, 0]] - padded = tf.pad(image, # pylint:disable=unexpected-keyword-arg,no-value-for-parameter - pad_sizes, - mode='REFLECT') - output = K.depthwise_conv2d(padded, kernels) - - if not magnitude: # direction of edges - # Reshape to [batch_size, h, w, d, num_kernels]. - shape = K.concatenate([image_shape, num_kernels], axis=0) - output = K.reshape(output, shape=shape) - output.set_shape(static_image_shape.concatenate(num_kernels)) - output = tf.atan(K.squeeze(output[:, :, :, :, 0] / output[:, :, :, :, 1], axis=None)) - # magnitude of edges -- unified x & y edges don't work well with Neural Networks - return output - - -class LossWrapper(): - """ A wrapper class for multiple keras losses to enable multiple masked weighted loss - functions on a single output. - - Notes - ----- - Whilst Keras does allow for applying multiple weighted loss functions, it does not allow - for an easy mechanism to add additional data (in our case masks) that are batch specific - but are not fed in to the model. - - This wrapper receives this additional mask data for the batch stacked onto the end of the - color channels of the received :param:`y_true` batch of images. These masks are then split - off the batch of images and applied to both the :param:`y_true` and :param:`y_pred` tensors - prior to feeding into the loss functions. - - For example, for an image of shape (4, 128, 128, 3) 3 additional masks may be stacked onto - the end of y_true, meaning we receive an input of shape (4, 128, 128, 6). This wrapper then - splits off (4, 128, 128, 3:6) from the end of the tensor, leaving the original y_true of - shape (4, 128, 128, 3) ready for masking and feeding through the loss functions. - """ - def __init__(self) -> None: - logger.debug("Initializing: %s", self.__class__.__name__) - self._loss_functions = [] - self._loss_weights = [] - self._mask_channels = [] - logger.debug("Initialized: %s", self.__class__.__name__) - - def add_loss(self, - function: tf.keras.losses.Loss, - weight: float = 1.0, - mask_channel: int = -1) -> None: - """ Add the given loss function with the given weight to the loss function chain. - - Parameters - ---------- - function: :class:`tf.keras.losses.Loss` - The loss function to add to the loss chain - weight: float, optional - The weighting to apply to the loss function. Default: `1.0` - mask_channel: int, optional - The channel in the `y_true` image that the mask exists in. Set to `-1` if there is no - mask for the given loss function. Default: `-1` - """ - logger.debug("Adding loss: (function: %s, weight: %s, mask_channel: %s)", - function, weight, mask_channel) - # Loss must be compiled inside LossContainer for keras to handle distibuted strategies - self._loss_functions.append(compile_utils.LossesContainer(function)) - self._loss_weights.append(weight) - self._mask_channels.append(mask_channel) - - def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: - """ Call the sub loss functions for the loss wrapper. - - Loss is returned as the weighted sum of the chosen losses. - - If masks are being applied to the loss function inputs, then they should be included as - additional channels at the end of :param:`y_true`, so that they can be split off and - applied to the actual inputs to the selected loss function(s). - - Parameters - ---------- - y_true: :class:`tensorflow.Tensor` - The ground truth batch of images, with any required masks stacked on the end - y_pred: :class:`tensorflow.Tensor` - The batch of model predictions - - Returns - ------- - :class:`tensorflow.Tensor` - The final weighted loss - """ - loss = 0.0 - for func, weight, mask_channel in zip(self._loss_functions, - self._loss_weights, - self._mask_channels): - logger.debug("Processing loss function: (func: %s, weight: %s, mask_channel: %s)", - func, weight, mask_channel) - n_true, n_pred = self._apply_mask(y_true, y_pred, mask_channel) - loss += (func(n_true, n_pred) * weight) - return loss - - @classmethod - def _apply_mask(cls, - y_true: tf.Tensor, - y_pred: tf.Tensor, - mask_channel: int, - mask_prop: float = 1.0) -> Tuple[tf.Tensor, tf.Tensor]: - """ Apply the mask to the input y_true and y_pred. If a mask is not required then - return the unmasked inputs. - - Parameters - ---------- - y_true: tensor or variable - The ground truth value - y_pred: tensor or variable - The predicted value - mask_channel: int - The channel within y_true that the required mask resides in - mask_prop: float, optional - The amount of mask propagation. Default: `1.0` - - Returns - ------- - tf.Tensor - The ground truth batch of images, with the required mask applied - tf.Tensor - The predicted batch of images with the required mask applied - """ - if mask_channel == -1: - logger.debug("No mask to apply") - return y_true[..., :3], y_pred[..., :3] - - logger.debug("Applying mask from channel %s", mask_channel) - mask = K.expand_dims(y_true[..., mask_channel], axis=-1) - mask_as_k_inv_prop = 1 - mask_prop - mask = (mask * mask_prop) + mask_as_k_inv_prop - - n_true = K.concatenate([y_true[..., i:i + 1] * mask for i in range(3)], axis=-1) - n_pred = K.concatenate([y_pred[..., i:i + 1] * mask for i in range(3)], axis=-1) - - return n_true, n_pred diff --git a/lib/model/networks/__init__.py b/lib/model/networks/__init__.py new file mode 100644 index 0000000000..e2be872d73 --- /dev/null +++ b/lib/model/networks/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python3 +""" Pre-defined networks for use in faceswap """ +from .simple_nets import AlexNet, SqueezeNet +from .clip import ViT, ViTConfig, TypeModels as TypeModelsViT diff --git a/lib/model/networks/clip.py b/lib/model/networks/clip.py new file mode 100644 index 0000000000..1b07e3fc76 --- /dev/null +++ b/lib/model/networks/clip.py @@ -0,0 +1,846 @@ +#!/usr/bin/env python3 +""" CLIP: https://github.com/openai/CLIP. This implementation only ports the visual transformer +part of the model. +""" +# TODO Fix Resnet. It is correct until final MHA +from __future__ import annotations +import inspect +import logging +import typing as T +import sys + +from dataclasses import dataclass + +import tensorflow as tf + +from lib.model.layers import QuickGELU +from lib.utils import GetModel + +keras = tf.keras +layers = tf.keras.layers +K = tf.keras.backend + +logger = logging.getLogger(__name__) + +TypeModels = T.Literal["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B-16", + "ViT-B-32", "ViT-L-14", "ViT-L-14-336px", "FaRL-B-16-16", "FaRL-B-16-64"] + + +@dataclass +class ViTConfig: + """ Configuration settings for ViT + + Parameters + ---------- + embed_dim: int + Dimensionality of the final shared embedding space + resolution: int + Spatial resolution of the input images + layer_conf: tuple[int, int, int, int] | int + Number of layers in the visual encoder, or a tuple of layer configurations for a custom + ResNet visual encoder + width: int + Width of the visual encoder layers + patch: int + Size of the patches to be extracted from the images. Only used for Visual encoder. + git_id: int, optional + The id of the model weights file stored in deepfakes_models repo if they exist. Default: 0 + """ + embed_dim: int + resolution: int + layer_conf: int | tuple[int, int, int, int] + width: int + patch: int + git_id: int = 0 + + def __post_init__(self): + """ Validate that patch_size is given correctly """ + assert (isinstance(self.layer_conf, (tuple, list)) and self.patch == 0) or ( + isinstance(self.layer_conf, int) and self.patch > 0) + + +ModelConfig: dict[TypeModels, ViTConfig] = { # Each model has a different set of parameters + "RN50": ViTConfig( + embed_dim=1024, resolution=224, layer_conf=(3, 4, 6, 3), width=64, patch=0, git_id=21), + "RN101": ViTConfig( + embed_dim=512, resolution=224, layer_conf=(3, 4, 23, 3), width=64, patch=0, git_id=22), + "RN50x4": ViTConfig( + embed_dim=640, resolution=288, layer_conf=(4, 6, 10, 6), width=80, patch=0, git_id=23), + "RN50x16": ViTConfig( + embed_dim=768, resolution=384, layer_conf=(6, 8, 18, 8), width=96, patch=0, git_id=24), + "RN50x64": ViTConfig( + embed_dim=1024, resolution=448, layer_conf=(3, 15, 36, 10), width=128, patch=0, git_id=25), + "ViT-B-16": ViTConfig( + embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=26), + "ViT-B-32": ViTConfig( + embed_dim=512, resolution=224, layer_conf=12, width=768, patch=32, git_id=27), + "ViT-L-14": ViTConfig( + embed_dim=768, resolution=224, layer_conf=24, width=1024, patch=14, git_id=28), + "ViT-L-14-336px": ViTConfig( + embed_dim=768, resolution=336, layer_conf=24, width=1024, patch=14, git_id=29), + "FaRL-B-16-16": ViTConfig( + embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=30), + "FaRL-B-16-64": ViTConfig( + embed_dim=512, resolution=224, layer_conf=12, width=768, patch=16, git_id=31)} + + +# ################## # +# VISUAL TRANSFORMER # +# ################## # + +class Transformer(): # pylint:disable=too-few-public-methods + """ A class representing a Transformer model with attention mechanism and residual connections. + + Parameters + ---------- + width: int + The dimension of the input and output vectors. + num_layers: int + The number of layers in the Transformer. + heads: int + The number of attention heads. + attn_mask: tf.Tensor, optional + The attention mask, by default None. + name: str, optional + The name of the Transformer model, by default "transformer". + + Methods + ------- + __call__() -> Model: + Calls the Transformer layers. + """ + _layer_names: dict[str, int] = {} + """ dict[str, int] for tracking unique layer names""" + + def __init__(self, + width: int, + num_layers: int, + heads: int, + attn_mask: tf.Tensor = None, + name: str = "transformer") -> None: + logger.debug("Initializing: %s (width: %s, num_layers: %s, heads: %s, attn_mask: %s, " + "name: %s)", + self.__class__.__name__, width, num_layers, heads, attn_mask, name) + self._width = width + self._num_layers = num_layers + self._heads = heads + self._attn_mask = attn_mask + self._name = name + logger.debug("Initialized: %s ", self.__class__.__name__) + + @classmethod + def _get_name(cls, name: str) -> str: + """ Return unique layer name for requested block. + + As blocks can be used multiple times, auto appends an integer to the end of the requested + name to keep all block names unique + + Parameters + ---------- + name: str + The requested name for the layer + + Returns + ------- + str + The unique name for this layer + """ + cls._layer_names[name] = cls._layer_names.setdefault(name, -1) + 1 + name = f"{name}.{cls._layer_names[name]}" + logger.debug("Generating block name: %s", name) + return name + + @classmethod + def _mlp(cls, inputs: tf.Tensor, key_dim: int, name: str) -> tf.Tensor: + """" Multilayer Perecptron for Block Ateention + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + The input to the MLP + key_dim: int + key dimension per head for MultiHeadAttention + name: str + The name to prefix on the layers + + Returns + ------- + :class:`tensorflow.Tensor` + The output from the MLP + """ + name = f"{name}.mlp" + var_x = layers.Dense(key_dim * 4, name=f"{name}.c_fc")(inputs) + var_x = QuickGELU(name=f"{name}.gelu")(var_x) + var_x = layers.Dense(key_dim, name=f"{name}.c_proj")(var_x) + return var_x + + def residual_attention_block(self, + inputs: tf.Tensor, + key_dim: int, + num_heads: int, + attn_mask: tf.Tensor, + name: str = "ResidualAttentionBlock") -> tf.Tensor: + """ Call the residual attention block + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input Tensor + key_dim: int + key dimension per head for MultiHeadAttention + num_heads: int + Number of heads for MultiHeadAttention + attn_mask: :class:`tensorflow.Tensor`, optional + Default: ``None`` + name: str, optional + The name for the layer. Default: "ResidualAttentionBlock" + + Returns + ------- + :class:`tf.Tensor` + The return Tensor + """ + name = self._get_name(name) + + var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{name}.ln_1")(inputs) + var_x = layers.MultiHeadAttention( + num_heads=num_heads, + key_dim=key_dim // num_heads, + name=f"{name}.attn")(var_x, var_x, var_x, attention_mask=attn_mask) + var_x = layers.Add()([inputs, var_x]) + var_y = var_x + var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{name}.ln_2")(var_x) + var_x = layers.Add()([var_y, self._mlp(var_x, key_dim, name)]) + return var_x + + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: + """ Call the Transformer layers + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input Tensor + + Returns + ------- + :class:`tf.Tensor` + The return Tensor + """ + logger.debug("Calling %s with input: %s", self.__class__.__name__, inputs.shape) + var_x = inputs + for _ in range(self._num_layers): + var_x = self.residual_attention_block(var_x, + self._width, + self._heads, + self._attn_mask, + name=f"{self._name}.resblocks") + return var_x + + +class EmbeddingLayer(tf.keras.layers.Layer): + """ Parent class for trainable embedding variables + + Parameters + ---------- + input_shape: tuple[int, ...] + The shape of the variable + scale: int + Amount to scale the random initialization by + name: str + The name of the layer + dtype: str, optional + The datatype for the layer. Mixed precision can mess up the embeddings. Default: "float32" + """ + def __init__(self, + input_shape: tuple[int, ...], + scale: int, + name: str, + *args, + dtype="float32", + **kwargs) -> None: + super().__init__(name=name, dtype=dtype, *args, **kwargs) + self._input_shape = input_shape + self._scale = scale + self._var: tf.Variable + + def build(self, input_shape: tuple[int, ...]) -> None: + """ Add the weights + + Parameters + ---------- + input_shape: tuple[int, ... + The input shape of the incoming tensor + """ + self._var = tf.Variable(self._scale * tf.random.normal(self._input_shape, + dtype=self.dtype), + trainable=True, + dtype=self.dtype) + super().build(input_shape) + + def get_config(self) -> dict[str, T.Any]: + """ Get the config dictionary for the layer + + Returns + ------- + dict[str, Any] + The config dictionary for the layer + """ + retval = super().get_config() + retval["input_shape"] = self._input_shape + retval["scale"] = self._scale + return retval + + +class ClassEmbedding(EmbeddingLayer): + """ Trainable Class Embedding layer """ + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """ Get the Class Embedding layer + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + Input tensor to the embedding layer + + Returns + ------- + :class:`tensorflow.Tensor` + The class embedding layer shaped for the input tensor + """ + return K.tile(self._var[None, None], [K.shape(inputs)[0], 1, 1]) + + +class PositionalEmbedding(EmbeddingLayer): + """ Trainable Positional Embedding layer """ + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """ Get the Positional Embedding layer + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + Input tensor to the embedding layer + + Returns + ------- + :class:`tensorflow.Tensor` + The positional embedding layer shaped for the input tensor + """ + return K.tile(self._var[None], [K.shape(inputs)[0], 1, 1]) + + +class Projection(EmbeddingLayer): + """ Trainable Projection Embedding Layer """ + def call(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor: + """ Get the Projection layer + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + Input tensor to the embedding layer + + Returns + ------- + :class:`tensorflow.Tensor` + The Projection layer expanded to the batch dimension and transposed for matmul + """ + return K.tile(K.transpose(self._var)[None], [K.shape(inputs)[0], 1, 1]) + + +class VisualTransformer(): # pylint:disable=too-few-public-methods + """ A class representing a Visual Transformer model for image classification tasks. + + Parameters + ---------- + input_resolution: int + The input resolution of the images. + patch_size: int + The size of the patches to be extracted from the images. + width: int + The dimension of the input and output vectors. + num_layers: int + The number of layers in the Transformer. + heads: int + The number of attention heads. + output_dim: int + The dimension of the output vector. + name: str, optional + The name of the Visual Transformer model, Default: "VisualTransformer". + + Methods + ------- + __call__() -> Model: + Builds and returns the Visual Transformer model. + """ + def __init__(self, + input_resolution: int, + patch_size: int, + width: int, + num_layers: int, + heads: int, + output_dim: int, + name: str = "VisualTransformer") -> None: + logger.debug("Initializing: %s (input_resolution: %s, patch_size: %s, width: %s, " + "layers: %s, heads: %s, output_dim: %s, name: %s)", + self.__class__.__name__, input_resolution, patch_size, width, num_layers, + heads, output_dim, name) + self._input_resolution = input_resolution + self._patch_size = patch_size + self._width = width + self._num_layers = num_layers + self._heads = heads + self._output_dim = output_dim + self._name = name + logger.debug("Initialized: %s", self.__class__.__name__) + + def __call__(self) -> tf.keras.models.Model: + """ Builds and returns the Visual Transformer model. + + Returns + ------- + Model + The Visual Transformer model. + """ + inputs = layers.Input([self._input_resolution, self._input_resolution, 3]) + var_x: tf.Tensor = layers.Conv2D(self._width, # shape = [*, grid, grid, width] + self._patch_size, + strides=self._patch_size, + use_bias=False, + name=f"{self._name}.conv1")(inputs) + + var_x = layers.Reshape((-1, self._width))(var_x) # shape = [*, grid ** 2, width] + + class_embed = ClassEmbedding((self._width, ), + self._width ** -0.5, + name=f"{self._name}.class_embedding")(var_x) + var_x = layers.Concatenate(axis=1)([class_embed, var_x]) + + pos_embed = PositionalEmbedding(((self._input_resolution // self._patch_size) ** 2 + 1, + self._width), + self._width ** -0.5, + name=f"{self._name}.positional_embedding")(var_x) + var_x = layers.Add()([var_x, pos_embed]) + var_x = layers.LayerNormalization(epsilon=1e-05, name=f"{self._name}.ln_pre")(var_x) + var_x = Transformer(self._width, + self._num_layers, + self._heads, + name=f"{self._name}.transformer")(var_x) + var_x = layers.LayerNormalization(epsilon=1e-05, + name=f"{self._name}.ln_post")(var_x[:, 0, :]) + proj = Projection((self._width, self._output_dim), + self._width ** -0.5, + name=f"{self._name}.proj")(var_x) + var_x = layers.Dot(axes=-1)([var_x, proj]) + return keras.models.Model(inputs=inputs, outputs=[var_x], name=self._name) + + +# ################ # +# MODIEFIED RESNET # +# ################ # +class Bottleneck(): # pylint:disable=too-few-public-methods + """ A ResNet bottleneck block that performs a sequence of convolutions, batch normalization, + and ReLU activation operations on an input tensor. + + Parameters + ---------- + inplanes: int + The number of input channels. + planes: int + The number of output channels. + stride: int, optional + The stride of the bottleneck block. Default: 1 + name: str, optional + The name of the bottleneck block. Default: "bottleneck" + """ + expansion = 4 + """ int: The factor by which the number of input channels is expanded to get the number of + output channels.""" + + def __init__(self, + inplanes: int, + planes: int, + stride: int = 1, + name: str = "bottleneck") -> None: + logger.debug("Initializing: %s (inplanes: %s, planes: %s, stride: %s, name: %s)", + self.__class__.__name__, inplanes, planes, stride, name) + self._inplanes = inplanes + self._planes = planes + self._stride = stride + self._name = name + logger.debug("Initialized: %s", self.__class__.__name__) + + def _downsample(self, inputs: tf.Tensor) -> tf.Tensor: + """ Perform downsample if required + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + The input the downsample + + Returns + ------- + :class:`tensorflow.Tensor` + The original tensor, if downsizing not required, otherwise the downsized tensor + """ + if self._stride <= 1 and self._inplanes == self._planes * self.expansion: + return inputs + + name = f"{self._name}.downsample" + out = layers.AveragePooling2D(self._stride, name=f"{name}.avgpool")(inputs) + out = layers.Conv2D(self._planes * self.expansion, + 1, + strides=1, + use_bias=False, + name=f"{name}.0")(out) + out = layers.BatchNormalization(name=f"{name}.1", epsilon=1e-5)(out) + return out + + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: + """ Performs the forward pass for a Bottleneck block. + + All conv layers have stride 1. an avgpool is performed after the second convolution when + stride > 1 + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + The input tensor to the Bottleneck block. + + Returns + ------- + :class:`tensorflow.Tensor` + The result of the forward pass through the Bottleneck block. + """ + out = layers.Conv2D(self._planes, 1, use_bias=False, name=f"{self._name}.conv1")(inputs) + out = layers.BatchNormalization(name=f"{self._name}.bn1", epsilon=1e-5)(out) + out = layers.ReLU()(out) + + out = layers.ZeroPadding2D(padding=((1, 1), (1, 1)))(out) + out = layers.Conv2D(self._planes, 3, use_bias=False, name=f"{self._name}.conv2")(out) + out = layers.BatchNormalization(name=f"{self._name}.bn2", epsilon=1e-5)(out) + out = layers.ReLU()(out) + + if self._stride > 1: + out = layers.AveragePooling2D(self._stride)(out) + + out = layers.Conv2D(self._planes * self.expansion, + 1, + use_bias=False, + name=f"{self._name}.conv3")(out) + out = layers.BatchNormalization(name=f"{self._name}.bn3", epsilon=1e-5)(out) + + identity = self._downsample(inputs) + + out += identity + out = layers.ReLU()(out) + return out + + +class AttentionPool2d(): # pylint:disable=too-few-public-methods + """ An Attention Pooling layer that applies a multi-head self-attention mechanism over a + spatial grid of features. + + Parameters + ---------- + spatial_dim: int + The dimensionality of the spatial grid of features. + embed_dim: int + The dimensionality of the feature embeddings. + num_heads: int + The number of attention heads. + output_dim: int + The output dimensionality of the attention layer. If None, it defaults to embed_dim. + name: str + The name of the layer. + """ + def __init__(self, + spatial_dim: int, + embed_dim: int, + num_heads: int, + output_dim: int | None = None, + name="AttentionPool2d"): + logger.debug("Initializing: %s (spatial_dim: %s, embed_dim: %s, num_heads: %s, " + "output_dim: %s, name: %s)", + self.__class__.__name__, spatial_dim, embed_dim, num_heads, output_dim, name) + + self._spatial_dim = spatial_dim + self._embed_dim = embed_dim + self._num_heads = num_heads + self._output_dim = output_dim + self._name = name + logger.debug("Initialized: %s", self.__class__.__name__) + + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: + """Performs the attention pooling operation on the input tensor. + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor`: + The input tensor of shape [batch_size, height, width, embed_dim]. + + Returns + ------- + :class:`tensorflow.Tensor`:: The result of the attention pooling operation + """ + var_x: tf.Tensor + var_x = layers.Reshape((-1, inputs.shape[-1]))(inputs) # NHWC -> N(HW)C + var_x = layers.Concatenate(axis=1)([K.mean(var_x, axis=1, # N(HW)C -> N(HW+1)C + keepdims=True), var_x]) + pos_embed = PositionalEmbedding((self._spatial_dim ** 2 + 1, self._embed_dim), # N(HW+1)C + self._embed_dim ** 0.5, + name=f"{self._name}.positional_embedding")(var_x) + var_x = layers.Add()([var_x, pos_embed]) + # TODO At this point torch + keras match. They mismatch after MHA + var_x = layers.MultiHeadAttention(num_heads=self._num_heads, + key_dim=self._embed_dim // self._num_heads, + output_shape=self._output_dim or self._embed_dim, + use_bias=True, + name=f"{self._name}.mha")(var_x[:, :1, ...], + var_x, + var_x) + # only return the first element in the sequence + return var_x[:, 0, ...] + + +class ModifiedResNet(): # pylint:disable=too-few-public-methods + """ A ResNet class that is similar to torchvision's but contains the following changes: + + - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max + pool. + - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions + with stride > 1 + - The final pooling layer is a QKV attention instead of an average pool + + Parameters + ---------- + input_resolution: int + The input resolution of the model. Default is 224. + width: int + The width of the model. Default is 64. + layer_config: list + A list containing the number of Bottleneck blocks for each layer. + output_dim: int + The output dimension of the model. + heads: int + The number of heads for the QKV attention. + name: str + The name of the model. Default is "ModifiedResNet". + """ + def __init__(self, + input_resolution: int, + width: int, + layer_config: tuple[int, int, int, int], + output_dim: int, + heads: int, + name="ModifiedResNet"): + self._input_resolution = input_resolution + self._width = width + self._layer_config = layer_config + self._heads = heads + self._output_dim = output_dim + self._name = name + + def _stem(self, inputs: tf.Tensor) -> tf.Tensor: + """ Applies the stem operation to the input tensor, which consists of 3 convolutional + layers with BatchNormalization and ReLU activation, followed by an average pooling + layer. + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + The input tensor + + Returns + ------- + :class:`tensorflow.Tensor` + The output tensor after applying the stem operation. + """ + var_x = inputs + for i in range(1, 4): + width = self._width if i == 3 else self._width // 2 + strides = 2 if i == 1 else 1 + var_x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)), name=f"conv{i}_padding")(var_x) + var_x = layers.Conv2D(width, + 3, + strides=strides, + use_bias=False, + name=f"conv{i}")(var_x) + var_x = layers.BatchNormalization(name=f"bn{i}", epsilon=1e-5)(var_x) + var_x = layers.ReLU()(var_x) + var_x = layers.AveragePooling2D(2, name="avgpool")(var_x) + return var_x + + def _bottleneck(self, + inputs: tf.Tensor, + planes: int, + blocks: int, + stride: int = 1, + name: str = "layer") -> tf.Tensor: + """ A private method that creates a sequential layer of Bottleneck blocks for the + ModifiedResNet model. + + Parameters + ---------- + inputs: :class:`tensorflow.Tensor` + The input tensor + planes: int + The number of output channels for the layer. + blocks: int + The number of Bottleneck blocks in the layer. + stride: int + The stride for the first Bottleneck block in the layer. Default is 1. + name: str + The name of the layer. Default is "layer". + + Returns + ------- + :class:`tensorflow.Tensor` + Sequential block of bottlenecks + """ + retval: tf.Tensor + retval = Bottleneck(planes, planes, stride, name=f"{name}.0")(inputs) + for i in range(1, blocks): + retval = Bottleneck(planes * Bottleneck.expansion, + planes, + name=f"{name}.{i}")(retval) + return retval + + def __call__(self) -> tf.keras.models.Model: + """ Implements the forward pass of the ModifiedResNet model. + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The modified resnet model. + """ + inputs = layers.Input((self._input_resolution, self._input_resolution, 3)) + var_x = self._stem(inputs) + + for i in range(4): + stride = 1 if i == 0 else 2 + var_x = self._bottleneck(var_x, + self._width * (2 ** i), + self._layer_config[i], + stride=stride, + name=f"{self._name}.layer{i + 1}") + + var_x = AttentionPool2d(self._input_resolution // 32, + self._width * 32, # the ResNet feature dimension + self._heads, + self._output_dim, + name=f"{self._name}.attnpool")(var_x) + return keras.models.Model(inputs, outputs=[var_x], name=self._name) + + +# ### # +# VIT # +# ### # +class ViT(): # pylint:disable=too-few-public-methods + """ Visiual Transform from CLIP + + A Convolutional Language-Image Pre-Training (CLIP) model that encodes images and text into a + shared latent space. + + Reference + --------- + https://arxiv.org/abs/2103.00020 + + Parameters + ---------- + name: ["RN50", "RN101", "RN50x4", "RN50x16", "RN50x64", "ViT-B-32", + "ViT-B-16", "ViT-L-14", "ViT-L-14-336px", "FaRL-B_16-64"] + The model configuration to use + input_size: int, optional + The required resolution size for the model. ``None`` for default preset size + load_weights: bool, optional + ``True`` to load pretrained weights. Default: ``False`` + """ + def __init__(self, + name: TypeModels, + input_size: int | None = None, + load_weights: bool = False) -> None: + logger.debug("Initializing: %s (name: %s, input_size: %s, load_weights: %s)", + self.__class__.__name__, name, input_size, load_weights) + assert name in ModelConfig, ("Name must be one of %s", list(ModelConfig)) + + self._name = name + self._load_weights = load_weights + + config = ModelConfig[name] + self._git_id = config.git_id + + res = input_size if input_size is not None else config.resolution + self._net = self._get_vision_net(config.layer_conf, + config.width, + config.embed_dim, + res, + config.patch) + logger.debug("Initialized: %s", self.__class__.__name__) + + def _get_vision_net(self, + layer_config: int | tuple[int, int, int, int], + width: int, + embed_dim: int, + resolution: int, + patch_size: int) -> tf.keras.models.Model: + """ Obtain the network for the vision layets + + Parameters + ---------- + layer_config: tuple[int, int, int, int] | int + Number of layers in the visual encoder, or a tuple of layer configurations for a custom + ResNet visual encoder. + width: int + Width of the visual encoder layers. + embed_dim: int + Dimensionality of the final shared embedding space. + resolution: int + Spatial resolution of the input images. + patch_size: int + Size of the patches to be extracted from the images. + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The :class:`ModifiedResNet` or :class:`VisualTransformer` vision model to use + """ + if isinstance(layer_config, (tuple, list)): + vision_heads = width * 32 // 64 + return ModifiedResNet(input_resolution=resolution, + width=width, + layer_config=layer_config, + output_dim=embed_dim, + heads=vision_heads, + name="visual") + vision_heads = width // 64 + return VisualTransformer(input_resolution=resolution, + width=width, + num_layers=layer_config, + output_dim=embed_dim, + heads=vision_heads, + patch_size=patch_size, + name="visual") + + def __call__(self) -> tf.keras.Model: + """ Get the configured ViT model + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The requested Visual Transformer model + """ + net: tf.keras.models.Model = self._net() + if self._load_weights and not self._git_id: + logger.warning("Trained weights are not available for '%s'", self._name) + return net + if self._load_weights: + model_path = GetModel(f"CLIPv_{self._name}_v1.h5", self._git_id).model_path + logger.info("Loading CLIPv trained weights for '%s'", self._name) + net.load_weights(model_path, by_name=True, skip_mismatch=True) + + return net + + +# Update layers into Keras custom objects +for name_, obj in inspect.getmembers(sys.modules[__name__]): + if (inspect.isclass(obj) and issubclass(obj, tf.keras.layers.Layer) + and obj.__module__ == __name__): + keras.utils.get_custom_objects().update({name_: obj}) diff --git a/lib/model/networks/simple_nets.py b/lib/model/networks/simple_nets.py new file mode 100644 index 0000000000..727161bcd8 --- /dev/null +++ b/lib/model/networks/simple_nets.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +""" Ports of existing NN Architecture for use in faceswap.py """ +from __future__ import annotations +import logging +import typing as T + +import tensorflow as tf + +# Fix intellisense/linting for tf.keras' thoroughly broken import system +keras = tf.keras +layers = keras.layers +Model = keras.models.Model + +if T.TYPE_CHECKING: + from tensorflow import Tensor + + +logger = logging.getLogger(__name__) + + +class _net(): # pylint:disable=too-few-public-methods + """ Base class for existing NeuralNet architecture + + Notes + ----- + All architectures assume channels_last format + + Parameters + ---------- + input_shape, Tuple, optional + The input shape for the model. Default: ``None`` + """ + def __init__(self, + input_shape: tuple[int, int, int] | None = None) -> None: + logger.debug("Initializing: %s (input_shape: %s)", self.__class__.__name__, input_shape) + self._input_shape = (None, None, 3) if input_shape is None else input_shape + assert len(self._input_shape) == 3 and self._input_shape[-1] == 3, ( + "Input shape must be in the format (height, width, channels) and the number of " + f"channels must equal 3. Received: {self._input_shape}") + logger.debug("Initialized: %s", self.__class__.__name__) + + +class AlexNet(_net): + """ AlexNet ported from torchvision version. + + Notes + ----- + This port only contains the features portion of the model. + + References + ---------- + https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf + + Parameters + ---------- + input_shape, Tuple, optional + The input shape for the model. Default: ``None`` + """ + def __init__(self, input_shape: tuple[int, int, int] | None = None) -> None: + super().__init__(input_shape) + self._feature_indices = [0, 3, 6, 8, 10] # For naming equivalent to PyTorch + self._filters = [64, 192, 384, 256, 256] # Filters at each block + + @classmethod + def _conv_block(cls, + inputs: Tensor, + padding: int, + filters: int, + kernel_size: int, + strides: int, + block_idx: int, + max_pool: bool) -> Tensor: + """ + The Convolutional block for AlexNet + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input tensor to the block + padding: int + The amount of zero paddin to apply prior to convolution + filters: int + The number of filters to apply during convolution + kernel_size: int + The kernel size of the convolution + strides: int + The number of strides for the convolution + block_idx: int + The index of the current block (for standardized naming convention) + max_pool: bool + ``True`` to apply a max pooling layer at the beginning of the block otherwise ``False`` + + Returns + ------- + :class:`tf.Tensor` + The output of the Convolutional block + """ + name = f"features.{block_idx}" + var_x = inputs + if max_pool: + var_x = layers.MaxPool2D(pool_size=3, strides=2, name=f"{name}.pool")(var_x) + var_x = layers.ZeroPadding2D(padding=padding, name=f"{name}.pad")(var_x) + var_x = layers.Conv2D(filters, + kernel_size=kernel_size, + strides=strides, + padding="valid", + activation="relu", + name=name)(var_x) + return var_x + + def __call__(self) -> tf.keras.models.Model: + """ Create the AlexNet Model + + Returns + ------- + :class:`keras.models.Model` + The compiled AlexNet model + """ + inputs = layers.Input(self._input_shape) + var_x = inputs + kernel_size = 11 + strides = 4 + + for idx, (filters, block_idx) in enumerate(zip(self._filters, self._feature_indices)): + padding = 2 if idx < 2 else 1 + do_max_pool = 0 < idx < 3 + var_x = self._conv_block(var_x, + padding, + filters, + kernel_size, + strides, + block_idx, + do_max_pool) + kernel_size = max(3, kernel_size // 2) + strides = 1 + return Model(inputs=inputs, outputs=[var_x]) + + +class SqueezeNet(_net): + """ SqueezeNet ported from torchvision version. + + Notes + ----- + This port only contains the features portion of the model. + + References + ---------- + https://arxiv.org/abs/1602.07360 + + Parameters + ---------- + input_shape, Tuple, optional + The input shape for the model. Default: ``None`` + """ + + @classmethod + def _fire(cls, + inputs: Tensor, + squeeze_planes: int, + expand_planes: int, + block_idx: int) -> Tensor: + """ The fire block for SqueezeNet. + + Parameters + ---------- + inputs: :class:`tf.Tensor` + The input to the fire block + squeeze_planes: int + The number of filters for the squeeze convolution + expand_planes: int + The number of filters for the expand convolutions + block_idx: int + The index of the current block (for standardized naming convention) + + Returns + ------- + :class:`tf.Tensor` + The output of the SqueezeNet fire block + """ + name = f"features.{block_idx}" + squeezed = layers.Conv2D(squeeze_planes, 1, + activation="relu", name=f"{name}.squeeze")(inputs) + expand1 = layers.Conv2D(expand_planes, 1, + activation="relu", name=f"{name}.expand1x1")(squeezed) + expand3 = layers.Conv2D(expand_planes, + 3, + activation="relu", + padding="same", + name=f"{name}.expand3x3")(squeezed) + return layers.Concatenate(axis=-1, name=name)([expand1, expand3]) + + def __call__(self) -> tf.keras.models.Model: + """ Create the SqueezeNet Model + + Returns + ------- + :class:`keras.models.Model` + The compiled SqueezeNet model + """ + inputs = layers.Input(self._input_shape) + var_x = layers.Conv2D(64, 3, strides=2, activation="relu", name="features.0")(inputs) + + block_idx = 2 + squeeze = 16 + expand = 64 + for idx in range(4): + if idx < 3: + var_x = layers.MaxPool2D(pool_size=3, strides=2)(var_x) + block_idx += 1 + var_x = self._fire(var_x, squeeze, expand, block_idx) + block_idx += 1 + var_x = self._fire(var_x, squeeze, expand, block_idx) + block_idx += 1 + squeeze += 16 + expand += 64 + return Model(inputs=inputs, outputs=[var_x]) diff --git a/lib/model/nn_blocks.py b/lib/model/nn_blocks.py index 341ae17ed5..63e431d33e 100644 --- a/lib/model/nn_blocks.py +++ b/lib/model/nn_blocks.py @@ -1,39 +1,29 @@ #!/usr/bin/env python3 """ Neural Network Blocks for faceswap.py. """ - +from __future__ import annotations import logging -from typing import Dict, Optional, Tuple, Union +import typing as T -from lib.utils import get_backend +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import ( # pylint:disable=import-error + Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D, Conv2DTranspose, + DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU, SeparableConv2D, UpSampling2D) +from tensorflow.keras.initializers import he_uniform, VarianceScaling # noqa:E501 # pylint:disable=import-error from .initializers import ICNR, ConvolutionAware from .layers import PixelShuffler, ReflectionPadding2D, Swish, KResizeImages from .normalization import InstanceNormalization -if get_backend() == "amd": - from keras.layers import ( - Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D, Conv2DTranspose, - DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU, SeparableConv2D, UpSampling2D) - from keras.initializers import he_uniform, VarianceScaling # pylint:disable=no-name-in-module - # type checking: - import keras - from plaidml.tile import Value as Tensor # pylint:disable=import-error -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import ( # noqa pylint:disable=no-name-in-module,import-error - Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D, Conv2DTranspose, - DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU, SeparableConv2D, UpSampling2D) - from tensorflow.keras.initializers import he_uniform, VarianceScaling # noqa pylint:disable=no-name-in-module,import-error - # type checking: +if T.TYPE_CHECKING: from tensorflow import keras from tensorflow import Tensor -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) _CONFIG: dict = {} -_NAMES: Dict[str, int] = {} +_NAMES: dict[str, int] = {} def set_config(configuration: dict) -> None: @@ -199,10 +189,10 @@ class Conv2DOutput(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int]], + kernel_size: int | tuple[int], activation: str = "sigmoid", padding: str = "same", **kwargs) -> None: - self._name = kwargs.pop("name") if "name" in kwargs else _get_name( + self._name = _get_name(kwargs.pop("name")) if "name" in kwargs else _get_name( f"conv_output_{filters}") self._filters = filters self._kernel_size = kernel_size @@ -275,11 +265,11 @@ class Conv2DBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 5, - strides: Union[int, Tuple[int, int]] = 2, + kernel_size: int | tuple[int, int] = 5, + strides: int | tuple[int, int] = 2, padding: str = "same", - normalization: Optional[str] = None, - activation: Optional[str] = "leakyrelu", + normalization: str | None = None, + activation: str | None = "leakyrelu", use_depthwise: bool = False, relu_alpha: float = 0.1, **kwargs) -> None: @@ -292,8 +282,9 @@ def __init__(self, self._use_reflect_padding = _CONFIG["reflect_padding"] + kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size self._args = (kernel_size, ) if use_depthwise else (filters, kernel_size) - self._strides = strides + self._strides = (strides, strides) if isinstance(strides, int) else strides self._padding = "valid" if self._use_reflect_padding else padding self._kwargs = kwargs self._normalization = None if not normalization else normalization.lower() @@ -324,10 +315,10 @@ def __call__(self, inputs: Tensor) -> Tensor: The output tensor from the Convolution 2D Layer """ if self._use_reflect_padding: - inputs = ReflectionPadding2D(stride=self._strides, - kernel_size=self._args[-1], + inputs = ReflectionPadding2D(stride=self._strides[0], + kernel_size=self._args[-1][0], # type:ignore[index] name=f"{self._name}_reflectionpadding2d")(inputs) - conv = DepthwiseConv2D if self._use_depthwise else Conv2D + conv: keras.layers.Layer = DepthwiseConv2D if self._use_depthwise else Conv2D var_x = conv(*self._args, strides=self._strides, padding=self._padding, @@ -372,8 +363,8 @@ class SeparableConv2DBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 5, - strides: Union[int, Tuple[int, int]] = 2, **kwargs) -> None: + kernel_size: int | tuple[int, int] = 5, + strides: int | tuple[int, int] = 2, **kwargs) -> None: self._name = _get_name(f"separableconv2d_{filters}") logger.debug("name: %s, filters: %s, kernel_size: %s, strides: %s, kwargs: %s)", self._name, filters, kernel_size, strides, kwargs) @@ -444,11 +435,11 @@ class UpscaleBlock(): # pylint:disable=too-few-public-methods def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: int | tuple[int, int] = 3, padding: str = "same", scale_factor: int = 2, - normalization: Optional[str] = None, - activation: Optional[str] = "leakyrelu", + normalization: str | None = None, + activation: str | None = "leakyrelu", **kwargs) -> None: self._name = _get_name(f"upscale_{filters}") logger.debug("name: %s. filters: %s, kernel_size: %s, padding: %s, scale_factor: %s, " @@ -531,9 +522,9 @@ class Upscale2xBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: int | tuple[int, int] = 3, padding: str = "same", - activation: Optional[str] = "leakyrelu", + activation: str | None = "leakyrelu", interpolation: str = "bilinear", sr_ratio: float = 0.5, scale_factor: int = 2, @@ -625,11 +616,11 @@ class UpscaleResizeImagesBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: int | tuple[int, int] = 3, padding: str = "same", - activation: Optional[str] = "leakyrelu", + activation: str | None = "leakyrelu", scale_factor: int = 2, - interpolation: str = "bilinear") -> None: + interpolation: T.Literal["nearest", "bilinear"] = "bilinear") -> None: self._name = _get_name(f"upscale_ri_{filters}") self._interpolation = interpolation self._size = scale_factor @@ -710,9 +701,9 @@ class UpscaleDNYBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: int | tuple[int, int] = 3, padding: str = "same", - activation: Optional[str] = "leakyrelu", + activation: str | None = "leakyrelu", size: int = 2, interpolation: str = "bilinear", **kwargs) -> None: @@ -767,7 +758,7 @@ class ResidualBlock(): # pylint:disable=too-few-public-methods """ def __init__(self, filters: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: int | tuple[int, int] = 3, padding: str = "same", **kwargs) -> None: self._name = _get_name(f"residual_{filters}") @@ -776,7 +767,8 @@ def __init__(self, self._use_reflect_padding = _CONFIG["reflect_padding"] self._filters = filters - self._kernel_size = kernel_size + self._kernel_size = (kernel_size, + kernel_size) if isinstance(kernel_size, int) else kernel_size self._padding = "valid" if self._use_reflect_padding else padding self._kwargs = kwargs @@ -796,7 +788,7 @@ def __call__(self, inputs: Tensor) -> Tensor: var_x = inputs if self._use_reflect_padding: var_x = ReflectionPadding2D(stride=1, - kernel_size=self._kernel_size, + kernel_size=self._kernel_size[0], name=f"{self._name}_reflectionpadding2d_0")(var_x) var_x = Conv2D(self._filters, kernel_size=self._kernel_size, @@ -806,7 +798,7 @@ def __call__(self, inputs: Tensor) -> Tensor: var_x = LeakyReLU(alpha=0.2, name=f"{self._name}_leakyrelu_1")(var_x) if self._use_reflect_padding: var_x = ReflectionPadding2D(stride=1, - kernel_size=self._kernel_size, + kernel_size=self._kernel_size[0], name=f"{self._name}_reflectionpadding2d_1")(var_x) kwargs = {key: val for key, val in self._kwargs.items() if key != "kernel_initializer"} diff --git a/lib/model/normalization/normalization_common.py b/lib/model/normalization.py similarity index 68% rename from lib/model/normalization/normalization_common.py rename to lib/model/normalization.py index 22e8419d5b..fcef640fe9 100644 --- a/lib/model/normalization/normalization_common.py +++ b/lib/model/normalization.py @@ -1,205 +1,18 @@ #!/usr/bin/env python3 -""" Normalization methods for faceswap.py common to both Plaid and Tensorflow Backends """ - -import sys +""" Normalization methods for faceswap.py specific to Tensorflow backend """ import inspect +import sys -from lib.utils import get_backend - -if get_backend() == "amd": - from keras.utils import get_custom_objects # pylint:disable=no-name-in-module - from keras.layers import Layer, InputSpec - from keras import initializers, regularizers, constraints, backend as K - from keras.backend import normalize_data_format # pylint:disable=no-name-in-module -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.utils import get_custom_objects # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.keras.layers import Layer, InputSpec # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.keras import initializers, regularizers, constraints, backend as K # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.python.keras.utils.conv_utils import normalize_data_format # noqa pylint:disable=no-name-in-module - - -class InstanceNormalization(Layer): - """Instance normalization layer (Lei Ba et al, 2016, Ulyanov et al., 2016). - - Normalize the activations of the previous layer at each step, i.e. applies a transformation - that maintains the mean activation close to 0 and the activation standard deviation close to 1. - - Parameters - ---------- - axis: int, optional - The axis that should be normalized (typically the features axis). For instance, after a - `Conv2D` layer with `data_format="channels_first"`, set `axis=1` in - :class:`InstanceNormalization`. Setting `axis=None` will normalize all values in each - instance of the batch. Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid - errors. Default: ``None`` - epsilon: float, optional - Small float added to variance to avoid dividing by zero. Default: `1e-3` - center: bool, optional - If ``True``, add offset of `beta` to normalized tensor. If ``False``, `beta` is ignored. - Default: ``True`` - scale: bool, optional - If ``True``, multiply by `gamma`. If ``False``, `gamma` is not used. When the next layer - is linear (also e.g. `relu`), this can be disabled since the scaling will be done by - the next layer. Default: ``True`` - beta_initializer: str, optional - Initializer for the beta weight. Default: `"zeros"` - gamma_initializer: str, optional - Initializer for the gamma weight. Default: `"ones"` - beta_regularizer: str, optional - Optional regularizer for the beta weight. Default: ``None`` - gamma_regularizer: str, optional - Optional regularizer for the gamma weight. Default: ``None`` - beta_constraint: float, optional - Optional constraint for the beta weight. Default: ``None`` - gamma_constraint: float, optional - Optional constraint for the gamma weight. Default: ``None`` - - References - ---------- - - Layer Normalization - https://arxiv.org/abs/1607.06450 - - - Instance Normalization: The Missing Ingredient for Fast Stylization - \ - https://arxiv.org/abs/1607.08022 - """ - # pylint:disable=too-many-instance-attributes,too-many-arguments - def __init__(self, - axis=None, - epsilon=1e-3, - center=True, - scale=True, - beta_initializer="zeros", - gamma_initializer="ones", - beta_regularizer=None, - gamma_regularizer=None, - beta_constraint=None, - gamma_constraint=None, - **kwargs): - self.beta = None - self.gamma = None - super().__init__(**kwargs) - self.supports_masking = True - self.axis = axis - self.epsilon = epsilon - self.center = center - self.scale = scale - self.beta_initializer = initializers.get(beta_initializer) - self.gamma_initializer = initializers.get(gamma_initializer) - self.beta_regularizer = regularizers.get(beta_regularizer) - self.gamma_regularizer = regularizers.get(gamma_regularizer) - self.beta_constraint = constraints.get(beta_constraint) - self.gamma_constraint = constraints.get(gamma_constraint) - - def build(self, input_shape): - """Creates the layer weights. - - Parameters - ---------- - input_shape: tensor - Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to - reference for weight shape computations. - """ - ndim = len(input_shape) - if self.axis == 0: - raise ValueError("Axis cannot be zero") - - if (self.axis is not None) and (ndim == 2): - raise ValueError("Cannot specify axis for rank 1 tensor") - - self.input_spec = InputSpec(ndim=ndim) # pylint:disable=attribute-defined-outside-init - - if self.axis is None: - shape = (1,) - else: - shape = (input_shape[self.axis],) - - if self.scale: - self.gamma = self.add_weight(shape=shape, - name="gamma", - initializer=self.gamma_initializer, - regularizer=self.gamma_regularizer, - constraint=self.gamma_constraint) - else: - self.gamma = None - if self.center: - self.beta = self.add_weight(shape=shape, - name="beta", - initializer=self.beta_initializer, - regularizer=self.beta_regularizer, - constraint=self.beta_constraint) - else: - self.beta = None - self.built = True # pylint:disable=attribute-defined-outside-init - - def call(self, inputs, training=None): # pylint:disable=arguments-differ,unused-argument - """This is where the layer's logic lives. - - Parameters - ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - - Returns - ------- - tensor - A tensor or list/tuple of tensors - """ - input_shape = K.int_shape(inputs) - reduction_axes = list(range(0, len(input_shape))) - - if self.axis is not None: - del reduction_axes[self.axis] - - del reduction_axes[0] - - mean = K.mean(inputs, reduction_axes, keepdims=True) - stddev = K.std(inputs, reduction_axes, keepdims=True) + self.epsilon - normed = (inputs - mean) / stddev - - broadcast_shape = [1] * len(input_shape) - if self.axis is not None: - broadcast_shape[self.axis] = input_shape[self.axis] - - if self.scale: - broadcast_gamma = K.reshape(self.gamma, broadcast_shape) - normed = normed * broadcast_gamma - if self.center: - broadcast_beta = K.reshape(self.beta, broadcast_shape) - normed = normed + broadcast_beta - return normed - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) containing the configuration of a - layer. The same layer can be reinstated later (without its trained weights) from this - configuration. - - The configuration of a layer does not include connectivity information, nor the layer - class name. These are handled by `Network` (one layer of abstraction above). +import tensorflow as tf - Returns - -------- - dict - A python dictionary containing the layer configuration - """ - config = { - "axis": self.axis, - "epsilon": self.epsilon, - "center": self.center, - "scale": self.scale, - "beta_initializer": initializers.serialize(self.beta_initializer), - "gamma_initializer": initializers.serialize(self.gamma_initializer), - "beta_regularizer": regularizers.serialize(self.beta_regularizer), - "gamma_regularizer": regularizers.serialize(self.gamma_regularizer), - "beta_constraint": constraints.serialize(self.beta_constraint), - "gamma_constraint": constraints.serialize(self.gamma_constraint) - } - base_config = super().get_config() - return dict(list(base_config.items()) + list(config.items())) +# Fix intellisense/linting for tf.keras' thoroughly broken import system +from tensorflow.python.keras.utils.conv_utils import normalize_data_format # noqa:E501 # pylint:disable=no-name-in-module +keras = tf.keras +layers = keras.layers +K = keras.backend -class AdaInstanceNormalization(Layer): +class AdaInstanceNormalization(layers.Layer): # type:ignore[name-defined] """ Adaptive Instance Normalization Layer for Keras. Parameters @@ -302,7 +115,7 @@ def get_config(self): base_config = super().get_config() return dict(list(base_config.items()) + list(config.items())) - def compute_output_shape(self, input_shape): # pylint:disable=no-self-use + def compute_output_shape(self, input_shape): """ Calculate the output shape from this layer. Parameters @@ -318,7 +131,7 @@ def compute_output_shape(self, input_shape): # pylint:disable=no-self-use return input_shape[0] -class GroupNormalization(Layer): +class GroupNormalization(layers.Layer): # type:ignore[name-defined] """ Group Normalization Parameters @@ -357,10 +170,10 @@ def __init__(self, axis=-1, gamma_init='one', beta_init='zero', gamma_regularize self.gamma = None super().__init__(**kwargs) self.axis = axis if isinstance(axis, (list, tuple)) else [axis] - self.gamma_init = initializers.get(gamma_init) - self.beta_init = initializers.get(beta_init) - self.gamma_regularizer = regularizers.get(gamma_regularizer) - self.beta_regularizer = regularizers.get(beta_regularizer) + self.gamma_init = keras.initializers.get(gamma_init) + self.beta_init = keras.initializers.get(beta_init) + self.gamma_regularizer = keras.regularizers.get(gamma_regularizer) + self.beta_regularizer = keras.regularizers.get(beta_regularizer) self.epsilon = epsilon self.group = group self.data_format = normalize_data_format(data_format) @@ -376,7 +189,7 @@ def build(self, input_shape): Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to reference for weight shape computations. """ - input_spec = [InputSpec(shape=input_shape)] + input_spec = [layers.InputSpec(shape=input_shape)] self.input_spec = input_spec # pylint:disable=attribute-defined-outside-init shape = [1 for _ in input_shape] if self.data_format == 'channels_last': @@ -397,7 +210,7 @@ def build(self, input_shape): name='beta') self.built = True # pylint:disable=attribute-defined-outside-init - def call(self, inputs, mask=None): # pylint:disable=unused-argument,arguments-differ + def call(self, inputs, *args, **kwargs): # noqa:C901 """This is where the layer's logic lives. Parameters @@ -486,16 +299,351 @@ def get_config(self): """ config = {'epsilon': self.epsilon, 'axis': self.axis, - 'gamma_init': initializers.serialize(self.gamma_init), - 'beta_init': initializers.serialize(self.beta_init), - 'gamma_regularizer': regularizers.serialize(self.gamma_regularizer), - 'beta_regularizer': regularizers.serialize(self.gamma_regularizer), + 'gamma_init': keras.initializers.serialize(self.gamma_init), + 'beta_init': keras.initializers.serialize(self.beta_init), + 'gamma_regularizer': keras.regularizers.serialize(self.gamma_regularizer), + 'beta_regularizer': keras.regularizers.serialize(self.gamma_regularizer), 'group': self.group} base_config = super().get_config() return dict(list(base_config.items()) + list(config.items())) +class InstanceNormalization(layers.Layer): # type:ignore[name-defined] + """Instance normalization layer (Lei Ba et al, 2016, Ulyanov et al., 2016). + + Normalize the activations of the previous layer at each step, i.e. applies a transformation + that maintains the mean activation close to 0 and the activation standard deviation close to 1. + + Parameters + ---------- + axis: int, optional + The axis that should be normalized (typically the features axis). For instance, after a + `Conv2D` layer with `data_format="channels_first"`, set `axis=1` in + :class:`InstanceNormalization`. Setting `axis=None` will normalize all values in each + instance of the batch. Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid + errors. Default: ``None`` + epsilon: float, optional + Small float added to variance to avoid dividing by zero. Default: `1e-3` + center: bool, optional + If ``True``, add offset of `beta` to normalized tensor. If ``False``, `beta` is ignored. + Default: ``True`` + scale: bool, optional + If ``True``, multiply by `gamma`. If ``False``, `gamma` is not used. When the next layer + is linear (also e.g. `relu`), this can be disabled since the scaling will be done by + the next layer. Default: ``True`` + beta_initializer: str, optional + Initializer for the beta weight. Default: `"zeros"` + gamma_initializer: str, optional + Initializer for the gamma weight. Default: `"ones"` + beta_regularizer: str, optional + Optional regularizer for the beta weight. Default: ``None`` + gamma_regularizer: str, optional + Optional regularizer for the gamma weight. Default: ``None`` + beta_constraint: float, optional + Optional constraint for the beta weight. Default: ``None`` + gamma_constraint: float, optional + Optional constraint for the gamma weight. Default: ``None`` + + References + ---------- + - Layer Normalization - https://arxiv.org/abs/1607.06450 + + - Instance Normalization: The Missing Ingredient for Fast Stylization - \ + https://arxiv.org/abs/1607.08022 + """ + # pylint:disable=too-many-instance-attributes,too-many-arguments + def __init__(self, + axis=None, + epsilon=1e-3, + center=True, + scale=True, + beta_initializer="zeros", + gamma_initializer="ones", + beta_regularizer=None, + gamma_regularizer=None, + beta_constraint=None, + gamma_constraint=None, + **kwargs): + self.beta = None + self.gamma = None + super().__init__(**kwargs) + self.supports_masking = True + self.axis = axis + self.epsilon = epsilon + self.center = center + self.scale = scale + self.beta_initializer = keras.initializers.get(beta_initializer) + self.gamma_initializer = keras.initializers.get(gamma_initializer) + self.beta_regularizer = keras.regularizers.get(beta_regularizer) + self.gamma_regularizer = keras.regularizers.get(gamma_regularizer) + self.beta_constraint = keras.constraints.get(beta_constraint) + self.gamma_constraint = keras.constraints.get(gamma_constraint) + + def build(self, input_shape): + """Creates the layer weights. + + Parameters + ---------- + input_shape: tensor + Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to + reference for weight shape computations. + """ + ndim = len(input_shape) + if self.axis == 0: + raise ValueError("Axis cannot be zero") + + if (self.axis is not None) and (ndim == 2): + raise ValueError("Cannot specify axis for rank 1 tensor") + + self.input_spec = layers.InputSpec(ndim=ndim) # noqa:E501 pylint:disable=attribute-defined-outside-init + + if self.axis is None: + shape = (1,) + else: + shape = (input_shape[self.axis],) + + if self.scale: + self.gamma = self.add_weight(shape=shape, + name="gamma", + initializer=self.gamma_initializer, + regularizer=self.gamma_regularizer, + constraint=self.gamma_constraint) + else: + self.gamma = None + if self.center: + self.beta = self.add_weight(shape=shape, + name="beta", + initializer=self.beta_initializer, + regularizer=self.beta_regularizer, + constraint=self.beta_constraint) + else: + self.beta = None + self.built = True # pylint:disable=attribute-defined-outside-init + + def call(self, inputs, training=None): # pylint:disable=arguments-differ,unused-argument + """This is where the layer's logic lives. + + Parameters + ---------- + inputs: tensor + Input tensor, or list/tuple of input tensors + + Returns + ------- + tensor + A tensor or list/tuple of tensors + """ + input_shape = K.int_shape(inputs) + reduction_axes = list(range(0, len(input_shape))) + + if self.axis is not None: + del reduction_axes[self.axis] + + del reduction_axes[0] + + mean = K.mean(inputs, reduction_axes, keepdims=True) + stddev = K.std(inputs, reduction_axes, keepdims=True) + self.epsilon + normed = (inputs - mean) / stddev + + broadcast_shape = [1] * len(input_shape) + if self.axis is not None: + broadcast_shape[self.axis] = input_shape[self.axis] + + if self.scale: + broadcast_gamma = K.reshape(self.gamma, broadcast_shape) + normed = normed * broadcast_gamma + if self.center: + broadcast_beta = K.reshape(self.beta, broadcast_shape) + normed = normed + broadcast_beta + return normed + + def get_config(self): + """Returns the config of the layer. + + A layer config is a Python dictionary (serializable) containing the configuration of a + layer. The same layer can be reinstated later (without its trained weights) from this + configuration. + + The configuration of a layer does not include connectivity information, nor the layer + class name. These are handled by `Network` (one layer of abstraction above). + + Returns + -------- + dict + A python dictionary containing the layer configuration + """ + config = { + "axis": self.axis, + "epsilon": self.epsilon, + "center": self.center, + "scale": self.scale, + "beta_initializer": keras.initializers.serialize(self.beta_initializer), + "gamma_initializer": keras.initializers.serialize(self.gamma_initializer), + "beta_regularizer": keras.regularizers.serialize(self.beta_regularizer), + "gamma_regularizer": keras.regularizers.serialize(self.gamma_regularizer), + "beta_constraint": keras.constraints.serialize(self.beta_constraint), + "gamma_constraint": keras.constraints.serialize(self.gamma_constraint) + } + base_config = super().get_config() + return dict(list(base_config.items()) + list(config.items())) + + +class RMSNormalization(layers.Layer): # type:ignore[name-defined] + """ Root Mean Square Layer Normalization (Biao Zhang, Rico Sennrich, 2019) + + RMSNorm is a simplification of the original layer normalization (LayerNorm). LayerNorm is a + regularization technique that might handle the internal covariate shift issue so as to + stabilize the layer activations and improve model convergence. It has been proved quite + successful in NLP-based model. In some cases, LayerNorm has become an essential component + to enable model optimization, such as in the SOTA NMT model Transformer. + + RMSNorm simplifies LayerNorm by removing the mean-centering operation, or normalizing layer + activations with RMS statistic. + + Parameters + ---------- + axis: int + The axis to normalize across. Typically this is the features axis. The left-out axes are + typically the batch axis/axes. This argument defaults to `-1`, the last dimension in the + input. + epsilon: float, optional + Small float added to variance to avoid dividing by zero. Default: `1e-8` + partial: float, optional + Partial multiplier for calculating pRMSNorm. Valid values are between `0.0` and `1.0`. + Setting to `0.0` or `1.0` disables. Default: `0.0` + bias: bool, optional + Whether to use a bias term for RMSNorm. Disabled by default because RMSNorm does not + enforce re-centering invariance. Default ``False`` + kwargs: dict + Standard keras layer kwargs + + References + ---------- + - RMS Normalization - https://arxiv.org/abs/1910.07467 + - Official implementation - https://github.com/bzhangGo/rmsnorm + """ + def __init__(self, axis=-1, epsilon=1e-8, partial=0.0, bias=False, **kwargs): + self.scale = None + self.offset = 0 + super().__init__(**kwargs) + + # Checks + if not isinstance(axis, int): + raise TypeError(f"Expected an int for the argument 'axis', but received: {axis}") + + if not 0.0 <= partial <= 1.0: + raise ValueError(f"partial must be between 0.0 and 1.0, but received {partial}") + + self.axis = axis + self.epsilon = epsilon + self.partial = partial + self.bias = bias + self.offset = 0. + + def build(self, input_shape): + """ Validate and populate :attr:`axis` + + Parameters + ---------- + input_shape: tensor + Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to + reference for weight shape computations. + """ + ndims = len(input_shape) + if ndims is None: + raise ValueError(f"Input shape {input_shape} has undefined rank.") + + # Resolve negative axis + if self.axis < 0: + self.axis += ndims + + # Validate axes + if self.axis < 0 or self.axis >= ndims: + raise ValueError(f"Invalid axis: {self.axis}") + + param_shape = [input_shape[self.axis]] + self.scale = self.add_weight( + name="scale", + shape=param_shape, + initializer="ones") + if self.bias: + self.offset = self.add_weight( + name="offset", + shape=param_shape, + initializer="zeros") + + self.built = True # pylint:disable=attribute-defined-outside-init + + def call(self, inputs, *args, **kwargs): + """ Call Root Mean Square Layer Normalization + + Parameters + ---------- + inputs: tensor + Input tensor, or list/tuple of input tensors + + Returns + ------- + tensor + A tensor or list/tuple of tensors + """ + # Compute the axes along which to reduce the mean / variance + input_shape = K.int_shape(inputs) + layer_size = input_shape[self.axis] + + if self.partial in (0.0, 1.0): + mean_square = K.mean(K.square(inputs), axis=self.axis, keepdims=True) + else: + partial_size = int(layer_size * self.partial) + partial_x, _ = tf.split( # pylint:disable=redundant-keyword-arg,no-value-for-parameter + inputs, + [partial_size, layer_size - partial_size], + axis=self.axis) + mean_square = K.mean(K.square(partial_x), axis=self.axis, keepdims=True) + + recip_square_root = tf.math.rsqrt(mean_square + self.epsilon) + output = self.scale * inputs * recip_square_root + self.offset + return output + + def compute_output_shape(self, input_shape): + """ The output shape of the layer is the same as the input shape. + + Parameters + ---------- + input_shape: tuple + The input shape to the layer + + Returns + ------- + tuple + The output shape to the layer + """ + return input_shape + + def get_config(self): + """Returns the config of the layer. + + A layer config is a Python dictionary (serializable) containing the configuration of a + layer. The same layer can be reinstated later (without its trained weights) from this + configuration. + + The configuration of a layer does not include connectivity information, nor the layer + class name. These are handled by `Network` (one layer of abstraction above). + + Returns + -------- + dict + A python dictionary containing the layer configuration + """ + base_config = super().get_config() + config = {"axis": self.axis, + "epsilon": self.epsilon, + "partial": self.partial, + "bias": self.bias} + return dict(list(base_config.items()) + list(config.items())) + + # Update normalization into Keras custom objects for name, obj in inspect.getmembers(sys.modules[__name__]): if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) + keras.utils.get_custom_objects().update({name: obj}) diff --git a/lib/model/normalization/__init__.py b/lib/model/normalization/__init__.py deleted file mode 100644 index a7f2bd759d..0000000000 --- a/lib/model/normalization/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -#!/usr/bin/env python3 -""" Conditional imports depending on whether the AMD version is installed or not """ - -from lib.utils import get_backend -from .normalization_common import AdaInstanceNormalization -from .normalization_common import GroupNormalization -from .normalization_common import InstanceNormalization - - -if get_backend() == "amd": - from .normalization_plaid import LayerNormalization, RMSNormalization -else: - from .normalization_tf import LayerNormalization, RMSNormalization diff --git a/lib/model/normalization/normalization_plaid.py b/lib/model/normalization/normalization_plaid.py deleted file mode 100644 index d9e28a5700..0000000000 --- a/lib/model/normalization/normalization_plaid.py +++ /dev/null @@ -1,384 +0,0 @@ -#!/usr/bin/env python3 -""" Normalization methods for faceswap.py. """ - -import sys -import inspect - -from plaidml.op import slice_tensor -from keras.layers import Layer -from keras import initializers, regularizers, constraints -from keras import backend as K -from keras.utils import get_custom_objects - - -class LayerNormalization(Layer): - """Instance normalization layer (Lei Ba et al, 2016). Implementation adapted from - tensorflow.keras implementation and https://github.com/CyberZHG/keras-layer-normalization - - Normalize the activations of the previous layer for each given example in a batch - independently, rather than across a batch like Batch Normalization. i.e. applies a - transformation that maintains the mean activation within each example close to 0 and the - activation standard deviation close to 1. - - Parameters - ---------- - axis: int or list/tuple - The axis or axes to normalize across. Typically this is the features axis/axes. - The left-out axes are typically the batch axis/axes. This argument defaults to `-1`, the - last dimension in the input. - epsilon: float, optional - Small float added to variance to avoid dividing by zero. Default: `1e-3` - center: bool, optional - If ``True``, add offset of `beta` to normalized tensor. If ``False``, `beta` is ignored. - Default: ``True`` - scale: bool, optional - If ``True``, multiply by `gamma`. If ``False``, `gamma` is not used. When the next layer - is linear (also e.g. `relu`), this can be disabled since the scaling will be done by - the next layer. Default: ``True`` - beta_initializer: str, optional - Initializer for the beta weight. Default: `"zeros"` - gamma_initializer: str, optional - Initializer for the gamma weight. Default: `"ones"` - beta_regularizer: str, optional - Optional regularizer for the beta weight. Default: ``None`` - gamma_regularizer: str, optional - Optional regularizer for the gamma weight. Default: ``None`` - beta_constraint: float, optional - Optional constraint for the beta weight. Default: ``None`` - gamma_constraint: float, optional - Optional constraint for the gamma weight. Default: ``None`` - kwargs: dict - Standard keras layer kwargs - - References - ---------- - - Layer Normalization - https://arxiv.org/abs/1607.06450 - - Keras implementation - https://github.com/CyberZHG/keras-layer-normalization - """ - def __init__(self, - axis=-1, - epsilon=1e-3, - center=True, - scale=True, - beta_initializer="zeros", - gamma_initializer="ones", - beta_regularizer=None, - gamma_regularizer=None, - beta_constraint=None, - gamma_constraint=None, - **kwargs): - - self.gamma = None - self.beta = None - super().__init__(**kwargs) - - if isinstance(axis, (list, tuple)): - self.axis = axis[:] - elif isinstance(axis, int): - self.axis = axis - else: - raise TypeError("Expected an int or a list/tuple of ints for the argument 'axis', " - f"but received: {axis}") - - self.epsilon = epsilon - self.center = center - self.scale = scale - self.beta_initializer = initializers.get(beta_initializer) - self.gamma_initializer = initializers.get(gamma_initializer) - self.beta_regularizer = regularizers.get(beta_regularizer) - self.gamma_regularizer = regularizers.get(gamma_regularizer) - self.beta_constraint = constraints.get(beta_constraint) - self.gamma_constraint = constraints.get(gamma_constraint) - self.supports_masking = True - - def build(self, input_shape): - """Creates the layer weights. - - Parameters - ---------- - input_shape: tensor - Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to - reference for weight shape computations. - """ - ndims = len(input_shape) - if ndims is None: - raise ValueError(f"Input shape {input_shape} has undefined rank.") - - # Convert axis to list and resolve negatives - if isinstance(self.axis, int): - self.axis = [self.axis] - elif isinstance(self.axis, tuple): - self.axis = list(self.axis) - for idx, axs in enumerate(self.axis): - if axs < 0: - self.axis[idx] = ndims + axs - - # Validate axes - for axs in self.axis: - if axs < 0 or axs >= ndims: - raise ValueError(f"Invalid axis: {axs}") - if len(self.axis) != len(set(self.axis)): - raise ValueError("Duplicate axis: {}".format(tuple(self.axis))) - - param_shape = [input_shape[dim] for dim in self.axis] - if self.scale: - self.gamma = self.add_weight( - name="gamma", - shape=param_shape, - initializer=self.gamma_initializer, - regularizer=self.gamma_regularizer, - constraint=self.gamma_constraint) - if self.center: - self.beta = self.add_weight( - name='beta', - shape=param_shape, - initializer=self.beta_initializer, - regularizer=self.beta_regularizer, - constraint=self.beta_constraint) - - self.built = True # pylint:disable=attribute-defined-outside-init - - def call(self, inputs, **kwargs): # pylint:disable=unused-argument - """This is where the layer's logic lives. - - Parameters - ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - - Returns - ------- - tensor - A tensor or list/tuple of tensors - """ - # Compute the axes along which to reduce the mean / variance - input_shape = K.int_shape(inputs) - ndims = len(input_shape) - - # Broadcasting only necessary for norm when the axis is not just the last dimension - broadcast_shape = [1] * ndims - for dim in self.axis: - broadcast_shape[dim] = input_shape[dim] - - def _broadcast(var): - if (var is not None and len(var.shape) != ndims and self.axis != [ndims - 1]): - return K.reshape(var, broadcast_shape) - return var - - # Calculate the moments on the last axis (layer activations). - mean = K.mean(inputs, self.axis, keepdims=True) - variance = K.mean(K.square(inputs - mean), axis=self.axis, keepdims=True) - std = K.sqrt(variance + self.epsilon) - outputs = (inputs - mean) / std - - scale, offset = _broadcast(self.gamma), _broadcast(self.beta) - if self.scale: - outputs *= scale - if self.center: - outputs *= offset - - return outputs - - def compute_output_shape(self, input_shape): # pylint:disable=no-self-use - """ The output shape of the layer is the same as the input shape. - - Parameters - ---------- - input_shape: tuple - The input shape to the layer - - Returns - ------- - tuple - The output shape to the layer - """ - return input_shape - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) containing the configuration of a - layer. The same layer can be reinstated later (without its trained weights) from this - configuration. - - The configuration of a layer does not include connectivity information, nor the layer - class name. These are handled by `Network` (one layer of abstraction above). - - Returns - -------- - dict - A python dictionary containing the layer configuration - """ - base_config = super().get_config() - config = dict(axis=self.axis, - epsilon=self.epsilon, - center=self.center, - scale=self.scale, - beta_initializer=initializers.serialize(self.beta_initializer), - gamma_initializer=initializers.serialize(self.gamma_initializer), - beta_regularizer=regularizers.serialize(self.beta_regularizer), - gamma_regularizer=regularizers.serialize(self.gamma_regularizer), - beta_constraint=constraints.serialize(self.beta_constraint), - gamma_constraint=constraints.serialize(self.gamma_constraint)) - return dict(list(base_config.items()) + list(config.items())) - - -class RMSNormalization(Layer): - """ Root Mean Square Layer Normalization (Biao Zhang, Rico Sennrich, 2019) - - RMSNorm is a simplification of the original layer normalization (LayerNorm). LayerNorm is a - regularization technique that might handle the internal covariate shift issue so as to - stabilize the layer activations and improve model convergence. It has been proved quite - successful in NLP-based model. In some cases, LayerNorm has become an essential component - to enable model optimization, such as in the SOTA NMT model Transformer. - - RMSNorm simplifies LayerNorm by removing the mean-centering operation, or normalizing layer - activations with RMS statistic. - - Parameters - ---------- - axis: int - The axis to normalize across. Typically this is the features axis. The left-out axes are - typically the batch axis/axes. This argument defaults to `-1`, the last dimension in the - input. - epsilon: float, optional - Small float added to variance to avoid dividing by zero. Default: `1e-8` - partial: float, optional - Partial multiplier for calculating pRMSNorm. Valid values are between `0.0` and `1.0`. - Setting to `0.0` or `1.0` disables. Default: `0.0` - bias: bool, optional - Whether to use a bias term for RMSNorm. Disabled by default because RMSNorm does not - enforce re-centering invariance. Default ``False`` - kwargs: dict - Standard keras layer kwargs - - References - ---------- - - RMS Normalization - https://arxiv.org/abs/1910.07467 - - Official implementation - https://github.com/bzhangGo/rmsnorm - """ - def __init__(self, axis=-1, epsilon=1e-8, partial=0.0, bias=False, **kwargs): - self.scale = None - self.offset = 0 - super().__init__(**kwargs) - - # Checks - if not isinstance(axis, int): - raise TypeError(f"Expected an int for the argument 'axis', but received: {axis}") - - if not 0.0 <= partial <= 1.0: - raise ValueError(f"partial must be between 0.0 and 1.0, but received {partial}") - - self.axis = axis - self.epsilon = epsilon - self.partial = partial - self.bias = bias - self.offset = 0. - - def build(self, input_shape): - """ Validate and populate :attr:`axis` - - Parameters - ---------- - input_shape: tensor - Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to - reference for weight shape computations. - """ - ndims = len(input_shape) - if ndims is None: - raise ValueError(f"Input shape {input_shape} has undefined rank.") - - # Resolve negative axis - if self.axis < 0: - self.axis += ndims - - # Validate axes - if self.axis < 0 or self.axis >= ndims: - raise ValueError(f"Invalid axis: {self.axis}") - - param_shape = [input_shape[self.axis]] - self.scale = self.add_weight( - name="scale", - shape=param_shape, - initializer="ones") - if self.bias: - self.offset = self.add_weight( - name="offset", - shape=param_shape, - initializer="zeros") - - self.built = True # pylint:disable=attribute-defined-outside-init - - def call(self, inputs, **kwargs): # pylint:disable=unused-argument - """ Call Root Mean Square Layer Normalization - - Parameters - ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - - Returns - ------- - tensor - A tensor or list/tuple of tensors - """ - # Compute the axes along which to reduce the mean / variance - input_shape = K.int_shape(inputs) - layer_size = input_shape[self.axis] - - if self.partial in (0.0, 1.0): - mean_square = K.mean(K.square(inputs), axis=self.axis, keepdims=True) - else: - partial_size = int(layer_size * self.partial) - partial_x = slice_tensor(inputs, - axes=[self.axis], - starts=[0], - ends=[partial_size]) - mean_square = K.mean(K.square(partial_x), axis=self.axis, keepdims=True) - - recip_square_root = 1. / K.sqrt(mean_square + self.epsilon) - output = self.scale * inputs * recip_square_root + self.offset - return output - - def compute_output_shape(self, input_shape): # pylint:disable=no-self-use - """ The output shape of the layer is the same as the input shape. - - Parameters - ---------- - input_shape: tuple - The input shape to the layer - - Returns - ------- - tuple - The output shape to the layer - """ - return input_shape - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) containing the configuration of a - layer. The same layer can be reinstated later (without its trained weights) from this - configuration. - - The configuration of a layer does not include connectivity information, nor the layer - class name. These are handled by `Network` (one layer of abstraction above). - - Returns - -------- - dict - A python dictionary containing the layer configuration - """ - base_config = super().get_config() - config = dict(axis=self.axis, - epsilon=self.epsilon, - partial=self.partial, - bias=self.bias) - return dict(list(base_config.items()) + list(config.items())) - - -# Update normalization into Keras custom objects -for name, obj in inspect.getmembers(sys.modules[__name__]): - if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) diff --git a/lib/model/normalization/normalization_tf.py b/lib/model/normalization/normalization_tf.py deleted file mode 100644 index b7a4abd028..0000000000 --- a/lib/model/normalization/normalization_tf.py +++ /dev/null @@ -1,171 +0,0 @@ -#!/usr/bin/env python3 -""" Normalization methods for faceswap.py specific to Tensorflow backend """ -import inspect -import sys - -import tensorflow as tf -# Ignore linting errors from Tensorflow's thoroughly broken import system -from tensorflow.keras import backend as K # pylint:disable=import-error -from tensorflow.keras.layers import Layer, LayerNormalization # noqa pylint:disable=no-name-in-module,unused-import,import-error -from tensorflow.keras.utils import get_custom_objects # noqa pylint:disable=no-name-in-module,import-error - - -class RMSNormalization(Layer): - """ Root Mean Square Layer Normalization (Biao Zhang, Rico Sennrich, 2019) - - RMSNorm is a simplification of the original layer normalization (LayerNorm). LayerNorm is a - regularization technique that might handle the internal covariate shift issue so as to - stabilize the layer activations and improve model convergence. It has been proved quite - successful in NLP-based model. In some cases, LayerNorm has become an essential component - to enable model optimization, such as in the SOTA NMT model Transformer. - - RMSNorm simplifies LayerNorm by removing the mean-centering operation, or normalizing layer - activations with RMS statistic. - - Parameters - ---------- - axis: int - The axis to normalize across. Typically this is the features axis. The left-out axes are - typically the batch axis/axes. This argument defaults to `-1`, the last dimension in the - input. - epsilon: float, optional - Small float added to variance to avoid dividing by zero. Default: `1e-8` - partial: float, optional - Partial multiplier for calculating pRMSNorm. Valid values are between `0.0` and `1.0`. - Setting to `0.0` or `1.0` disables. Default: `0.0` - bias: bool, optional - Whether to use a bias term for RMSNorm. Disabled by default because RMSNorm does not - enforce re-centering invariance. Default ``False`` - kwargs: dict - Standard keras layer kwargs - - References - ---------- - - RMS Normalization - https://arxiv.org/abs/1910.07467 - - Official implementation - https://github.com/bzhangGo/rmsnorm - """ - def __init__(self, axis=-1, epsilon=1e-8, partial=0.0, bias=False, **kwargs): - self.scale = None - self.offset = 0 - super().__init__(**kwargs) - - # Checks - if not isinstance(axis, int): - raise TypeError(f"Expected an int for the argument 'axis', but received: {axis}") - - if not 0.0 <= partial <= 1.0: - raise ValueError(f"partial must be between 0.0 and 1.0, but received {partial}") - - self.axis = axis - self.epsilon = epsilon - self.partial = partial - self.bias = bias - self.offset = 0. - - def build(self, input_shape): - """ Validate and populate :attr:`axis` - - Parameters - ---------- - input_shape: tensor - Keras tensor (future input to layer) or ``list``/``tuple`` of Keras tensors to - reference for weight shape computations. - """ - ndims = len(input_shape) - if ndims is None: - raise ValueError(f"Input shape {input_shape} has undefined rank.") - - # Resolve negative axis - if self.axis < 0: - self.axis += ndims - - # Validate axes - if self.axis < 0 or self.axis >= ndims: - raise ValueError(f"Invalid axis: {self.axis}") - - param_shape = [input_shape[self.axis]] - self.scale = self.add_weight( - name="scale", - shape=param_shape, - initializer="ones") - if self.bias: - self.offset = self.add_weight( - name="offset", - shape=param_shape, - initializer="zeros") - - self.built = True # pylint:disable=attribute-defined-outside-init - - def call(self, inputs, **kwargs): # pylint:disable=unused-argument - """ Call Root Mean Square Layer Normalization - - Parameters - ---------- - inputs: tensor - Input tensor, or list/tuple of input tensors - - Returns - ------- - tensor - A tensor or list/tuple of tensors - """ - # Compute the axes along which to reduce the mean / variance - input_shape = K.int_shape(inputs) - layer_size = input_shape[self.axis] - - if self.partial in (0.0, 1.0): - mean_square = K.mean(K.square(inputs), axis=self.axis, keepdims=True) - else: - partial_size = int(layer_size * self.partial) - partial_x, _ = tf.split( # pylint:disable=redundant-keyword-arg,no-value-for-parameter - inputs, - [partial_size, layer_size - partial_size], - axis=self.axis) - mean_square = K.mean(K.square(partial_x), axis=self.axis, keepdims=True) - - recip_square_root = tf.math.rsqrt(mean_square + self.epsilon) - output = self.scale * inputs * recip_square_root + self.offset - return output - - def compute_output_shape(self, input_shape): # pylint:disable=no-self-use - """ The output shape of the layer is the same as the input shape. - - Parameters - ---------- - input_shape: tuple - The input shape to the layer - - Returns - ------- - tuple - The output shape to the layer - """ - return input_shape - - def get_config(self): - """Returns the config of the layer. - - A layer config is a Python dictionary (serializable) containing the configuration of a - layer. The same layer can be reinstated later (without its trained weights) from this - configuration. - - The configuration of a layer does not include connectivity information, nor the layer - class name. These are handled by `Network` (one layer of abstraction above). - - Returns - -------- - dict - A python dictionary containing the layer configuration - """ - base_config = super().get_config() - config = dict(axis=self.axis, - epsilon=self.epsilon, - partial=self.partial, - bias=self.bias) - return dict(list(base_config.items()) + list(config.items())) - - -# Update normalization into Keras custom objects -for name, obj in inspect.getmembers(sys.modules[__name__]): - if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) diff --git a/lib/model/optimizers_tf.py b/lib/model/optimizers.py similarity index 92% rename from lib/model/optimizers_tf.py rename to lib/model/optimizers.py index 9b028a55fa..33efe7c8bd 100644 --- a/lib/model/optimizers_tf.py +++ b/lib/model/optimizers.py @@ -1,8 +1,5 @@ #!/usr/bin/env python3 """ Custom Optimizers for TensorFlow 2.x/tf.keras """ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function import inspect import sys @@ -10,8 +7,8 @@ import tensorflow as tf # Ignore linting errors from Tensorflow's thoroughly broken import system -from tensorflow.keras.optimizers import (Adam, Nadam, RMSprop) # noqa pylint:disable=no-name-in-module,unused-import,import-error -from tensorflow.keras.utils import get_custom_objects # noqa pylint:disable=no-name-in-module,import-error +from tensorflow.keras.optimizers import Adam, Nadam, RMSprop # noqa:E501,F401 pylint:disable=import-error,unused-import +keras = tf.keras class AdaBelief(tf.keras.optimizers.Optimizer): @@ -381,22 +378,22 @@ def get_config(self): The optimizer configuration. """ config = super().get_config() - config.update(dict(learning_rate=self._serialize_hyperparameter("learning_rate"), - beta_1=self._serialize_hyperparameter("beta_1"), - beta_2=self._serialize_hyperparameter("beta_2"), - decay=self._serialize_hyperparameter("decay"), - weight_decay=self._serialize_hyperparameter("weight_decay"), - sma_threshold=self._serialize_hyperparameter("sma_threshold"), - epsilon=self.epsilon, - amsgrad=self.amsgrad, - rectify=self.rectify, - total_steps=self._serialize_hyperparameter("total_steps"), - warmup_proportion=self._serialize_hyperparameter("warmup_proportion"), - min_lr=self._serialize_hyperparameter("min_lr"))) + config.update({"learning_rate": self._serialize_hyperparameter("learning_rate"), + "beta_1": self._serialize_hyperparameter("beta_1"), + "beta_2": self._serialize_hyperparameter("beta_2"), + "decay": self._serialize_hyperparameter("decay"), + "weight_decay": self._serialize_hyperparameter("weight_decay"), + "sma_threshold": self._serialize_hyperparameter("sma_threshold"), + "epsilon": self.epsilon, + "amsgrad": self.amsgrad, + "rectify": self.rectify, + "total_steps": self._serialize_hyperparameter("total_steps"), + "warmup_proportion": self._serialize_hyperparameter("warmup_proportion"), + "min_lr": self._serialize_hyperparameter("min_lr")}) return config # Update layers into Keras custom objects for _name, obj in inspect.getmembers(sys.modules[__name__]): if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({_name: obj}) + keras.utils.get_custom_objects().update({_name: obj}) diff --git a/lib/model/optimizers_plaid.py b/lib/model/optimizers_plaid.py deleted file mode 100644 index 848cabff3b..0000000000 --- a/lib/model/optimizers_plaid.py +++ /dev/null @@ -1,156 +0,0 @@ -#!/usr/bin/env python3 -""" Custom Optimizers for PlaidML/Keras 2.2. """ -import inspect -import sys - -from keras import backend as K -from keras.optimizers import Optimizer, Adam, Nadam, RMSprop # noqa pylint:disable=unused-import -from keras.utils import get_custom_objects - - -class AdaBelief(Optimizer): - """AdaBelief optimizer. - - Default parameters follow those provided in the original paper. - - Parameters - ---------- - learning_rate: float - The learning rate. - beta_1: float - The exponential decay rate for the 1st moment estimates. - beta_2: float - The exponential decay rate for the 2nd moment estimates. - epsilon: float, optional - A small constant for numerical stability. Default: `K.epsilon()`. - amsgrad: bool - Whether to apply AMSGrad variant of this algorithm from the paper "On the Convergence - of Adam and beyond". - - References - ---------- - AdaBelief - A Method for Stochastic Optimization - https://arxiv.org/abs/1412.6980v8 - On the Convergence of AdaBelief and Beyond - https://openreview.net/forum?id=ryQu7f-RZ - - Adapted from https://github.com/liaoxuanzhi/adabelief - - BSD 2-Clause License - - Copyright (c) 2021, Juntang Zhuang - All rights reserved. - - Redistribution and use in source and binary forms, with or without - modification, are permitted provided that the following conditions are met: - - 1. Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. - - 2. Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - - THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - - """ - - def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, - epsilon=None, decay=0., weight_decay=0.0, **kwargs): - super().__init__(**kwargs) - with K.name_scope(self.__class__.__name__): - self.iterations = K.variable(0, dtype='int64', name='iterations') - self.lr = K.variable(lr, name='lr') - self.beta_1 = K.variable(beta_1, name='beta_1') - self.beta_2 = K.variable(beta_2, name='beta_2') - self.decay = K.variable(decay, name='decay') - if epsilon is None: - epsilon = K.epsilon() - self.epsilon = float(epsilon) - self.initial_decay = decay - self.weight_decay = float(weight_decay) - - def get_updates(self, loss, params): # pylint:disable=too-many-locals - """ Get the weight updates - - Parameters - ---------- - loss: list - The loss to update - params: list - The variables - """ - grads = self.get_gradients(loss, params) - self.updates = [K.update_add(self.iterations, 1)] - - l_r = self.lr - if self.initial_decay > 0: - l_r = l_r * (1. / (1. + self.decay * K.cast(self.iterations, - K.dtype(self.decay)))) - - var_t = K.cast(self.iterations, K.floatx()) + 1 - # bias correction - bias_correction1 = 1. - K.pow(self.beta_1, var_t) - bias_correction2 = 1. - K.pow(self.beta_2, var_t) - - m_s = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params] - v_s = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params] - - self.weights = [self.iterations] + m_s + v_s - - for param, grad, var_m, var_v in zip(params, grads, m_s, v_s): - if self.weight_decay != 0.: - grad += self.weight_decay * K.stop_gradient(param) - - m_t = (self.beta_1 * var_m) + (1. - self.beta_1) * grad - m_corr_t = m_t / bias_correction1 - - v_t = (self.beta_2 * var_v) + (1. - self.beta_2) * K.square(grad - m_t) + self.epsilon - v_corr_t = K.sqrt(v_t / bias_correction2) - - p_t = param - l_r * m_corr_t / (v_corr_t + self.epsilon) - - self.updates.append(K.update(var_m, m_t)) - self.updates.append(K.update(var_v, v_t)) - new_param = p_t - - # Apply constraints. - if getattr(param, 'constraint', None) is not None: - new_param = param.constraint(new_param) - - self.updates.append(K.update(param, new_param)) - return self.updates - - def get_config(self): - """ Returns the config of the optimizer. - - An optimizer config is a Python dictionary (serializable) containing the configuration of - an optimizer. The same optimizer can be re-instantiated later (without any saved state) - from this configuration. - - Returns - ------- - dict - The optimizer configuration. - """ - config = dict(lr=float(K.get_value(self.lr)), - beta_1=float(K.get_value(self.beta_1)), - beta_2=float(K.get_value(self.beta_2)), - decay=float(K.get_value(self.decay)), - epsilon=self.epsilon, - weight_decay=self.weight_decay) - base_config = super().get_config() - return dict(list(base_config.items()) + list(config.items())) - - -# Update layers into Keras custom objects -for name, obj in inspect.getmembers(sys.modules[__name__]): - if inspect.isclass(obj) and obj.__module__ == __name__: - get_custom_objects().update({name: obj}) diff --git a/lib/model/session.py b/lib/model/session.py index ac84048a1e..a400b2fde6 100644 --- a/lib/model/session.py +++ b/lib/model/session.py @@ -1,22 +1,23 @@ #!/usr/bin python3 """ Settings manager for Keras Backend """ - +from __future__ import annotations +from contextlib import nullcontext import logging +import typing as T import numpy as np import tensorflow as tf +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Activation # pylint:disable=import-error +from tensorflow.keras.models import load_model as k_load_model, Model # noqa:E501 # pylint:disable=import-error + from lib.utils import get_backend -if get_backend() == "amd": - from keras.layers import Activation - from keras.models import load_model as k_load_model, Model -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Activation # noqa pylint:disable=no-name-in-module,import-error - from tensorflow.keras.models import load_model as k_load_model, Model # noqa pylint:disable=no-name-in-module,import-error +if T.TYPE_CHECKING: + from collections.abc import Callable -logger = logging.getLogger(__name__) # pylint:disable=invalid-name +logger = logging.getLogger(__name__) class KSession(): @@ -26,8 +27,7 @@ class KSession(): actions performed on a model are handled consistently and can be performed in parallel in separate threads. - This is an early implementation of this class, and should be expanded out over time - with relevant `AMD`, `CPU` and `NVIDIA` backend methods. + This is an early implementation of this class, and should be expanded out over time. Notes ----- @@ -49,74 +49,61 @@ class KSession(): exclude_gpus: list, optional A list of indices correlating to connected GPUs that Tensorflow should not use. Pass ``None`` to not exclude any GPUs. Default: ``None`` - + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` """ - def __init__(self, name, model_path, model_kwargs=None, allow_growth=False, exclude_gpus=None): - logger.trace("Initializing: %s (name: %s, model_path: %s, model_kwargs: %s, " - "allow_growth: %s, exclude_gpus: %s)", self.__class__.__name__, name, - model_path, model_kwargs, allow_growth, exclude_gpus) + def __init__(self, + name: str, + model_path: str, + model_kwargs: dict | None = None, + allow_growth: bool = False, + exclude_gpus: list[int] | None = None, + cpu_mode: bool = False) -> None: + logger.trace("Initializing: %s (name: %s, model_path: %s, " # type:ignore + "model_kwargs: %s, allow_growth: %s, exclude_gpus: %s, cpu_mode: %s)", + self.__class__.__name__, name, model_path, model_kwargs, allow_growth, + exclude_gpus, cpu_mode) self._name = name self._backend = get_backend() - self._set_session(allow_growth, exclude_gpus) + self._context = self._set_session(allow_growth, + [] if exclude_gpus is None else exclude_gpus, + cpu_mode) self._model_path = model_path self._model_kwargs = {} if not model_kwargs else model_kwargs - self._model = None - logger.trace("Initialized: %s", self.__class__.__name__,) + self._model: Model | None = None + logger.trace("Initialized: %s", self.__class__.__name__,) # type:ignore - def predict(self, feed, batch_size=None): + def predict(self, + feed: list[np.ndarray] | np.ndarray, + batch_size: int | None = None) -> list[np.ndarray] | np.ndarray: """ Get predictions from the model. This method is a wrapper for :func:`keras.predict()` function. For Tensorflow backends - this is a straight call to the predict function. For PlaidML backends, this attempts - to optimize the inference batch sizes to reduce the number of kernels that need to be - compiled. + this is a straight call to the predict function. Parameters ---------- feed: numpy.ndarray or list The feed to be provided to the model as input. This should be a :class:`numpy.ndarray` for single inputs or a `list` of :class:`numpy.ndarray` objects for multiple inputs. - """ - if self._backend == "amd" and batch_size is not None: - return self._amd_predict_with_optimized_batchsizes(feed, batch_size) - return self._model.predict(feed, batch_size=batch_size) - - def _amd_predict_with_optimized_batchsizes(self, feed, batch_size): - """ Minimizes the amount of kernels to be compiled when using the ``amd`` backend with - varying batch sizes while trying to keep the batchsize as high as possible. + batchsize: int, optional + The batch size to run prediction at. Default ``None`` - Parameters - ---------- - feed: numpy.ndarray or list - The feed to be provided to the model as input. This should be a ``numpy.ndarray`` - for single inputs or a ``list`` of ``numpy.ndarray`` objects for multiple inputs. - batch_size: int - The upper batchsize to use. + Returns + ------- + :class:`numpy.ndarray` + The predictions from the model """ - if isinstance(feed, np.ndarray): - feed = [feed] - items = feed[0].shape[0] - done_items = 0 - results = [] - while done_items < items: - if batch_size < 4: # Not much difference in BS < 4 - batch_size = 1 - batch_items = ((items - done_items) // batch_size) * batch_size - if batch_items: - pred_data = [x[done_items:done_items + batch_items] for x in feed] - pred = self._model.predict(pred_data, batch_size=batch_size) - done_items += batch_items - results.append(pred) - batch_size //= 2 - if isinstance(results[0], np.ndarray): - return np.concatenate(results) - return [np.concatenate(x) for x in zip(*results)] - - def _set_session(self, allow_growth, exclude_gpus): + assert self._model is not None + with self._context: + return self._model.predict(feed, verbose=0, batch_size=batch_size) + + def _set_session(self, + allow_growth: bool, + exclude_gpus: list, + cpu_mode: bool) -> T.ContextManager: """ Sets the backend session options. - For AMD backend this does nothing. - For CPU backends, this hides any GPUs from Tensorflow. For Nvidia backends, this hides any GPUs that Tensorflow should not use and applies @@ -124,20 +111,21 @@ def _set_session(self, allow_growth, exclude_gpus): Parameters ---------- - allow_growth: bool, optional + allow_growth: bool Enable the Tensorflow GPU allow_growth configuration option. This option prevents Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation - and slower performance. Default: False - exclude_gpus: list, optional + and slower performance + exclude_gpus: list A list of indices correlating to connected GPUs that Tensorflow should not use. Pass - ``None`` to not exclude any GPUs. Default: ``None`` + ``None`` to not exclude any GPUs + cpu_mode: bool + ``True`` run the model on CPU. Default: ``False`` """ - if self._backend == "amd": - return + retval = nullcontext() if self._backend == "cpu": - logger.verbose("Hiding GPUs from Tensorflow") + logger.verbose("Hiding GPUs from Tensorflow") # type:ignore tf.config.set_visible_devices([], "GPU") - return + return retval gpus = tf.config.list_physical_devices('GPU') if exclude_gpus: @@ -145,12 +133,16 @@ def _set_session(self, allow_growth, exclude_gpus): logger.debug("Filtering devices to: %s", gpus) tf.config.set_visible_devices(gpus, "GPU") - if allow_growth: + if allow_growth and self._backend == "nvidia": for gpu in gpus: logger.info("Setting allow growth for GPU: %s", gpu) tf.config.experimental.set_memory_growth(gpu, True) - def load_model(self): + if cpu_mode: + retval = tf.device("/device:cpu:0") + return retval + + def load_model(self) -> None: """ Loads a model. This method is a wrapper for :func:`keras.models.load_model()`. Loads a model and its @@ -161,12 +153,12 @@ def load_model(self): For Tensorflow backends, the `make_predict_function` method is called on the model to make it thread safe. """ - logger.verbose("Initializing plugin model: %s", self._name) - self._model = k_load_model(self._model_path, compile=False, **self._model_kwargs) - if self._backend != "amd": + logger.verbose("Initializing plugin model: %s", self._name) # type:ignore + with self._context: + self._model = k_load_model(self._model_path, compile=False, **self._model_kwargs) self._model.make_predict_function() - def define_model(self, function): + def define_model(self, function: Callable) -> None: """ Defines a model from the given function. This method acts as a wrapper for :class:`keras.models.Model()`. @@ -178,9 +170,10 @@ def define_model(self, function): ``outputs``. The function that generates these results should be passed in, NOT the results themselves, as the function needs to be executed within the correct context. """ - self._model = Model(*function()) + with self._context: + self._model = Model(*function()) - def load_model_weights(self): + def load_model_weights(self) -> None: """ Load model weights for a defined model inside the correct session. This method is a wrapper for :class:`keras.load_weights()`. Once a model has been defined @@ -190,12 +183,13 @@ def load_model_weights(self): For Tensorflow backends, the `make_predict_function` method is called on the model to make it thread safe. """ - logger.verbose("Initializing plugin model: %s", self._name) - self._model.load_weights(self._model_path) - if self._backend != "amd": + logger.verbose("Initializing plugin model: %s", self._name) # type:ignore + assert self._model is not None + with self._context: + self._model.load_weights(self._model_path) self._model.make_predict_function() - def append_softmax_activation(self, layer_index=-1): + def append_softmax_activation(self, layer_index: int = -1) -> None: """ Append a softmax activation layer to a model Occasionally a softmax activation layer needs to be added to a model's output. @@ -208,5 +202,7 @@ def append_softmax_activation(self, layer_index=-1): softmax activation layer. Default: `-1` (The final layer of the model) """ logger.debug("Appending Softmax Activation to model: (layer_index: %s)", layer_index) - softmax = Activation("softmax", name="softmax")(self._model.layers[layer_index].output) - self._model = Model(inputs=self._model.input, outputs=[softmax]) + assert self._model is not None + with self._context: + softmax = Activation("softmax", name="softmax")(self._model.layers[layer_index].output) + self._model = Model(inputs=self._model.input, outputs=[softmax]) diff --git a/lib/multithreading.py b/lib/multithreading.py index 58f72ebf4a..a2c4300d44 100644 --- a/lib/multithreading.py +++ b/lib/multithreading.py @@ -1,14 +1,23 @@ #!/usr/bin/env python3 """ Multithreading/processing utils for faceswap """ - +from __future__ import annotations import logging +import typing as T from multiprocessing import cpu_count import queue as Queue import sys import threading +from types import TracebackType + +if T.TYPE_CHECKING: + from collections.abc import Callable, Generator -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) +_ErrorType: T.TypeAlias = tuple[type[BaseException], + BaseException, + TracebackType] | tuple[T.Any, T.Any, T.Any] | None +_THREAD_NAMES: set[str] = set() def total_cpus(): @@ -16,26 +25,80 @@ def total_cpus(): return cpu_count() +def _get_name(name: str) -> str: + """ Obtain a unique name for a thread + + Parameters + ---------- + name: str + The requested name + + Returns + ------- + str + The request name with "_#" appended (# being an integer) making the name unique + """ + idx = 0 + real_name = name + while True: + if real_name in _THREAD_NAMES: + real_name = f"{name}_{idx}" + idx += 1 + continue + _THREAD_NAMES.add(real_name) + return real_name + + class FSThread(threading.Thread): - """ Subclass of thread that passes errors back to parent """ - def __init__(self, group=None, target=None, name=None, # pylint: disable=too-many-arguments - args=(), kwargs=None, *, daemon=None): - super().__init__(group=group, target=target, name=name, - args=args, kwargs=kwargs, daemon=daemon) - self.err = None - - def check_and_raise_error(self): - """ Checks for errors in thread and raises them in caller """ + """ Subclass of thread that passes errors back to parent + + Parameters + ---------- + target: callable object, Optional + The callable object to be invoked by the run() method. If ``None`` nothing is called. + Default: ``None`` + name: str, optional + The thread name. if ``None`` a unique name is constructed of the form "Thread-N" where N + is a small decimal number. Default: ``None`` + args: tuple + The argument tuple for the target invocation. Default: (). + kwargs: dict + keyword arguments for the target invocation. Default: {}. + """ + _target: Callable + _args: tuple + _kwargs: dict[str, T.Any] + _name: str + + def __init__(self, + target: Callable | None = None, + name: str | None = None, + args: tuple = (), + kwargs: dict[str, T.Any] | None = None, + *, + daemon: bool | None = None) -> None: + super().__init__(target=target, name=name, args=args, kwargs=kwargs, daemon=daemon) + self.err: _ErrorType = None + + def check_and_raise_error(self) -> None: + """ Checks for errors in thread and raises them in caller. + + Raises + ------ + Error + Re-raised error from within the thread + """ if not self.err: return logger.debug("Thread error caught: %s", self.err) raise self.err[1].with_traceback(self.err[2]) - def run(self): + def run(self) -> None: + """ Runs the target, reraising any errors from within the thread in the caller. """ try: - if self._target: + if self._target is not None: self._target(*self._args, **self._kwargs) - except Exception as err: # pylint: disable=broad-except + except Exception as err: # pylint:disable=broad-except self.err = sys.exc_info() logger.debug("Error in thread (%s): %s", self._name, str(err)) finally: @@ -45,53 +108,85 @@ def run(self): class MultiThread(): - """ Threading for IO heavy ops - Catches errors in thread and rethrows to parent """ - def __init__(self, target, *args, thread_count=1, name=None, **kwargs): - self._name = name if name else target.__name__ + """ Threading for IO heavy ops. Catches errors in thread and rethrows to parent. + + Parameters + ---------- + target: callable object + The callable object to be invoked by the run() method. + args: tuple + The argument tuple for the target invocation. Default: (). + thread_count: int, optional + The number of threads to use. Default: 1 + name: str, optional + The thread name. if ``None`` a unique name is constructed of the form {target.__name__}_N + where N is an incrementing integer. Default: ``None`` + kwargs: dict + keyword arguments for the target invocation. Default: {}. + """ + def __init__(self, + target: Callable, + *args, + thread_count: int = 1, + name: str | None = None, + **kwargs) -> None: + self._name = _get_name(name if name else target.__name__) logger.debug("Initializing %s: (target: '%s', thread_count: %s)", self.__class__.__name__, self._name, thread_count) - logger.trace("args: %s, kwargs: %s", args, kwargs) + logger.trace("args: %s, kwargs: %s", args, kwargs) # type:ignore self.daemon = True self._thread_count = thread_count - self._threads = list() + self._threads: list[FSThread] = [] self._target = target self._args = args self._kwargs = kwargs logger.debug("Initialized %s: '%s'", self.__class__.__name__, self._name) @property - def has_error(self): - """ Return true if a thread has errored, otherwise false """ + def has_error(self) -> bool: + """ bool: ``True`` if a thread has errored, otherwise ``False`` """ return any(thread.err for thread in self._threads) @property - def errors(self): - """ Return a list of thread errors """ + def errors(self) -> list[_ErrorType]: + """ list: List of thread error values """ return [thread.err for thread in self._threads if thread.err] @property - def name(self): - """ Return thread name """ + def name(self) -> str: + """ :str: The name of the thread """ return self._name - def check_and_raise_error(self): - """ Checks for errors in thread and raises them in caller """ + def check_and_raise_error(self) -> None: + """ Checks for errors in thread and raises them in caller. + + Raises + ------ + Error + Re-raised error from within the thread + """ if not self.has_error: return logger.debug("Thread error caught: %s", self.errors) error = self.errors[0] + assert error is not None raise error[1].with_traceback(error[2]) - def is_alive(self): - """ Return true if any thread is alive else false """ + def is_alive(self) -> bool: + """ Check if any threads are still alive + + Returns + ------- + bool + ``True`` if any threads are alive. ``False`` if no threads are alive + """ return any(thread.is_alive() for thread in self._threads) - def start(self): - """ Start a thread with the given method and args """ + def start(self) -> None: + """ Start all the threads for the given method, args and kwargs """ logger.debug("Starting thread(s): '%s'", self._name) for idx in range(self._thread_count): - name = "{}_{}".format(self._name, idx) + name = self._name if self._thread_count == 1 else f"{self._name}_{idx}" logger.debug("Starting thread %s of %s: '%s'", idx + 1, self._thread_count, name) thread = FSThread(name=name, @@ -103,44 +198,83 @@ def start(self): self._threads.append(thread) logger.debug("Started all threads '%s': %s", self._name, len(self._threads)) - def completed(self): - """ Return False if there are any alive threads else True """ + def completed(self) -> bool: + """ Check if all threads have completed + + Returns + ------- + ``True`` if all threads have completed otherwise ``False`` + """ retval = all(not thread.is_alive() for thread in self._threads) logger.debug(retval) return retval - def join(self): - """ Join the running threads, catching and re-raising any errors """ + def join(self) -> None: + """ Join the running threads, catching and re-raising any errors + + Clear the list of threads for class instance re-use + """ logger.debug("Joining Threads: '%s'", self._name) for thread in self._threads: - logger.debug("Joining Thread: '%s'", thread._name) # pylint: disable=protected-access + logger.debug("Joining Thread: '%s'", thread._name) # pylint:disable=protected-access thread.join() if thread.err: logger.error("Caught exception in thread: '%s'", - thread._name) # pylint: disable=protected-access + thread._name) # pylint:disable=protected-access raise thread.err[1].with_traceback(thread.err[2]) + del self._threads + self._threads = [] logger.debug("Joined all Threads: '%s'", self._name) class BackgroundGenerator(MultiThread): - """ Run a queue in the background. From: - https://stackoverflow.com/questions/7323664/ """ - # See below why prefetch count is flawed - def __init__(self, generator, prefetch=1, thread_count=2, - queue=None, args=None, kwargs=None): - # pylint:disable=too-many-arguments - super().__init__(target=self._run, thread_count=thread_count) - self.queue = queue or Queue.Queue(prefetch) + """ Run a task in the background background and queue data for consumption + + Parameters + ---------- + generator: iterable + The generator to run in the background + prefetch, int, optional + The number of items to pre-fetch from the generator before blocking (see Notes). Default: 1 + name: str, optional + The thread name. if ``None`` a unique name is constructed of the form + {generator.__name__}_N where N is an incrementing integer. Default: ``None`` + args: tuple, Optional + The argument tuple for generator invocation. Default: ``None``. + kwargs: dict, Optional + keyword arguments for the generator invocation. Default: ``None``. + + Notes + ----- + Putting to the internal queue only blocks if put is called while queue has already + reached max size. Therefore this means prefetch is actually 1 more than the parameter + supplied (N in the queue, one waiting for insertion) + + References + ---------- + https://stackoverflow.com/questions/7323664/ + """ + def __init__(self, + generator: Callable, + prefetch: int = 1, + name: str | None = None, + args: tuple | None = None, + kwargs: dict[str, T.Any] | None = None) -> None: + super().__init__(name=name, target=self._run) + self.queue: Queue.Queue = Queue.Queue(prefetch) self.generator = generator self._gen_args = args or tuple() - self._gen_kwargs = kwargs or dict() + self._gen_kwargs = kwargs or {} self.start() - def _run(self): - """ Put until queue size is reached. - Note: put blocks only if put is called while queue has already - reached max size => this makes prefetch + thread_count prefetched items! - N in the the queue, one waiting for insertion per thread! """ + def _run(self) -> None: + """ Run the :attr:`_generator` and put into the queue until until queue size is reached. + + Raises + ------ + Exception + If there is a failure to run the generator and put to the queue + """ try: for item in self.generator(*self._gen_args, **self._gen_kwargs): self.queue.put(item) @@ -149,8 +283,14 @@ def _run(self): self.queue.put(None) raise - def iterator(self): - """ Iterate items out of the queue """ + def iterator(self) -> Generator: + """ Iterate items out of the queue + + Yields + ------ + Any + The items from the generator + """ while True: next_item = self.queue.get() self.check_and_raise_error() diff --git a/lib/plaidml_utils.py b/lib/plaidml_utils.py deleted file mode 100644 index 706a9f6fbf..0000000000 --- a/lib/plaidml_utils.py +++ /dev/null @@ -1,32 +0,0 @@ -''' -Multiple plaidml implementation. -''' - -import plaidml - - -def pad(data, paddings, mode="CONSTANT", name=None, constant_value=0): - """ PlaidML Pad """ - # TODO: use / implement other padding method when required - # CONSTANT -> SpatialPadding ? | Doesn't support first and last axis + - # no support for constant_value - # SYMMETRIC -> Requires implement ? - if mode.upper() != "REFLECT": - raise NotImplementedError("pad only supports mode == 'REFLECT'") - if constant_value != 0: - raise NotImplementedError("pad does not support constant_value != 0") - return plaidml.op.reflection_padding(data, paddings) - - -def is_plaidml_error(error): - """ Test whether the given exception is a plaidml Exception. - - error: :class:`Exception` - The generated error - - Returns - ------- - bool - ``True`` if the given error has been generated from plaidML otherwise ``False`` - """ - return isinstance(error, plaidml.exceptions.PlaidMLError) diff --git a/lib/queue_manager.py b/lib/queue_manager.py index f9fb83e63b..9fd5122aa4 100644 --- a/lib/queue_manager.py +++ b/lib/queue_manager.py @@ -7,28 +7,48 @@ import logging import threading -from queue import Queue, Empty as QueueEmpty # pylint: disable=unused-import; # noqa +from queue import Queue, Empty as QueueEmpty # pylint:disable=unused-import; # noqa from time import sleep -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -class QueueManager(): - """ Manage queues for availabilty across processes - Don't import this class directly, instead - import the variable: queue_manager """ - def __init__(self): +class EventQueue(Queue): + """ Standard Queue object with a separate global shutdown parameter indicating that the main + process, and by extension this queue, should be shut down. + + Parameters + ---------- + shutdown_event: :class:`threading.Event` + The global shutdown event common to all managed queues + maxsize: int, Optional + Upperbound limit on the number of items that can be placed in the queue. Default: `0` + """ + def __init__(self, shutdown_event: threading.Event, maxsize: int = 0) -> None: + super().__init__(maxsize=maxsize) + self._shutdown = shutdown_event + + @property + def shutdown(self) -> threading.Event: + """ :class:`threading.Event`: The global shutdown event """ + return self._shutdown + + +class _QueueManager(): + """ Manage :class:`EventQueue` objects for availabilty across processes. + + Notes + ----- + Don't import this class directly, instead import via :func:`queue_manager` """ + def __init__(self) -> None: logger.debug("Initializing %s", self.__class__.__name__) self.shutdown = threading.Event() - self.queues = dict() + self.queues: dict[str, EventQueue] = {} logger.debug("Initialized %s", self.__class__.__name__) - def add_queue(self, name, maxsize=0, create_new=False): - """ Add a queue to the manager. - - Adds an event "shutdown" to the queue that can be used to indicate to a process that any - activity on the queue should cease. + def add_queue(self, name: str, maxsize: int = 0, create_new: bool = False) -> str: + """ Add a :class:`EventQueue` to the manager. Parameters ---------- @@ -50,77 +70,108 @@ def add_queue(self, name, maxsize=0, create_new=False): logger.debug("QueueManager adding: (name: '%s', maxsize: %s, create_new: %s)", name, maxsize, create_new) if not create_new and name in self.queues: - raise ValueError("Queue '{}' already exists.".format(name)) + raise ValueError(f"Queue '{name}' already exists.") if create_new and name in self.queues: i = 0 while name in self.queues: name = f"{name}{i}" logger.debug("Duplicate queue name. Updated to: '%s'", name) - queue = Queue(maxsize=maxsize) - - setattr(queue, "shutdown", self.shutdown) - self.queues[name] = queue + self.queues[name] = EventQueue(self.shutdown, maxsize=maxsize) logger.debug("QueueManager added: (name: '%s')", name) return name - def del_queue(self, name): - """ remove a queue from the manager """ + def del_queue(self, name: str) -> None: + """ Remove a queue from the manager + + Parameters + ---------- + name: str + The name of the queue to be deleted. Must exist within the queue manager. + """ logger.debug("QueueManager deleting: '%s'", name) del self.queues[name] logger.debug("QueueManager deleted: '%s'", name) - def get_queue(self, name, maxsize=0): - """ Return a queue from the manager - If it doesn't exist, create it """ + def get_queue(self, name: str, maxsize: int = 0) -> EventQueue: + """ Return a :class:`EventQueue` from the manager. If it doesn't exist, create it. + + Parameters + ---------- + name: str + The name of the queue to obtain + maxsize: int, Optional + The maximum queue size. Set to `0` for unlimited. Only used if the requested queue + does not already exist. Default: `0` + """ logger.debug("QueueManager getting: '%s'", name) - queue = self.queues.get(name, None) + queue = self.queues.get(name) if not queue: self.add_queue(name, maxsize) queue = self.queues[name] logger.debug("QueueManager got: '%s'", name) return queue - def terminate_queues(self): - """ Set shutdown event, clear and send EOF to all queues - To be called if there is an error """ + def terminate_queues(self) -> None: + """ Terminates all managed queues. + + Sets the global shutdown event, clears and send EOF to all queues. To be called if there + is an error """ logger.debug("QueueManager terminating all queues") self.shutdown.set() - self.flush_queues() + self._flush_queues() for q_name, queue in self.queues.items(): logger.debug("QueueManager terminating: '%s'", q_name) queue.put("EOF") logger.debug("QueueManager terminated all queues") - def flush_queues(self): - """ Empty out all queues """ + def _flush_queues(self): + """ Empty out the contents of every managed queue. """ for q_name in self.queues: self.flush_queue(q_name) logger.debug("QueueManager flushed all queues") - def flush_queue(self, q_name): - """ Empty out a specific queue """ - logger.debug("QueueManager flushing: '%s'", q_name) - queue = self.queues[q_name] + def flush_queue(self, name: str) -> None: + """ Flush the contents from a managed queue. + + Parameters + ---------- + name: str + The name of the managed :class:`EventQueue` to flush + """ + logger.debug("QueueManager flushing: '%s'", name) + queue = self.queues[name] while not queue.empty(): queue.get(True, 1) - def debug_monitor(self, update_secs=2): - """ Debug tool for monitoring queues """ - thread = threading.Thread(target=self.debug_queue_sizes, - args=(update_secs, )) + def debug_monitor(self, update_interval: int = 2) -> None: + """ A debug tool for monitoring managed :class:`EventQueues`. + + Prints queue sizes to the console for all managed queues. + + Parameters + ---------- + update_interval: int, Optional + The number of seconds between printing information to the console. Default: 2 + """ + thread = threading.Thread(target=self._debug_queue_sizes, + args=(update_interval, )) thread.daemon = True thread.start() - def debug_queue_sizes(self, update_secs): - """ Output the queue sizes - logged to INFO so it also displays in console + def _debug_queue_sizes(self, update_interval) -> None: + """ Print the queue size for each managed queue to console. + + Parameters + ---------- + update_interval: int + The number of seconds between printing information to the console """ while True: logger.info("====================================================") for name in sorted(self.queues.keys()): logger.info("%s: %s", name, self.queues[name].qsize()) - sleep(update_secs) + sleep(update_interval) -queue_manager = QueueManager() # pylint: disable=invalid-name +queue_manager = _QueueManager() # pylint:disable=invalid-name diff --git a/lib/serializer.py b/lib/serializer.py index db6b85d94d..ab48ec129f 100644 --- a/lib/serializer.py +++ b/lib/serializer.py @@ -17,10 +17,11 @@ try: import yaml + _HAS_YAML = True except ImportError: - yaml = None + _HAS_YAML = False -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Serializer(): @@ -72,13 +73,13 @@ def save(self, filename, data): with open(filename, self._write_option) as s_file: s_file.write(self.marshal(data)) except IOError as err: - msg = "Error writing to '{}': {}".format(filename, err.strerror) + msg = f"Error writing to '{filename}': {err.strerror}" raise FaceswapError(msg) from err def _check_extension(self, filename): """ Check the filename has an extension. If not add the correct one for the serializer """ extension = os.path.splitext(filename)[1] - retval = filename if extension else "{}.{}".format(filename, self.file_extension) + retval = filename if extension else f"{filename}.{self.file_extension}" logger.debug("Original filename: '%s', final filename: '%s'", filename, retval) return retval @@ -109,7 +110,7 @@ def load(self, filename): retval = self.unmarshal(data) except IOError as err: - msg = "Error reading from '{}': {}".format(filename, err.strerror) + msg = f"Error reading from '{filename}': {err.strerror}" raise FaceswapError(msg) from err logger.debug("data type: %s", type(retval)) return retval @@ -137,7 +138,7 @@ def marshal(self, data): try: retval = self._marshal(data) except Exception as err: - msg = "Error serializing data for type {}: {}".format(type(data), str(err)) + msg = f"Error serializing data for type {type(data)}: {str(err)}" raise FaceswapError(msg) from err logger.debug("returned data type: %s", type(retval)) return retval @@ -165,19 +166,16 @@ def unmarshal(self, serialized_data): try: retval = self._unmarshal(serialized_data) except Exception as err: - msg = "Error unserializing data for type {}: {}".format(type(serialized_data), - str(err)) + msg = f"Error unserializing data for type {type(serialized_data)}: {str(err)}" raise FaceswapError(msg) from err logger.debug("returned data type: %s", type(retval)) return retval - @classmethod - def _marshal(cls, data): + def _marshal(self, data): """ Override for serializer specific marshalling """ raise NotImplementedError() - @classmethod - def _unmarshal(cls, data): + def _unmarshal(self, data): """ Override for serializer specific unmarshalling """ raise NotImplementedError() @@ -188,13 +186,11 @@ def __init__(self): super().__init__() self._file_extension = "yml" - @classmethod - def _marshal(cls, data): + def _marshal(self, data): return yaml.dump(data, default_flow_style=False).encode("utf-8") - @classmethod - def _unmarshal(cls, data): - return yaml.load(data.decode("utf-8"), Loader=yaml.FullLoader) + def _unmarshal(self, data): + return yaml.load(data.decode("utf-8", errors="replace"), Loader=yaml.FullLoader) class _JSONSerializer(Serializer): @@ -203,13 +199,11 @@ def __init__(self): super().__init__() self._file_extension = "json" - @classmethod - def _marshal(cls, data): + def _marshal(self, data): return json.dumps(data, indent=2).encode("utf-8") - @classmethod - def _unmarshal(cls, data): - return json.loads(data.decode("utf-8")) + def _unmarshal(self, data): + return json.loads(data.decode("utf-8", errors="replace")) class _PickleSerializer(Serializer): @@ -218,12 +212,10 @@ def __init__(self): super().__init__() self._file_extension = "pickle" - @classmethod - def _marshal(cls, data): + def _marshal(self, data): return pickle.dumps(data) - @classmethod - def _unmarshal(cls, data): + def _unmarshal(self, data): return pickle.loads(data) @@ -260,13 +252,13 @@ def __init__(self): def _marshal(self, data): """ Pickle and compress data """ - data = self._child._marshal(data) # pylint: disable=protected-access + data = self._child._marshal(data) # pylint:disable=protected-access return zlib.compress(data) def _unmarshal(self, data): """ Decompress and unpicke data """ data = zlib.decompress(data) - return self._child._unmarshal(data) # pylint: disable=protected-access + return self._child._unmarshal(data) # pylint:disable=protected-access def get_serializer(serializer): @@ -294,9 +286,9 @@ def get_serializer(serializer): retval = _JSONSerializer() elif serializer.lower() == "pickle": retval = _PickleSerializer() - elif serializer.lower() == "yaml" and yaml is not None: + elif serializer.lower() == "yaml" and _HAS_YAML: retval = _YAMLSerializer() - elif serializer.lower() == "yaml" and yaml is None: + elif serializer.lower() == "yaml": logger.warning("You must have PyYAML installed to use YAML as the serializer." "Switching to JSON as the serializer.") retval = _JSONSerializer @@ -336,9 +328,9 @@ def get_serializer_from_filename(filename): retval = _NPYSerializer() elif extension == ".fsa": retval = _CompressedSerializer() - elif extension in (".yaml", ".yml") and yaml is not None: + elif extension in (".yaml", ".yml") and _HAS_YAML: retval = _YAMLSerializer() - elif extension in (".yaml", ".yml") and yaml is None: + elif extension in (".yaml", ".yml"): logger.warning("You must have PyYAML installed to use YAML as the serializer.\n" "Switching to JSON as the serializer.") retval = _JSONSerializer() diff --git a/lib/sysinfo.py b/lib/sysinfo.py index e0ef06fbfe..7eda070361 100644 --- a/lib/sysinfo.py +++ b/lib/sysinfo.py @@ -6,147 +6,136 @@ import os import platform import sys + from subprocess import PIPE, Popen import psutil -from lib.gpu_stats import GPUStats +from lib.git import git +from lib.gpu_stats import GPUStats, GPUInfo +from lib.utils import get_backend from setup import CudaCheck -class _SysInfo(): # pylint:disable=too-few-public-methods +class _SysInfo(): """ Obtain information about the System, Python and GPU """ - def __init__(self): + def __init__(self) -> None: self._state_file = _State().state_file self._configs = _Configs().configs - self._system = dict(platform=platform.platform(), - system=platform.system(), - machine=platform.machine(), - release=platform.release(), - processor=platform.processor(), - cpu_count=os.cpu_count()) - self._python = dict(implementation=platform.python_implementation(), - version=platform.python_version()) - self._gpu = GPUStats(log=False).sys_info + self._system = {"platform": platform.platform(), + "system": platform.system().lower(), + "machine": platform.machine(), + "release": platform.release(), + "processor": platform.processor(), + "cpu_count": os.cpu_count()} + self._python = {"implementation": platform.python_implementation(), + "version": platform.python_version()} + self._gpu = self._get_gpu_info() self._cuda_check = CudaCheck() @property - def _encoding(self): + def _encoding(self) -> str: """ str: The system preferred encoding """ return locale.getpreferredencoding() @property - def _is_conda(self): + def _is_conda(self) -> bool: """ bool: `True` if running in a Conda environment otherwise ``False``. """ return ("conda" in sys.version.lower() or os.path.exists(os.path.join(sys.prefix, 'conda-meta'))) @property - def _is_linux(self): + def _is_linux(self) -> bool: """ bool: `True` if running on a Linux system otherwise ``False``. """ - return self._system["system"].lower() == "linux" + return self._system["system"] == "linux" @property - def _is_macos(self): + def _is_macos(self) -> bool: """ bool: `True` if running on a macOS system otherwise ``False``. """ - return self._system["system"].lower() == "darwin" + return self._system["system"] == "darwin" @property - def _is_windows(self): + def _is_windows(self) -> bool: """ bool: `True` if running on a Windows system otherwise ``False``. """ - return self._system["system"].lower() == "windows" + return self._system["system"] == "windows" @property - def _is_virtual_env(self): + def _is_virtual_env(self) -> bool: """ bool: `True` if running inside a virtual environment otherwise ``False``. """ if not self._is_conda: retval = (hasattr(sys, "real_prefix") or (hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix)) else: prefix = os.path.dirname(sys.prefix) - retval = (os.path.basename(prefix) == "envs") + retval = os.path.basename(prefix) == "envs" return retval @property - def _ram_free(self): + def _ram_free(self) -> int: """ int: The amount of free RAM in bytes. """ return psutil.virtual_memory().free @property - def _ram_total(self): + def _ram_total(self) -> int: """ int: The amount of total RAM in bytes. """ return psutil.virtual_memory().total @property - def _ram_available(self): + def _ram_available(self) -> int: """ int: The amount of available RAM in bytes. """ return psutil.virtual_memory().available @property - def _ram_used(self): + def _ram_used(self) -> int: """ int: The amount of used RAM in bytes. """ return psutil.virtual_memory().used @property - def _fs_command(self): + def _fs_command(self) -> str: """ str: The command line command used to execute faceswap. """ return " ".join(sys.argv) @property - def _installed_pip(self): + def _installed_pip(self) -> str: """ str: The list of installed pip packages within Faceswap's scope. """ - pip = Popen("{} -m pip freeze".format(sys.executable), - shell=True, stdout=PIPE) - installed = pip.communicate()[0].decode().splitlines() + with Popen(f"{sys.executable} -m pip freeze", shell=True, stdout=PIPE) as pip: + installed = pip.communicate()[0].decode(self._encoding, errors="replace").splitlines() return "\n".join(installed) @property - def _installed_conda(self): + def _installed_conda(self) -> str: """ str: The list of installed Conda packages within Faceswap's scope. """ if not self._is_conda: - return None - conda = Popen("conda list", shell=True, stdout=PIPE, stderr=PIPE) - stdout, stderr = conda.communicate() + return "" + with Popen("conda list", shell=True, stdout=PIPE, stderr=PIPE) as conda: + stdout, stderr = conda.communicate() if stderr: return "Could not get package list" - installed = stdout.decode().splitlines() + installed = stdout.decode(self._encoding, errors="replace").splitlines() return "\n".join(installed) @property - def _conda_version(self): + def _conda_version(self) -> str: """ str: The installed version of Conda, or `N/A` if Conda is not installed. """ if not self._is_conda: return "N/A" - conda = Popen("conda --version", shell=True, stdout=PIPE, stderr=PIPE) - stdout, stderr = conda.communicate() + with Popen("conda --version", shell=True, stdout=PIPE, stderr=PIPE) as conda: + stdout, stderr = conda.communicate() if stderr: return "Conda is used, but version not found" - version = stdout.decode().splitlines() + version = stdout.decode(self._encoding, errors="replace").splitlines() return "\n".join(version) @property - def _git_branch(self): - """ str: The git branch that is currently being used to execute Faceswap. """ - git = Popen("git status", shell=True, stdout=PIPE, stderr=PIPE) - stdout, stderr = git.communicate() - if stderr: - return "Not Found" - branch = stdout.decode().splitlines()[0].replace("On branch ", "") - return branch - - @property - def _git_commits(self): + def _git_commits(self) -> str: """ str: The last 5 git commits for the currently running Faceswap. """ - git = Popen("git log --pretty=oneline --abbrev-commit -n 5", - shell=True, stdout=PIPE, stderr=PIPE) - stdout, stderr = git.communicate() - if stderr: + commits = git.get_commits(3) + if not commits: return "Not Found" - commits = stdout.decode().splitlines() - return ". ".join(commits) + return " | ".join(commits) @property - def _cuda_version(self): + def _cuda_version(self) -> str: """ str: The installed CUDA version. """ # TODO Handle multiple CUDA installs retval = self._cuda_check.cuda_version @@ -157,7 +146,7 @@ def _cuda_version(self): return retval @property - def _cudnn_version(self): + def _cudnn_version(self) -> str: """ str: The installed cuDNN version. """ retval = self._cuda_check.cudnn_version if not retval: @@ -166,7 +155,26 @@ def _cudnn_version(self): retval += ". Check Conda packages for Conda cuDNN" return retval - def full_info(self): + def _get_gpu_info(self) -> GPUInfo: + """ Obtain GPU Stats. If an error is raised, swallow the error, and add to GPUInfo output + + Returns + ------- + :class:`~lib.gpu_stats.GPUInfo` + The information on connected GPUs + """ + try: + retval = GPUStats(log=False).sys_info + except Exception as err: # pylint:disable=broad-except + err_string = f"{type(err)}: {err}" + retval = GPUInfo(vram=[], + vram_free=[], + driver="N/A", + devices=[f"Error obtaining GPU Stats: '{err_string}'"], + devices_active=[]) + return retval + + def full_info(self) -> str: """ Obtain extensive system information stats, formatted into a human readable format. Returns @@ -176,7 +184,8 @@ def full_info(self): console or a log file. """ retval = "\n============ System Information ============\n" - sys_info = {"os_platform": self._system["platform"], + sys_info = {"backend": get_backend(), + "os_platform": self._system["platform"], "os_machine": self._system["machine"], "os_release": self._system["release"], "py_conda_version": self._conda_version, @@ -188,19 +197,21 @@ def full_info(self): "sys_processor": self._system["processor"], "sys_ram": self._format_ram(), "encoding": self._encoding, - "git_branch": self._git_branch, + "git_branch": git.branch, "git_commits": self._git_commits, "gpu_cuda": self._cuda_version, "gpu_cudnn": self._cudnn_version, - "gpu_driver": self._gpu["driver"], - "gpu_devices": ", ".join(["GPU_{}: {}".format(idx, device) - for idx, device in enumerate(self._gpu["devices"])]), - "gpu_vram": ", ".join(["GPU_{}: {}MB".format(idx, int(vram)) - for idx, vram in enumerate(self._gpu["vram"])]), - "gpu_devices_active": ", ".join(["GPU_{}".format(idx) - for idx in self._gpu["devices_active"]])} + "gpu_driver": self._gpu.driver, + "gpu_devices": ", ".join([f"GPU_{idx}: {device}" + for idx, device in enumerate(self._gpu.devices)]), + "gpu_vram": ", ".join( + f"GPU_{idx}: {int(vram)}MB ({int(vram_free)}MB free)" + for idx, (vram, vram_free) in enumerate(zip(self._gpu.vram, + self._gpu.vram_free))), + "gpu_devices_active": ", ".join([f"GPU_{idx}" + for idx in self._gpu.devices_active])} for key in sorted(sys_info.keys()): - retval += ("{0: <20} {1}\n".format(key + ":", sys_info[key])) + retval += (f"{key + ':':<20} {sys_info[key]}\n") retval += "\n=============== Pip Packages ===============\n" retval += self._installed_pip if self._is_conda: @@ -211,7 +222,7 @@ def full_info(self): retval += self._configs return retval - def _format_ram(self): + def _format_ram(self) -> str: """ Format the RAM stats into Megabytes to make it more readable. Returns @@ -219,15 +230,15 @@ def _format_ram(self): str The total, available, used and free RAM displayed in Megabytes """ - retval = list() + retval = [] for name in ("total", "available", "used", "free"): - value = getattr(self, "_ram_{}".format(name)) + value = getattr(self, f"_ram_{name}") value = int(value / (1024 * 1024)) - retval.append("{}: {}MB".format(name.capitalize(), value)) + retval.append(f"{name.capitalize()}: {value}MB") return ", ".join(retval) -def get_sysinfo(): +def get_sysinfo() -> str: """ Obtain extensive system information stats, formatted into a human readable format. If an error occurs obtaining the system information, then the error message is returned instead. @@ -240,8 +251,9 @@ def get_sysinfo(): """ try: retval = _SysInfo().full_info() - except Exception as err: # pylint: disable=broad-except - retval = "Exception occured trying to retrieve sysinfo: {}".format(err) + except Exception as err: # pylint:disable=broad-except + retval = f"Exception occured trying to retrieve sysinfo: {str(err)}" + raise return retval @@ -249,11 +261,11 @@ class _Configs(): # pylint:disable=too-few-public-methods """ Parses the config files in /faceswap/config and outputs the information stored within them in a human readable format. """ - def __init__(self): + def __init__(self) -> None: self.config_dir = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), "config") self.configs = self._get_configs() - def _get_configs(self): + def _get_configs(self) -> str: """ Obtain the formatted configurations from the config folder. Returns @@ -261,13 +273,16 @@ def _get_configs(self): str The current configuration in the config files formatted in a human readable format """ - config_files = [os.path.join(self.config_dir, cfile) - for cfile in os.listdir(self.config_dir) - if os.path.basename(cfile) == ".faceswap" - or os.path.splitext(cfile)[1] == ".ini"] - return self._parse_configs(config_files) + try: + config_files = [os.path.join(self.config_dir, cfile) + for cfile in os.listdir(self.config_dir) + if os.path.basename(cfile) == ".faceswap" + or os.path.splitext(cfile)[1] == ".ini"] + return self._parse_configs(config_files) + except FileNotFoundError: + return "" - def _parse_configs(self, config_files): + def _parse_configs(self, config_files: list[str]) -> str: """ Parse the given list of config files into a human readable format. Parameters @@ -284,14 +299,14 @@ def _parse_configs(self, config_files): for cfile in config_files: fname = os.path.basename(cfile) ext = os.path.splitext(cfile)[1] - formatted += "\n--------- {} ---------\n".format(fname) + formatted += f"\n--------- {fname} ---------\n" if ext == ".ini": formatted += self._parse_ini(cfile) elif fname == ".faceswap": formatted += self._parse_json(cfile) return formatted - def _parse_ini(self, config_file): + def _parse_ini(self, config_file: str) -> str: """ Parse an ``.ini`` formatted config file into a human readable format. Parameters @@ -305,20 +320,20 @@ def _parse_ini(self, config_file): The current configuration in the config file formatted in a human readable format """ formatted = "" - with open(config_file, "r") as cfile: + with open(config_file, "r", encoding="utf-8", errors="replace") as cfile: for line in cfile.readlines(): line = line.strip() if line.startswith("#") or not line: continue item = line.split("=") if len(item) == 1: - formatted += "\n{}\n".format(item[0].strip()) + formatted += f"\n{item[0].strip()}\n" else: formatted += self._format_text(item[0], item[1]) return formatted - def _parse_json(self, config_file): - """ Parse an ``.json`` formatted config file into a python dictionary. + def _parse_json(self, config_file: str) -> str: + """ Parse an ``.json`` formatted config file into a formatted string. Parameters ---------- @@ -330,15 +345,15 @@ def _parse_json(self, config_file): dict The current configuration in the config file formatted as a python dictionary """ - formatted = "" - with open(config_file, "r") as cfile: + formatted: str = "" + with open(config_file, "r", encoding="utf-8", errors="replace") as cfile: conf_dict = json.load(cfile) for key in sorted(conf_dict.keys()): formatted += self._format_text(key, conf_dict[key]) return formatted @staticmethod - def _format_text(key, value): + def _format_text(key: str, value: str) -> str: """Format a key value pair into a consistently spaced string output for display. Parameters @@ -353,25 +368,25 @@ def _format_text(key, value): str The formatted key value pair for display """ - return "{0: <25} {1}\n".format(key.strip() + ":", value.strip()) + return f"{key.strip() + ':':<25} {value.strip()}\n" class _State(): # pylint:disable=too-few-public-methods """ Parses the state file in the current model directory, if the model is training, and formats the content into a human readable format. """ - def __init__(self): + def __init__(self) -> None: self._model_dir = self._get_arg("-m", "--model-dir") self._trainer = self._get_arg("-t", "--trainer") self.state_file = self._get_state_file() @property - def _is_training(self): + def _is_training(self) -> bool: """ bool: ``True`` if this function has been called during a training session otherwise ``False``. """ return len(sys.argv) > 1 and sys.argv[1].lower() == "train" @staticmethod - def _get_arg(*args): + def _get_arg(*args: str) -> str | None: """ Obtain the value for a given command line option from sys.argv. Returns @@ -385,7 +400,7 @@ def _get_arg(*args): return cmd[cmd.index(opt) + 1] return None - def _get_state_file(self): + def _get_state_file(self) -> str: """ Parses the model's state file and compiles the contents into a human readable string. Returns @@ -395,14 +410,14 @@ def _get_state_file(self): """ if not self._is_training or self._model_dir is None or self._trainer is None: return "" - fname = os.path.join(self._model_dir, "{}_state.json".format(self._trainer)) + fname = os.path.join(self._model_dir, f"{self._trainer}_state.json") if not os.path.isfile(fname): return "" retval = "\n\n=============== State File =================\n" - with open(fname, "r") as sfile: + with open(fname, "r", encoding="utf-8", errors="replace") as sfile: retval += sfile.read() return retval -sysinfo = get_sysinfo() # pylint: disable=invalid-name +sysinfo = get_sysinfo() # pylint:disable=invalid-name diff --git a/lib/training/__init__.py b/lib/training/__init__.py index a478362969..0e30c90340 100644 --- a/lib/training/__init__.py +++ b/lib/training/__init__.py @@ -1,6 +1,19 @@ #!/usr/bin/env python3 """ Package for handling alignments files, detected faces and aligned faces along with their associated objects. """ +from __future__ import annotations +import typing as T from .augmentation import ImageAugmentation -from .generator import TrainingDataGenerator +from .generator import Feeder +from .lr_finder import LearningRateFinder +from .preview_cv import PreviewBuffer, TriggerType + +if T.TYPE_CHECKING: + from .preview_cv import PreviewBase + Preview: type[PreviewBase] + +try: + from .preview_tk import PreviewTk as Preview +except ImportError: + from .preview_cv import PreviewCV as Preview diff --git a/lib/training/augmentation.py b/lib/training/augmentation.py index 5b0eff2309..8fd911969b 100644 --- a/lib/training/augmentation.py +++ b/lib/training/augmentation.py @@ -1,208 +1,207 @@ #!/usr/bin/env python3 """ Processes the augmentation of images for feeding into a Faceswap model. """ +from __future__ import annotations import logging +import typing as T import cv2 +import numexpr as ne import numpy as np from scipy.interpolate import griddata from lib.image import batch_convert_color +from lib.logger import parse_class_init -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from lib.config import ConfigValueType +logger = logging.getLogger(__name__) -class ImageAugmentation(): - """ Performs augmentation on batches of training images. - Parameters - ---------- - batchsize: int - The number of images that will be fed through the augmentation functions at once. - is_display: bool - Whether the images being fed through will be used for Preview or Time-lapse. Disables - the "warp" augmentation for these images. - input_size: int - The expected input size for the model. It is assumed that the input to the model is always - a square image. This is the size, in pixels, of the `width` and the `height` of the input - to the model. - output_shapes: list - A list of tuples defining the output shapes from the model, in the order that the outputs - are returned. The tuples should be in (`height`, `width`, `channels`) format. - coverage_ratio: float - The ratio of the training image to be trained on. Dictates how much of the image will be - cropped out. E.G: a coverage ratio of 0.625 will result in cropping a 160px box from a - 256px image (:math:`256 * 0.625 = 160`) - config: dict - The configuration `dict` generated from :file:`config.train.ini` containing the trainer - plugin configuration options. +class AugConstants: # pylint:disable=too-many-instance-attributes,too-few-public-methods + """ Dataclass for holding constants for Image Augmentation. - Attributes + Parameters ---------- - initialized: bool - Flag to indicate whether :class:`ImageAugmentation` has been initialized with the training - image size in order to cache certain augmentation operations (see :func:`initialize`) - is_display: bool - Flag to indicate whether these augmentations are for time-lapses/preview images (``True``) - or standard training data (``False``) + config: dict[str, ConfigValueType] + The user training configuration options + processing_size: int: + The size of image to augment the data for + batch_size: int + The batch size that augmented data is being prepared for """ - def __init__(self, batchsize, is_display, input_size, output_shapes, coverage_ratio, config): - logger.debug("Initializing %s: (batchsize: %s, is_display: %s, input_size: %s, " - "output_shapes: %s, coverage_ratio: %s, config: %s)", - self.__class__.__name__, batchsize, is_display, input_size, output_shapes, - coverage_ratio, config) + def __init__(self, + config: dict[str, ConfigValueType], + processing_size: int, + batch_size: int) -> None: + logger.debug(parse_class_init(locals())) + self.clahe_base_contrast: int = 0 + """int: The base number for Contrast Limited Adaptive Histogram Equalization""" + self.clahe_chance: float = 0.0 + """float: Probability to perform Contrast Limited Adaptive Histogram Equilization""" + self.clahe_max_size: int = 0 + """int: Maximum clahe window size""" + + self.lab_adjust: np.ndarray + """:class:`numpy.ndarray`: Adjustment amounts for L*A*B augmentation""" + self.transform_rotation: int = 0 + """int: Rotation range for transformations""" + self.transform_zoom: float = 0.0 + """float: Zoom range for transformations""" + self.transform_shift: float = 0.0 + """float: Shift range for transformations""" + self.warp_maps: np.ndarray + """:class:`numpy.ndarray`The stacked (x, y) mappings for image warping""" + self.warp_pad: tuple[int, int] = (0, 0) + """:tuple[int, int]: The padding to apply for image warping""" + self.warp_slices: slice + """:slice: The slices for extracting a warped image""" + self.warp_lm_edge_anchors: np.ndarray + """::class:`numpy.ndarray`: The edge anchors for landmark based warping""" + self.warp_lm_grids: np.ndarray + """::class:`numpy.ndarray`: The grids for landmark based warping""" - self.initialized = False - self.is_display = is_display + self._config = config + self._size = processing_size + self._load_config(batch_size) + logger.debug("Initialized: %s", self.__class__.__name__) + + def _load_clahe(self) -> None: + """ Load the CLAHE constants from user config """ + color_clahe_chance = self._config.get("color_clahe_chance", 50) + color_clahe_max_size = self._config.get("color_clahe_max_size", 4) + assert isinstance(color_clahe_chance, int) + assert isinstance(color_clahe_max_size, int) + + self.clahe_base_contrast = max(2, self._size // 128) + self.clahe_chance = color_clahe_chance / 100 + self.clahe_max_size = color_clahe_max_size + logger.debug("clahe_base_contrast: %s, clahe_chance: %s, clahe_max_size: %s", + self.clahe_base_contrast, self.clahe_chance, self.clahe_max_size) + + def _load_lab(self) -> None: + """ Load the random L*A*B augmentation constants """ + color_lightness = self._config.get("color_lightness", 30) + color_ab = self._config.get("color_ab", 8) + assert isinstance(color_lightness, int) + assert isinstance(color_ab, int) + + amount_l = int(color_lightness) / 100 + amount_ab = int(color_ab) / 100 + + self.lab_adjust = np.array([amount_l, amount_ab, amount_ab], dtype="float32") + logger.debug("lab_adjust: %s", self.lab_adjust) + + def _load_transform(self) -> None: + """ Load the random transform constants """ + shift_range = self._config.get("shift_range", 5) + rotation_range = self._config.get("rotation_range", 10) + zoom_amount = self._config.get("zoom_amount", 5) + assert isinstance(shift_range, int) + assert isinstance(rotation_range, int) + assert isinstance(zoom_amount, int) - # Set on first image load from initialize - self._training_size = 0 - self._constants = None + self.transform_shift = (shift_range / 100) * self._size + self.transform_rotation = rotation_range + self.transform_zoom = zoom_amount / 100 + logger.debug("transform_shift: %s, transform_rotation: %s, transform_zoom: %s", + self.transform_shift, self.transform_rotation, self.transform_zoom) - self._batchsize = batchsize - self._config = config - # Transform and Warp args - self._input_size = input_size - self._output_sizes = [shape[1] for shape in output_shapes if shape[2] == 3] - logger.debug("Output sizes: %s", self._output_sizes) - # Warp args - self._coverage_ratio = coverage_ratio - self._scale = 5 # Normal random variable scale + def _load_warp(self, batch_size: int) -> None: + """ Load the warp augmentation constants - logger.debug("Initialized %s", self.__class__.__name__) + Parameters + ---------- + batch_size: int + The batch size that augmented data is being prepared for + """ + warp_range = np.linspace(0, self._size, 5, dtype='float32') + warp_mapx = np.broadcast_to(warp_range, (batch_size, 5, 5)).astype("float32") + warp_mapy = np.broadcast_to(warp_mapx[0].T, (batch_size, 5, 5)).astype("float32") + warp_pad = int(1.25 * self._size) - def initialize(self, training_size): - """ Initializes the caching of constants for use in various image augmentations. + self.warp_maps = np.stack((warp_mapx, warp_mapy), axis=1) + self.warp_pad = (warp_pad, warp_pad) + self.warp_slices = slice(warp_pad // 10, -warp_pad // 10) + logger.debug("warp_maps: (%s, %s), warp_pad: %s, warp_slices: %s", + self.warp_maps.shape, self.warp_maps.dtype, + self.warp_pad, self.warp_slices) - The training image size is not known prior to loading the images from disk and commencing - training, so it cannot be set in the :func:`__init__` method. When the first training batch - is loaded this function should be called to initialize the class and perform various - calculations based on this input size to cache certain constants for image augmentation - calculations. + def _load_warp_to_landmarks(self, batch_size: int) -> None: + """ Load the warp-to-landmarks augmentation constants Parameters ---------- - training_size: int - The size of the training images stored on disk that are to be fed into - :class:`ImageAugmentation`. The training images should always be square and of the - same size. This is the size, in pixels, of the `width` and the `height` of the - training images. - """ - logger.debug("Initializing constants. training_size: %s", training_size) - self._training_size = training_size - coverage = int(self._training_size * self._coverage_ratio // 2) * 2 - - # Color Aug - clahe_base_contrast = training_size // 128 - # Target Images - tgt_slices = slice(self._training_size // 2 - coverage // 2, - self._training_size // 2 + coverage // 2) - - # Random Warp - warp_range_ = np.linspace(self._training_size // 2 - coverage // 2, - self._training_size // 2 + coverage // 2, 5, dtype='float32') - warp_mapx = np.broadcast_to(warp_range_, (self._batchsize, 5, 5)).astype("float32") - warp_mapy = np.broadcast_to(warp_mapx[0].T, (self._batchsize, 5, 5)).astype("float32") - - warp_pad = int(1.25 * self._input_size) - warp_slices = slice(warp_pad // 10, -warp_pad // 10) - - # Random Warp Landmarks - p_mx = self._training_size - 1 - p_hf = (self._training_size // 2) - 1 + batch_size: int + The batch size that augmented data is being prepared for + """ + p_mx = self._size - 1 + p_hf = (self._size // 2) - 1 edge_anchors = np.array([(0, 0), (0, p_mx), (p_mx, p_mx), (p_mx, 0), (p_hf, 0), (p_hf, p_mx), (p_mx, p_hf), (0, p_hf)]).astype("int32") - edge_anchors = np.broadcast_to(edge_anchors, (self._batchsize, 8, 2)) - grids = np.mgrid[0:p_mx:complex(self._training_size), 0:p_mx:complex(self._training_size)] - - self._constants = dict(clahe_base_contrast=clahe_base_contrast, - tgt_slices=tgt_slices, - warp_mapx=warp_mapx, - warp_mapy=warp_mapy, - warp_pad=warp_pad, - warp_slices=warp_slices, - warp_lm_edge_anchors=edge_anchors, - warp_lm_grids=grids) - self.initialized = True - logger.debug("Initialized constants: %s", {k: str(v) if isinstance(v, np.ndarray) else v - for k, v in self._constants.items()}) - - # <<< TARGET IMAGES >>> # - def get_targets(self, batch): - """ Returns the target images, and masks, if required. + edge_anchors = np.broadcast_to(edge_anchors, (batch_size, 8, 2)) + grids = np.mgrid[0: p_mx: complex(self._size), # type:ignore[misc] + 0: p_mx: complex(self._size)] # type:ignore[misc] + + self.warp_lm_edge_anchors = edge_anchors + self.warp_lm_grids = grids + logger.debug("warp_lm_edge_anchors: (%s, %s), warp_lm_grids: (%s, %s)", + self.warp_lm_edge_anchors.shape, self.warp_lm_edge_anchors.dtype, + self.warp_lm_grids.shape, self.warp_lm_grids.dtype) + + def _load_config(self, batch_size: int) -> None: + """ Load the constants into the class from user config Parameters ---------- - batch: :class:`numpy.ndarray` - This should be a 4+-dimensional array of training images in the format (`batchsize`, - `height`, `width`, `channels`). Targets should be requested after performing image - transformations but prior to performing warps. + batch_size: int + The batch size that augmented data is being prepared for + """ + logger.debug("Loading augmentation constants") + self._load_clahe() + self._load_lab() + self._load_transform() + self._load_warp(batch_size) + self._load_warp_to_landmarks(batch_size) + logger.debug("Loaded augmentation constants") - The 4th channel should be the mask. Any channels above the 4th should be any additional - masks that are requested. - Returns - ------- - dict - The following keys will be within the returned dictionary: - - * **targets** (`list`) - A list of 4-dimensional :class:`numpy.ndarray` s in the \ - order and size of each output of the model as defined in :attr:`output_shapes`. The \ - format of these arrays will be (`batchsize`, `height`, `width`, `3`). **NB:** \ - masks are not included in the `targets` list. If masks are to be included in the \ - output they will be returned as their own item from the `masks` key. - - * **masks** (:class:`numpy.ndarray`) - A 4-dimensional array containing the target \ - masks in the format (`batchsize`, `height`, `width`, `1`). - """ - logger.trace("Compiling targets: batch shape: %s", batch.shape) - slices = self._constants["tgt_slices"] - target_batch = [np.array([cv2.resize(image[slices, slices, :], - (size, size), - cv2.INTER_AREA) - for image in batch], dtype='float32') / 255. - for size in self._output_sizes] - logger.trace("Target image shapes: %s", - [tgt_images.shape for tgt_images in target_batch]) - - retval = self._separate_target_mask(target_batch) - logger.trace("Final targets: %s", - {k: v.shape if isinstance(v, np.ndarray) else [img.shape for img in v] - for k, v in retval.items()}) - return retval - - @staticmethod - def _separate_target_mask(target_batch): - """ Return the batch and the batch of final masks +class ImageAugmentation(): + """ Performs augmentation on batches of training images. - Parameters - ---------- - target_batch: list - List of 4 dimension :class:`numpy.ndarray` objects resized the model outputs. - The 4th channel of the array contains the face mask, any additional channels after - this are additional masks (e.g. eye mask and mouth mask) + Parameters + ---------- + batch_size: int + The number of images that will be fed through the augmentation functions at once. + processing_size: int + The largest input or output size of the model. This is the size that images are processed + at. + config: dict + The configuration `dict` generated from :file:`config.train.ini` containing the trainer + plugin configuration options. + """ + def __init__(self, + batch_size: int, + processing_size: int, + config: dict[str, ConfigValueType]) -> None: + logger.debug(parse_class_init(locals())) + self._processing_size = processing_size + self._batch_size = batch_size + + # flip_args + flip_chance = config.get("random_flip", 50) + assert isinstance(flip_chance, int) + self._flip_chance = flip_chance - Returns - ------- - dict: - The targets and the masks separated into their own items. The targets are a list of - 3 channel, 4 dimensional :class:`numpy.ndarray` objects sized for each output from the - model. The masks are a :class:`numpy.ndarray` of the final output size. Any additional - masks(e.g. eye and mouth masks) will be collated together into a :class:`numpy.ndarray` - of the final output size. The number of channels will be the number of additional - masks available - """ - logger.trace("target_batch shapes: %s", [tgt.shape for tgt in target_batch]) - retval = dict(targets=[batch[..., :3] for batch in target_batch], - masks=target_batch[-1][..., 3][..., None]) - if target_batch[-1].shape[-1] > 4: - retval["additional_masks"] = target_batch[-1][..., 4:] - logger.trace("returning: %s", {k: v.shape if isinstance(v, np.ndarray) else [tgt.shape - for tgt in v] - for k, v in retval.items()}) - return retval + # Warp args + self._warp_scale = 5 / 256 * self._processing_size # Normal random variable scale + self._warp_lm_scale = 2 / 256 * self._processing_size # Normal random variable scale + + self._constants = AugConstants(config, processing_size, batch_size) + logger.debug("Initialized %s", self.__class__.__name__) # <<< COLOR AUGMENTATION >>> # - def color_adjust(self, batch): + def color_adjust(self, batch: np.ndarray) -> np.ndarray: """ Perform color augmentation on the passed in batch. The color adjustment parameters are set in :file:`config.train.ini` @@ -219,49 +218,44 @@ def color_adjust(self, batch): A 4-dimensional array of the same shape as :attr:`batch` with color augmentation applied. """ - if not self.is_display: - logger.trace("Augmenting color") - batch = batch_convert_color(batch, "BGR2LAB") - batch = self._random_clahe(batch) - batch = self._random_lab(batch) - batch = batch_convert_color(batch, "LAB2BGR") + logger.trace("Augmenting color") # type:ignore[attr-defined] + batch = batch_convert_color(batch, "BGR2LAB") + self._random_lab(batch) + self._random_clahe(batch) + batch = batch_convert_color(batch, "LAB2BGR") return batch - def _random_clahe(self, batch): + def _random_clahe(self, batch: np.ndarray) -> None: """ Randomly perform Contrast Limited Adaptive Histogram Equalization on a batch of images """ - base_contrast = self._constants["clahe_base_contrast"] + base_contrast = self._constants.clahe_base_contrast - batch_random = np.random.rand(self._batchsize) - indices = np.where(batch_random < self._config.get("color_clahe_chance", 50) / 100)[0] + batch_random = np.random.rand(self._batch_size) + indices = np.where(batch_random < self._constants.clahe_chance)[0] if not np.any(indices): - return batch - - grid_bases = np.rint(np.random.uniform(0, - self._config.get("color_clahe_max_size", 4), - size=indices.shape[0])).astype("uint8") - contrast_adjustment = (grid_bases * (base_contrast // 2)) - grid_sizes = contrast_adjustment + base_contrast - logger.trace("Adjusting Contrast. Grid Sizes: %s", grid_sizes) - - clahes = [cv2.createCLAHE(clipLimit=2.0, # pylint: disable=no-member + return + grid_bases = np.random.randint(self._constants.clahe_max_size + 1, + size=indices.shape[0], + dtype="uint8") + grid_sizes = (grid_bases * (base_contrast // 2)) + base_contrast + logger.trace("Adjusting Contrast. Grid Sizes: %s", grid_sizes) # type:ignore[attr-defined] + + clahes = [cv2.createCLAHE(clipLimit=2.0, tileGridSize=(grid_size, grid_size)) for grid_size in grid_sizes] for idx, clahe in zip(indices, clahes): - batch[idx, :, :, 0] = clahe.apply(batch[idx, :, :, 0]) - return batch + batch[idx, :, :, 0] = clahe.apply(batch[idx, :, :, 0], ) - def _random_lab(self, batch): + def _random_lab(self, batch: np.ndarray) -> None: """ Perform random color/lightness adjustment in L*a*b* color space on a batch of images """ - amount_l = self._config.get("color_lightness", 30) / 100 - amount_ab = self._config.get("color_ab", 8) / 100 - adjust = np.array([amount_l, amount_ab, amount_ab], dtype="float32") - randoms = ( - (np.random.rand(self._batchsize, 1, 1, 3).astype("float32") * (adjust * 2)) - adjust) - logger.trace("Random LAB adjustments: %s", randoms) - + randoms = np.random.uniform(-self._constants.lab_adjust, + self._constants.lab_adjust, + size=(self._batch_size, 1, 1, 3)).astype("float32") + logger.trace("Random LAB adjustments: %s", randoms) # type:ignore[attr-defined] + # Iterating through the images and channels is much faster than numpy.where and slightly + # faster than numexpr.where. for image, rand in zip(batch, randoms): for idx in range(rand.shape[-1]): adjustment = rand[:, :, idx] @@ -269,10 +263,9 @@ def _random_lab(self, batch): image[:, :, idx] = ((255 - image[:, :, idx]) * adjustment) + image[:, :, idx] else: image[:, :, idx] = image[:, :, idx] * (1 + adjustment) - return batch # <<< IMAGE AUGMENTATION >>> # - def transform(self, batch): + def transform(self, batch: np.ndarray): """ Perform random transformation on the passed in batch. The transformation parameters are set in :file:`config.train.ini` @@ -282,47 +275,36 @@ def transform(self, batch): batch: :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `channels`) and in `BGR` format. - - Returns - ---------- - :class:`numpy.ndarray` - A 4-dimensional array of the same shape as :attr:`batch` with transformation applied. """ - if self.is_display: - return batch - logger.trace("Randomly transforming image") - rotation_range = self._config.get("rotation_range", 10) - zoom_range = self._config.get("zoom_amount", 5) / 100 - shift_range = self._config.get("shift_range", 5) / 100 - - rotation = np.random.uniform(-rotation_range, - rotation_range, - size=self._batchsize).astype("float32") - scale = np.random.uniform(1 - zoom_range, - 1 + zoom_range, - size=self._batchsize).astype("float32") - tform = np.random.uniform( - -shift_range, - shift_range, - size=(self._batchsize, 2)).astype("float32") * self._training_size - + logger.trace("Randomly transforming image") # type:ignore[attr-defined] + + rotation = np.random.uniform(-self._constants.transform_rotation, + self._constants.transform_rotation, + size=self._batch_size).astype("float32") + scale = np.random.uniform(1 - self._constants.transform_zoom, + 1 + self._constants.transform_zoom, + size=self._batch_size).astype("float32") + + tform = np.random.uniform(-self._constants.transform_shift, + self._constants.transform_shift, + size=(self._batch_size, 2)).astype("float32") mats = np.array( - [cv2.getRotationMatrix2D((self._training_size // 2, self._training_size // 2), + [cv2.getRotationMatrix2D((self._processing_size // 2, self._processing_size // 2), rot, scl) for rot, scl in zip(rotation, scale)]).astype("float32") mats[..., 2] += tform - batch = np.array([cv2.warpAffine(image, - mat, - (self._training_size, self._training_size), - borderMode=cv2.BORDER_REPLICATE) - for image, mat in zip(batch, mats)]) + for image, mat in zip(batch, mats): + cv2.warpAffine(image, + mat, + (self._processing_size, self._processing_size), + dst=image, + borderMode=cv2.BORDER_REPLICATE) - logger.trace("Randomly transformed image") - return batch + logger.trace("Randomly transformed image") # type:ignore[attr-defined] - def random_flip(self, batch): + def random_flip(self, batch: np.ndarray): """ Perform random horizontal flipping on the passed in batch. The probability of flipping an image is set in :file:`config.train.ini` @@ -332,21 +314,15 @@ def random_flip(self, batch): batch: :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `channels`) and in `BGR` format. - - Returns - ---------- - :class:`numpy.ndarray` - A 4-dimensional array of the same shape as :attr:`batch` with transformation applied. """ - if not self.is_display: - logger.trace("Randomly flipping image") - randoms = np.random.rand(self._batchsize) - indices = np.where(randoms > self._config.get("random_flip", 50) / 100)[0] - batch[indices] = batch[indices, :, ::-1] - logger.trace("Randomly flipped %s images of %s", len(indices), self._batchsize) - return batch - - def warp(self, batch, to_landmarks=False, **kwargs): + logger.trace("Randomly flipping image") # type:ignore[attr-defined] + randoms = np.random.rand(self._batch_size) + indices = np.where(randoms <= self._flip_chance / 100)[0] + batch[indices] = batch[indices, :, ::-1] + logger.trace("Randomly flipped %s images of %s", # type:ignore[attr-defined] + len(indices), self._batch_size) + + def warp(self, batch: np.ndarray, to_landmarks: bool = False, **kwargs) -> np.ndarray: """ Perform random warping on the passed in batch by one of two methods. Parameters @@ -367,43 +343,71 @@ def warp(self, batch, to_landmarks=False, **kwargs): * **batch_dst_points** (:class:`numpy.ndarray`) - A batch of randomly chosen closest \ match destination faces landmarks. This is a 3-dimensional array in the shape \ (`batchsize`, `68`, `2`). + Returns ---------- :class:`numpy.ndarray` A 4-dimensional array of the same shape as :attr:`batch` with warping applied. """ if to_landmarks: - return self._random_warp_landmarks(batch, **kwargs).astype("float32") / 255.0 - return self._random_warp(batch).astype("float32") / 255.0 - - def _random_warp(self, batch): - """ Randomly warp the input batch """ - logger.trace("Randomly warping batch") - mapx = self._constants["warp_mapx"] - mapy = self._constants["warp_mapy"] - pad = self._constants["warp_pad"] - slices = self._constants["warp_slices"] - - rands = np.random.normal(size=(self._batchsize, 2, 5, 5), - scale=self._scale).astype("float32") - batch_maps = np.stack((mapx, mapy), axis=1) + rands - batch_interp = np.array([[cv2.resize(map_, (pad, pad))[slices, slices] for map_ in maps] + return self._random_warp_landmarks(batch, **kwargs) + return self._random_warp(batch) + + def _random_warp(self, batch: np.ndarray) -> np.ndarray: + """ Randomly warp the input batch + + Parameters + ---------- + batch: :class:`numpy.ndarray` + The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, + `3`) and in `BGR` format. + + Returns + ---------- + :class:`numpy.ndarray` + A 4-dimensional array of the same shape as :attr:`batch` with warping applied. + """ + logger.trace("Randomly warping batch") # type:ignore[attr-defined] + slices = self._constants.warp_slices + rands = np.random.normal(size=(self._batch_size, 2, 5, 5), + scale=self._warp_scale).astype("float32") + batch_maps = ne.evaluate("m + r", local_dict={"m": self._constants.warp_maps, "r": rands}) + batch_interp = np.array([[cv2.resize(map_, self._constants.warp_pad)[slices, slices] + for map_ in maps] for maps in batch_maps]) warped_batch = np.array([cv2.remap(image, interp[0], interp[1], cv2.INTER_LINEAR) for image, interp in zip(batch, batch_interp)]) - logger.trace("Warped image shape: %s", warped_batch.shape) + logger.trace("Warped image shape: %s", warped_batch.shape) # type:ignore[attr-defined] return warped_batch - def _random_warp_landmarks(self, batch, batch_src_points, batch_dst_points): - """ From dfaker. Warp the image to a similar set of landmarks from the opposite side """ - logger.trace("Randomly warping landmarks") - edge_anchors = self._constants["warp_lm_edge_anchors"] - grids = self._constants["warp_lm_grids"] - slices = self._constants["tgt_slices"] + def _random_warp_landmarks(self, + batch: np.ndarray, + batch_src_points: np.ndarray, + batch_dst_points: np.ndarray) -> np.ndarray: + """ From dfaker. Warp the image to a similar set of landmarks from the opposite side + + batch: :class:`numpy.ndarray` + The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, + `3`) and in `BGR` format. + batch_src_points :class:`numpy.ndarray` + A batch of 68 point landmarks for the source faces. This is a 3-dimensional array in + the shape (`batchsize`, `68`, `2`). + batch_dst_points :class:`numpy.ndarray` + A batch of randomly chosen closest match destination faces landmarks. This is a + 3-dimensional array in the shape (`batchsize`, `68`, `2`). + + Returns + ---------- + :class:`numpy.ndarray` + A 4-dimensional array of the same shape as :attr:`batch` with warping applied. + """ + logger.trace("Randomly warping landmarks") # type:ignore[attr-defined] + edge_anchors = self._constants.warp_lm_edge_anchors + grids = self._constants.warp_lm_grids batch_dst = (batch_dst_points + np.random.normal(size=batch_dst_points.shape, - scale=2.0)) + scale=self._warp_lm_scale)) face_cores = [cv2.convexHull(np.concatenate([src[17:], dst[17:]], axis=0)) for src, dst in zip(batch_src_points.astype("int32"), @@ -418,48 +422,21 @@ def _random_warp_landmarks(self, batch, batch_src_points, batch_dst_points): for src, dst, face_core in zip(batch_src[:, :18, :], batch_dst[:, :18, :], face_cores)] - batch_src = [np.delete(src, idxs, axis=0) for idxs, src in zip(rem_indices, batch_src)] - batch_dst = [np.delete(dst, idxs, axis=0) for idxs, dst in zip(rem_indices, batch_dst)] + lbatch_src = [np.delete(src, idxs, axis=0) for idxs, src in zip(rem_indices, batch_src)] + lbatch_dst = [np.delete(dst, idxs, axis=0) for idxs, dst in zip(rem_indices, batch_dst)] grid_z = np.array([griddata(dst, src, (grids[0], grids[1]), method="linear") - for src, dst in zip(batch_src, batch_dst)]) - maps = grid_z.reshape((self._batchsize, - self._training_size, - self._training_size, + for src, dst in zip(lbatch_src, lbatch_dst)]) + maps = grid_z.reshape((self._batch_size, + self._processing_size, + self._processing_size, 2)).astype("float32") + warped_batch = np.array([cv2.remap(image, map_[..., 1], map_[..., 0], cv2.INTER_LINEAR, - cv2.BORDER_TRANSPARENT) + borderMode=cv2.BORDER_TRANSPARENT) for image, map_ in zip(batch, maps)]) - warped_batch = np.array([cv2.resize(image[slices, slices, :], - (self._input_size, self._input_size), - cv2.INTER_AREA) - for image in warped_batch]) - logger.trace("Warped batch shape: %s", warped_batch.shape) + logger.trace("Warped batch shape: %s", warped_batch.shape) # type:ignore[attr-defined] return warped_batch - - def skip_warp(self, batch): - """ Returns the images resized and cropped for feeding the model, if warping has been - disabled. - - Parameters - ---------- - batch: :class:`numpy.ndarray` - The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, - `3`) and in `BGR` format. - - Returns - ------- - :class:`numpy.ndarray` - The given batch cropped and resized for feeding the model - """ - logger.trace("Compiling skip warp images: batch shape: %s", batch.shape) - slices = self._constants["tgt_slices"] - retval = np.array([cv2.resize(image[slices, slices, :], - (self._input_size, self._input_size), - cv2.INTER_AREA) - for image in batch], dtype='float32') / 255. - logger.trace("feed batch shape: %s", retval.shape) - return retval diff --git a/lib/training/cache.py b/lib/training/cache.py new file mode 100644 index 0000000000..8915f79ec6 --- /dev/null +++ b/lib/training/cache.py @@ -0,0 +1,535 @@ +#!/usr/bin/env python3 +""" Holds the data cache for training data generators """ +from __future__ import annotations +import logging +import os +import typing as T + +from threading import Lock + +import cv2 +import numpy as np +from tqdm import tqdm + +from lib.align import CenteringType, DetectedFace, LandmarkType +from lib.image import read_image_batch, read_image_meta_batch +from lib.utils import FaceswapError + +if T.TYPE_CHECKING: + from lib.align.alignments import PNGHeaderAlignmentsDict, PNGHeaderDict + from lib.config import ConfigValueType + +logger = logging.getLogger(__name__) +_FACE_CACHES: dict[str, "_Cache"] = {} + + +def get_cache(side: T.Literal["a", "b"], + filenames: list[str] | None = None, + config: dict[str, ConfigValueType] | None = None, + size: int | None = None, + coverage_ratio: float | None = None) -> "_Cache": + """ Obtain a :class:`_Cache` object for the given side. If the object does not pre-exist then + create it. + + Parameters + ---------- + side: str + `"a"` or `"b"`. The side of the model to obtain the cache for + filenames: list + The filenames of all the images. This can either be the full path or the base name. If the + full paths are passed in, they are stripped to base name for use as the cache key. Must be + passed for the first call of this function for each side. For subsequent calls this + parameter is ignored. Default: ``None`` + config: dict, optional + The user selected training configuration options. Must be passed for the first call of this + function for each side. For subsequent calls this parameter is ignored. Default: ``None`` + size: int, optional + The largest output size of the model. Must be passed for the first call of this function + for each side. For subsequent calls this parameter is ignored. Default: ``None`` + coverage_ratio: float: optional + The coverage ratio that the model is using. Must be passed for the first call of this + function for each side. For subsequent calls this parameter is ignored. Default: ``None`` + + Returns + ------- + :class:`_Cache` + The face meta information cache for the requested side + """ + if not _FACE_CACHES.get(side): + assert config is not None, ("config must be provided for first call to cache") + assert filenames is not None, ("filenames must be provided for first call to cache") + assert size is not None, ("size must be provided for first call to cache") + assert coverage_ratio is not None, ("coverage_ratio must be provided for first call to " + "cache") + logger.debug("Creating cache. side: %s, size: %s, coverage_ratio: %s", + side, size, coverage_ratio) + _FACE_CACHES[side] = _Cache(filenames, config, size, coverage_ratio) + return _FACE_CACHES[side] + + +def _check_reset(face_cache: "_Cache") -> bool: + """ Check whether a given cache needs to be reset because a face centering change has been + detected in the other cache. + + Parameters + ---------- + face_cache: :class:`_Cache` + The cache object that is checking whether it should reset + + Returns + ------- + bool + ``True`` if the given object should reset the cache, otherwise ``False`` + """ + check_cache = next((cache for cache in _FACE_CACHES.values() if cache != face_cache), None) + retval = False if check_cache is None else check_cache.check_reset() + return retval + + +class _Cache(): + """ A thread safe mechanism for collecting and holding face meta information (masks, " + "alignments data etc.) for multiple :class:`TrainingDataGenerator`s. + + Each side may have up to 3 generators (training, preview and time-lapse). To conserve VRAM + these need to share access to the same face information for the images they are processing. + + As the cache is populated at run-time, thread safe writes are required for the first epoch. + Following that, the cache is only used for reads, which is thread safe intrinsically. + + It would probably be quicker to set locks on each individual face, but for code complexity + reasons, and the fact that the lock is only taken up during cache population, and it should + only be being read multiple times on save iterations, we lock the whole cache during writes. + + Parameters + ---------- + filenames: list + The filenames of all the images. This can either be the full path or the base name. If the + full paths are passed in, they are stripped to base name for use as the cache key. + config: dict + The user selected training configuration options + size: int + The largest output size of the model + coverage_ratio: float + The coverage ratio that the model is using. + """ + def __init__(self, + filenames: list[str], + config: dict[str, ConfigValueType], + size: int, + coverage_ratio: float) -> None: + logger.debug("Initializing: %s (filenames: %s, size: %s, coverage_ratio: %s)", + self.__class__.__name__, len(filenames), size, coverage_ratio) + self._lock = Lock() + self._cache_info = {"cache_full": False, "has_reset": False} + self._partially_loaded: list[str] = [] + + self._image_count = len(filenames) + self._cache: dict[str, DetectedFace] = {} + self._aligned_landmarks: dict[str, np.ndarray] = {} + self._extract_version = 0.0 + self._size = size + + assert config["centering"] in T.get_args(CenteringType) + self._centering: CenteringType = T.cast(CenteringType, config["centering"]) + self._config = config + self._coverage_ratio = coverage_ratio + + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def cache_full(self) -> bool: + """bool: ``True`` if the cache has been fully populated. ``False`` if there are items still + to be cached. """ + if self._cache_info["cache_full"]: + return self._cache_info["cache_full"] + with self._lock: + return self._cache_info["cache_full"] + + @property + def aligned_landmarks(self) -> dict[str, np.ndarray]: + """ dict: The filename as key, aligned landmarks as value. """ + # Note: Aligned landmarks are only used for warp-to-landmarks, so this can safely populate + # all of the aligned landmarks for the entire cache. + if not self._aligned_landmarks: + with self._lock: + # For Warp-To-Landmarks a race condition can occur where this is referenced from + # the opposite side prior to it being populated, so block on a lock. + self._aligned_landmarks = {key: face.aligned.landmarks + for key, face in self._cache.items()} + return self._aligned_landmarks + + @property + def size(self) -> int: + """ int: The pixel size of the cropped aligned face """ + return self._size + + def check_reset(self) -> bool: + """ Check whether this cache has been reset due to a face centering change, and reset the + flag if it has. + + Returns + ------- + bool + ``True`` if the cache has been reset because of a face centering change due to + legacy alignments, otherwise ``False``. """ + retval = self._cache_info["has_reset"] + if retval: + logger.debug("Resetting 'has_reset' flag") + self._cache_info["has_reset"] = False + return retval + + def get_items(self, filenames: list[str]) -> list[DetectedFace]: + """ Obtain the cached items for a list of filenames. The returned list is in the same order + as the provided filenames. + + Parameters + ---------- + filenames: list + A list of image filenames to obtain the cached data for + + Returns + ------- + list + List of DetectedFace objects holding the cached metadata. The list returns in the same + order as the filenames received + """ + return [self._cache[os.path.basename(filename)] for filename in filenames] + + def cache_metadata(self, filenames: list[str]) -> np.ndarray: + """ Obtain the batch with metadata for items that need caching and cache DetectedFace + objects to :attr:`_cache`. + + Parameters + ---------- + filenames: list + List of full paths to image file names + + Returns + ------- + :class:`numpy.ndarray` + The batch of face images loaded from disk + """ + keys = [os.path.basename(filename) for filename in filenames] + with self._lock: + if _check_reset(self): + self._reset_cache(False) + + needs_cache = [filename + for filename, key in zip(filenames, keys) + if key not in self._cache or key in self._partially_loaded] + logger.trace("Needs cache: %s", needs_cache) # type: ignore + + if not needs_cache: + # Don't bother reading the metadata if no images in this batch need caching + logger.debug("All metadata already cached for: %s", keys) + return read_image_batch(filenames) + + try: + batch, metadata = read_image_batch(filenames, with_metadata=True) + except ValueError as err: + if "inhomogeneous" in str(err): + raise FaceswapError( + "There was an error loading a batch of images. This is most likely due to " + "non-faceswap extracted faces in your training folder." + "\nAll training images should be Faceswap extracted faces." + "\nAll training images should be the same size." + f"\nThe files that caused this error are: {filenames}") from err + raise + if len(batch.shape) == 1: + folder = os.path.dirname(filenames[0]) + details = [ + f"{key} ({f'{img.shape[1]}px' if isinstance(img, np.ndarray) else type(img)})" + for key, img in zip(keys, batch)] + msg = (f"There are mismatched image sizes in the folder '{folder}'. All training " + "images for each side must have the same dimensions.\nThe batch that " + f"failed contains the following files:\n{details}.") + raise FaceswapError(msg) + + # Populate items into cache + for filename in needs_cache: + key = os.path.basename(filename) + meta = metadata[filenames.index(filename)] + + # Version Check + self._validate_version(meta, filename) + if self._partially_loaded: # Faces already loaded for Warp-to-landmarks + self._partially_loaded.remove(key) + detected_face = self._cache[key] + else: + detected_face = self._load_detected_face(filename, meta["alignments"]) + + self._prepare_masks(filename, detected_face) + self._cache[key] = detected_face + + # Update the :attr:`cache_full` attribute + cache_full = not self._partially_loaded and len(self._cache) == self._image_count + if cache_full: + logger.verbose("Cache filled: '%s'", os.path.dirname(filenames[0])) # type: ignore + self._cache_info["cache_full"] = cache_full + + return batch + + def pre_fill(self, filenames: list[str], side: T.Literal["a", "b"]) -> None: + """ When warp to landmarks is enabled, the cache must be pre-filled, as each side needs + access to the other side's alignments. + + Parameters + ---------- + filenames: list + The list of full paths to the images to load the metadata from + side: str + `"a"` or `"b"`. The side of the model being cached. Used for info output + + Raises + ------ + FaceSwapError + If unsupported landmark type exists + """ + with self._lock: + for filename, meta in tqdm(read_image_meta_batch(filenames), + desc=f"WTL: Caching Landmarks ({side.upper()})", + total=len(filenames), + leave=False): + if "itxt" not in meta or "alignments" not in meta["itxt"]: + raise FaceswapError(f"Invalid face image found. Aborting: '{filename}'") + + meta = meta["itxt"] + key = os.path.basename(filename) + # Version Check + self._validate_version(meta, filename) + detected_face = self._load_detected_face(filename, meta["alignments"]) + + aligned = detected_face.aligned + assert aligned is not None + if aligned.landmark_type != LandmarkType.LM_2D_68: + raise FaceswapError("68 Point facial Landmarks are required for Warp-to-" + f"landmarks. The face that failed was: '{filename}'") + + self._cache[key] = detected_face + self._partially_loaded.append(key) + + def _validate_version(self, png_meta: PNGHeaderDict, filename: str) -> None: + """ Validate that there are not a mix of v1.0 extracted faces and v2.x faces. + + Parameters + ---------- + png_meta: dict + The information held within the Faceswap PNG Header + filename: str + The full path to the file being validated + + Raises + ------ + FaceswapError + If a version 1.0 face appears in a 2.x set or vice versa + """ + alignment_version = png_meta["source"]["alignments_version"] + + if not self._extract_version: + logger.debug("Setting initial extract version: %s", alignment_version) + self._extract_version = alignment_version + if alignment_version == 1.0 and self._centering != "legacy": + self._reset_cache(True) + return + + if (self._extract_version == 1.0 and alignment_version > 1.0) or ( + alignment_version == 1.0 and self._extract_version > 1.0): + raise FaceswapError("Mixing legacy and full head extracted facesets is not supported. " + "The following folder contains a mix of extracted face types: " + f"'{os.path.dirname(filename)}'") + + self._extract_version = min(alignment_version, self._extract_version) + + def _reset_cache(self, set_flag: bool) -> None: + """ In the event that a legacy extracted face has been seen, and centering is not legacy + the cache will need to be reset for legacy centering. + + Parameters + ---------- + set_flag: bool + ``True`` if the flag should be set to indicate that the cache is being reset because of + a legacy face set/centering mismatch. ``False`` if the cache is being reset because it + has detected a reset flag from the opposite cache. + """ + if set_flag: + logger.warning("You are using legacy extracted faces but have selected '%s' centering " + "which is incompatible. Switching centering to 'legacy'", + self._centering) + self._config["centering"] = "legacy" + self._centering = "legacy" + self._cache = {} + self._cache_info["cache_full"] = False + if set_flag: + self._cache_info["has_reset"] = True + + def _load_detected_face(self, + filename: str, + alignments: PNGHeaderAlignmentsDict) -> DetectedFace: + """ Load a :class:`DetectedFace` object and load its associated `aligned` property. + + Parameters + ---------- + filename: str + The file path for the current image + alignments: dict + The alignments for a single face, extracted from a PNG header + + Returns + ------- + :class:`lib.align.DetectedFace` + The loaded Detected Face object + """ + detected_face = DetectedFace() + detected_face.from_png_meta(alignments) + detected_face.load_aligned(None, + size=self._size, + centering=self._centering, + coverage_ratio=self._coverage_ratio, + is_aligned=True, + is_legacy=self._extract_version == 1.0) + logger.trace("Cached aligned face for: %s", filename) # type: ignore + return detected_face + + def _prepare_masks(self, filename: str, detected_face: DetectedFace) -> None: + """ Prepare the masks required from training, and compile into a single compressed array + + Parameters + ---------- + filename: str + The file path for the current image + detected_face: :class:`lib.align.DetectedFace` + The detected face object that holds the masks + """ + masks = [(self._get_face_mask(filename, detected_face))] + for area in T.get_args(T.Literal["eye", "mouth"]): + masks.append(self._get_localized_mask(filename, detected_face, area)) + + detected_face.store_training_masks(masks, delete_masks=True) + logger.trace("Stored masks for filename: %s)", filename) # type: ignore + + def _get_face_mask(self, filename: str, detected_face: DetectedFace) -> np.ndarray | None: + """ Obtain the training sized face mask from the :class:`DetectedFace` for the requested + mask type. + + Parameters + ---------- + filename: str + The file path for the current image + detected_face: :class:`lib.align.DetectedFace` + The detected face object that holds the masks + + Raises + ------ + FaceswapError + If the requested mask type is not available an error is returned along with a list + of available masks + """ + if not self._config["penalized_mask_loss"] and not self._config["learn_mask"]: + return None + + if not self._config["mask_type"]: + logger.debug("No mask selected. Not validating") + return None + + if self._config["mask_type"] not in detected_face.mask: + exist_masks = list(detected_face.mask) + msg = "No masks exist for this face" + if exist_masks: + msg = f"The masks that exist for this face are: {exist_masks}" + raise FaceswapError( + f"You have selected the mask type '{self._config['mask_type']}' but at least one " + "face does not contain the selected mask.\n" + f"The face that failed was: '{filename}'\n{msg}") + + mask = detected_face.mask[str(self._config["mask_type"])] + assert isinstance(self._config["mask_dilation"], float) + assert isinstance(self._config["mask_blur_kernel"], int) + assert isinstance(self._config["mask_threshold"], int) + mask.set_dilation(self._config["mask_dilation"]) + mask.set_blur_and_threshold(blur_kernel=self._config["mask_blur_kernel"], + threshold=self._config["mask_threshold"]) + + pose = detected_face.aligned.pose + mask.set_sub_crop(pose.offset[mask.stored_centering], + pose.offset[self._centering], + self._centering, + self._coverage_ratio) + face_mask = mask.mask + if self._size != face_mask.shape[0]: + interpolator = cv2.INTER_CUBIC if mask.stored_size < self._size else cv2.INTER_AREA + face_mask = cv2.resize(face_mask, + (self._size, self._size), + interpolation=interpolator)[..., None] + + logger.trace("Obtained face mask for: %s %s", filename, face_mask.shape) # type: ignore + return face_mask + + def _get_localized_mask(self, + filename: str, + detected_face: DetectedFace, + area: T.Literal["eye", "mouth"]) -> np.ndarray | None: + """ Obtain a localized mask for the given area if it is required for training. + + Parameters + ---------- + filename: str + The file path for the current image + detected_face: :class:`lib.align.DetectedFace` + The detected face object that holds the masks + area: str + `"eye"` or `"mouth"`. The area of the face to obtain the mask for + """ + multiplier = self._config[f"{area}_multiplier"] + assert isinstance(multiplier, int) + if not self._config["penalized_mask_loss"] or multiplier <= 1: + return None + try: + mask = detected_face.get_landmark_mask(area, self._size // 16, 2.5) + except FaceswapError as err: + logger.error(str(err)) + raise FaceswapError("Eye/Mouth multiplier masks could not be generated due to missing " + f"landmark data. The file that failed was: '{filename}'") from err + logger.trace("Caching localized '%s' mask for: %s %s", # type: ignore + area, filename, mask.shape) + return mask + + +class RingBuffer(): + """ Rolling buffer for holding training/preview batches + + Parameters + ---------- + batch_size: int + The batch size to create the buffer for + image_shape: tuple + The height/width/channels shape of a single image in the batch + buffer_size: int, optional + The number of arrays to hold in the rolling buffer. Default: `2` + dtype: str, optional + The datatype to create the buffer as. Default: `"uint8"` + """ + def __init__(self, + batch_size: int, + image_shape: tuple[int, int, int], + buffer_size: int = 2, + dtype: str = "uint8") -> None: + logger.debug("Initializing: %s (batch_size: %s, image_shape: %s, buffer_size: %s, " + "dtype: %s", self.__class__.__name__, batch_size, image_shape, buffer_size, + dtype) + self._max_index = buffer_size - 1 + self._index = 0 + self._buffer = [np.empty((batch_size, *image_shape), dtype=dtype) + for _ in range(buffer_size)] + logger.debug("Initialized: %s", self.__class__.__name__) # type: ignore + + def __call__(self) -> np.ndarray: + """ Obtain the next array from the ring buffer + + Returns + ------- + :class:`np.ndarray` + A pre-allocated numpy array from the buffer + """ + retval = self._buffer[self._index] + self._index += 1 if self._index < self._max_index else -self._max_index + return retval diff --git a/lib/training/generator.py b/lib/training/generator.py index c8bbf2f56c..0624246e0a 100644 --- a/lib/training/generator.py +++ b/lib/training/generator.py @@ -1,432 +1,397 @@ #!/usr/bin/env python3 """ Handles Data Augmentation for feeding Faceswap Models """ - +from __future__ import annotations import logging import os +import typing as T +from concurrent import futures from random import shuffle, choice -from threading import Lock -from zlib import decompress -import numpy as np import cv2 -from tqdm import tqdm -from lib.align import AlignedFace, DetectedFace, get_centered_size -from lib.image import read_image_batch, read_image_meta_batch +import numpy as np +import numexpr as ne +from lib.align import AlignedFace, DetectedFace +from lib.align.aligned_face import CenteringType +from lib.image import read_image_batch from lib.multithreading import BackgroundGenerator from lib.utils import FaceswapError from . import ImageAugmentation +from .cache import get_cache, RingBuffer -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from collections.abc import Generator + from lib.config import ConfigValueType + from plugins.train.model._base import ModelBase + from .cache import _Cache -_FACE_CACHES = dict() +logger = logging.getLogger(__name__) +BatchType = tuple[np.ndarray, list[np.ndarray]] -def _get_cache(side, filenames, config): - """ Obtain a :class:`_Cache` object for the given side. If the object does not pre-exist then - create it. +class DataGenerator(): + """ Parent class for Training and Preview Data Generators. + + This class is called from :mod:`plugins.train.trainer._base` and launches a background + iterator that compiles augmented data, target data and sample data. Parameters ---------- - side: str - `"a"` or `"b"`. The side of the model to obtain the cache for - filenames: list - The filenames of all the images. This can either be the full path or the base name. If the - full paths are passed in, they are stripped to base name for use as the cache key. + model: :class:`~plugins.train.model.ModelBase` + The model that this data generator is feeding config: dict - The user selected training configuration options - - Returns - ------- - :class:`_Cache` - The face meta information cache for the requested side + The configuration `dict` generated from :file:`config.train.ini` containing the trainer + plugin configuration options. + side: {'a' or 'b'} + The side of the model that this iterator is for. + images: list + A list of image paths that will be used to compile the final augmented data from. + batch_size: int + The batch size for this iterator. Images will be returned in :class:`numpy.ndarray` + objects of this size from the iterator. """ - if not _FACE_CACHES.get(side): - logger.debug("Creating cache. Side: %s", side) - _FACE_CACHES[side] = _Cache(filenames, config) - return _FACE_CACHES[side] - - -def _check_reset(face_cache): - """ Check whether a given cache needs to be reset because a face centering change has been - detected in the other cache. + def __init__(self, + config: dict[str, ConfigValueType], + model: ModelBase, + side: T.Literal["a", "b"], + images: list[str], + batch_size: int) -> None: + logger.debug("Initializing %s: (model: %s, side: %s, images: %s , " + "batch_size: %s, config: %s)", self.__class__.__name__, model.name, side, + len(images), batch_size, config) + self._config = config + self._side = side + self._images = images + self._batch_size = batch_size + + self._process_size = max(img[1] for img in model.input_shapes + model.output_shapes) + self._output_sizes = self._get_output_sizes(model) + self._model_input_size = max(img[1] for img in model.input_shapes) + + self._coverage_ratio = model.coverage_ratio + self._color_order = model.color_order.lower() + self._use_mask = self._config["mask_type"] and (self._config["penalized_mask_loss"] or + self._config["learn_mask"]) + + self._validate_samples() + self._buffer = RingBuffer(batch_size, + (self._process_size, self._process_size, self._total_channels), + dtype="uint8") + self._face_cache: _Cache = get_cache(side, + filenames=images, + config=self._config, + size=self._process_size, + coverage_ratio=self._coverage_ratio) + logger.debug("Initialized %s", self.__class__.__name__) - Parameters - ---------- - face_cache: :class:`_Cache` - The cache object that is checking whether it should reset + @property + def _total_channels(self) -> int: + """int: The total number of channels, including mask channels that the target image + should hold. """ + channels = 3 + if self._config["mask_type"] and (self._config["learn_mask"] or + self._config["penalized_mask_loss"]): + channels += 1 + + mults = [area for area in ["eye", "mouth"] + if T.cast(int, self._config[f"{area}_multiplier"]) > 1] + if self._config["penalized_mask_loss"] and mults: + channels += len(mults) + return channels + + def _get_output_sizes(self, model: ModelBase) -> list[int]: + """ Obtain the size of each output tensor for the model. - Returns - ------- - bool - ``True`` if the given object should reset the cache, otherwise ``False`` - """ - check_cache = next((cache for cache in _FACE_CACHES.values() if cache != face_cache), None) - retval = check_cache if check_cache is None else check_cache.check_reset() - return retval + Parameters + ---------- + model: :class:`~plugins.train.model.ModelBase` + The model that this data generator is feeding + Returns + ------- + list + A list of integers for the model output size for the current side + """ + out_shapes = model.output_shapes + split = len(out_shapes) // 2 + side_out = out_shapes[:split] if self._side == "a" else out_shapes[split:] + retval = [shape[1] for shape in side_out if shape[-1] != 1] + logger.debug("side: %s, model output shapes: %s, output sizes: %s", + self._side, model.output_shapes, retval) + return retval -class _Cache(): - """ A thread safe mechanism for collecting and holding face meta information (masks, " - "alignments data etc.) for multiple :class:`TrainingDataGenerator`s. + def minibatch_ab(self, do_shuffle: bool = True) -> Generator[BatchType, None, None]: + """ A Background iterator to return augmented images, samples and targets. - Each side may have up to 3 generators (training, preview and time-lapse). To conserve VRAM - these need to share access to the same face information for the images they are processing. + The exit point from this class and the sole attribute that should be referenced. Called + from :mod:`plugins.train.trainer._base`. Returns an iterator that yields images for + training, preview and time-lapses. - As the cache is populated at run-time, thread safe writes are required for the first epoch. - Following that, the cache is only used for reads, which is thread safe intrinsically. + Parameters + ---------- + do_shuffle: bool, optional + Whether data should be shuffled prior to loading from disk. If true, each time the full + list of filenames are processed, the data will be reshuffled to make sure they are not + returned in the same order. Default: ``True`` - It would probably be quicker to set locks on each individual face, but for code complexity - reasons, and the fact that the lock is only taken up during cache population, and it should - only be being read multiple times on save iterations, we lock the whole cache during writes. + Yields + ------ + feed: list + 4-dimensional array of faces to feed the training the model (:attr:`x` parameter for + :func:`keras.models.model.train_on_batch`.). The array returned is in the format + (`batch size`, `height`, `width`, `channels`). + targets: list + List of 4-dimensional :class:`numpy.ndarray` objects in the order and size of each + output of the model. The format of these arrays will be (`batch size`, `height`, + `width`, `x`). This is the :attr:`y` parameter for + :func:`keras.models.model.train_on_batch`. The number of channels here will vary. + The first 3 channels are (rgb/bgr). The 4th channel is the face mask. Any subsequent + channels are area masks (e.g. eye/mouth masks) + """ + logger.debug("do_shuffle: %s", do_shuffle) + args = (do_shuffle, ) + batcher = BackgroundGenerator(self._minibatch, args=args) + return batcher.iterator() - Parameters - ---------- - filenames: list - The filenames of all the images. This can either be the full path or the base name. If the - full paths are passed in, they are stripped to base name for use as the cache key. - config: dict - The user selected training configuration options - """ - def __init__(self, filenames, config): - self._lock = Lock() - self._cache = {os.path.basename(filename): dict(cached=False) for filename in filenames} - self._aligned_landmarks = None - self._partial_load = False - self._cache_full = False - self._extract_version = None - self._has_reset = False - self._size = None - - self._centering = config["centering"] - self._config = config + # << INTERNAL METHODS >> # + def _validate_samples(self) -> None: + """ Ensures that the total number of images within :attr:`images` is greater or equal to + the selected :attr:`batch_size`. - @property - def cache_full(self): - """bool: ``True`` if the cache has been fully populated. ``False`` if there are items still - to be cached. """ - if self._cache_full: - return self._cache_full - with self._lock: - return self._cache_full + Raises + ------ + :class:`FaceswapError` + If the number of images loaded is smaller than the selected batch size + """ + length = len(self._images) + msg = ("Number of images is lower than batch-size (Note that too few images may lead to " + f"bad training). # images: {length}, batch-size: {self._batch_size}") + try: + assert length >= self._batch_size, msg + except AssertionError as err: + msg += ("\nYou should increase the number of images in your training set or lower " + "your batch-size.") + raise FaceswapError(msg) from err - @property - def partially_loaded(self): - """ bool: ``True`` if the cache has been partially loaded for Warp To Landmarks otherwise - ``False`` """ - if self._partial_load: - return self._partial_load - with self._lock: - return self._partial_load + def _minibatch(self, do_shuffle: bool) -> Generator[BatchType, None, None]: + """ A generator function that yields the augmented, target and sample images for the + current batch on the current side. - @property - def extract_version(self): - """ float: The alignments file version used to extract the faces. """ - return self._extract_version + Parameters + ---------- + do_shuffle: bool, optional + Whether data should be shuffled prior to loading from disk. If true, each time the full + list of filenames are processed, the data will be reshuffled to make sure they are not + returned in the same order. Default: ``True`` - @property - def aligned_landmarks(self): - """ dict: The filename as key, aligned landmarks as value """ - if self._aligned_landmarks is None: - with self._lock: - # For Warp-To-Landmarks a race condition can occur where this is referenced from - # the opposite side prior to it being populated, so block on a lock. - self._aligned_landmarks = {key: val["aligned_face"].landmarks - for key, val in self._cache.items()} - return self._aligned_landmarks + Yields + ------ + feed: list + 4-dimensional array of faces to feed the training the model (:attr:`x` parameter for + :func:`keras.models.model.train_on_batch`.). The array returned is in the format + (`batch size`, `height`, `width`, `channels`). + targets: list + List of 4-dimensional :class:`numpy.ndarray` objects in the order and size of each + output of the model. The format of these arrays will be (`batch size`, `height`, + `width`, `x`). This is the :attr:`y` parameter for + :func:`keras.models.model.train_on_batch`. The number of channels here will vary. + The first 3 channels are (rgb/bgr). The 4th channel is the face mask. Any subsequent + channels are area masks (e.g. eye/mouth masks) + """ + logger.debug("Loading minibatch generator: (image_count: %s, do_shuffle: %s)", + len(self._images), do_shuffle) - @property - def crop_size(self): - """ int: The pixel size of the cropped aligned face """ - return self._size + def _img_iter(imgs): + """ Infinite iterator for recursing through image list and reshuffling at each epoch""" + while True: + if do_shuffle: + shuffle(imgs) + for img in imgs: + yield img - def check_reset(self): - """ Check whether this cache has been reset due to a face centering change, and reset the - flag if it has. + img_iter = _img_iter(self._images[:]) + while True: + img_paths = [next(img_iter) # pylint:disable=stop-iteration-return + for _ in range(self._batch_size)] + retval = self._process_batch(img_paths) + yield retval - Returns - ------- - bool - ``True`` if the cache has been reset because of a face centering change due to - legacy alignments, otherwise ``False``. """ - retval = self._has_reset - if retval: - logger.debug("Resetting 'has_reset' flag") - self._has_reset = False - return retval + def _get_images_with_meta(self, filenames: list[str]) -> tuple[np.ndarray, list[DetectedFace]]: + """ Obtain the raw face images with associated :class:`DetectedFace` objects for this + batch. - def get_items(self, filenames): - """ Obtain the cached items for a list of filenames. The returned list is in the same order - as the provided filenames. + If this is the first time a face has been loaded, then it's meta data is extracted + from the png header and added to :attr:`_face_cache`. Parameters ---------- filenames: list - A list of image filenames to obtain the cached data for + List of full paths to image file names Returns ------- + raw_faces: :class:`numpy.ndarray` + The full sized batch of training images for the given filenames list - List of dictionaries containing the cached metadata. The list returns in the same order - as the filenames received + Batch of :class:`~lib.align.DetectedFace` objects for the given filename including the + aligned face objects for the model output size """ - return [self._cache[os.path.basename(filename)] for filename in filenames] - - def cache_metadata(self, filenames): - """ Obtain the batch with metadata for items that need caching and cache them to - :attr:`_cache`. + if not self._face_cache.cache_full: + raw_faces = self._face_cache.cache_metadata(filenames) + else: + raw_faces = read_image_batch(filenames) + + detected_faces = self._face_cache.get_items(filenames) + logger.trace( # type:ignore[attr-defined] + "filenames: %s, raw_faces: '%s', detected_faces: %s", + filenames, raw_faces.shape, len(detected_faces)) + return raw_faces, detected_faces + + def _crop_to_coverage(self, + filenames: list[str], + images: np.ndarray, + detected_faces: list[DetectedFace], + batch: np.ndarray) -> None: + """ Crops the training image out of the full extract image based on the centering and + coveage used in the user's configuration settings. - Parameters - ---------- - filenames: list - List of full paths to image file names + If legacy extract images are being used then this just returns the extracted batch with + their corresponding landmarks. - Returns - ------- - :class:`numpy.ndarray` - The batch of face images loaded from disk - """ - keys = [os.path.basename(filename) for filename in filenames] - with self._lock: - if _check_reset(self): - self._reset_cache(False) - - needs_cache = [filename - for filename, key in zip(filenames, keys) - if not self._cache[key]["cached"]] - logger.trace("Needs cache: %s", needs_cache) - - if not needs_cache: - # Don't bother reading the metadata if no images in this batch need caching - logger.debug("All metadata already cached for: %s", keys) - return read_image_batch(filenames) - - batch, metadata = read_image_batch(filenames, with_metadata=True) - - if len(batch.shape) == 1: - folder = os.path.dirname(filenames[0]) - details = [ - "{0} ({1})".format( - key, f"{img.shape[1]}px" if isinstance(img, np.ndarray) else type(img)) - for key, img in zip(keys, batch)] - msg = (f"There are mismatched image sizes in the folder '{folder}'. All training " - "images for each side must have the same dimensions.\nThe batch that " - f"failed contains the following files:\n{details}.") - raise FaceswapError(msg) - - # Populate items into cache - for filename in needs_cache: - key = os.path.basename(filename) - meta = metadata[filenames.index(filename)] - - # Version Check - self._validate_version(meta, filename) - if self._partial_load: # Faces already loaded for Warp-to-landmarks - detected_face = self._cache[key]["detected_face"] - else: - detected_face = self._add_aligned_face(filename, - meta["alignments"], - batch.shape[1]) - - self._add_mask(filename, detected_face) - for area in ("eye", "mouth"): - self._add_localized_mask(filename, detected_face, area) - - self._cache[key]["cached"] = True - # Update the :attr:`cache_full` attribute - cache_full = all(item["cached"] for item in self._cache.values()) - if cache_full: - logger.verbose("Cache filled: '%s'", os.path.dirname(filenames[0])) - self._cache_full = cache_full - - return batch - - def pre_fill(self, filenames, side): - """ When warp to landmarks is enabled, the cache must be pre-filled, as each side needs - access to the other side's alignments. + Uses thread pool execution for about a 33% speed increase @ 64 batch size Parameters ---------- filenames: list - The list of full paths to the images to load the metadata from - side: str - `"a"` or `"b"`. The side of the model being cached. Used for info output + The list of filenames that correspond to this batch + images: :class:`numpy.ndarray` + The batch of faces that have been loaded from disk + detected_faces: list + The list of :class:`lib.align.DetectedFace` items corresponding to the batch + batch: :class:`np.ndarray` + The pre-allocated array to hold this batch """ - with self._lock: - for filename, meta in tqdm(read_image_meta_batch(filenames), - desc="WTL: Caching Landmarks ({})".format(side.upper()), - total=len(filenames), - leave=False): - if "itxt" not in meta or "alignments" not in meta["itxt"]: - raise FaceswapError(f"Invalid face image found. Aborting: '{filename}'") - - size = meta["width"] - meta = meta["itxt"] - # Version Check - self._validate_version(meta, filename) - detected_face = self._add_aligned_face(filename, meta["alignments"], size) - self._cache[os.path.basename(filename)]["detected_face"] = detected_face - self._partial_load = True - - def _validate_version(self, png_meta, filename): - """ Validate that there are not a mix of v1.0 extracted faces and v2.x faces. + logger.trace( # type:ignore[attr-defined] + "Cropping training images info: (filenames: %s, side: '%s')", filenames, self._side) - Parameters - ---------- - png_meta: dict - The information held within the Faceswap PNG Header - filename: str - The full path to the file being validated - - Raises - ------ - FaceswapError - If a version 1.0 face appears in a 2.x set or vice versa - """ - alignment_version = png_meta["source"]["alignments_version"] + with futures.ThreadPoolExecutor() as executor: + proc = {executor.submit(face.aligned.extract_face, img): idx + for idx, (face, img) in enumerate(zip(detected_faces, images))} - if not self._extract_version: - logger.debug("Setting initial extract version: %s", alignment_version) - self._extract_version = alignment_version - if alignment_version == 1.0 and self._centering != "legacy": - self._reset_cache(True) - return + for future in futures.as_completed(proc): + batch[proc[future], ..., :3] = future.result() - if (self._extract_version == 1.0 and alignment_version > 1.0) or ( - alignment_version == 1.0 and self._extract_version > 1.0): - raise FaceswapError("Mixing legacy and full head extracted facesets is not supported. " - "The following folder contains a mix of extracted face types: " - "{}".format(os.path.dirname(filename))) + def _apply_mask(self, detected_faces: list[DetectedFace], batch: np.ndarray) -> None: + """ Applies the masks to the 4th channel of the batch. - self._extract_version = min(alignment_version, self._extract_version) + If the configuration options `eye_multiplier` and/or `mouth_multiplier` are greater than 1 + then these masks are applied to the final channels of the batch respectively. - def _reset_cache(self, set_flag): - """ In the event that a legacy extracted face has been seen, and centering is not legacy - the cache will need to be reset for legacy centering. + If masks are not being used then this function returns having done nothing Parameters ---------- - set_flag: bool - ``True`` if the flag should be set to indicate that the cache is being reset because of - a legacy face set/centering mismatch. ``False`` if the cache is being reset because it - has detected a reset flag from the opposite cache. + detected_face: list + The list of :class:`~lib.align.DetectedFace` objects corresponding to the batch + batch: :class:`numpy.ndarray` + The preallocated array to apply masks to + side: str + '"a"' or '"b"' the side that is being processed """ - if set_flag: - logger.warning("You are using legacy extracted faces but have selected '%s' centering " - "which is incompatible. Switching centering to 'legacy'", - self._centering) - self._config["centering"] = "legacy" - self._centering = "legacy" - self._cache = {key: dict(cached=False) for key in self._cache} - self._cache_full = False - self._size = None - if set_flag: - self._has_reset = True - - def _add_aligned_face(self, filename, alignments, image_size): - """ Add a :class:`lib.align.AlignedFace` object to the cache. + if not self._use_mask: + return + + masks = np.array([face.get_training_masks() for face in detected_faces]) + batch[..., 3:] = masks + + logger.trace("side: %s, masks: %s, batch: %s", # type:ignore[attr-defined] + self._side, masks.shape, batch.shape) + + def _process_batch(self, filenames: list[str]) -> BatchType: + """ Prepares data for feeding through subclassed methods. + + If this is the first time a face has been loaded, then it's meta data is extracted from the + png header and added to :attr:`_face_cache` Parameters ---------- - filename: str - The file path for the current image - alignments: dict - The alignments for a single face, extracted from a PNG header - image_size: int - The pixel size of the image loaded from disk + filenames: list + List of full paths to image file names for a single batch Returns ------- - :class:`lib.align.DetectedFace` - The Detected Face object that was used to create the Aligned Face + :class:`numpy.ndarray` + 4-dimensional array of faces to feed the training the model. + list + List of 4-dimensional :class:`numpy.ndarray`. The number of channels here will vary. + The first 3 channels are (rgb/bgr). The 4th channel is the face mask. Any subsequent + channels are area masks (e.g. eye/mouth masks) """ - if self._size is None: - self._size = get_centered_size("legacy" if self._extract_version == 1.0 else "head", - self._centering, - image_size) + raw_faces, detected_faces = self._get_images_with_meta(filenames) + batch = self._buffer() + self._crop_to_coverage(filenames, raw_faces, detected_faces, batch) + self._apply_mask(detected_faces, batch) + feed, targets = self.process_batch(filenames, raw_faces, detected_faces, batch) - detected_face = DetectedFace() - detected_face.from_png_meta(alignments) + logger.trace( # type:ignore[attr-defined] + "Processed %s batch side %s. (filenames: %s, feed: %s, targets: %s)", + self.__class__.__name__, self._side, filenames, feed.shape, [t.shape for t in targets]) - aligned_face = AlignedFace(detected_face.landmarks_xy, - centering=self._centering, - size=self._size, - is_aligned=True) - logger.trace("Caching aligned face for: %s", filename) - self._cache[os.path.basename(filename)]["aligned_face"] = aligned_face - return detected_face + return feed, targets - def _add_mask(self, filename, detected_face): - """ Load the mask to the cache if a mask is required for training. + def process_batch(self, + filenames: list[str], + images: np.ndarray, + detected_faces: list[DetectedFace], + batch: np.ndarray) -> BatchType: + """ Override for processing the batch for the current generator. Parameters ---------- - filename: str - The file path for the current image - detected_face: :class:`lib.align.DetectedFace` - The detected face object that holds the masks + filenames: list + List of full paths to image file names for a single batch + images: :class:`numpy.ndarray` + The batch of faces corresponding to the filenames + detected_faces: list + List of :class:`~lib.align.DetectedFace` objects with aligned data and masks loaded for + the current batch + batch: :class:`numpy.ndarray` + The pre-allocated batch with images and masks populated for the selected coverage and + centering - Raises - ------ - FaceswapError - If the requested mask type is not available an error is returned along with a list - of available masks + Returns + ------- + list + 4-dimensional array of faces to feed the training the model. + list + List of 4-dimensional :class:`numpy.ndarray`. The number of channels here will vary. + The first 3 channels are (rgb/bgr). The 4th channel is the face mask. Any subsequent + channels are area masks (e.g. eye/mouth masks) """ - if not self._config["penalized_mask_loss"] and not self._config["learn_mask"]: - return - - if not self._config["mask_type"]: - logger.debug("No mask selected. Not validating") - return - - if self._config["mask_type"] not in detected_face.mask: - raise FaceswapError( - "You have selected the mask type '{}' but at least one face does not contain the " - "selected mask.\nThe face that failed was: '{}'\nThe masks that exist for this " - "face are: {}".format( - self._config["mask_type"], filename, list(detected_face.mask))) + raise NotImplementedError() - key = os.path.basename(filename) - mask = detected_face.mask[self._config["mask_type"]] - mask.set_blur_and_threshold(blur_kernel=self._config["mask_blur_kernel"], - threshold=self._config["mask_threshold"]) + def _set_color_order(self, batch) -> None: + """ Set the color order correctly for the model's input type. - pose = self._cache[key]["aligned_face"].pose - mask.set_sub_crop(pose.offset[self._centering] - pose.offset[mask.stored_centering], - self._centering) - - logger.trace("Caching mask for: %s", filename) - self._cache[key]["mask"] = mask + batch: :class:`numpy.ndarray` + The pre-allocated batch with images in the first 3 channels in BGR order + """ + if self._color_order == "rgb": + batch[..., :3] = batch[..., [2, 1, 0]] - def _add_localized_mask(self, filename, detected_face, area): - """ Load a localized mask to the cache for the given area if it is required for training. + def _to_float32(self, in_array: np.ndarray) -> np.ndarray: + """ Cast an UINT8 array in 0-255 range to float32 in 0.0-1.0 range. - Parameters - ---------- - filename: str - The file path for the current image - detected_face: :class:`lib.align.DetectedFace` - The detected face object that holds the masks - area: str - `"eye"` or `"mouth"`. The area of the face to obtain the mask for + in_array: :class:`numpy.ndarray` + The input uint8 array """ - if not self._config["penalized_mask_loss"] or self._config[f"{area}_multiplier"] <= 1: - return - key = "eyes" if area == "eye" else area + return ne.evaluate("x / c", + local_dict={"x": in_array, "c": np.float32(255)}, + casting="unsafe") - logger.trace("Caching localized '%s' mask for: %s", key, filename) - self._cache[os.path.basename(filename)][f"mask_{key}"] = detected_face.get_landmark_mask( - self._size, - key, - aligned=True, - centering=self._centering, - dilation=self._size // 32, - blur_kernel=self._size // 16, - as_zip=True) - -class TrainingDataGenerator(): # pylint:disable=too-few-public-methods +class TrainingDataGenerator(DataGenerator): """ A Training Data Generator for compiling data for feeding to a model. This class is called from :mod:`plugins.train.trainer._base` and launches a background @@ -434,406 +399,570 @@ class TrainingDataGenerator(): # pylint:disable=too-few-public-methods Parameters ---------- - model_input_size: int - The expected input size for the model. It is assumed that the input to the model is always - a square image. This is the size, in pixels, of the `width` and the `height` of the input - to the model. - model_output_shapes: list - A list of tuples defining the output shapes from the model, in the order that the outputs - are returned. The tuples should be in (`height`, `width`, `channels`) format. - coverage_ratio: float - The ratio of the training image to be trained on. Dictates how much of the image will be - cropped out. E.G: a coverage ratio of 0.625 will result in cropping a 160px box from a - 256px image (:math:`256 * 0.625 = 160`). - color_order: ["rgb", "bgr"] - The color order that the model expects as input - augment_color: bool - ``True`` if color is to be augmented, otherwise ``False`` - no_flip: bool - ``True`` if the image shouldn't be randomly flipped as part of augmentation, otherwise - ``False`` - no_warp: bool - ``True`` if the image shouldn't be warped as part of augmentation, otherwise ``False`` - warp_to_landmarks: bool - ``True`` if the random warp method should warp to similar landmarks from the other side, - ``False`` if the standard random warp method should be used. - face_cache: dict - A thread safe dictionary containing a cache of information relating to all faces being - trained on + model: :class:`~plugins.train.model.ModelBase` + The model that this data generator is feeding config: dict The configuration `dict` generated from :file:`config.train.ini` containing the trainer plugin configuration options. + side: {'a' or 'b'} + The side of the model that this iterator is for. + images: list + A list of image paths that will be used to compile the final augmented data from. + batch_size: int + The batch size for this iterator. Images will be returned in :class:`numpy.ndarray` + objects of this size from the iterator. """ - def __init__(self, model_input_size, model_output_shapes, coverage_ratio, color_order, - augment_color, no_flip, no_warp, warp_to_landmarks, config): - logger.debug("Initializing %s: (model_input_size: %s, model_output_shapes: %s, " - "coverage_ratio: %s, color_order: %s, augment_color: %s, no_flip: %s, " - "no_warp: %s, warp_to_landmarks: %s, config: %s)", - self.__class__.__name__, model_input_size, model_output_shapes, - coverage_ratio, color_order, augment_color, no_flip, no_warp, - warp_to_landmarks, config) - self._config = config - self._model_input_size = model_input_size - self._model_output_shapes = model_output_shapes - self._coverage_ratio = coverage_ratio - self._color_order = color_order.lower() - self._augment_color = augment_color - self._no_flip = no_flip - self._warp_to_landmarks = warp_to_landmarks - self._no_warp = no_warp - - # Batchsize and processing class are set when this class is called by a feeder - # from lib.training_data - self._batchsize = 0 - self._face_cache = None - self._nearest_landmarks = dict() - self._processing = None - logger.debug("Initialized %s", self.__class__.__name__) + def __init__(self, + config: dict[str, ConfigValueType], + model: ModelBase, + side: T.Literal["a", "b"], + images: list[str], + batch_size: int) -> None: + super().__init__(config, model, side, images, batch_size) + self._augment_color = not model.command_line_arguments.no_augment_color + self._no_flip = model.command_line_arguments.no_flip + self._no_warp = model.command_line_arguments.no_warp + self._warp_to_landmarks = (not self._no_warp + and model.command_line_arguments.warp_to_landmarks) - def minibatch_ab(self, images, batchsize, side, - do_shuffle=True, is_preview=False, is_timelapse=False): - """ A Background iterator to return augmented images, samples and targets. + if self._warp_to_landmarks: + self._face_cache.pre_fill(images, side) + self._processing = ImageAugmentation(batch_size, + self._process_size, + self._config) + self._nearest_landmarks: dict[str, tuple[str, ...]] = {} + logger.debug("Initialized %s", self.__class__.__name__) - The exit point from this class and the sole attribute that should be referenced. Called - from :mod:`plugins.train.trainer._base`. Returns an iterator that yields images for - training, preview and time-lapses. + def _create_targets(self, batch: np.ndarray) -> list[np.ndarray]: + """ Compile target images, with masks, for the model output sizes. Parameters ---------- - images: list - A list of image paths that will be used to compile the final augmented data from. - batchsize: int - The batchsize for this iterator. Images will be returned in :class:`numpy.ndarray` - objects of this size from the iterator. - side: {'a' or 'b'} - The side of the model that this iterator is for. - do_shuffle: bool, optional - Whether data should be shuffled prior to loading from disk. If true, each time the full - list of filenames are processed, the data will be reshuffled to make sure they are not - returned in the same order. Default: ``True`` - is_preview: bool, optional - Indicates whether this iterator is generating preview images. If ``True`` then certain - augmentations will not be performed. Default: ``False`` - is_timelapse: bool optional - Indicates whether this iterator is generating time-lapse images. If ``True``, then - certain augmentations will not be performed. Default: ``False`` + batch: :class:`numpy.ndarray` + This should be a 4-dimensional array of training images in the format (`batch size`, + `height`, `width`, `channels`). Targets should be requested after performing image + transformations but prior to performing warps. The 4th channel should be the mask. + Any channels above the 4th should be any additional area masks (e.g. eye/mouth) that + are required. - Yields - ------ - dict - The following items are contained in each `dict` yielded from this iterator: - - * **feed** (:class:`numpy.ndarray`) - The feed for the model. The array returned is \ - in the format (`batchsize`, `height`, `width`, `channels`). This is the :attr:`x` \ - parameter for :func:`keras.models.model.train_on_batch`. - - * **targets** (`list`) - A list of 4-dimensional :class:`numpy.ndarray` objects in \ - the order and size of each output of the model as defined in \ - :attr:`model_output_shapes`. the format of these arrays will be (`batchsize`, \ - `height`, `width`, `3`). This is the :attr:`y` parameter for \ - :func:`keras.models.model.train_on_batch` **NB:** masks are not included in the \ - `targets` list. If required for feeding into the Keras model, they will need to be \ - added to this list in :mod:`plugins.train.trainer._base` from the `masks` key. - - * **masks** (:class:`numpy.ndarray`) - A 4-dimensional array containing the target \ - masks in the format (`batchsize`, `height`, `width`, `1`). - - * **samples** (:class:`numpy.ndarray`) - A 4-dimensional array containing the samples \ - for feeding to the model's predict function for generating preview and time-lapse \ - samples. The array will be in the format (`batchsize`, `height`, `width`, \ - `channels`). **NB:** This item will only exist in the `dict` if :attr:`is_preview` \ - or :attr:`is_timelapse` is ``True`` + Returns + ------- + list + List of 4-dimensional target images, at all model output sizes, with masks compiled + into channels 4+ for each output size """ - logger.debug("Queue batches: (image_count: %s, batchsize: %s, side: '%s', do_shuffle: %s, " - "is_preview, %s, is_timelapse: %s)", len(images), batchsize, side, do_shuffle, - is_preview, is_timelapse) - self._batchsize = batchsize - self._face_cache = _get_cache(side, images, self._config) - self._processing = ImageAugmentation(batchsize, - is_preview or is_timelapse, - self._model_input_size, - self._model_output_shapes, - self._coverage_ratio, - self._config) + logger.trace("Compiling targets: batch shape: %s", # type:ignore[attr-defined] + batch.shape) + if len(self._output_sizes) == 1 and self._output_sizes[0] == self._process_size: + # Rolling buffer here makes next to no difference, so just create array on the fly + retval = [self._to_float32(batch)] + else: + retval = [self._to_float32(np.array([cv2.resize(image, + (size, size), + interpolation=cv2.INTER_AREA) + for image in batch])) + for size in self._output_sizes] + logger.trace("Processed targets: %s", # type:ignore[attr-defined] + [t.shape for t in retval]) + return retval - if self._warp_to_landmarks and not self._face_cache.partially_loaded: - self._face_cache.pre_fill(images, side) + def process_batch(self, + filenames: list[str], + images: np.ndarray, + detected_faces: list[DetectedFace], + batch: np.ndarray) -> BatchType: + """ Performs the augmentation and compiles target images and samples. - args = (images, side, do_shuffle, batchsize) - batcher = BackgroundGenerator(self._minibatch, thread_count=2, args=args) - return batcher.iterator() + Parameters + ---------- + filenames: list + List of full paths to image file names for a single batch + images: :class:`numpy.ndarray` + The batch of faces corresponding to the filenames + detected_faces: list + List of :class:`~lib.align.DetectedFace` objects with aligned data and masks loaded for + the current batch + batch: :class:`numpy.ndarray` + The pre-allocated batch with images and masks populated for the selected coverage and + centering - # << INTERNAL METHODS >> # - def _validate_samples(self, data): - """ Ensures that the total number of images within :attr:`images` is greater or equal to - the selected :attr:`batchsize`. Raises an exception if this is not the case. """ - length = len(data) - msg = ("Number of images is lower than batch-size (Note that too few " - "images may lead to bad training). # images: {}, " - "batch-size: {}".format(length, self._batchsize)) - try: - assert length >= self._batchsize, msg - except AssertionError as err: - msg += ("\nYou should increase the number of images in your training set or lower " - "your batch-size.") - raise FaceswapError(msg) from err + Returns + ------- + feed: :class:`numpy.ndarray` + 4-dimensional array of faces to feed the training the model (:attr:`x` parameter for + :func:`keras.models.model.train_on_batch`.). The array returned is in the format + (`batch size`, `height`, `width`, `channels`). + targets: list + List of 4-dimensional :class:`numpy.ndarray` objects in the order and size of each + output of the model. The format of these arrays will be (`batch size`, `height`, + `width`, `x`). This is the :attr:`y` parameter for + :func:`keras.models.model.train_on_batch`. The number of channels here will vary. + The first 3 channels are (rgb/bgr). The 4th channel is the face mask. Any subsequent + channels are area masks (e.g. eye/mouth masks) + """ + logger.trace("Process training: (side: '%s', filenames: '%s', images: %s, " # type:ignore + "batch: %s, detected_faces: %s)", self._side, filenames, images.shape, + batch.shape, len(detected_faces)) - def _minibatch(self, images, side, do_shuffle, batchsize): - """ A generator function that yields the augmented, target and sample images. - see :func:`minibatch_ab` for more details on the output. """ - logger.debug("Loading minibatch generator: (image_count: %s, side: '%s', do_shuffle: %s)", - len(images), side, do_shuffle) - self._validate_samples(images) + # Color Augmentation of the image only + if self._augment_color: + batch[..., :3] = self._processing.color_adjust(batch[..., :3]) - def _img_iter(imgs): - while True: - if do_shuffle: - shuffle(imgs) - for img in imgs: - yield img + # Random Transform and flip + self._processing.transform(batch) - img_iter = _img_iter(images) - while True: - img_paths = [next(img_iter) for _ in range(batchsize)] - yield self._process_batch(img_paths, side) + if not self._no_flip: + self._processing.random_flip(batch) - logger.debug("Finished minibatch generator: (side: '%s')", side) + # Switch color order for RGB models + self._set_color_order(batch) - def _process_batch(self, filenames, side): - """ Performs the augmentation and compiles target images and samples. + # Get Targets + targets = self._create_targets(batch) - If this is the first time a face has been loaded, then it's meta data is extracted from the - png header and added to :attr:`_face_cache` + # TODO Look at potential for applying mask on input + # Random Warp + if self._warp_to_landmarks: + landmarks = np.array([face.aligned.landmarks for face in detected_faces]) + batch_dst_pts = self._get_closest_match(filenames, landmarks) + warp_kwargs = {"batch_src_points": landmarks, "batch_dst_points": batch_dst_pts} + else: + warp_kwargs = {} + + warped = batch[..., :3] if self._no_warp else self._processing.warp( + batch[..., :3], + self._warp_to_landmarks, + **warp_kwargs) + + if self._model_input_size != self._process_size: + feed = self._to_float32(np.array([cv2.resize(image, + (self._model_input_size, + self._model_input_size), + interpolation=cv2.INTER_AREA) + for image in warped])) + else: + feed = self._to_float32(warped) + + return feed, targets - See - :func:`minibatch_ab` for more details on the output. + def _get_closest_match(self, filenames: list[str], batch_src_points: np.ndarray) -> np.ndarray: + """ Only called if the :attr:`_warp_to_landmarks` is ``True``. Gets the closest + matched 68 point landmarks from the opposite training set. Parameters ---------- filenames: list - List of full paths to image file names - side: str - The side of the model being trained on (`a` or `b`) - """ - logger.trace("Process batch: (filenames: '%s', side: '%s')", filenames, side) + Filenames for current batch + batch_src_points: :class:`np.ndarray` + The source landmarks for the current batch - if not self._face_cache.cache_full: - batch = self._face_cache.cache_metadata(filenames) - else: - batch = read_image_batch(filenames) + Returns + ------- + :class:`np.ndarray` + Randomly selected closest matches from the other side's landmarks + """ + logger.trace( # type:ignore[attr-defined] + "Retrieving closest matched landmarks: (filenames: '%s', src_points: '%s')", + filenames, batch_src_points) + lm_side: T.Literal["a", "b"] = "a" if self._side == "b" else "b" + other_cache = get_cache(lm_side) + landmarks = other_cache.aligned_landmarks - cache = self._face_cache.get_items(filenames) - batch, landmarks = self._crop_to_center(filenames, cache, batch, side) - batch = self._apply_mask(filenames, cache, batch, side) - processed = dict() + try: + closest_matches = [self._nearest_landmarks[os.path.basename(filename)] + for filename in filenames] + except KeyError: + # Resize mismatched training image size landmarks + sizes = {side: cache.size for side, cache in zip((self._side, lm_side), + (self._face_cache, other_cache))} + if len(set(sizes.values())) > 1: + scale = sizes[self._side] / sizes[lm_side] + landmarks = {key: lms * scale for key, lms in landmarks.items()} + closest_matches = self._cache_closest_matches(filenames, batch_src_points, landmarks) - # Initialize processing training size on first image - if not self._processing.initialized: - self._processing.initialize(batch.shape[1]) + batch_dst_points = np.array([landmarks[choice(fname)] for fname in closest_matches]) + logger.trace("Returning: (batch_dst_points: %s)", # type:ignore[attr-defined] + batch_dst_points.shape) + return batch_dst_points - # Get Landmarks prior to manipulating the image - if self._warp_to_landmarks: - batch_dst_pts = self._get_closest_match(filenames, side, landmarks) - warp_kwargs = dict(batch_src_points=landmarks, batch_dst_points=batch_dst_pts) - else: - warp_kwargs = dict() + def _cache_closest_matches(self, + filenames: list[str], + batch_src_points: np.ndarray, + landmarks: dict[str, np.ndarray]) -> list[tuple[str, ...]]: + """ Cache the nearest landmarks for this batch - # Color Augmentation of the image only - if self._augment_color: - batch[..., :3] = self._processing.color_adjust(batch[..., :3]) + Parameters + ---------- + filenames: list + Filenames for current batch + batch_src_points: :class:`np.ndarray` + The source landmarks for the current batch + landmarks: dict + The destination landmarks with associated filenames - # Random Transform and flip - batch = self._processing.transform(batch) - if not self._no_flip: - batch = self._processing.random_flip(batch) + """ + logger.trace("Caching closest matches") # type:ignore + dst_landmarks = list(landmarks.items()) + dst_points = np.array([lm[1] for lm in dst_landmarks]) + batch_closest_matches: list[tuple[str, ...]] = [] - # Switch color order for RGB models - if self._color_order == "rgb": - batch[..., :3] = batch[..., [2, 1, 0]] + for filename, src_points in zip(filenames, batch_src_points): + closest = (np.mean(np.square(src_points - dst_points), axis=(1, 2))).argsort()[:10] + closest_matches = tuple(dst_landmarks[i][0] for i in closest) + self._nearest_landmarks[os.path.basename(filename)] = closest_matches + batch_closest_matches.append(closest_matches) + logger.trace("Cached closest matches") # type:ignore + return batch_closest_matches - # Add samples to output if this is for display - if self._processing.is_display: - processed["samples"] = batch[..., :3].astype("float32") / 255.0 - # Get Targets - processed.update(self._processing.get_targets(batch)) +class PreviewDataGenerator(DataGenerator): + """ Generator for compiling images for generating previews. - # Random Warp # TODO change masks to have a input mask and a warped target mask - if self._no_warp: - processed["feed"] = [self._processing.skip_warp(batch[..., :3])] - else: - processed["feed"] = [self._processing.warp(batch[..., :3], - self._warp_to_landmarks, - **warp_kwargs)] + This class is called from :mod:`plugins.train.trainer._base` and launches a background + iterator that compiles sample preview data for feeding the model's predict function and for + display. - logger.trace("Processed batch: (filenames: %s, side: '%s', processed: %s)", - filenames, - side, - {k: v.shape if isinstance(v, np.ndarray) else[i.shape for i in v] - for k, v in processed.items()}) - return processed + Parameters + ---------- + model: :class:`~plugins.train.model.ModelBase` + The model that this data generator is feeding + config: dict + The configuration `dict` generated from :file:`config.train.ini` containing the trainer + plugin configuration options. + side: {'a' or 'b'} + The side of the model that this iterator is for. + images: list + A list of image paths that will be used to compile the final images. + batch_size: int + The batch size for this iterator. Images will be returned in :class:`numpy.ndarray` + objects of this size from the iterator. + """ + def _create_samples(self, + images: np.ndarray, + detected_faces: list[DetectedFace]) -> list[np.ndarray]: + """ Compile the 'sample' images. These are the 100% coverage images which hold the model + output in the preview window. - def _crop_to_center(self, filenames, cache, batch, side): - """ Crops the training image out of the full extract image based on the centering used in - the user's configuration settings. + Parameters + ---------- + images: :class:`numpy.ndarray` + The original batch of images as loaded from disk. + detected_faces: list + List of :class:`~lib.align.DetectedFace` for the current batch - If legacy extract images are being used then this just returns the extracted batch with - their corresponding landmarks. + Returns + ------- + list + List of 4-dimensional target images, at final model output size + """ + logger.trace( # type:ignore[attr-defined] + "Compiling samples: images shape: %s, detected_faces: %s ", + images.shape, len(detected_faces)) + output_size = self._output_sizes[-1] + full_size = 2 * int(np.rint((output_size / self._coverage_ratio) / 2)) + + assert self._config["centering"] in T.get_args(CenteringType) + retval = np.empty((full_size, full_size, 3), dtype="float32") + retval = self._to_float32(np.array([ + AlignedFace(face.landmarks_xy, + image=images[idx], + centering=T.cast(CenteringType, + self._config["centering"]), + size=full_size, + dtype="uint8", + is_aligned=True).face + for idx, face in enumerate(detected_faces)])) + + logger.trace("Processed samples: %s", retval.shape) # type:ignore[attr-defined] + return [retval] + + def process_batch(self, + filenames: list[str], + images: np.ndarray, + detected_faces: list[DetectedFace], + batch: np.ndarray) -> BatchType: + """ Creates the full size preview images and the sub-cropped images for feeding the model's + predict function. Parameters ---------- filenames: list - The list of filenames that correspond to this batch - cache: list - The list of cached items (aligned faces, masks etc.) corresponding to the batch + List of full paths to image file names for a single batch + images: :class:`numpy.ndarray` + The batch of faces corresponding to the filenames + detected_faces: list + List of :class:`~lib.align.DetectedFace` objects with aligned data and masks loaded for + the current batch batch: :class:`numpy.ndarray` - The batch of faces that have been loaded from disk - side: str - '"a"' or '"b"' the side that is being processed + The pre-allocated batch with images and masks populated for the selected coverage and + centering Returns ------- - batch: :class:`numpy.ndarray` - The centered faces cropped out of the loaded batch - landmarks: :class:`numpy.ndarray` - The aligned landmarks for this batch. NB: The aligned landmarks do not directly - correspond to the size of the extracted face. They are scaled to the source training - image, not the sub-image. - - Raises - ------ - FaceswapError - If Alignment information is not available for any of the images being loaded in - the batch + feed: :class:`numpy.ndarray` + List of 4-dimensional :class:`numpy.ndarray` objects at model output size for feeding + the model's predict function. The first 3 channels are (rgb/bgr). The 4th channel is + the face mask. + samples: list + 4-dimensional array containing the 100% coverage images at the model's centering for + for generating previews. The array returned is in the format + (`batch size`, `height`, `width`, `channels`). """ - logger.trace("Cropping training images info: (filenames: %s, side: '%s')", filenames, side) - aligned = [item["aligned_face"] for item in cache] + logger.trace("Process preview: (side: '%s', filenames: '%s', images: %s, " # type:ignore + "batch: %s, detected_faces: %s)", self._side, filenames, images.shape, + batch.shape, len(detected_faces)) - if self._face_cache.extract_version == 1.0: - # Legacy extract. Don't crop, just return batch with landmarks - return batch, np.array([face.landmarks for face in aligned]) + # Switch color order for RGB models + self._set_color_order(batch) + self._set_color_order(images) - landmarks = np.array([face.landmarks for face in aligned]) - cropped = np.array([align.extract_face(img) for align, img in zip(aligned, batch)]) - return cropped, landmarks + if not self._use_mask: + mask = np.zeros_like(batch[..., 0])[..., None] + 255 + batch = np.concatenate([batch, mask], axis=-1) - def _apply_mask(self, filenames, cache, batch, side): - """ Applies the mask to the 4th channel of the image. If masks are not being used - applies a dummy all ones mask. + feed = self._to_float32(batch[..., :4]) # Don't resize here: we want masks at output res. - If the configuration options `eye_multiplier` and/or `mouth_multiplier` are greater than 1 - then these masks are applied to the final channels of the batch respectively. + # If user sets model input size as larger than output size, the preview will error, so + # resize in these rare instances + out_size = max(self._output_sizes) + if self._process_size > out_size: + feed = np.array([cv2.resize(img, (out_size, out_size), interpolation=cv2.INTER_AREA) + for img in feed]) + + samples = self._create_samples(images, detected_faces) + + return feed, samples + + +class Feeder(): + """ Handles the processing of a Batch for training the model and generating samples. + + Parameters + ---------- + images: dict + The list of full paths to the training images for this :class:`_Feeder` for each side + model: plugin from :mod:`plugins.train.model` + The selected model that will be running this trainer + batch_size: int + The size of the batch to be processed for each side at each iteration + config: dict + The configuration for this trainer + include_preview: bool, optional + ``True`` to create a feeder for generating previews. Default: ``True`` + """ + def __init__(self, + images: dict[T.Literal["a", "b"], list[str]], + model: ModelBase, + batch_size: int, + config: dict[str, ConfigValueType], + include_preview: bool = True) -> None: + logger.debug("Initializing %s: num_images: %s, batch_size: %s, config: %s, " + "include_preview: %s)", self.__class__.__name__, + {k: len(v) for k, v in images.items()}, batch_size, config, include_preview) + self._model = model + self._images = images + self._batch_size = batch_size + self._config = config + self._feeds = { + side: self._load_generator(side, False).minibatch_ab() + for side in T.get_args(T.Literal["a", "b"])} + + self._display_feeds = {"preview": self._set_preview_feed() if include_preview else {}, + "timelapse": {}} + logger.debug("Initialized %s:", self.__class__.__name__) + + def _load_generator(self, + side: T.Literal["a", "b"], + is_display: bool, + batch_size: int | None = None, + images: list[str] | None = None) -> DataGenerator: + """ Load the :class:`~lib.training_data.TrainingDataGenerator` for this feeder. Parameters ---------- - filenames: list - The list of filenames that correspond to this batch - cache: list - The list of cached items (aligned faces, masks etc.) corresponding to the batch - batch: :class:`numpy.ndarray` - The batch of faces that have been loaded from disk - side: str - '"a"' or '"b"' the side that is being processed + side: ["a", "b"] + The side of the model to load the generator for + is_display: bool + ``True`` if the generator is for creating preview/time-lapse images. ``False`` if it is + for creating training images + batch_size: int, optional + If ``None`` then the batch size selected in command line arguments is used, otherwise + the batch size provided here is used. + images: list, optional. Default: ``None`` + If provided then this will be used as the list of images for the generator. If ``None`` + then the training folder images for the side will be used. Default: ``None`` Returns ------- - :class:`numpy.ndarray` - The batch with masks applied to the final channels + :class:`~lib.training_data.TrainingDataGenerator` + The training data generator """ - logger.trace("Input filenames: %s, batch shape: %s, side: %s", - filenames, batch.shape, side) - size = batch.shape[1] - - for key in ("mask", "mask_eyes", "mask_mouth"): - lookup = cache[0].get(key) - if lookup is None and key != "mask": - continue - - if lookup is None and key == "mask": - logger.trace("Creating dummy masks. side: %s", side) - masks = np.ones_like(batch[..., :1], dtype=batch.dtype) - else: - logger.trace("Obtaining masks for batch. (key: %s side: %s)", key, side) - - masks = np.array([self._get_mask(item[key], size) - for item in cache], dtype=batch.dtype) - masks = self._resize_masks(size, masks) - logger.trace("masks: (key: %s, shape: %s)", key, masks.shape) - batch = np.concatenate((batch, masks), axis=-1) - logger.trace("Output batch shape: %s, side: %s", batch.shape, side) - return batch - - @classmethod - def _get_mask(cls, item, size): - """ Decompress zipped eye and mouth masks, or return the stored mask + logger.debug("Loading generator, side: %s, is_display: %s, batch_size: %s", + side, is_display, batch_size) + generator = PreviewDataGenerator if is_display else TrainingDataGenerator + retval = generator(self._config, + self._model, + side, + self._images[side] if images is None else images, + self._batch_size if batch_size is None else batch_size) + return retval + + def _set_preview_feed(self) -> dict[T.Literal["a", "b"], Generator[BatchType, None, None]]: + """ Set the preview feed for this feeder. + + Creates a generator from :class:`lib.training_data.PreviewDataGenerator` specifically + for previews for the feeder. + + Returns + ------- + dict + The side ("a" or "b") as key, :class:`~lib.training_data.PreviewDataGenerator` as + value. + """ + retval: dict[T.Literal["a", "b"], Generator[BatchType, None, None]] = {} + num_images = self._config.get("preview_images", 14) + assert isinstance(num_images, int) + for side in T.get_args(T.Literal["a", "b"]): + logger.debug("Setting preview feed: (side: '%s')", side) + preview_images = min(max(num_images, 2), 16) + batchsize = min(len(self._images[side]), preview_images) + retval[side] = self._load_generator(side, + True, + batch_size=batchsize).minibatch_ab() + return retval + + def get_batch(self) -> tuple[list[list[np.ndarray]], ...]: + """ Get the feed data and the targets for each training side for feeding into the model's + train function. + + Returns + ------- + model_inputs: list + The inputs to the model for each side A and B + model_targets: list + The targets for the model for each side A and B + """ + model_inputs: list[list[np.ndarray]] = [] + model_targets: list[list[np.ndarray]] = [] + for side in ("a", "b"): + side_feed, side_targets = next(self._feeds[side]) + if self._model.config["learn_mask"]: # Add the face mask as it's own target + side_targets += [side_targets[-1][..., 3][..., None]] + logger.trace( # type:ignore[attr-defined] + "side: %s, input_shapes: %s, target_shapes: %s", + side, side_feed.shape, [i.shape for i in side_targets]) + model_inputs.append([side_feed]) + model_targets.append(side_targets) + + return model_inputs, model_targets + + def generate_preview(self, is_timelapse: bool = False + ) -> dict[T.Literal["a", "b"], list[np.ndarray]]: + """ Generate the images for preview window or timelapse Parameters ---------- - item: :class:`lib.align.Mask` or `bytes` - Either a stored face mask object or a zipped eye or mouth mask - size: int - The size of the stored eye or mouth mask for reshaping + is_timelapse, bool, optional + ``True`` if preview is to be generated for a Timelapse otherwise ``False``. + Default: ``False`` Returns ------- - class:`numpy.ndarray` - The decompressed mask + dict + Dictionary for side A and B of list of numpy arrays corresponding to the + samples, targets and masks for this preview """ - if isinstance(item, bytes): - retval = np.frombuffer(decompress(item), dtype="uint8").reshape(size, size, 1) - else: - retval = item.mask - return retval + logger.debug("Generating preview (is_timelapse: %s)", is_timelapse) + + batchsizes: list[int] = [] + feed: dict[T.Literal["a", "b"], np.ndarray] = {} + samples: dict[T.Literal["a", "b"], np.ndarray] = {} + masks: dict[T.Literal["a", "b"], np.ndarray] = {} + + # MyPy can't recurse into nested dicts to get the type :( + iterator = T.cast(dict[T.Literal["a", "b"], "Generator[BatchType, None, None]"], + self._display_feeds["timelapse" if is_timelapse else "preview"]) + for side in T.get_args(T.Literal["a", "b"]): + side_feed, side_samples = next(iterator[side]) + batchsizes.append(len(side_samples[0])) + samples[side] = side_samples[0] + feed[side] = side_feed[..., :3] + masks[side] = side_feed[..., 3][..., None] + + logger.debug("Generated samples: is_timelapse: %s, images: %s", is_timelapse, + {key: {k: v.shape for k, v in item.items()} + for key, item + in zip(("feed", "samples", "sides"), (feed, samples, masks))}) + return self.compile_sample(min(batchsizes), feed, samples, masks) + + def compile_sample(self, + image_count: int, + feed: dict[T.Literal["a", "b"], np.ndarray], + samples: dict[T.Literal["a", "b"], np.ndarray], + masks: dict[T.Literal["a", "b"], np.ndarray] + ) -> dict[T.Literal["a", "b"], list[np.ndarray]]: + """ Compile the preview samples for display. - @classmethod - def _resize_masks(cls, target_size, masks): - """ Resize the masks to the target size """ - logger.trace("target size: %s, masks shape: %s", target_size, masks.shape) - mask_size = masks.shape[1] - if target_size == mask_size: - logger.trace("Mask and targets the same size. Not resizing") - return masks - interpolator = cv2.INTER_CUBIC if mask_size < target_size else cv2.INTER_AREA - masks = np.array([cv2.resize(mask, - (target_size, target_size), - interpolation=interpolator)[..., None] - for mask in masks]) - logger.trace("Resized masks: %s", masks.shape) - return masks - - def _get_closest_match(self, filenames, side, batch_src_points): - """ Only called if the :attr:`_warp_to_landmarks` is ``True``. Gets the closest - matched 68 point landmarks from the opposite training set. """ - logger.trace("Retrieving closest matched landmarks: (filenames: '%s', src_points: '%s'", - filenames, batch_src_points) - lm_side = "a" if side == "b" else "b" - landmarks = _FACE_CACHES[lm_side].aligned_landmarks - - closest_matches = [self._nearest_landmarks.get(os.path.basename(filename)) - for filename in filenames] - if None in closest_matches: - # Resize mismatched training image size landmarks - sizes = {side: cache.crop_size for side, cache in _FACE_CACHES.items()} - if len(set(sizes.values())) > 1: - scale = sizes[side] / sizes[lm_side] - landmarks = {key: lms * scale for key, lms in landmarks.items()} - closest_matches = self._cache_closest_matches(filenames, batch_src_points, landmarks) + Parameters + ---------- + image_count: int + The number of images to limit the sample output to. + feed: dict + Dictionary for side "a", "b" of :class:`numpy.ndarray`. The images that should be fed + into the model for obtaining a prediction + samples: dict + Dictionary for side "a", "b" of :class:`numpy.ndarray`. The 100% coverage target images + that should be used for creating the preview. + masks: dict + Dictionary for side "a", "b" of :class:`numpy.ndarray`. The masks that should be used + for creating the preview. - batch_dst_points = np.array([landmarks[choice(fname)] for fname in closest_matches]) - logger.trace("Returning: (batch_dst_points: %s)", batch_dst_points.shape) - return batch_dst_points + Returns + ------- + list + The list of samples, targets and masks as :class:`numpy.ndarrays` for creating a + preview image + """ + num_images = self._config.get("preview_images", 14) + assert isinstance(num_images, int) + num_images = min(image_count, num_images) + retval: dict[T.Literal["a", "b"], list[np.ndarray]] = {} + for side in T.get_args(T.Literal["a", "b"]): + logger.debug("Compiling samples: (side: '%s', samples: %s)", side, num_images) + retval[side] = [feed[side][0:num_images], + samples[side][0:num_images], + masks[side][0:num_images]] + logger.debug("Compiled Samples: %s", {k: [i.shape for i in v] for k, v in retval.items()}) + return retval - def _cache_closest_matches(self, filenames, batch_src_points, landmarks): - """ Cache the nearest landmarks for this batch """ - logger.trace("Caching closest matches") - dst_landmarks = list(landmarks.items()) - dst_points = np.array([lm[1] for lm in dst_landmarks]) - batch_closest_matches = list() + def set_timelapse_feed(self, + images: dict[T.Literal["a", "b"], list[str]], + batch_size: int) -> None: + """ Set the time-lapse feed for this feeder. - for filename, src_points in zip(filenames, batch_src_points): - closest = (np.mean(np.square(src_points - dst_points), axis=(1, 2))).argsort()[:10] - closest_matches = tuple(dst_landmarks[i][0] for i in closest) - self._nearest_landmarks[os.path.basename(filename)] = closest_matches - batch_closest_matches.append(closest_matches) - logger.trace("Cached closest matches") - return batch_closest_matches + Creates a generator from :class:`lib.training_data.PreviewDataGenerator` specifically + for generating time-lapse previews for the feeder. + + Parameters + ---------- + images: dict + The list of full paths to the images for creating the time-lapse for each side + batch_size: int + The number of images to be used to create the time-lapse preview. + """ + logger.debug("Setting time-lapse feed: (input_images: '%s', batch_size: %s)", + images, batch_size) + + # MyPy can't recurse into nested dicts to get the type :( + iterator = T.cast(dict[T.Literal["a", "b"], "Generator[BatchType, None, None]"], + self._display_feeds["timelapse"]) + + for side in T.get_args(T.Literal["a", "b"]): + imgs = images[side] + logger.debug("Setting preview feed: (side: '%s', images: %s)", side, len(imgs)) + + iterator[side] = self._load_generator(side, + True, + batch_size=batch_size, + images=imgs).minibatch_ab(do_shuffle=False) + logger.debug("Set time-lapse feed: %s", self._display_feeds["timelapse"]) diff --git a/lib/training/lr_finder.py b/lib/training/lr_finder.py new file mode 100644 index 0000000000..d7c3298c69 --- /dev/null +++ b/lib/training/lr_finder.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +""" Learning Rate Finder for faceswap.py. """ +from __future__ import annotations +import logging +import os +import shutil +import typing as T +from datetime import datetime +from enum import Enum + +import tensorflow as tf +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +from tqdm import tqdm + +if T.TYPE_CHECKING: + from lib.config import ConfigValueType + from lib.training import Feeder + from plugins.train.model._base import ModelBase + +keras = tf.keras +K = keras.backend + +logger = logging.getLogger(__name__) + + +class LRStrength(Enum): + """ Enum for how aggressively to set the optimal learning rate """ + DEFAULT = 10 + AGGRESSIVE = 5 + EXTREME = 2.5 + + +class LearningRateFinder: + """ Learning Rate Finder + + Parameters + ---------- + model: :class:`tensorflow.keras.models.Model` + The keras model to find the optimal learning rate for + config: dict + The configuration options for the model + feeder: :class:`~lib.training.generator.Feeder` + The feeder for training the model + stop_factor: int + When to stop finding the optimal learning rate + beta: float + Amount to smooth loss by, for graphing purposes + """ + def __init__(self, + model: ModelBase, + config: dict[str, ConfigValueType], + feeder: Feeder, + stop_factor: int = 4, + beta: float = 0.98) -> None: + logger.debug("Initializing %s: (model: %s, config: %s, feeder: %s, stop_factor: %s, " + "beta: %s)", + self.__class__.__name__, model, config, feeder, stop_factor, beta) + + self._iterations = T.cast(int, config["lr_finder_iterations"]) + self._save_graph = config["lr_finder_mode"] in ("graph_and_set", "graph_and_exit") + self._strength = LRStrength[T.cast(str, config["lr_finder_strength"]).upper()].value + self._config = config + + self._start_lr = 1e-10 + end_lr = 1e+1 + + self._model = model + self._feeder = feeder + self._stop_factor = stop_factor + self._beta = beta + self._lr_multiplier: float = (end_lr / self._start_lr) ** (1.0 / self._iterations) + + self._metrics: dict[T.Literal["learning_rates", "losses"], list[float]] = { + "learning_rates": [], + "losses": []} + self._loss: dict[T.Literal["avg", "best"], float] = {"avg": 0.0, "best": 1e9} + + logger.debug("Initialized %s", self.__class__.__name__) + + def _on_batch_end(self, iteration: int, loss: float) -> None: + """ Learning rate actions to perform at the end of a batch + + Parameters + ---------- + iteration: int + The current iteration + loss: float + The loss value for the current batch + """ + learning_rate = K.get_value(self._model.model.optimizer.lr) + self._metrics["learning_rates"].append(learning_rate) + + self._loss["avg"] = (self._beta * self._loss["avg"]) + ((1 - self._beta) * loss) + smoothed = self._loss["avg"] / (1 - (self._beta ** iteration)) + self._metrics["losses"].append(smoothed) + + stop_loss = self._stop_factor * self._loss["best"] + + if iteration > 1 and smoothed > stop_loss: + self._model.model.stop_training = True + return + + if iteration == 1 or smoothed < self._loss["best"]: + self._loss["best"] = smoothed + + learning_rate *= self._lr_multiplier + + K.set_value(self._model.model.optimizer.lr, learning_rate) + + def _update_description(self, progress_bar: tqdm) -> None: + """ Update the description of the progress bar for the current iteration + + Parameters + ---------- + progress_bar: :class:`tqdm.tqdm` + The learning rate finder progress bar to update + """ + current = self._metrics['learning_rates'][-1] + best_idx = self._metrics["losses"].index(self._loss["best"]) + best = self._metrics["learning_rates"][best_idx] / self._strength + progress_bar.set_description(f"Current: {current:.1e} Best: {best:.1e}") + + def _train(self) -> None: + """ Train the model for the given number of iterations to find the optimal + learning rate and show progress""" + logger.info("Finding optimal learning rate...") + pbar = tqdm(range(1, self._iterations + 1), + desc="Current: N/A Best: N/A ", + leave=False) + for idx in pbar: + model_inputs, model_targets = self._feeder.get_batch() + loss: list[float] = self._model.model.train_on_batch(model_inputs, y=model_targets) + if np.isnan(loss[0]): + break + self._on_batch_end(idx, loss[0]) + self._update_description(pbar) + + def _reset_model(self, original_lr: float, new_lr: float) -> None: + """ Reset the model's weights to initial values, reset the model's optimizer and set the + learning rate + + Parameters + ---------- + original_lr: float + The model's original learning rate + new_lr: float + The discovered optimal learning rate + """ + self._model.state.update_session_config("learning_rate", new_lr) + self._model.state.save() + + logger.debug("Loading initial weights") + self._model.model.load_weights(self._model.io.filename) + + if self._config["lr_finder_mode"] == "graph_and_exit": + return + + opt_conf = self._model.model.optimizer.get_config() + logger.debug("Recompiling model to reset optimizer state. Optimizer config: %s", opt_conf) + new_opt = self._model.model.optimizer.__class__(**opt_conf) + self._model.model.compile(optimizer=new_opt, loss=self._model.model.loss) + + logger.info("Updating Learning Rate from %s to %s", f"{original_lr:.1e}", f"{new_lr:.1e}") + K.set_value(self._model.model.optimizer.lr, new_lr) + + def find(self) -> bool: + """ Find the optimal learning rate + + Returns + ------- + bool + ``True`` if the learning rate was succesfully discovered otherwise ``False`` + """ + if not self._model.io.model_exists: + self._model.io.save() + + original_lr = K.get_value(self._model.model.optimizer.lr) + K.set_value(self._model.model.optimizer.lr, self._start_lr) + + self._train() + print() + + best_idx = self._metrics["losses"].index(self._loss["best"]) + new_lr = self._metrics["learning_rates"][best_idx] / self._strength + if new_lr < 1e-9: + logger.error("The optimal learning rate could not be found. This is most likely " + "because you did not run the finder for enough iterations.") + shutil.rmtree(self._model.io.model_dir) + return False + + self._plot_loss() + self._reset_model(original_lr, new_lr) + return True + + def _plot_loss(self, skip_begin: int = 10, skip_end: int = 1) -> None: + """ Plot a graph of loss vs learning rate and save to the training folder + + Parameters + ---------- + skip_begin: int, optional + Number of iterations to skip at the start. Default: `10` + skip_end: int, optional + Number of iterations to skip at the end. Default: `1` + """ + if not self._save_graph: + return + + matplotlib.use("Agg") + lrs = self._metrics["learning_rates"][skip_begin:-skip_end] + losses = self._metrics["losses"][skip_begin:-skip_end] + plt.plot(lrs, losses, label="Learning Rate") + best_idx = self._metrics["losses"].index(self._loss["best"]) + best_lr = self._metrics["learning_rates"][best_idx] + for val, color in zip(LRStrength, ("g", "y", "r")): + l_r = best_lr / val.value + idx = lrs.index(next(r for r in lrs if r >= l_r)) + plt.plot(l_r, losses[idx], + f"{color}o", + label=f"{val.name.title()}: {l_r:.1e}") + + plt.xscale("log") + plt.xlabel("Learning Rate (Log Scale)") + plt.ylabel("Loss") + plt.title("Learning Rate Finder") + plt.legend() + + now = datetime.now().strftime("%Y-%m-%d_%H.%M.%S") + output = os.path.join(self._model.io.model_dir, f"learning_rate_finder_{now}.png") + logger.info("Saving Learning Rate Finder graph to: '%s'", output) + plt.savefig(output) diff --git a/lib/training/preview_cv.py b/lib/training/preview_cv.py new file mode 100644 index 0000000000..c0a6458af2 --- /dev/null +++ b/lib/training/preview_cv.py @@ -0,0 +1,189 @@ +#!/usr/bin/python +""" The pop up preview window for Faceswap. + +If Tkinter is installed, then this will be used to manage the preview image, otherwise we +fallback to opencv's imshow +""" +from __future__ import annotations +import logging +import typing as T + +from threading import Event, Lock +from time import sleep + +import cv2 + +if T.TYPE_CHECKING: + from collections.abc import Generator + import numpy as np + +logger = logging.getLogger(__name__) +TriggerType = dict[T.Literal["toggle_mask", "refresh", "save", "quit", "shutdown"], Event] +TriggerKeysType = T.Literal["m", "r", "s", "enter"] +TriggerNamesType = T.Literal["toggle_mask", "refresh", "save", "quit"] + + +class PreviewBuffer(): + """ A thread safe class for holding preview images """ + def __init__(self) -> None: + logger.debug("Initializing: %s", self.__class__.__name__) + self._images: dict[str, np.ndarray] = {} + self._lock = Lock() + self._updated = Event() + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def is_updated(self) -> bool: + """ bool: ``True`` when new images have been loaded into the preview buffer """ + return self._updated.is_set() + + def add_image(self, name: str, image: np.ndarray) -> None: + """ Add an image to the preview buffer in a thread safe way """ + logger.debug("Adding image: (name: '%s', shape: %s)", name, image.shape) + with self._lock: + self._images[name] = image + logger.debug("Added images: %s", list(self._images)) + self._updated.set() + + def get_images(self) -> Generator[tuple[str, np.ndarray], None, None]: + """ Get the latest images from the preview buffer. When iterator is exhausted clears the + :attr:`updated` event. + + Yields + ------ + name: str + The name of the image + :class:`numpy.ndarray` + The image in BGR format + """ + logger.debug("Retrieving images: %s", list(self._images)) + with self._lock: + for name, image in self._images.items(): + logger.debug("Yielding: '%s' (%s)", name, image.shape) + yield name, image + if self.is_updated: + logger.debug("Clearing updated event") + self._updated.clear() + logger.debug("Retrieved images") + + +class PreviewBase(): # pylint:disable=too-few-public-methods + """ Parent class for OpenCV and Tkinter Preview Windows + + Parameters + ---------- + preview_buffer: :class:`PreviewBuffer` + The thread safe object holding the preview images + triggers: dict, optional + Dictionary of event triggers for pop-up preview. Not required when running inside the GUI. + Default: `None` + """ + def __init__(self, + preview_buffer: PreviewBuffer, + triggers: TriggerType | None = None) -> None: + logger.debug("Initializing %s parent (triggers: %s)", + self.__class__.__name__, triggers) + self._triggers = triggers + self._buffer = preview_buffer + self._keymaps: dict[TriggerKeysType, TriggerNamesType] = {"m": "toggle_mask", + "r": "refresh", + "s": "save", + "enter": "quit"} + self._title = "" + logger.debug("Initialized %s parent", self.__class__.__name__) + + @property + def _should_shutdown(self) -> bool: + """ bool: ``True`` if the preview has received an external signal to shutdown otherwise + ``False`` """ + if self._triggers is None or not self._triggers["shutdown"].is_set(): + return False + logger.debug("Shutdown signal received") + return True + + def _launch(self) -> None: + """ Wait until an image is loaded into the preview buffer and call the child's + :func:`_display_preview` function """ + logger.debug("Launching %s", self.__class__.__name__) + while True: + if not self._buffer.is_updated: + logger.debug("Waiting for preview image") + sleep(1) + continue + break + logger.debug("Launching preview") + self._display_preview() + + def _display_preview(self) -> None: + """ Override for preview viewer's display loop """ + raise NotImplementedError() + + +class PreviewCV(PreviewBase): # pylint:disable=too-few-public-methods + """ Simple fall back preview viewer using OpenCV for when TKinter is not available + + Parameters + ---------- + preview_buffer: :class:`PreviewBuffer` + The thread safe object holding the preview images + triggers: dict + Dictionary of event triggers for pop-up preview. + """ + def __init__(self, + preview_buffer: PreviewBuffer, + triggers: TriggerType) -> None: + logger.debug("Unable to import Tkinter. Falling back to OpenCV") + super().__init__(preview_buffer, triggers=triggers) + self._triggers: TriggerType = self._triggers + self._windows: list[str] = [] + + self._lookup = {ord(key): val + for key, val in self._keymaps.items() if key != "enter"} + self._lookup[ord("\n")] = self._keymaps["enter"] + self._lookup[ord("\r")] = self._keymaps["enter"] + + self._launch() + + @property + def _window_closed(self) -> bool: + """ bool: ``True`` if any window has been closed otherwise ``False`` """ + retval = any(cv2.getWindowProperty(win, cv2.WND_PROP_VISIBLE) < 1 for win in self._windows) + if retval: + logger.debug("Window closed detected") + return retval + + def _check_keypress(self, key: int): + """ Check whether we have received a valid key press from OpenCV window and handle + accordingly. + + Parameters + ---------- + key_press: int + The key press received from OpenCV + """ + if not key or key == -1 or key not in self._lookup: + return + + if key == ord("r"): + print("") # Let log print on different line from loss output + logger.info("Refresh preview requested...") + + self._triggers[self._lookup[key]].set() + logger.debug("Processed keypress '%s'. Set event for '%s'", key, self._lookup[key]) + + def _display_preview(self): + """ Handle the displaying of the images currently in :attr:`_preview_buffer`""" + while True: + if self._buffer.is_updated or self._window_closed: + for name, image in self._buffer.get_images(): + logger.debug("showing image: '%s' (%s)", name, image.shape) + cv2.imshow(name, image) + self._windows.append(name) + + key = cv2.waitKey(1000) + self._check_keypress(key) + + if self._triggers["shutdown"].is_set(): + logger.debug("Shutdown received") + break + logger.debug("%s shutdown", self.__class__.__name__) diff --git a/lib/training/preview_tk.py b/lib/training/preview_tk.py new file mode 100644 index 0000000000..633ead57e6 --- /dev/null +++ b/lib/training/preview_tk.py @@ -0,0 +1,943 @@ +#!/usr/bin/python +""" The pop up preview window for Faceswap. + +If Tkinter is installed, then this will be used to manage the preview image, otherwise we +fallback to opencv's imshow +""" +from __future__ import annotations +import logging +import os +import sys +import tkinter as tk +import typing as T + +from datetime import datetime +from platform import system +from tkinter import ttk +from math import ceil, floor + +from PIL import Image, ImageTk + +import cv2 + +from .preview_cv import PreviewBase, TriggerKeysType + +if T.TYPE_CHECKING: + import numpy as np + from .preview_cv import PreviewBuffer, TriggerType + +logger = logging.getLogger(__name__) + + +class _Taskbar(): + """ Taskbar at bottom of Preview window + + Parameters + ---------- + parent: :class:`tkinter.Frame` + The parent frame that holds the canvas and taskbar + taskbar: :class:`tkinter.ttk.Frame` or ``None`` + None if preview is a pop-up window otherwise ttk.Frame if taskbar is managed by the GUI + """ + def __init__(self, parent: tk.Frame, taskbar: ttk.Frame | None) -> None: + logger.debug("Initializing %s (parent: '%s', taskbar: %s)", + self.__class__.__name__, parent, taskbar) + self._is_standalone = taskbar is None + self._gui_mapped: list[tk.Widget] = [] + self._frame = tk.Frame(parent) if taskbar is None else taskbar + + self._min_max_scales = (20, 400) + self._vars = {"save": tk.BooleanVar(), + "scale": tk.StringVar(), + "slider": tk.IntVar(), + "interpolator": tk.IntVar()} + self._interpolators = [("nearest_neighbour", cv2.INTER_NEAREST), + ("bicubic", cv2.INTER_CUBIC)] + self._scale = self._add_scale_combo() + self._slider = self._add_scale_slider() + self._add_interpolator_radio() + + if self._is_standalone: + self._add_save_button() + self._frame.pack(side=tk.BOTTOM, fill=tk.X, padx=2, pady=2) + + logger.debug("Initialized %s ('%s')", self.__class__.__name__, self) + + @property + def min_scale(self) -> int: + """ int: The minimum allowed scale """ + return self._min_max_scales[0] + + @property + def max_scale(self) -> int: + """ int: The maximum allowed scale """ + return self._min_max_scales[1] + + @property + def save_var(self) -> tk.BooleanVar: + """:class:`tkinter.IntVar`: Variable which is set to ``True`` when the save button has + been. pressed """ + retval = self._vars["save"] + assert isinstance(retval, tk.BooleanVar) + return retval + + @property + def scale_var(self) -> tk.StringVar: + """:class:`tkinter.StringVar`: The variable holding the currently selected "##%" formatted + percentage scaling amount displayed in the Combobox. """ + retval = self._vars["scale"] + assert isinstance(retval, tk.StringVar) + return retval + + @property + def slider_var(self) -> tk.IntVar: + """:class:`tkinter.IntVar`: The variable holding the currently selected percentage scaling + amount in the slider. """ + retval = self._vars["slider"] + assert isinstance(retval, tk.IntVar) + return retval + + @property + def interpolator_var(self) -> tk.IntVar: + """:class:`tkinter.IntVar`: The variable holding the CV2 Interpolator Enum. """ + retval = self._vars["interpolator"] + assert isinstance(retval, tk.IntVar) + return retval + + def _track_widget(self, widget: tk.Widget) -> None: + """ If running embedded in the GUI track the widgets so that they can be destroyed if + the preview is disabled """ + if self._is_standalone: + return + logger.debug("Tracking option bar widget for GUI: %s", widget) + self._gui_mapped.append(widget) + + def _add_scale_combo(self) -> ttk.Combobox: + """ Add a scale combo for selecting zoom amount. + + Returns + ------- + :class:`tkinter.ttk.Combobox` + The Combobox widget + """ + logger.debug("Adding scale combo") + self.scale_var.set("100%") + scale = ttk.Combobox(self._frame, + textvariable=self.scale_var, + values=["Fit"], + state="readonly", + width=10) + scale.pack(side=tk.RIGHT) + scale.bind("", self._clear_combo_focus) # Remove auto-focus on widget text box + self._track_widget(scale) + logger.debug("Added scale combo: '%s'", scale) + return scale + + def _clear_combo_focus(self, *args) -> None: # pylint:disable=unused-argument + """ Remove the highlighting and stealing of focus that the combobox annoyingly + implements. """ + logger.debug("Clearing scale combo focus") + self._scale.selection_clear() + self._scale.winfo_toplevel().focus_set() + logger.debug("Cleared scale combo focus") + + def _add_scale_slider(self) -> tk.Scale: + """ Add a scale slider for zooming the image. + + Returns + ------- + :class:`tkinter.Scale` + The scale widget + """ + logger.debug("Adding scale slider") + self.slider_var.set(100) + slider = tk.Scale(self._frame, + orient=tk.HORIZONTAL, + to=self.max_scale, + showvalue=False, + variable=self.slider_var, + command=self._on_slider_update) + slider.pack(side=tk.RIGHT) + self._track_widget(slider) + logger.debug("Added scale slider: '%s'", slider) + return slider + + def _add_interpolator_radio(self) -> None: + """ Add a radio box to choose interpolator """ + frame = tk.Frame(self._frame) + for text, mode in self._interpolators: + logger.debug("Adding %s radio button", text) + radio = tk.Radiobutton(frame, text=text, value=mode, variable=self.interpolator_var) + radio.pack(side=tk.LEFT, anchor=tk.W) + self._track_widget(radio) + + logger.debug("Added %s radio button", radio) + self.interpolator_var.set(cv2.INTER_NEAREST) + frame.pack(side=tk.RIGHT) + self._track_widget(frame) + + def _add_save_button(self) -> None: + """ Add a save button for saving out original preview """ + logger.debug("Adding save button") + button = tk.Button(self._frame, + text="Save", + cursor="hand2", + command=lambda: self.save_var.set(True)) + button.pack(side=tk.LEFT) + logger.debug("Added save burron: '%s'", button) + + def _on_slider_update(self, value) -> None: + """ Callback for when the scale slider is adjusted. Adjusts the combo box display to the + current slider value. + + Parameters + ---------- + value: int + The value that the slider has been set to + """ + self.scale_var.set(f"{value}%") + + def set_min_max_scale(self, min_scale: int, max_scale: int) -> None: + """ Set the minimum and maximum value that we allow an image to be scaled down to. This + impacts the slider and combo box min/max values: + + Parameters + ---------- + min_scale: int + The minimum percentage scale that is permitted + max_scale: int + The maximum percentage scale that is permitted + """ + logger.debug("Setting min/max scales: (min: %s, max: %s)", min_scale, max_scale) + self._min_max_scales = (min_scale, max_scale) + self._slider.config(from_=self.min_scale, to=max_scale) + scales = [10, 25, 50, 75, 100, 200, 300, 400, 800] + if min_scale not in scales: + scales.insert(0, min_scale) + if max_scale not in scales: + scales.append(max_scale) + choices = ["Fit", *[f"{x}%" for x in scales if self.max_scale >= x >= self.min_scale]] + self._scale.config(values=choices) + logger.debug("Set min/max scale. min_max_scales: %s, scale combo choices: %s", + self._min_max_scales, choices) + + def cycle_interpolators(self, *args) -> None: # pylint:disable=unused-argument + """ Cycle interpolators on a keypress callback """ + current = next(i for i in self._interpolators if i[1] == self.interpolator_var.get()) + next_idx = self._interpolators.index(current) + 1 + next_idx = 0 if next_idx == len(self._interpolators) else next_idx + self.interpolator_var.set(self._interpolators[next_idx][1]) + + def destroy_widgets(self) -> None: + """ Remove the taskbar widgets when the preview within the GUI has been disabled """ + if self._is_standalone: + return + + for widget in self._gui_mapped: + if widget.winfo_ismapped(): + logger.debug("Removing widget: %s", widget) + widget.pack_forget() + widget.destroy() + del widget + + for var in list(self._vars): + logger.debug("Deleting tk variable: %s", var) + del self._vars[var] + + +class _PreviewCanvas(tk.Canvas): # pylint:disable=too-many-ancestors + """ The canvas that holds the preview image + + Parameters + ---------- + parent: :class:`tkinter.Frame` + The parent frame that will hold the Canvas and taskbar + scale_var: :class:`tkinter.StringVar` + The variable that holds the value from the scale combo box + screen_dimensions: tuple + The (`width`, `height`) of the displaying monitor + is_standalone: bool + ``True`` if the preview is standalone, ``False`` if it is in the GUI + """ + def __init__(self, + parent: tk.Frame, + scale_var: tk.StringVar, + screen_dimensions: tuple[int, int], + is_standalone: bool) -> None: + logger.debug("Initializing %s (parent: '%s', scale_var: %s, screen_dimensions: %s)", + self.__class__.__name__, parent, scale_var, screen_dimensions) + frame = tk.Frame(parent) + super().__init__(frame) + + self._is_standalone = is_standalone + self._screen_dimensions = screen_dimensions + self._var_scale = scale_var + self._configure_scrollbars(frame) + self._image: ImageTk.PhotoImage | None = None + self._image_id = self.create_image(self.width / 2, + self.height / 2, + anchor=tk.CENTER, + image=self._image) + self.pack(fill=tk.BOTH, expand=True) + self.bind("", self._resize) + frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + logger.debug("Initialized %s ('%s')", self.__class__.__name__, self) + + @property + def image_id(self) -> int: + """ int: The ID of the preview image item within the canvas """ + return self._image_id + + @property + def width(self) -> int: + """int: The pixel width of canvas""" + return self.winfo_width() + + @property + def height(self) -> int: + """int: The pixel width of the canvas""" + return self.winfo_height() + + def _configure_scrollbars(self, frame: tk.Frame) -> None: + """ Add X and Y scrollbars to the frame and set to scroll the canvas. + + Parameters + ---------- + frame: :class:`tkinter.Frame` + The parent frame to the canvas + """ + logger.debug("Configuring scrollbars") + x_scrollbar = tk.Scrollbar(frame, orient="horizontal", command=self.xview) + x_scrollbar.pack(side=tk.BOTTOM, fill=tk.X) + + y_scrollbar = tk.Scrollbar(frame, command=self.yview) + y_scrollbar.pack(side=tk.RIGHT, fill=tk.Y) + + self.configure(xscrollcommand=x_scrollbar.set, yscrollcommand=y_scrollbar.set) + logger.debug("Configured scrollbars. x: '%s', y: '%s'", x_scrollbar, y_scrollbar) + + def _resize(self, event: tk.Event) -> None: # pylint:disable=unused-argument + """ Place the image in center of canvas on resize event and move to top left + + Parameters + ---------- + event: :class:`tkinter.Event` + The canvas resize event. Unused. + """ + if self._var_scale.get() == "Fit": # Trigger an update to resize image + logger.debug("Triggering redraw for 'Fit' Scaling") + self._var_scale.set("Fit") + return + + self.configure(scrollregion=self.bbox("all")) + self.update_idletasks() + + assert self._image is not None + self._center_image(self.width / 2, self.height / 2) + + # Move to top left when resizing into screen dimensions (initial startup) + if self.width > self._screen_dimensions[0]: + logger.debug("Moving image to left edge") + self.xview_moveto(0.0) + if self.height > self._screen_dimensions[1]: + logger.debug("Moving image to top edge") + self.yview_moveto(0.0) + + def _center_image(self, point_x: float, point_y: float) -> None: + """ Center the image on the canvas on a resize or image update. + + Parameters + ---------- + point_x: int + The x point to center on + point_y: int + The y point to center on + """ + canvas_location = (self.canvasx(point_x), self.canvasy(point_y)) + logger.debug("Centering canvas for size (%s, %s). New image coordinates: %s", + point_x, point_y, canvas_location) + self.coords(self.image_id, canvas_location) + + def set_image(self, + image: ImageTk.PhotoImage, + center_image: bool = False) -> None: + """ Update the canvas with the given image and update area/scrollbars accordingly + + Parameters + ---------- + image: :class:`ImageTK.PhotoImage` + The preview image to display in the canvas + bool, optional + ``True`` if the image should be re-centered. Default ``True`` + """ + logger.debug("Setting canvas image. ID: %s, size: %s for canvas size: %s (recenter: %s)", + self.image_id, (image.width(), image.height()), (self.width, self.height), + center_image) + self._image = image + self.itemconfig(self.image_id, image=self._image) + + if self._is_standalone: # canvas size should not be updated inside GUI + self.config(width=self._image.width(), height=self._image.height()) + + self.update_idletasks() + if center_image: + self._center_image(self.width / 2, self.height / 2) + self.configure(scrollregion=self.bbox("all")) + logger.debug("set canvas image. Canvas size: %s", (self.width, self.height)) + + +class _Image(): + """ Holds the source image and the resized display image for the canvas + + Parameters + ---------- + save_variable: :class:`tkinter.BooleanVar` + Variable that indicates a save preview has been requested in standalone mode + is_standalone: bool + ``True`` if the preview is running in standalone mode. ``False`` if it is running in the + GUI + """ + def __init__(self, save_variable: tk.BooleanVar, is_standalone: bool) -> None: + logger.debug("Initializing %s: (save_variable: %s, is_standalone: %s)", + self.__class__.__name__, save_variable, is_standalone) + self._is_standalone = is_standalone + self._source: np.ndarray | None = None + self._display: ImageTk.PhotoImage | None = None + self._scale = 1.0 + self._interpolation = cv2.INTER_NEAREST + + self._save_var = save_variable + self._save_var.trace("w", self.save_preview) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def display_image(self) -> ImageTk.PhotoImage: + """ :class:`PIL.ImageTk.PhotoImage`: The current display image """ + assert self._display is not None + return self._display + + @property + def source(self) -> np.ndarray: + """ :class:`PIL.Image.Image`: The current source preview image """ + assert self._source is not None + return self._source + + @property + def scale(self) -> int: + """int: The current display scale as a percentage of original image size """ + return int(self._scale * 100) + + def set_source_image(self, name: str, image: np.ndarray) -> None: + """ Set the source image to :attr:`source` + + Parameters + ---------- + name: str + The name of the preview image to load + image: :class:`numpy.ndarray` + The image to use in RGB format + """ + logger.debug("Setting source image. name: '%s', shape: %s", name, image.shape) + self._source = image + + def set_display_image(self) -> None: + """ Obtain the scaled image and set to :attr:`display_image` """ + logger.debug("Setting display image. Scale: %s", self._scale) + image = self.source[..., 2::-1] # TO RGB + if self._scale not in (0.0, 1.0): # Scale will be 0,0 on initial load in GUI + interp = self._interpolation if self._scale > 1.0 else cv2.INTER_NEAREST + dims = (int(round(self.source.shape[1] * self._scale, 0)), + int(round(self.source.shape[0] * self._scale, 0))) + image = cv2.resize(image, dims, interpolation=interp) + self._display = ImageTk.PhotoImage(Image.fromarray(image)) + logger.debug("Set display image. Size: %s", + (self._display.width(), self._display.height())) + + def set_scale(self, scale: float) -> bool: + """ Set the display scale to the given value. + + Parameters + ---------- + scale: float + The value to set scaling to + + Returns + ------- + bool + ``True`` if the scale has been changed otherwise ``False`` + """ + if self._scale == scale: + return False + logger.debug("Setting scale: %s", scale) + self._scale = scale + return True + + def set_interpolation(self, interpolation: int) -> bool: + """ Set the interpolation enum to the given value. + + Parameters + ---------- + interpolation: int + The value to set interpolation to + + Returns + ------- + bool + ``True`` if the interpolation has been changed otherwise ``False`` + """ + if self._interpolation == interpolation: + return False + logger.debug("Setting interpolation: %s") + self._interpolation = interpolation + return True + + def save_preview(self, *args) -> None: + """ Save out the full size preview to the faceswap folder on a save button press + + Parameters + ---------- + args: tuple + Tuple containing either the key press event (Ctrl+s shortcut), the tk variable + arguments (standalone save button press) or the folder location (GUI save button press) + """ + if self._is_standalone and not self._save_var.get() and not isinstance(args[0], tk.Event): + return + + if self._is_standalone: + root_path = os.path.join(os.path.realpath(os.path.dirname(sys.argv[0]))) + else: + root_path = args[0] + + now = datetime.now().strftime("%Y-%m-%d_%H.%M.%S") + filename = os.path.join(root_path, f"preview_{now}.png") + cv2.imwrite(filename, self.source) + print("") + logger.info("Saved preview to: '%s'", filename) + + if self._is_standalone: + self._save_var.set(False) + + +class _Bindings(): # pylint:disable=too-few-public-methods + """ Handle Mouse and Keyboard bindings for the canvas. + + Parameters + ---------- + canvas: :class:`_PreviewCanvas` + The canvas that holds the preview image + taskbar: :class:`_Taskbar` + The taskbar widget which holds the scaling variables + image: :class:`_Image` + The object which holds the source and display version of the preview image + is_standalone: bool + ``True`` if the preview is standalone, ``False`` if it is embedded in the GUI + """ + def __init__(self, + canvas: _PreviewCanvas, + taskbar: _Taskbar, + image: _Image, + is_standalone: bool) -> None: + logger.debug("Initializing %s (canvas: '%s', taskbar: '%s', image: '%s')", + self.__class__.__name__, canvas, taskbar, image) + self._canvas = canvas + self._taskbar = taskbar + self._image = image + + self._drag_data: list[float] = [0., 0.] + self._set_mouse_bindings() + self._set_key_bindings(is_standalone) + logger.debug("Initialized %s", self.__class__.__name__,) + + def _on_bound_zoom(self, event: tk.Event) -> None: + """ Action to perform on a valid zoom key press or mouse wheel action + + Parameters + ---------- + event: :class:`tkinter.Event` + The key press or mouse wheel event + """ + if event.keysym in ("KP_Add", "plus") or event.num == 4 or event.delta > 0: + scale = min(self._taskbar.max_scale, self._image.scale + 25) + else: + scale = max(self._taskbar.min_scale, self._image.scale - 25) + logger.trace("Bound zoom action: (event: %s, scale: %s)", event, scale) # type: ignore + self._taskbar.scale_var.set(f"{scale}%") + + def _on_mouse_click(self, event: tk.Event) -> None: + """ log initial click coordinates for mouse click + drag action + + Parameters + ---------- + event: :class:`tkinter.Event` + The mouse event + """ + self._drag_data = [event.x / self._image.display_image.width(), + event.y / self._image.display_image.height()] + logger.trace("Mouse click action: (event: %s, drag_data: %s)", # type: ignore + event, self._drag_data) + + def _on_mouse_drag(self, event: tk.Event) -> None: + """ Drag image left, right, up or down + + Parameters + ---------- + event: :class:`tkinter.Event` + The mouse event + """ + location_x = event.x / self._image.display_image.width() + location_y = event.y / self._image.display_image.height() + + if self._canvas.xview() != (0.0, 1.0): + to_x = min(1.0, max(0.0, self._drag_data[0] - location_x + self._canvas.xview()[0])) + self._canvas.xview_moveto(to_x) + if self._canvas.yview() != (0.0, 1.0): + to_y = min(1.0, max(0.0, self._drag_data[1] - location_y + self._canvas.yview()[0])) + self._canvas.yview_moveto(to_y) + + self._drag_data = [location_x, location_y] + + def _on_key_move(self, event: tk.Event) -> None: + """ Action to perform on a valid move key press + + Parameters + ---------- + event: :class:`tkinter.Event` + The key press event + """ + move_axis = self._canvas.xview if event.keysym in ("Left", "Right") else self._canvas.yview + visible = move_axis()[1] - move_axis()[0] + amount = -visible / 25 if event.keysym in ("Up", "Left") else visible / 25 + logger.trace("Key move event: (event: %s, move_axis: %s, visible: %s, " # type: ignore + "amount: %s)", move_axis, visible, amount) + move_axis(tk.MOVETO, min(1.0, max(0.0, move_axis()[0] + amount))) + + def _set_mouse_bindings(self) -> None: + """ Set the mouse bindings for interacting with the preview image + + Mousewheel: Zoom in and out + Mouse click: Move image + """ + logger.debug("Binding mouse events") + if system() == "Linux": + self._canvas.tag_bind(self._canvas.image_id, "", self._on_bound_zoom) + self._canvas.tag_bind(self._canvas.image_id, "", self._on_bound_zoom) + else: + self._canvas.bind("", self._on_bound_zoom) + + self._canvas.tag_bind(self._canvas.image_id, "", self._on_mouse_click) + self._canvas.tag_bind(self._canvas.image_id, "", self._on_mouse_drag) + logger.debug("Bound mouse events") + + def _set_key_bindings(self, is_standalone: bool) -> None: + """ Set the keyboard bindings. + + Up/Down/Left/Right: Moves image + +/-: Zooms image + ctrl+s: Save + i: Cycle interpolators + + Parameters + ---------- + ``True`` if the preview is standalone, ``False`` if it is embedded in the GUI + """ + if not is_standalone: + # Don't bind keys for GUI as it adds complication + return + logger.debug("Binding key events") + root = self._canvas.winfo_toplevel() + for key in ("Left", "Right", "Up", "Down"): + root.bind(f"<{key}>", self._on_key_move) + for key in ("Key-plus", "Key-minus", "Key-KP_Add", "Key-KP_Subtract"): + root.bind(f"<{key}>", self._on_bound_zoom) + root.bind("", self._image.save_preview) + root.bind("", self._taskbar.cycle_interpolators) + logger.debug("Bound key events") + + +class PreviewTk(PreviewBase): + """ Holds a preview window for displaying the pop out preview. + + Parameters + ---------- + preview_buffer: :class:`PreviewBuffer` + The thread safe object holding the preview images + parent: tkinter widget, optional + If this viewer is being called from the GUI the parent widget should be passed in here. + If this is a standalone pop-up window then pass ``None``. Default: ``None`` + taskbar: :class:`tkinter.ttk.Frame`, optional + If this viewer is being called from the GUI the parent's option frame should be passed in + here. If this is a standalone pop-up window then pass ``None``. Default: ``None`` + triggers: dict, optional + Dictionary of event triggers for pop-up preview. Not required when running inside the GUI. + Default: `None` + """ + def __init__(self, + preview_buffer: PreviewBuffer, + parent: tk.Widget | None = None, + taskbar: ttk.Frame | None = None, + triggers: TriggerType | None = None) -> None: + logger.debug("Initializing %s (parent: '%s')", self.__class__.__name__, parent) + super().__init__(preview_buffer, triggers=triggers) + self._is_standalone = parent is None + self._initialized = False + self._root = parent if parent is not None else tk.Tk() + self._master_frame = tk.Frame(self._root) + + self._taskbar = _Taskbar(self._master_frame, taskbar) + + self._screen_dimensions = self._get_geometry() + self._canvas = _PreviewCanvas(self._master_frame, + self._taskbar.scale_var, + self._screen_dimensions, + self._is_standalone) + + self._image = _Image(self._taskbar.save_var, self._is_standalone) + + _Bindings(self._canvas, self._taskbar, self._image, self._is_standalone) + + self._taskbar.scale_var.trace("w", self._set_scale) + self._taskbar.interpolator_var.trace("w", self._set_interpolation) + + self._process_triggers() + + if self._is_standalone: + self.pack(fill=tk.BOTH, expand=True) + + self._output_helptext() + + logger.debug("Initialized %s", self.__class__.__name__) + + self._launch() + + @property + def master_frame(self) -> tk.Frame: + """ :class:`tkinter.Frame`: The master frame that holds the preview window """ + return self._master_frame + + def pack(self, *args, **kwargs): + """ Redirect calls to pack the widget to pack the actual :attr:`_master_frame`. + + Takes standard :class:`tkinter.Frame` pack arguments + """ + logger.debug("Packing master frame: (args: %s, kwargs: %s)", args, kwargs) + self._master_frame.pack(*args, **kwargs) + + def save(self, location: str) -> None: + """ Save action to be performed when save button pressed from the GUI. + + location: str + Full path to the folder to save the preview image to + """ + self._image.save_preview(location) + + def remove_option_controls(self) -> None: + """ Remove the taskbar options controls when the preview is disabled in the GUI """ + self._taskbar.destroy_widgets() + + def _output_helptext(self) -> None: + """ Output the keybindings to Console. """ + if not self._is_standalone: + return + logger.info("---------------------------------------------------") + logger.info(" Preview key bindings:") + logger.info(" Zoom: +/-") + logger.info(" Toggle Zoom Mode: i") + logger.info(" Move: arrow keys") + logger.info(" Save Preview: Ctrl+s") + logger.info("---------------------------------------------------") + + def _get_geometry(self) -> tuple[int, int]: + """ Obtain the geometry of the current screen (standalone) or the dimensions of the widget + holding the preview window (GUI). + + Just pulling screen width and height does not account for multiple monitors, so dummy in a + window to pull actual dimensions before hiding it again. + + Returns + ------- + Tuple + The (`width`, `height`) of the current monitor's display + """ + if not self._is_standalone: + root = self._root.winfo_toplevel() # Get dims of whole GUI + retval = root.winfo_width(), root.winfo_height() + logger.debug("Obtained frame geometry: %s", retval) + return retval + + assert isinstance(self._root, tk.Tk) + logger.debug("Obtaining screen geometry") + self._root.update_idletasks() + self._root.attributes("-fullscreen", True) + self._root.state("iconic") + retval = self._root.winfo_width(), self._root.winfo_height() + self._root.attributes("-fullscreen", False) + self._root.state("withdraw") + logger.debug("Obtained screen geometry: %s", retval) + return retval + + def _set_min_max_scales(self) -> None: + """ Set the minimum and maximum area that we allow to scale image to. """ + logger.debug("Calculating minimum scale for screen dimensions %s", self._screen_dimensions) + half_screen = tuple(x // 2 for x in self._screen_dimensions) + min_scales = (half_screen[0] / self._image.source.shape[1], + half_screen[1] / self._image.source.shape[0]) + min_scale = min(1.0, *min_scales) + min_scale = (ceil(min_scale * 10)) * 10 + + eight_screen = tuple(x * 8 for x in self._screen_dimensions) + max_scales = (eight_screen[0] / self._image.source.shape[1], + eight_screen[1] / self._image.source.shape[0]) + max_scale = min(8.0, max(1.0, min(max_scales))) + max_scale = (floor(max_scale * 10)) * 10 + + logger.debug("Calculated minimum scale: %s, maximum_scale: %s", min_scale, max_scale) + self._taskbar.set_min_max_scale(min_scale, max_scale) + + def _initialize_window(self) -> None: + """ Initialize the window to fit into the current screen """ + logger.debug("Initializing window") + assert isinstance(self._root, tk.Tk) + width = min(self._master_frame.winfo_reqwidth(), self._screen_dimensions[0]) + height = min(self._master_frame.winfo_reqheight(), self._screen_dimensions[1]) + self._set_min_max_scales() + self._root.state("normal") + self._root.geometry(f"{width}x{height}") + self._root.protocol("WM_DELETE_WINDOW", lambda: None) # Intercept close window + self._initialized = True + logger.debug("Initialized window: (width: %s, height: %s)", width, height) + + def _update_image(self, center_image: bool = False) -> None: + """ Update the image displayed in the canvas and set the canvas size and scroll region + accordingly + + center_image: bool = ``True`` + ``True`` if the image in the canvas should be recentered. Defaul:``True`` + """ + logger.debug("Updating image (center_image: %s)", center_image) + self._image.set_display_image() + self._canvas.set_image(self._image.display_image, center_image) + logger.debug("Updated image") + + def _convert_fit_scale(self) -> str: + """ Convert "Fit" scale to the actual scaling amount + + Returns + ------- + str + The fit scaling in '##%' format + """ + logger.debug("Converting 'Fit' scaling") + width_scale = self._canvas.width / self._image.source.shape[1] + height_scale = self._canvas.height / self._image.source.shape[0] + scale = min(width_scale, height_scale) * 100 + retval = f"{floor(scale)}%" + logger.debug("Converted 'Fit' scaling: (width_scale: %s, height_scale: %s, scale: %s, " + "retval: '%s'", width_scale, height_scale, scale, retval) + return retval + + def _set_scale(self, *args) -> None: # pylint:disable=unused-argument + """ Update the image on a scale request """ + txtscale = self._taskbar.scale_var.get() + logger.debug("Setting scale: '%s'", txtscale) + txtscale = self._convert_fit_scale() if txtscale == "Fit" else txtscale + scale = int(txtscale[:-1]) # Strip percentage and convert to int + logger.debug("Got scale: %s", scale) + + if self._image.set_scale(scale / 100): + logger.debug("Updating for new scale") + self._taskbar.slider_var.set(scale) + self._update_image(center_image=True) + + def _set_interpolation(self, *args) -> None: # pylint:disable=unused-argument + """ Callback for when the interpolator is change""" + interp = self._taskbar.interpolator_var.get() + if not self._image.set_interpolation(interp) or self._image.scale <= 1.0: + return + self._update_image(center_image=False) + + def _process_triggers(self) -> None: + """ Process the standard faceswap key press triggers: + + m = toggle_mask + r = refresh + s = save + enter = quit + """ + if self._triggers is None: # Don't need triggers for GUI + return + logger.debug("Processing triggers") + root = self._canvas.winfo_toplevel() + for key in self._keymaps: + bindkey = "Return" if key == "enter" else key + logger.debug("Adding trigger for key: '%s'", bindkey) + + root.bind(f"<{bindkey}>", self._on_keypress) + logger.debug("Processed triggers") + + def _on_keypress(self, event: tk.Event) -> None: + """ Update the triggers on a keypress event for picking up by main faceswap process. + + Parameters + ---------- + event: :class:`tkinter.Event` + The valid preview trigger keypress + """ + if self._triggers is None: # Don't need triggers for GUI + return + keypress = "enter" if event.keysym == "Return" else event.keysym + key = T.cast(TriggerKeysType, keypress) + logger.debug("Processing keypress '%s'", key) + if key == "r": + print("") # Let log print on different line from loss output + logger.info("Refresh preview requested...") + + self._triggers[self._keymaps[key]].set() + logger.debug("Processed keypress '%s'. Set event for '%s'", key, self._keymaps[key]) + + def _display_preview(self) -> None: + """ Handle the displaying of the images currently in :attr:`_preview_buffer`""" + if self._should_shutdown: + self._root.destroy() + + if not self._buffer.is_updated: + self._root.after(1000, self._display_preview) + return + + for name, image in self._buffer.get_images(): + logger.debug("Updating image: (name: '%s', shape: %s)", name, image.shape) + if self._is_standalone and not self._title: + assert isinstance(self._root, tk.Tk) + self._title = name + logger.debug("Setting title: '%s;", self._title) + self._root.title(self._title) + self._image.set_source_image(name, image) + self._update_image(center_image=not self._initialized) + + self._root.after(1000, self._display_preview) + + if not self._initialized and self._is_standalone: + self._initialize_window() + self._root.mainloop() + if not self._initialized: # Set initialized to True for GUI + self._set_min_max_scales() + self._taskbar.scale_var.set("Fit") + self._initialized = True + + +def main(): + """ Load image from first given argument and display + + python -m lib.training.preview_tk + """ + from lib.logger import log_setup # pylint:disable=import-outside-toplevel + from .preview_cv import PreviewBuffer # pylint:disable=import-outside-toplevel + log_setup("DEBUG", "faceswap_preview.log", "Test", False) + + img = cv2.imread(sys.argv[-1], cv2.IMREAD_UNCHANGED) + buff = PreviewBuffer() # pylint:disable=used-before-assignment + buff.add_image("test_image", img) + PreviewTk(buff) + + +if __name__ == "__main__": + main() diff --git a/lib/utils.py b/lib/utils.py index 7f1f47458a..be7588fdfd 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -1,27 +1,35 @@ #!/usr/bin python3 """ Utilities available across all scripts """ - +from __future__ import annotations import json import logging import os import sys -import urllib +import tkinter as tk +import typing as T import warnings import zipfile -from re import finditer from multiprocessing import current_process +from re import finditer from socket import timeout as socket_timeout, error as socket_error +from threading import get_ident +from time import time +from urllib import request, error as urlliberror +import numpy as np from tqdm import tqdm +if T.TYPE_CHECKING: + from argparse import Namespace + from http.client import HTTPResponse + # Global variables -_image_extensions = [ # pylint:disable=invalid-name - ".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff"] -_video_extensions = [ # pylint:disable=invalid-name - ".avi", ".flv", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm", ".wmv", - ".ts", ".vob"] -_TF_VERS = None +IMAGE_EXTENSIONS = [".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff"] +VIDEO_EXTENSIONS = [".avi", ".flv", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm", ".wmv", + ".ts", ".vob"] +_TF_VERS: tuple[int, int] | None = None +ValidBackends = T.Literal["nvidia", "cpu", "apple_silicon", "directml", "rocm"] class _Backend(): # pylint:disable=too-few-public-methods @@ -29,13 +37,18 @@ class _Backend(): # pylint:disable=too-few-public-methods Variable. If file doesn't exist and a variable hasn't been set, create the config file. """ - def __init__(self): - self._backends = {"1": "amd", "2": "cpu", "3": "nvidia", "4": "apple_silicon"} + def __init__(self) -> None: + self._backends: dict[str, ValidBackends] = {"1": "cpu", + "2": "directml", + "3": "nvidia", + "4": "apple_silicon", + "5": "rocm"} + self._valid_backends = list(self._backends.values()) self._config_file = self._get_config_file() self.backend = self._get_backend() @classmethod - def _get_config_file(cls): + def _get_config_file(cls) -> str: """ Obtain the location of the main Faceswap configuration file. Returns @@ -47,7 +60,7 @@ def _get_config_file(cls): config_file = os.path.join(pypath, "config", ".faceswap") return config_file - def _get_backend(self): + def _get_backend(self) -> ValidBackends: """ Return the backend from either the `FACESWAP_BACKEND` Environment Variable or from the :file:`config/.faceswap` configuration file. If neither of these exist, prompt the user to select a backend. @@ -59,7 +72,9 @@ def _get_backend(self): """ # Check if environment variable is set, if so use that if "FACESWAP_BACKEND" in os.environ: - fs_backend = os.environ["FACESWAP_BACKEND"].lower() + fs_backend = T.cast(ValidBackends, os.environ["FACESWAP_BACKEND"].lower()) + assert fs_backend in T.get_args(ValidBackends), ( + f"Faceswap backend must be one of {T.get_args(ValidBackends)}") print(f"Setting Faceswap backend from environment variable to {fs_backend.upper()}") return fs_backend # Intercept for sphinx docs build @@ -75,14 +90,14 @@ def _get_backend(self): except json.decoder.JSONDecodeError: self._configure_backend() continue - fs_backend = config.get("backend", None) - if fs_backend is None or fs_backend.lower() not in self._backends.values(): + fs_backend = config.get("backend", "").lower() + if not fs_backend or fs_backend not in self._backends.values(): fs_backend = self._configure_backend() if current_process().name == "MainProcess": print(f"Setting Faceswap backend to {fs_backend.upper()}") - return fs_backend.lower() + return fs_backend - def _configure_backend(self): + def _configure_backend(self) -> ValidBackends: """ Get user input to select the backend that Faceswap should use. Returns @@ -92,12 +107,14 @@ def _configure_backend(self): """ print("First time configuration. Please select the required backend") while True: - selection = input("1: AMD, 2: CPU, 3: NVIDIA, 4: APPLE SILICON: ") - if selection not in ("1", "2", "3", "4"): + txt = ", ".join([": ".join([key, val.upper().replace("_", " ")]) + for key, val in self._backends.items()]) + selection = input(f"{txt}: ") + if selection not in self._backends: print(f"'{selection}' is not a valid selection. Please try again") continue break - fs_backend = self._backends[selection].lower() + fs_backend = self._backends[selection] config = {"backend": fs_backend} with open(self._config_file, "w", encoding="utf8") as cnf: json.dump(config, cnf) @@ -105,48 +122,68 @@ def _configure_backend(self): return fs_backend -_FS_BACKEND = _Backend().backend +_FS_BACKEND: ValidBackends = _Backend().backend -def get_backend(): +def get_backend() -> ValidBackends: """ Get the backend that Faceswap is currently configured to use. Returns ------- str - The backend configuration in use by Faceswap + The backend configuration in use by Faceswap. One of ["cpu", "directml", "nvidia", "rocm", + "apple_silicon"] + + Example + ------- + >>> from lib.utils import get_backend + >>> get_backend() + 'nvidia' """ return _FS_BACKEND -def set_backend(backend): +def set_backend(backend: str) -> None: """ Override the configured backend with the given backend. Parameters ---------- - backend: ["amd", "cpu", "nvidia", "apple_silicon"] + backend: ["cpu", "directml", "nvidia", "rocm", "apple_silicon"] The backend to set faceswap to + + Example + ------- + >>> from lib.utils import set_backend + >>> set_backend("nvidia") """ global _FS_BACKEND # pylint:disable=global-statement - _FS_BACKEND = backend.lower() + backend = T.cast(ValidBackends, backend.lower()) + _FS_BACKEND = backend -def get_tf_version(): - """ Obtain the major.minor version of currently installed Tensorflow. +def get_tf_version() -> tuple[int, int]: + """ Obtain the major. minor version of currently installed Tensorflow. Returns ------- - float - The currently installed tensorflow version + tuple[int, int] + A tuple of the form (major, minor) representing the version of TensorFlow that is installed + + Example + ------- + >>> from lib.utils import get_tf_version + >>> get_tf_version() + (2, 10) """ global _TF_VERS # pylint:disable=global-statement if _TF_VERS is None: import tensorflow as tf # pylint:disable=import-outside-toplevel - _TF_VERS = float(".".join(tf.__version__.split(".")[:2])) # pylint:disable=no-member + split = tf.__version__.split(".")[:2] + _TF_VERS = (int(split[0]), int(split[1])) return _TF_VERS -def get_folder(path, make_folder=True): +def get_folder(path: str, make_folder: bool = True) -> str: """ Return a path to a folder, creating it if it doesn't exist Parameters @@ -162,34 +199,56 @@ def get_folder(path, make_folder=True): str or `None` The path to the requested folder. If `make_folder` is set to ``False`` and the requested path does not exist, then ``None`` is returned + + Example + ------- + >>> from lib.utils import get_folder + >>> get_folder('/tmp/myfolder') + '/tmp/myfolder' + + >>> get_folder('/tmp/myfolder', make_folder=False) + '' """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name + logger = logging.getLogger(__name__) logger.debug("Requested path: '%s'", path) if not make_folder and not os.path.isdir(path): logger.debug("%s does not exist", path) - return None + return "" os.makedirs(path, exist_ok=True) logger.debug("Returning: '%s'", path) return path -def get_image_paths(directory, extension=None): - """ Obtain a list of full paths that reside within a folder. +def get_image_paths(directory: str, extension: str | None = None) -> list[str]: + """ Gets the image paths from a given directory. + + The function searches for files with the specified extension(s) in the given directory, and + returns a list of their paths. If no extension is provided, the function will search for files + with any of the following extensions: '.bmp', '.jpeg', '.jpg', '.png', '.tif', '.tiff' Parameters ---------- directory: str - The folder that contains the images to be returned + The directory to search in extension: str - The specific image extensions that should be returned + The file extension to search for. If not provided, all image file types will be searched + for Returns ------- - list + list[str] The list of full paths to the images contained within the given folder + + Example + ------- + >>> from lib.utils import get_image_paths + >>> get_image_paths('/path/to/directory') + ['/path/to/directory/image1.jpg', '/path/to/directory/image2.png'] + >>> get_image_paths('/path/to/directory', '.jpg') + ['/path/to/directory/image1.jpg'] """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name - image_extensions = _image_extensions if extension is None else [extension] + logger = logging.getLogger(__name__) + image_extensions = IMAGE_EXTENSIONS if extension is None else [extension] dir_contents = [] if not os.path.exists(directory): @@ -198,32 +257,68 @@ def get_image_paths(directory, extension=None): dir_scanned = sorted(os.scandir(directory), key=lambda x: x.name) logger.debug("Scanned Folder contains %s files", len(dir_scanned)) - logger.trace("Scanned Folder Contents: %s", dir_scanned) + logger.trace("Scanned Folder Contents: %s", dir_scanned) # type:ignore[attr-defined] for chkfile in dir_scanned: if any(chkfile.name.lower().endswith(ext) for ext in image_extensions): - logger.trace("Adding '%s' to image list", chkfile.path) + logger.trace("Adding '%s' to image list", chkfile.path) # type:ignore[attr-defined] dir_contents.append(chkfile.path) logger.debug("Returning %s images", len(dir_contents)) return dir_contents -def convert_to_secs(*args): - """ Convert a time to seconds. +def get_dpi() -> float | None: + """ Gets the DPI (dots per inch) of the display screen. + + Returns + ------- + float or ``None`` + The DPI of the display screen or ``None`` if the dpi couldn't be obtained (ie: if the + function is called on a headless system) + + Example + ------- + >>> from lib.utils import get_dpi + >>> get_dpi() + 96.0 + """ + logger = logging.getLogger(__name__) + try: + root = tk.Tk() + dpi = root.winfo_fpixels('1i') + except tk.TclError: + logger.warning("Display not detected. Could not obtain DPI") + return None + + return float(dpi) + + +def convert_to_secs(*args: int) -> int: + """ Convert time in hours, minutes, and seconds to seconds. Parameters ---------- - args: tuple - 2 or 3 ints. If 2 ints are supplied, then (`minutes`, `seconds`) is implied. If 3 ints are - supplied then (`hours`, `minutes`, `seconds`) is implied. + *args: int + 1, 2 or 3 ints. If 2 ints are supplied, then (`minutes`, `seconds`) is implied. If 3 ints + are supplied then (`hours`, `minutes`, `seconds`) is implied. Returns ------- int The given time converted to seconds + + Example + ------- + >>> from lib.utils import convert_to_secs + >>> convert_to_secs(1, 30, 0) + 5400 + >>> convert_to_secs(0, 15, 30) + 930 + >>> convert_to_secs(0, 0, 45) + 45 """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name + logger = logging.getLogger(__name__) logger.debug("from time: %s", args) retval = 0.0 if len(args) == 1: @@ -232,12 +327,13 @@ def convert_to_secs(*args): retval = 60 * float(args[0]) + float(args[1]) elif len(args) == 3: retval = 3600 * float(args[0]) + 60 * float(args[1]) + float(args[2]) + retval = int(retval) logger.debug("to secs: %s", retval) return retval -def full_path_split(path): - """ Split a full path to a location into all of it's separate components. +def full_path_split(path: str) -> list[str]: + """ Split a file path into all of its parts. Parameters ---------- @@ -251,12 +347,14 @@ def full_path_split(path): Example ------- - >>> path = "/foo/baz/bar" - >>> full_path_split(path) - >>> ["foo", "baz", "bar"] + >>> from lib.utils import full_path_split + >>> full_path_split("/usr/local/bin/python") + ['usr', 'local', 'bin', 'python'] + >>> full_path_split("relative/path/to/file.txt") + ['relative', 'path', 'to', 'file.txt']] """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name - allparts = [] + logger = logging.getLogger(__name__) + allparts: list[str] = [] while True: parts = os.path.split(path) if parts[0] == path: # sentinel for absolute paths @@ -267,28 +365,35 @@ def full_path_split(path): break path = parts[0] allparts.insert(0, parts[1]) - logger.trace("path: %s, allparts: %s", path, allparts) + logger.trace("path: %s, allparts: %s", path, allparts) # type:ignore[attr-defined] + # Remove any empty strings which may have got inserted + allparts = [part for part in allparts if part] return allparts -def set_system_verbosity(log_level): +def set_system_verbosity(log_level: str): """ Set the verbosity level of tensorflow and suppresses future and deprecation warnings from - any modules + any modules. + + This function sets the `TF_CPP_MIN_LOG_LEVEL` environment variable to control the verbosity of + TensorFlow output, as well as filters certain warning types to be ignored. The log level is + determined based on the input string `log_level`. Parameters ---------- log_level: str - The requested Faceswap log level + The requested Faceswap log level. References ---------- https://stackoverflow.com/questions/35911252/disable-tensorflow-debugging-information - Can be set to: - 0: all logs shown. 1: filter out INFO logs. 2: filter out WARNING logs. 3: filter out ERROR - logs. - """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name + Example + ------- + >>> from lib.utils import set_system_verbosity + >>> set_system_verbosity('warning') + """ + logger = logging.getLogger(__name__) from lib.logger import get_loglevel # pylint:disable=import-outside-toplevel numeric_level = get_loglevel(log_level) log_level = "3" if numeric_level > 15 else "0" @@ -299,41 +404,91 @@ def set_system_verbosity(log_level): warnings.simplefilter(action='ignore', category=warncat) -def deprecation_warning(function, additional_info=None): - """ Log at warning level that a function will be removed in a future update. +def deprecation_warning(function: str, additional_info: str | None = None) -> None: + """ Log a deprecation warning message. + + This function logs a warning message to indicate that the specified function has been + deprecated and will be removed in future. An optional additional message can also be included. Parameters ---------- function: str - The function that will be deprecated. + The name of the function that will be deprecated. additional_info: str, optional Any additional information to display with the deprecation message. Default: ``None`` + + Example + ------- + >>> from lib.utils import deprecation_warning + >>> deprecation_warning('old_function', 'Use new_function instead.') """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name + logger = logging.getLogger(__name__) logger.debug("func_name: %s, additional_info: %s", function, additional_info) - msg = f"{function} has been deprecated and will be removed from a future update." + msg = f"{function} has been deprecated and will be removed from a future update." if additional_info is not None: msg += f" {additional_info}" logger.warning(msg) -def camel_case_split(identifier): - """ Split a camel case name +def handle_deprecated_cliopts(arguments: Namespace) -> Namespace: + """ Handle deprecated command line arguments and update to correct argument. + + Deprecated cli opts will be provided in the following format: + `"depr___"` Parameters ---------- - identifier: str - The camel case text to be split + arguments: :class:`argpares.Namespace` + The passed in faceswap cli arguments Returns ------- - list - A list of the given identifier split into it's constituent parts + :class:`argpares.Namespace` + The cli arguments with deprecated values mapped to the correct entry + """ + logger = logging.getLogger(__name__) + + for key, selected in vars(arguments).items(): + if not key.startswith("depr_") or key.startswith("depr_") and selected is None: + continue # Not a deprecated opt + if isinstance(selected, bool) and not selected: + continue # store-true opt with default value + + opt, old, new = key.replace("depr_", "").rsplit("_", maxsplit=2) + deprecation_warning(f"Command line option '-{old}'", f"Use '-{new}, --{opt}' instead") + + exist = getattr(arguments, opt) + if exist == selected: + logger.debug("Keeping existing '%s' value of '%s'", opt, exist) + else: + logger.debug("Updating arg '%s' from '%s' to '%s' from deprecated opt", + opt, exist, selected) + + return arguments +def camel_case_split(identifier: str) -> list[str]: + """ Split a camelCase string into a list of its individual parts + + Parameters + ---------- + identifier: str + The camelCase text to be split + + Returns + ------- + list[str] + A list of the individual parts of the camelCase string. + References ---------- https://stackoverflow.com/questions/29916065 + + Example + ------- + >>> from lib.utils import camel_case_split + >>> camel_case_split('camelCaseExample') + ['camel', 'Case', 'Example'] """ matches = finditer( ".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)", @@ -341,16 +496,25 @@ def camel_case_split(identifier): return [m.group(0) for m in matches] -def safe_shutdown(got_error=False): - """ Close all tracked queues and threads in event of crash or on shut down. +def safe_shutdown(got_error: bool = False) -> None: + """ Safely shut down the system. + + This function terminates the queue manager and exits the program in a clean and orderly manner. + An optional boolean parameter can be used to indicate whether an error occurred during the + program's execution. Parameters ---------- got_error: bool, optional - ``True`` if this function is being called as the result of raised error, otherwise - ``False``. Default: ``False`` + ``True`` if this function is being called as the result of raised error. Default: ``False`` + + Example + ------- + >>> from lib.utils import safe_shutdown + >>> safe_shutdown() + >>> safe_shutdown(True) """ - logger = logging.getLogger(__name__) # pylint:disable=invalid-name + logger = logging.getLogger(__name__) logger.debug("Safely shutting down") from lib.queue_manager import queue_manager # pylint:disable=import-outside-toplevel queue_manager.terminate_queues() @@ -365,12 +529,21 @@ class FaceswapError(Exception): ------ FaceswapError on a captured error + + Example + ------- + >>> from lib.utils import FaceswapError + >>> try: + ... # Some code that may raise an error + ... except SomeError: + ... raise FaceswapError("There was an error while running the code") + FaceswapError: There was an error while running the code """ pass # pylint:disable=unnecessary-pass -class GetModel(): # pylint:disable=too-few-public-methods - """ Check for models in their cache path. +class GetModel(): + """ Check for models in the cache path. If available, return the path, if not available, get, unzip and install model @@ -378,9 +551,6 @@ class GetModel(): # pylint:disable=too-few-public-methods ---------- model_filename: str or list The name of the model to be loaded (see notes below) - cache_dir: str - The model cache folder of the current plugin calling this class. IE: The folder that holds - the model to be loaded. git_model_id: int The second digit in the github tag that identifies this model. See https://github.com/deepfakes-models/faceswap-models for more information @@ -395,14 +565,19 @@ class GetModel(): # pylint:disable=too-few-public-methods number: `_v.` (eg: `["mtcnn_det_v1.1.py", "mtcnn_det_v1.2.py", "mtcnn_det_v1.3.py"]`, `["resnet_ssd_v1.caffemodel" ,"resnet_ssd_v1.prototext"]` + + Example + ------- + >>> from lib.utils import GetModel + >>> model_downloader = GetModel("s3fd_keras_v2.h5", 11) """ - def __init__(self, model_filename, cache_dir, git_model_id): + def __init__(self, model_filename: str | list[str], git_model_id: int) -> None: self.logger = logging.getLogger(__name__) if not isinstance(model_filename, list): model_filename = [model_filename] self._model_filename = model_filename - self._cache_dir = cache_dir + self._cache_dir = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), ".fs_cache") self._git_model_id = git_model_id self._url_base = "https://github.com/deepfakes-models/faceswap-models/releases/download" self._chunk_size = 1024 # Chunk size for downloading and unzipping @@ -410,86 +585,77 @@ def __init__(self, model_filename, cache_dir, git_model_id): self._get() @property - def _model_full_name(self): + def _model_full_name(self) -> str: """ str: The full model name from the filename(s). """ common_prefix = os.path.commonprefix(self._model_filename) retval = os.path.splitext(common_prefix)[0] - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def _model_name(self): + def _model_name(self) -> str: """ str: The model name from the model's full name. """ retval = self._model_full_name[:self._model_full_name.rfind("_")] - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def _model_version(self): + def _model_version(self) -> int: """ int: The model's version number from the model full name. """ retval = int(self._model_full_name[self._model_full_name.rfind("_") + 2:]) - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def model_path(self): - """ str: The model path(s) in the cache folder. """ - retval = [os.path.join(self._cache_dir, fname) for fname in self._model_filename] - retval = retval[0] if len(retval) == 1 else retval - self.logger.trace(retval) + def model_path(self) -> str | list[str]: + """ str or list[str]: The model path(s) in the cache folder. + + Example + ------- + >>> from lib.utils import GetModel + >>> model_downloader = GetModel("s3fd_keras_v2.h5", 11) + >>> model_downloader.model_path + '/path/to/s3fd_keras_v2.h5' + """ + paths = [os.path.join(self._cache_dir, fname) for fname in self._model_filename] + retval: str | list[str] = paths[0] if len(paths) == 1 else paths + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def _model_zip_path(self): + def _model_zip_path(self) -> str: """ str: The full path to downloaded zip file. """ retval = os.path.join(self._cache_dir, f"{self._model_full_name}.zip") - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def _model_exists(self): + def _model_exists(self) -> bool: """ bool: ``True`` if the model exists in the cache folder otherwise ``False``. """ if isinstance(self.model_path, list): retval = all(os.path.exists(pth) for pth in self.model_path) else: retval = os.path.exists(self.model_path) - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval @property - def _plugin_section(self): - """ str: The plugin section from the config_dir """ - path = os.path.normpath(self._cache_dir) - split = path.split(os.sep) - retval = split[split.index("plugins") + 1] - self.logger.trace(retval) - return retval - - @property - def _url_section(self): - """ int: The section ID in github for this plugin type. """ - sections = dict(extract=1, train=2, convert=3) - retval = sections[self._plugin_section] - self.logger.trace(retval) - return retval - - @property - def _url_download(self): + def _url_download(self) -> str: """ strL Base download URL for models. """ - tag = f"v{self._url_section}.{self._git_model_id}.{self._model_version}" + tag = f"v{self._git_model_id}.{self._model_version}" retval = f"{self._url_base}/{tag}/{self._model_full_name}.zip" - self.logger.trace("Download url: %s", retval) + self.logger.trace("Download url: %s", retval) # type:ignore[attr-defined] return retval @property - def _url_partial_size(self): - """ float: How many bytes have already been downloaded. """ + def _url_partial_size(self) -> int: + """ int: How many bytes have already been downloaded. """ zip_file = self._model_zip_path retval = os.path.getsize(zip_file) if os.path.exists(zip_file) else 0 - self.logger.trace(retval) + self.logger.trace(retval) # type:ignore[attr-defined] return retval - def _get(self): + def _get(self) -> None: """ Check the model exists, if not, download the model, unzip it and place it in the model's cache folder. """ if self._model_exists: @@ -499,22 +665,22 @@ def _get(self): self._unzip_model() os.remove(self._model_zip_path) - def _download_model(self): + def _download_model(self) -> None: """ Download the model zip from github to the cache folder. """ self.logger.info("Downloading model: '%s' from: %s", self._model_name, self._url_download) for attempt in range(self._retries): try: downloaded_size = self._url_partial_size - req = urllib.request.Request(self._url_download) + req = request.Request(self._url_download) if downloaded_size != 0: req.add_header("Range", f"bytes={downloaded_size}-") - with urllib.request.urlopen(req, timeout=10) as response: + with request.urlopen(req, timeout=10) as response: self.logger.debug("header info: {%s}", response.info()) self.logger.debug("Return Code: %s", response.getcode()) self._write_zipfile(response, downloaded_size) break except (socket_error, socket_timeout, - urllib.error.HTTPError, urllib.error.URLError) as err: + urlliberror.HTTPError, urlliberror.URLError) as err: if attempt + 1 < self._retries: self.logger.warning("Error downloading model (%s). Retrying %s of %s...", str(err), attempt + 2, self._retries) @@ -527,17 +693,19 @@ def _download_model(self): self._url_download, self._cache_dir) sys.exit(1) - def _write_zipfile(self, response, downloaded_size): + def _write_zipfile(self, response: HTTPResponse, downloaded_size: int) -> None: """ Write the model zip file to disk. Parameters ---------- - response: :class:`urllib.request.urlopen` + response: :class:`http.client.HTTPResponse` The response from the model download task downloaded_size: int The amount of bytes downloaded so far """ - length = int(response.getheader("content-length")) + downloaded_size + content_length = response.getheader("content-length") + content_length = "0" if content_length is None else content_length + length = int(content_length) + downloaded_size if length == downloaded_size: self.logger.info("Zip already exists. Skipping download") return @@ -558,7 +726,7 @@ def _write_zipfile(self, response, downloaded_size): out_file.write(buffer) pbar.close() - def _unzip_model(self): + def _unzip_model(self) -> None: """ Unzip the model file to the cache folder """ self.logger.info("Extracting: '%s'", self._model_name) try: @@ -568,12 +736,12 @@ def _unzip_model(self): self.logger.error("Unable to extract model file: %s", str(err)) sys.exit(1) - def _write_model(self, zip_file): + def _write_model(self, zip_file: zipfile.ZipFile) -> None: """ Extract files from zip file and write, with progress bar. Parameters ---------- - zip_file: str + zip_file: :class:`zipfile.ZipFile` The downloaded model zip file """ length = sum(f.file_size for f in zip_file.infolist()) @@ -595,5 +763,161 @@ def _write_model(self, zip_file): break pbar.update(len(buffer)) out_file.write(buffer) - zip_file.close() pbar.close() + + +class DebugTimes(): + """ A simple tool to help debug timings. + + Parameters + ---------- + min: bool, Optional + Display minimum time taken in summary stats. Default: ``True`` + mean: bool, Optional + Display mean time taken in summary stats. Default: ``True`` + max: bool, Optional + Display maximum time taken in summary stats. Default: ``True`` + + Example + ------- + >>> from lib.utils import DebugTimes + >>> debug_times = DebugTimes() + >>> debug_times.step_start("step 1") + >>> # do something here + >>> debug_times.step_end("step 1") + >>> debug_times.summary() + ---------------------------------- + Step Count Min + ---------------------------------- + step 1 1 0.000000 + """ + def __init__(self, + show_min: bool = True, show_mean: bool = True, show_max: bool = True) -> None: + self._times: dict[str, list[float]] = {} + self._steps: dict[str, float] = {} + self._interval = 1 + self._display = {"min": show_min, "mean": show_mean, "max": show_max} + + def step_start(self, name: str, record: bool = True) -> None: + """ Start the timer for the given step name. + + Parameters + ---------- + name: str + The name of the step to start the timer for + record: bool, optional + ``True`` to record the step time, ``False`` to not record it. + Used for when you have conditional code to time, but do not want to insert if/else + statements in the code. Default: `True` + + Example + ------- + >>> from lib.util import DebugTimes + >>> debug_times = DebugTimes() + >>> debug_times.step_start("Example Step") + >>> # do something here + >>> debug_times.step_end("Example Step") + """ + if not record: + return + storename = name + str(get_ident()) + self._steps[storename] = time() + + def step_end(self, name: str, record: bool = True) -> None: + """ Stop the timer and record elapsed time for the given step name. + + Parameters + ---------- + name: str + The name of the step to end the timer for + record: bool, optional + ``True`` to record the step time, ``False`` to not record it. + Used for when you have conditional code to time, but do not want to insert if/else + statements in the code. Default: `True` + + Example + ------- + >>> from lib.util import DebugTimes + >>> debug_times = DebugTimes() + >>> debug_times.step_start("Example Step") + >>> # do something here + >>> debug_times.step_end("Example Step") + """ + if not record: + return + storename = name + str(get_ident()) + self._times.setdefault(name, []).append(time() - self._steps.pop(storename)) + + @classmethod + def _format_column(cls, text: str, width: int) -> str: + """ Pad the given text to be aligned to the given width. + + Parameters + ---------- + text: str + The text to be formatted + width: int + The size of the column to insert the text into + + Returns + ------- + str + The text with the correct amount of padding applied + """ + return f"{text}{' ' * (width - len(text))}" + + def summary(self, decimal_places: int = 6, interval: int = 1) -> None: + """ Print a summary of step times. + + Parameters + ---------- + decimal_places: int, optional + The number of decimal places to display the summary elapsed times to. Default: 6 + interval: int, optional + How many times summary must be called before printing to console. Default: 1 + + Example + ------- + >>> from lib.utils import DebugTimes + >>> debug = DebugTimes() + >>> debug.step_start("test") + >>> time.sleep(0.5) + >>> debug.step_end("test") + >>> debug.summary() + ---------------------------------- + Step Count Min + ---------------------------------- + test 1 0.500000 + """ + interval = max(1, interval) + if interval != self._interval: + self._interval += 1 + return + + name_col = max(len(key) for key in self._times) + 4 + items_col = 8 + time_col = (decimal_places + 4) * sum(1 for v in self._display.values() if v) + separator = "-" * (name_col + items_col + time_col) + print("") + print(separator) + header = (f"{self._format_column('Step', name_col)}" + f"{self._format_column('Count', items_col)}") + header += f"{self._format_column('Min', time_col)}" if self._display["min"] else "" + header += f"{self._format_column('Avg', time_col)}" if self._display["mean"] else "" + header += f"{self._format_column('Max', time_col)}" if self._display["max"] else "" + print(header) + print(separator) + for key, val in self._times.items(): + num = str(len(val)) + contents = f"{self._format_column(key, name_col)}{self._format_column(num, items_col)}" + if self._display["min"]: + _min = f"{np.min(val):.{decimal_places}f}" + contents += f"{self._format_column(_min, time_col)}" + if self._display["mean"]: + avg = f"{np.mean(val):.{decimal_places}f}" + contents += f"{self._format_column(avg, time_col)}" + if self._display["max"]: + _max = f"{np.max(val):.{decimal_places}f}" + contents += f"{self._format_column(_max, time_col)}" + print(contents) + self._interval = 1 diff --git a/lib/vgg_face.py b/lib/vgg_face.py index be917dba1a..d10c957df9 100644 --- a/lib/vgg_face.py +++ b/lib/vgg_face.py @@ -7,8 +7,6 @@ """ import logging -import sys -import os import cv2 import numpy as np @@ -16,7 +14,7 @@ from lib.utils import GetModel -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class VGGFace(): @@ -37,9 +35,7 @@ def __init__(self, backend="CPU"): # <<< GET MODEL >>> # def get_model(self, git_model_id, model_filename, backend): """ Check if model is available, if not, download and unzip it """ - root_path = os.path.abspath(os.path.dirname(sys.argv[0])) - cache_path = os.path.join(root_path, "plugins", "extract", "recognition", ".cache") - model = GetModel(model_filename, cache_path, git_model_id).model_path + model = GetModel(model_filename, git_model_id).model_path model = cv2.dnn.readNetFromCaffe(model[1], model[0]) model.setPreferableTarget(self.get_backend(backend)) return model @@ -50,7 +46,7 @@ def get_backend(backend): if backend == "OPENCL": logger.info("Using OpenCL backend. If the process runs, you can safely ignore any of " "the failure messages.") - retval = getattr(cv2.dnn, "DNN_TARGET_{}".format(backend)) + retval = getattr(cv2.dnn, f"DNN_TARGET_{backend}") return retval def predict(self, face): diff --git a/locales/es/LC_MESSAGES/gui.menu.mo b/locales/es/LC_MESSAGES/gui.menu.mo new file mode 100644 index 0000000000..f31697bbbc Binary files /dev/null and b/locales/es/LC_MESSAGES/gui.menu.mo differ diff --git a/locales/es/LC_MESSAGES/gui.menu.po b/locales/es/LC_MESSAGES/gui.menu.po new file mode 100644 index 0000000000..ba02769135 --- /dev/null +++ b/locales/es/LC_MESSAGES/gui.menu.po @@ -0,0 +1,155 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: faceswap.spanish\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-07 13:54+0100\n" +"PO-Revision-Date: 2023-06-07 14:11+0100\n" +"Last-Translator: \n" +"Language-Team: tokafondo\n" +"Language: es_ES\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/gui/menu.py:37 +msgid "faceswap.dev - Guides and Forum" +msgstr "faceswap.dev - Guías y foro" + +#: lib/gui/menu.py:38 +msgid "Patreon - Support this project" +msgstr "Patreon - Apoya este proyecto" + +#: lib/gui/menu.py:39 +msgid "Discord - The FaceSwap Discord server" +msgstr "Discord - El servidor de Discord de FaceSwap" + +#: lib/gui/menu.py:40 +msgid "Github - Our Source Code" +msgstr "Github - Nuestro código fuente" + +#: lib/gui/menu.py:60 +msgid "File" +msgstr "" + +#: lib/gui/menu.py:61 +msgid "Settings" +msgstr "" + +#: lib/gui/menu.py:62 +msgid "Help" +msgstr "" + +#: lib/gui/menu.py:85 +msgid "Configure Settings..." +msgstr "" + +#: lib/gui/menu.py:116 +msgid "New Project..." +msgstr "" + +#: lib/gui/menu.py:121 +msgid "Open Project..." +msgstr "" + +#: lib/gui/menu.py:126 +msgid "Save Project" +msgstr "" + +#: lib/gui/menu.py:131 +msgid "Save Project as..." +msgstr "" + +#: lib/gui/menu.py:136 +msgid "Reload Project from Disk" +msgstr "" + +#: lib/gui/menu.py:141 +msgid "Close Project" +msgstr "" + +#: lib/gui/menu.py:147 +msgid "Open Task..." +msgstr "" + +#: lib/gui/menu.py:154 +msgid "Open recent" +msgstr "" + +#: lib/gui/menu.py:156 +msgid "Quit" +msgstr "" + +#: lib/gui/menu.py:211 +msgid "{} Task" +msgstr "" + +#: lib/gui/menu.py:223 +msgid "Clear recent files" +msgstr "" + +#: lib/gui/menu.py:391 +msgid "Check for updates..." +msgstr "" + +#: lib/gui/menu.py:394 +msgid "Update Faceswap..." +msgstr "" + +#: lib/gui/menu.py:398 +msgid "Switch Branch" +msgstr "" + +#: lib/gui/menu.py:401 +msgid "Resources" +msgstr "" + +#: lib/gui/menu.py:404 +msgid "Output System Information" +msgstr "" + +#: lib/gui/menu.py:589 +msgid "currently selected Task" +msgstr "tarea actualmente seleccionada" + +#: lib/gui/menu.py:589 +msgid "Project" +msgstr "Proyecto" + +#: lib/gui/menu.py:591 +msgid "Reload {} from disk" +msgstr "Recargar {} del disco" + +#: lib/gui/menu.py:593 +msgid "Create a new {}..." +msgstr "Crear un nuevo {}..." + +#: lib/gui/menu.py:595 +msgid "Reset {} to default" +msgstr "Reiniciar {} a los ajustes por defecto" + +#: lib/gui/menu.py:597 +msgid "Save {}" +msgstr "Guardar {}" + +#: lib/gui/menu.py:599 +msgid "Save {} as..." +msgstr "Guardar {} como..." + +#: lib/gui/menu.py:603 +msgid " from a task or project file" +msgstr " de un archivo de tarea o proyecto" + +#: lib/gui/menu.py:604 +msgid "Load {}..." +msgstr "Cargar {}..." + +#: lib/gui/menu.py:659 +msgid "Configure {} settings..." +msgstr "Configurar los ajustes de {}..." diff --git a/locales/es/LC_MESSAGES/gui.tooltips.mo b/locales/es/LC_MESSAGES/gui.tooltips.mo index c1fdf1a808..9df3225181 100644 Binary files a/locales/es/LC_MESSAGES/gui.tooltips.mo and b/locales/es/LC_MESSAGES/gui.tooltips.mo differ diff --git a/locales/es/LC_MESSAGES/gui.tooltips.po b/locales/es/LC_MESSAGES/gui.tooltips.po index 82b9c0c53e..ab23f031fc 100644 --- a/locales/es/LC_MESSAGES/gui.tooltips.po +++ b/locales/es/LC_MESSAGES/gui.tooltips.po @@ -6,16 +6,16 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" "POT-Creation-Date: 2021-03-22 18:37+0000\n" -"PO-Revision-Date: 2021-03-22 18:39+0000\n" +"PO-Revision-Date: 2023-06-07 14:12+0100\n" "Last-Translator: \n" "Language-Team: tokafondo\n" "Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.2\n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" #: lib/gui/command.py:184 msgid "Output command line options to the console" @@ -129,62 +129,6 @@ msgstr "Grabar {} a un fichero" msgid "Enable or disable {} display" msgstr "Activar o desactivar la muestra de {}" -#: lib/gui/menu.py:32 -msgid "faceswap.dev - Guides and Forum" -msgstr "faceswap.dev - Guías y foro" - -#: lib/gui/menu.py:33 -msgid "Patreon - Support this project" -msgstr "Patreon - Apoya este proyecto" - -#: lib/gui/menu.py:34 -msgid "Discord - The FaceSwap Discord server" -msgstr "Discord - El servidor de Discord de FaceSwap" - -#: lib/gui/menu.py:35 -msgid "Github - Our Source Code" -msgstr "Github - Nuestro código fuente" - -#: lib/gui/menu.py:527 -msgid "Configure {} settings..." -msgstr "Configurar los ajustes de {}..." - -#: lib/gui/menu.py:535 -msgid "Project" -msgstr "Proyecto" - -#: lib/gui/menu.py:535 -msgid "currently selected Task" -msgstr "tarea actualmente seleccionada" - -#: lib/gui/menu.py:537 -msgid "Reload {} from disk" -msgstr "Recargar {} del disco" - -#: lib/gui/menu.py:539 -msgid "Create a new {}..." -msgstr "Crear un nuevo {}..." - -#: lib/gui/menu.py:541 -msgid "Reset {} to default" -msgstr "Reiniciar {} a los ajustes por defecto" - -#: lib/gui/menu.py:543 -msgid "Save {}" -msgstr "Guardar {}" - -#: lib/gui/menu.py:545 -msgid "Save {} as..." -msgstr "Guardar {} como..." - -#: lib/gui/menu.py:549 -msgid " from a task or project file" -msgstr " de un archivo de tarea o proyecto" - -#: lib/gui/menu.py:550 -msgid "Load {}..." -msgstr "Cargar {}..." - #: lib/gui/popup_configure.py:209 msgid "Close without saving" msgstr "Cerrar sin guardar" diff --git a/locales/es/LC_MESSAGES/lib.cli.args.mo b/locales/es/LC_MESSAGES/lib.cli.args.mo index 78be3fdb01..914c7bf3b8 100644 Binary files a/locales/es/LC_MESSAGES/lib.cli.args.mo and b/locales/es/LC_MESSAGES/lib.cli.args.mo differ diff --git a/locales/es/LC_MESSAGES/lib.cli.args.po b/locales/es/LC_MESSAGES/lib.cli.args.po old mode 100644 new mode 100755 index 6662ea056e..d6b64bef41 --- a/locales/es/LC_MESSAGES/lib.cli.args.po +++ b/locales/es/LC_MESSAGES/lib.cli.args.po @@ -5,24 +5,25 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-05-17 18:04+0100\n" -"PO-Revision-Date: 2021-05-17 18:18+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:06+0000\n" +"PO-Revision-Date: 2024-03-28 18:14+0000\n" "Last-Translator: \n" "Language-Team: tokafondo\n" "Language: es\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.3\n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" -#: lib/cli/args.py:177 lib/cli/args.py:187 lib/cli/args.py:195 -#: lib/cli/args.py:205 +#: lib/cli/args.py:188 lib/cli/args.py:199 lib/cli/args.py:208 +#: lib/cli/args.py:219 msgid "Global Options" msgstr "Opciones Globales" -#: lib/cli/args.py:178 +#: lib/cli/args.py:190 msgid "" "R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond " "to any GPU(s) that you do not wish to be made available to Faceswap. " @@ -34,12 +35,12 @@ msgstr "" "con Faceswap. Marcar todas las GPUs forzará a Faceswap a usar sólo la CPU,\n" "L|{}" -#: lib/cli/args.py:188 +#: lib/cli/args.py:201 msgid "" "Optionally overide the saved config with the path to a custom config file." msgstr "Usar un fichero alternativo de configuración, almacenado en esta ruta." -#: lib/cli/args.py:196 +#: lib/cli/args.py:210 msgid "" "Log level. Stick with INFO or VERBOSE unless you need to file an error " "report. Be careful with TRACE as it will generate a lot of data" @@ -47,953 +48,12 @@ msgstr "" "Nivel de registro. Dejarlo en INFO o VERBOSE, a menos que necesite informar " "de un error. Tenga en cuenta que TRACE generará muchísima información" -#: lib/cli/args.py:206 +#: lib/cli/args.py:220 msgid "Path to store the logfile. Leave blank to store in the faceswap folder" msgstr "" "Ruta para almacenar el fichero de registro. Dejarlo en blanco para " "almacenarlo en la carpeta pde instalación de faceswap" -#: lib/cli/args.py:299 lib/cli/args.py:308 lib/cli/args.py:316 -#: lib/cli/args.py:630 lib/cli/args.py:639 -msgid "Data" -msgstr "Datos" - -#: lib/cli/args.py:300 -msgid "" -"Input directory or video. Either a directory containing the image files you " -"wish to process or path to a video file. NB: This should be the source video/" -"frames NOT the source faces." -msgstr "" -"Directorio o vídeo de entrada. Un directorio que contenga los archivos de " -"imagen que desea procesar o la ruta a un archivo de vídeo. NB: Debe ser el " -"vídeo/los fotogramas de origen, NO las caras de origen." - -#: lib/cli/args.py:309 -msgid "Output directory. This is where the converted files will be saved." -msgstr "" -"Directorio de salida. Aquí es donde se guardarán los archivos convertidos." - -#: lib/cli/args.py:317 -msgid "" -"Optional path to an alignments file. Leave blank if the alignments file is " -"at the default location." -msgstr "" -"Ruta opcional a un archivo de alineaciones. Dejar en blanco si el archivo de " -"alineaciones está en la ubicación por defecto." - -#: lib/cli/args.py:340 -msgid "" -"Extract faces from image or video sources.\n" -"Extraction plugins can be configured in the 'Settings' Menu" -msgstr "" -"Extrae caras de fuentes de imagen o video.\n" -"Los plugins de extracción pueden ser configuradas en el menú de 'Ajustes'" - -#: lib/cli/args.py:365 lib/cli/args.py:381 lib/cli/args.py:393 -#: lib/cli/args.py:428 lib/cli/args.py:446 lib/cli/args.py:458 -#: lib/cli/args.py:649 lib/cli/args.py:676 lib/cli/args.py:712 -msgid "Plugins" -msgstr "Extensiones" - -#: lib/cli/args.py:366 -msgid "" -"R|Detector to use. Some of these have configurable settings in '/config/" -"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" -"L|cv2-dnn: A CPU only extractor which is the least reliable and least " -"resource intensive. Use this if not using a GPU and time is important.\n" -"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " -"than other GPU detectors but can often return more false positives.\n" -"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " -"fewer false positives than other GPU detectors, but is a lot more resource " -"intensive." -msgstr "" -"R|Detector de caras a usar. Algunos tienen ajustes configurables en '/config/" -"extract.ini' o 'Ajustes > Configurar Extensiones de Extracción:\n" -"L|cv2-dnn: Extractor que usa sólo la CPU. Es el menos fiable y el que menos " -"recursos usa. Elegir este si necesita rapidez y no usar la GPU.\n" -"L|mtcnn: Buen detector. Rápido en la CPU y más rápido en la GPU. Usa menos " -"recursos que otros detectores basados en GPU, pero puede devolver más falsos " -"positivos.\n" -"L|s3fd: El mejor detector. Lento en la CPU, y más rápido en la GPU. Puede " -"detectar más caras y tiene menos falsos positivos que otros detectores " -"basados en GPU, pero uso muchos más recursos." - -#: lib/cli/args.py:382 -msgid "" -"R|Aligner to use.\n" -"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " -"but less accurate. Only use this if not using a GPU and time is important.\n" -"L|fan: Best aligner. Fast on GPU, slow on CPU." -msgstr "" -"R|Alineador a usar.\n" -"L|cv2-dnn: Detector que usa sólo la CPU. Más rápido, usa menos recursos, " -"pero es menos preciso. Elegir este si necesita rapidez y no usar la GPU.\n" -"L|fan: El mejor alineador. Rápido en la GPU, y lento en la CPU." - -#: lib/cli/args.py:394 -msgid "" -"R|Additional Masker(s) to use. The masks generated here will all take up GPU " -"RAM. You can select none, one or multiple masks, but the extraction may take " -"longer the more you select. NB: The Extended and Components (landmark based) " -"masks are automatically generated on extraction.\n" -"L|bisenet-fp: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked including full head masking " -"(configurable in mask settings).\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " -"faces clear of obstructions. Profile faces and obstructions may result in " -"sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " -"frontal faces. The mask model has been specifically trained to recognize " -"some facial obstructions (hands and eyeglasses). Profile faces may result in " -"sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " -"faces. The mask model has been trained by community members and will need " -"testing for further description. Profile faces may result in sub-par " -"performance.\n" -"The auto generated masks are as follows:\n" -"L|components: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks and the mask is extended upwards onto the " -"forehead.\n" -"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" -msgstr "" -"R|Enmascarador(es) adicional(es) a usar. Las máscaras generadas aquí usarán " -"todas RAM de la GPU. Puede seleccionar una, varias o ninguna máscaras, pero " -"la extracción tardará más cuanto más marque. Las máscaras Extended y " -"Components son siempre generadas durante la extracción.\n" -"L|bisenet-fp: Máscara relativamente ligera basada en NN que proporciona un " -"control más refinado sobre el área a enmascarar, incluido el enmascaramiento " -"completo de la cabeza (configurable en la configuración de la máscara).\n" -"L|vgg-clear: Máscara diseñada para proporcionar una segmentación inteligente " -"de rostros principalmente frontales y libres de obstrucciones. Los rostros " -"de perfil y las obstrucciones pueden dar lugar a un rendimiento inferior.\n" -"L|vgg-obstructed: Máscara diseñada para proporcionar una segmentación " -"inteligente de rostros principalmente frontales. El modelo de la máscara ha " -"sido entrenado específicamente para reconocer algunas obstrucciones faciales " -"(manos y gafas). Los rostros de perfil pueden dar lugar a un rendimiento " -"inferior.\n" -"L|unet-dfl: Máscara diseñada para proporcionar una segmentación inteligente " -"de rostros principalmente frontales. El modelo de máscara ha sido entrenado " -"por los miembros de la comunidad y necesitará ser probado para una mayor " -"descripción. Los rostros de perfil pueden dar lugar a un rendimiento " -"inferior.\n" -"Las máscaras que siempre se generan son:\n" -"L|components: Máscara diseñada para proporcionar una segmentación facial " -"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " -"Se construye un casco convexo alrededor del exterior de los puntos de " -"referencia para crear una máscara.\n" -"L|extended: Máscara diseñada para proporcionar una segmentación facial " -"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " -"Se construye un casco convexo alrededor del exterior de los puntos de " -"referencia y la máscara se extiende hacia arriba en la frente.\n" -"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" - -#: lib/cli/args.py:429 -msgid "" -"R|Performing normalization can help the aligner better align faces with " -"difficult lighting conditions at an extraction speed cost. Different methods " -"will yield different results on different sets. NB: This does not impact the " -"output face, just the input to the aligner.\n" -"L|none: Don't perform normalization on the face.\n" -"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " -"face.\n" -"L|hist: Equalize the histograms on the RGB channels.\n" -"L|mean: Normalize the face colors to the mean." -msgstr "" -"R|Realizar la normalización puede ayudar al alineador a alinear mejor las " -"caras con condiciones de iluminación difíciles a un coste de velocidad de " -"extracción. Diferentes métodos darán diferentes resultados en diferentes " -"conjuntos. NB: Esto no afecta a la cara de salida, sólo a la entrada del " -"alineador.\n" -"L|none: No realice la normalización en la cara.\n" -"L|clahe: Realice la ecualización adaptativa del histograma con contraste " -"limitado en el rostro.\n" -"L|hist: Iguala los histogramas de los canales RGB.\n" -"L|mean: Normalizar los colores de la cara a la media." - -#: lib/cli/args.py:447 -msgid "" -"The number of times to re-feed the detected face into the aligner. Each time " -"the face is re-fed into the aligner the bounding box is adjusted by a small " -"amount. The final landmarks are then averaged from each iteration. Helps to " -"remove 'micro-jitter' but at the cost of slower extraction speed. The more " -"times the face is re-fed into the aligner, the less micro-jitter should " -"occur but the longer extraction will take." -msgstr "" -"El número de veces que hay que volver a introducir la cara detectada en el " -"alineador. Cada vez que la cara se vuelve a introducir en el alineador, el " -"cuadro delimitador se ajusta en una pequeña cantidad. Los puntos de " -"referencia finales se promedian en cada iteración. Esto ayuda a eliminar el " -"'micro-jitter', pero a costa de una menor velocidad de extracción. Cuantas " -"más veces se vuelva a introducir la cara en el alineador, menos " -"microfluctuaciones se producirán, pero la extracción será más larga." - -#: lib/cli/args.py:459 -msgid "" -"If a face isn't found, rotate the images to try to find a face. Can find " -"more faces at the cost of extraction speed. Pass in a single number to use " -"increments of that size up to 360, or pass in a list of numbers to enumerate " -"exactly what angles to check." -msgstr "" -"Si no se encuentra una cara, gira las imágenes para intentar encontrar una " -"cara. Puede encontrar más caras a costa de la velocidad de extracción. Pase " -"un solo número para usar incrementos de ese tamaño hasta 360, o pase una " -"lista de números para enumerar exactamente qué ángulos comprobar." - -#: lib/cli/args.py:471 lib/cli/args.py:481 lib/cli/args.py:494 -#: lib/cli/args.py:508 lib/cli/args.py:749 lib/cli/args.py:763 -#: lib/cli/args.py:776 lib/cli/args.py:790 -msgid "Face Processing" -msgstr "Proceso de Caras" - -#: lib/cli/args.py:472 -msgid "" -"Filters out faces detected below this size. Length, in pixels across the " -"diagonal of the bounding box. Set to 0 for off" -msgstr "" -"Filtra las caras detectadas por debajo de este tamaño. Longitud, en píxeles " -"a lo largo de la diagonal del cuadro delimitador. Establecer a 0 para " -"desactivar" - -#: lib/cli/args.py:482 lib/cli/args.py:764 -msgid "" -"Optionally filter out people who you do not wish to process by passing in an " -"image of that person. Should be a front portrait with a single person in the " -"image. Multiple images can be added space separated. NB: Using face filter " -"will significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Opcionalmente, puede filtrar las personas que no desea procesar pasando una " -"imagen de esa persona. Debe ser un retrato frontal con una sola persona en " -"la imagen. Se pueden añadir varias imágenes separadas por espacios. NB: El " -"uso del filtro de caras disminuirá significativamente la velocidad de " -"extracción y no se puede garantizar su precisión." - -#: lib/cli/args.py:495 lib/cli/args.py:777 -msgid "" -"Optionally select people you wish to process by passing in an image of that " -"person. Should be a front portrait with a single person in the image. " -"Multiple images can be added space separated. NB: Using face filter will " -"significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Opcionalmente, seleccione las personas que desea procesar pasando una imagen " -"de esa persona. Debe ser un retrato frontal con una sola persona en la " -"imagen. Se pueden añadir varias imágenes separadas por espacios. NB: El uso " -"del filtro facial disminuirá significativamente la velocidad de extracción y " -"no se puede garantizar su precisión." - -#: lib/cli/args.py:509 lib/cli/args.py:791 -msgid "" -"For use with the optional nfilter/filter files. Threshold for positive face " -"recognition. Lower values are stricter. NB: Using face filter will " -"significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Para usar con los archivos opcionales nfilter/filter. Umbral para el " -"reconocimiento positivo de caras. Los valores más bajos son más estrictos. " -"NB: El uso del filtro facial disminuirá significativamente la velocidad de " -"extracción y no se puede garantizar su precisión." - -#: lib/cli/args.py:520 lib/cli/args.py:532 lib/cli/args.py:544 -#: lib/cli/args.py:556 -msgid "output" -msgstr "salida" - -#: lib/cli/args.py:521 -msgid "" -"The output size of extracted faces. Make sure that the model you intend to " -"train supports your required size. This will only need to be changed for hi-" -"res models." -msgstr "" -"El tamaño de salida de las caras extraídas. Asegúrese de que el modelo que " -"pretende entrenar admite el tamaño deseado. Esto sólo tendrá que ser " -"cambiado para los modelos de alta resolución." - -#: lib/cli/args.py:533 -msgid "" -"Extract every 'nth' frame. This option will skip frames when extracting " -"faces. For example a value of 1 will extract faces from every frame, a value " -"of 10 will extract faces from every 10th frame." -msgstr "" -"Extraer cada 'enésimo' fotograma. Esta opción omitirá los fotogramas al " -"extraer las caras. Por ejemplo, un valor de 1 extraerá las caras de cada " -"fotograma, un valor de 10 extraerá las caras de cada 10 fotogramas." - -#: lib/cli/args.py:545 -msgid "" -"Automatically save the alignments file after a set amount of frames. By " -"default the alignments file is only saved at the end of the extraction " -"process. NB: If extracting in 2 passes then the alignments file will only " -"start to be saved out during the second pass. WARNING: Don't interrupt the " -"script when writing the file because it might get corrupted. Set to 0 to " -"turn off" -msgstr "" -"Guardar automáticamente el archivo de alineaciones después de una cantidad " -"determinada de cuadros. Por defecto, el archivo de alineaciones sólo se " -"guarda al final del proceso de extracción. Nota: Si se extrae en 2 pases, el " -"archivo de alineaciones sólo se empezará a guardar durante el segundo pase. " -"ADVERTENCIA: No interrumpa el script al escribir el archivo porque podría " -"corromperse. Poner a 0 para desactivar" - -#: lib/cli/args.py:557 -msgid "Draw landmarks on the ouput faces for debugging purposes." -msgstr "" -"Dibujar puntos de referencia en las caras de salida para fines de depuración." - -#: lib/cli/args.py:563 lib/cli/args.py:572 lib/cli/args.py:580 -#: lib/cli/args.py:587 lib/cli/args.py:803 lib/cli/args.py:814 -#: lib/cli/args.py:822 lib/cli/args.py:841 lib/cli/args.py:847 -msgid "settings" -msgstr "ajustes" - -#: lib/cli/args.py:564 -msgid "" -"Don't run extraction in parallel. Will run each part of the extraction " -"process separately (one after the other) rather than all at the smae time. " -"Useful if VRAM is at a premium." -msgstr "" -"No ejecute la extracción en paralelo. Ejecutará cada parte del proceso de " -"extracción por separado (una tras otra) en lugar de hacerlo todo al mismo " -"tiempo. Útil si la VRAM es escasa." - -#: lib/cli/args.py:573 -msgid "" -"Skips frames that have already been extracted and exist in the alignments " -"file" -msgstr "" -"Omite los fotogramas que ya han sido extraídos y que existen en el archivo " -"de alineaciones" - -#: lib/cli/args.py:581 -msgid "Skip frames that already have detected faces in the alignments file" -msgstr "" -"Omitir los fotogramas que ya tienen caras detectadas en el archivo de " -"alineaciones" - -#: lib/cli/args.py:588 -msgid "Skip saving the detected faces to disk. Just create an alignments file" -msgstr "" -"No guardar las caras detectadas en el disco. Crear sólo un archivo de " -"alineaciones" - -#: lib/cli/args.py:610 -msgid "" -"Swap the original faces in a source video/images to your final faces.\n" -"Conversion plugins can be configured in the 'Settings' Menu" -msgstr "" -"Cambia las caras originales de un vídeo/imágenes de origen por las caras " -"finales.\n" -"Los plugins de conversión pueden ser configurados en el menú \"Configuración" -"\"" - -#: lib/cli/args.py:631 -msgid "" -"Only required if converting from images to video. Provide The original video " -"that the source frames were extracted from (for extracting the fps and " -"audio)." -msgstr "" -"Sólo es necesario si se convierte de imágenes a vídeo. Proporcione el vídeo " -"original del que se extrajeron los fotogramas de origen (para extraer los " -"fps y el audio)." - -#: lib/cli/args.py:640 -msgid "" -"Model directory. The directory containing the trained model you wish to use " -"for conversion." -msgstr "" -"Directorio del modelo. El directorio que contiene el modelo entrenado que " -"desea utilizar para la conversión." - -#: lib/cli/args.py:650 -msgid "" -"R|Performs color adjustment to the swapped face. Some of these options have " -"configurable settings in '/config/convert.ini' or 'Settings > Configure " -"Convert Plugins':\n" -"L|avg-color: Adjust the mean of each color channel in the swapped " -"reconstruction to equal the mean of the masked area in the original image.\n" -"L|color-transfer: Transfers the color distribution from the source to the " -"target image using the mean and standard deviations of the L*a*b* color " -"space.\n" -"L|manual-balance: Manually adjust the balance of the image in a variety of " -"color spaces. Best used with the Preview tool to set correct values.\n" -"L|match-hist: Adjust the histogram of each color channel in the swapped " -"reconstruction to equal the histogram of the masked area in the original " -"image.\n" -"L|seamless-clone: Use cv2's seamless clone function to remove extreme " -"gradients at the mask seam by smoothing colors. Generally does not give very " -"satisfactory results.\n" -"L|none: Don't perform color adjustment." -msgstr "" -"R|Realiza un ajuste de color a la cara intercambiada. Algunas de estas " -"opciones tienen ajustes configurables en '/config/convert.ini' o 'Ajustes > " -"Configurar Extensiones de Conversión':\n" -"L|avg-color: Ajuste la media de cada canal de color en la reconstrucción " -"intercambiada para igualar la media del área enmascarada en la imagen " -"original.\n" -"L|color-transfer: Transfiere la distribución del color de la imagen de " -"origen a la de destino utilizando la media y las desviaciones estándar del " -"espacio de color L*a*b*.\n" -"L|manual-balance: Ajuste manualmente el equilibrio de la imagen en una " -"variedad de espacios de color. Se utiliza mejor con la herramienta de vista " -"previa para establecer los valores correctos.\n" -"L|match-hist: Ajuste el histograma de cada canal de color en la " -"reconstrucción intercambiada para igualar el histograma del área enmascarada " -"en la imagen original.\n" -"L|seamless-clone: Utilice la función de clonación sin costuras de cv2 para " -"eliminar los gradientes extremos en la costura de la máscara, suavizando los " -"colores. Generalmente no da resultados muy satisfactorios.\n" -"L|none: No realice el ajuste de color." - -#: lib/cli/args.py:677 -msgid "" -"R|Masker to use. NB: The mask you require must exist within the alignments " -"file. You can add additional masks with the Mask Tool.\n" -"L|none: Don't use a mask.\n" -"L|bisenet-fp-face: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked (configurable in mask settings). " -"Use this version of bisenet-fp if your model is trained with 'face' or " -"'legacy' centering.\n" -"L|bisenet-fp-head: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked (configurable in mask settings). " -"Use this version of bisenet-fp if your model is trained with 'head' " -"centering.\n" -"L|components: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks and the mask is extended upwards onto the " -"forehead.\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " -"faces clear of obstructions. Profile faces and obstructions may result in " -"sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " -"frontal faces. The mask model has been specifically trained to recognize " -"some facial obstructions (hands and eyeglasses). Profile faces may result in " -"sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " -"faces. The mask model has been trained by community members and will need " -"testing for further description. Profile faces may result in sub-par " -"performance.\n" -"L|predicted: If the 'Learn Mask' option was enabled during training, this " -"will use the mask that was created by the trained model." -msgstr "" -"R|Máscara a utilizar. NB: La máscara que necesita debe existir en el archivo " -"de alineaciones. Puede añadir máscaras adicionales con la herramienta de " -"máscaras.\n" -"L|none: No utilizar una máscara.\n" -"L|bisenet-fp-face: Máscara relativamente ligera basada en NN que proporciona " -"un control más refinado sobre el área a enmascarar (configurable en la " -"configuración de la máscara). Utilice esta versión de bisenet-fp si su " -"modelo está entrenado con centrado 'face' o 'legacy'.\n" -"L|bisenet-fp-head: Máscara relativamente ligera basada en NN que proporciona " -"un control más refinado sobre el área a enmascarar (configurable en la " -"configuración de la máscara). Utilice esta versión de bisenet-fp si su " -"modelo está entrenado con centrado de 'cabeza'.\n" -"L|components: Máscara diseñada para proporcionar una segmentación facial " -"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " -"Se construye un casco convexo alrededor del exterior de los puntos de " -"referencia para crear una máscara.\n" -"L|extended: Máscara diseñada para proporcionar una segmentación facial " -"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " -"Se construye un casco convexo alrededor del exterior de los puntos de " -"referencia y la máscara se extiende hacia arriba en la frente.\n" -"L|vgg-clear: Máscara diseñada para proporcionar una segmentación inteligente " -"de rostros principalmente frontales y libres de obstrucciones. Los rostros " -"de perfil y las obstrucciones pueden dar lugar a un rendimiento inferior.\n" -"L|vgg-obstructed: Máscara diseñada para proporcionar una segmentación " -"inteligente de rostros principalmente frontales. El modelo de la máscara ha " -"sido entrenado específicamente para reconocer algunas obstrucciones faciales " -"(manos y gafas). Los rostros de perfil pueden dar lugar a un rendimiento " -"inferior.\n" -"L|unet-dfl: Máscara diseñada para proporcionar una segmentación inteligente " -"de rostros principalmente frontales. El modelo de máscara ha sido entrenado " -"por los miembros de la comunidad y necesitará ser probado para una mayor " -"descripción. Los rostros de perfil pueden dar lugar a un rendimiento " -"inferior.\n" -"L|predicted: Si la opción 'Learn Mask' se habilitó durante el entrenamiento, " -"esto usará la máscara que fue creada por el modelo entrenado." - -#: lib/cli/args.py:713 -msgid "" -"R|The plugin to use to output the converted images. The writers are " -"configurable in '/config/convert.ini' or 'Settings > Configure Convert " -"Plugins:'\n" -"L|ffmpeg: [video] Writes out the convert straight to video. When the input " -"is a series of images then the '-ref' (--reference-video) parameter must be " -"set.\n" -"L|gif: [animated image] Create an animated gif.\n" -"L|opencv: [images] The fastest image writer, but less options and formats " -"than other plugins.\n" -"L|pillow: [images] Slower than opencv, but has more options and supports " -"more formats." -msgstr "" -"R|El plugin a utilizar para dar salida a las imágenes convertidas. Los " -"escritores son configurables en '/config/convert.ini' o 'Ajustes > " -"Configurar Extensiones de Conversión:'\n" -"L|ffmpeg: [video] Escribe la conversión directamente en vídeo. Cuando la " -"entrada es una serie de imágenes, el parámetro '-ref' (--reference-video) " -"debe ser establecido.\n" -"L|gif: [imagen animada] Crea un gif animado.\n" -"L|opencv: [images] El escritor de imágenes más rápido, pero con menos " -"opciones y formatos que otros plugins.\n" -"L|pillow: [images] Más lento que opencv, pero tiene más opciones y soporta " -"más formatos." - -#: lib/cli/args.py:732 lib/cli/args.py:739 lib/cli/args.py:833 -msgid "Frame Processing" -msgstr "Proceso de fotogramas" - -#: lib/cli/args.py:733 -msgid "" -"Scale the final output frames by this amount. 100%% will output the frames " -"at source dimensions. 50%% at half size 200%% at double size" -msgstr "" -"Escala los fotogramas finales de salida en esta cantidad. 100%% dará salida " -"a los fotogramas a las dimensiones de origen. 50%% a la mitad de tamaño. 200%" -"% al doble de tamaño" - -#: lib/cli/args.py:740 -msgid "" -"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " -"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " -"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " -"converting from images, then the filenames must end with the frame-number!" -msgstr "" -"Rangos de fotogramas a los que aplicar la transferencia, por ejemplo, para " -"los fotogramas de 10 a 50 y de 90 a 100 utilice --frame-ranges 10-50 90-100. " -"Los fotogramas que queden fuera del rango seleccionado se descartarán a " -"menos que se seleccione '-k' (--keep-unchanged). Nota: Si está convirtiendo " -"imágenes, ¡los nombres de los archivos deben terminar con el número de " -"fotograma!" - -#: lib/cli/args.py:750 -msgid "" -"If you have not cleansed your alignments file, then you can filter out faces " -"by defining a folder here that contains the faces extracted from your input " -"files/video. If this folder is defined, then only faces that exist within " -"your alignments file and also exist within the specified folder will be " -"converted. Leaving this blank will convert all faces that exist within the " -"alignments file." -msgstr "" -"Si no ha limpiado su archivo de alineaciones, puede filtrar las caras " -"definiendo aquí una carpeta que contenga las caras extraídas de sus archivos/" -"vídeos de entrada. Si se define esta carpeta, sólo se convertirán las caras " -"que existan en el archivo de alineaciones y también en la carpeta " -"especificada. Si se deja en blanco, se convertirán todas las caras que " -"existan en el archivo de alineaciones." - -#: lib/cli/args.py:804 -msgid "" -"The maximum number of parallel processes for performing conversion. " -"Converting images is system RAM heavy so it is possible to run out of memory " -"if you have a lot of processes and not enough RAM to accommodate them all. " -"Setting this to 0 will use the maximum available. No matter what you set " -"this to, it will never attempt to use more processes than are available on " -"your system. If singleprocess is enabled this setting will be ignored." -msgstr "" -"El número máximo de procesos paralelos para realizar la conversión. La " -"conversión de imágenes requiere mucha RAM del sistema, por lo que es posible " -"que se agote la memoria si tiene muchos procesos y no hay suficiente RAM " -"para acomodarlos a todos. Si se ajusta a 0, se utilizará el máximo " -"disponible. No importa lo que establezca, nunca intentará utilizar más " -"procesos que los disponibles en su sistema. Si 'singleprocess' está " -"habilitado, este ajuste será ignorado." - -#: lib/cli/args.py:815 -msgid "" -"[LEGACY] This only needs to be selected if a legacy model is being loaded or " -"if there are multiple models in the model folder" -msgstr "" -"[LEGACY] Sólo es necesario seleccionar esta opción si se está cargando un " -"modelo heredado si hay varios modelos en la carpeta de modelos" - -#: lib/cli/args.py:823 -msgid "" -"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " -"alignments file for your destination video. However, if you wish you can " -"generate the alignments on-the-fly by enabling this option. This will use an " -"inferior extraction pipeline and will lead to substandard results. If an " -"alignments file is found, this option will be ignored." -msgstr "" -"Activar la conversión sobre la marcha. NO se recomienda. Debe generar un " -"archivo de alineación limpio para su vídeo de destino. Sin embargo, si lo " -"desea, puede generar las alineaciones sobre la marcha activando esta opción. " -"Esto utilizará una tubería de extracción inferior y conducirá a resultados " -"de baja calidad. Si se encuentra un archivo de alineaciones, esta opción " -"será ignorada." - -#: lib/cli/args.py:834 -msgid "" -"When used with --frame-ranges outputs the unchanged frames that are not " -"processed instead of discarding them." -msgstr "" -"Cuando se usa con --frame-ranges, la salida incluye los fotogramas no " -"procesados en vez de descartarlos." - -#: lib/cli/args.py:842 -msgid "Swap the model. Instead converting from of A -> B, converts B -> A" -msgstr "" -"Intercambiar el modelo. En vez de convertir de A a B, convierte de B a A" - -#: lib/cli/args.py:848 -msgid "Disable multiprocessing. Slower but less resource intensive." -msgstr "Desactiva el multiproceso. Es más lento, pero usa menos recursos." - -#: lib/cli/args.py:864 -msgid "" -"Train a model on extracted original (A) and swap (B) faces.\n" -"Training models can take a long time. Anything from 24hrs to over a week\n" -"Model plugins can be configured in the 'Settings' Menu" -msgstr "" -"Entrene un modelo con las caras originales (A) e intercambiadas (B) " -"extraídas.\n" -"El entrenamiento de los modelos puede llevar mucho tiempo. Desde 24 horas " -"hasta más de una semana.\n" -"Los plugins de los modelos pueden configurarse en el menú \"Ajustes\"" - -#: lib/cli/args.py:883 lib/cli/args.py:892 -msgid "faces" -msgstr "caras" - -#: lib/cli/args.py:884 -msgid "" -"Input directory. A directory containing training images for face A. This is " -"the original face, i.e. the face that you want to remove and replace with " -"face B." -msgstr "" -"Directorio de entrada. Un directorio que contiene imágenes de entrenamiento " -"para la cara A. Esta es la cara original, es decir, la cara que se quiere " -"eliminar y sustituir por la cara B." - -#: lib/cli/args.py:893 -msgid "" -"Input directory. A directory containing training images for face B. This is " -"the swap face, i.e. the face that you want to place onto the head of person " -"A." -msgstr "" -"Directorio de entrada. Un directorio que contiene imágenes de entrenamiento " -"para la cara B. Esta es la cara de intercambio, es decir, la cara que se " -"quiere colocar en la cabeza de la persona A." - -#: lib/cli/args.py:901 lib/cli/args.py:913 lib/cli/args.py:929 -#: lib/cli/args.py:954 lib/cli/args.py:964 -msgid "model" -msgstr "modelo" - -#: lib/cli/args.py:902 -msgid "" -"Model directory. This is where the training data will be stored. You should " -"always specify a new folder for new models. If starting a new model, select " -"either an empty folder, or a folder which does not exist (which will be " -"created). If continuing to train an existing model, specify the location of " -"the existing model." -msgstr "" -"Directorio del modelo. Aquí es donde se almacenarán los datos de " -"entrenamiento. Siempre debe especificar una nueva carpeta para los nuevos " -"modelos. Si se inicia un nuevo modelo, seleccione una carpeta vacía o una " -"carpeta que no exista (que se creará). Si continúa entrenando un modelo " -"existente, especifique la ubicación del modelo existente." - -#: lib/cli/args.py:914 -msgid "" -"R|Load the weights from a pre-existing model into a newly created model. For " -"most models this will load weights from the Encoder of the given model into " -"the encoder of the newly created model. Some plugins may have specific " -"configuration options allowing you to load weights from other layers. " -"Weights will only be loaded when creating a new model. This option will be " -"ignored if you are resuming an existing model. Generally you will also want " -"to 'freeze-weights' whilst the rest of your model catches up with your " -"Encoder.\n" -"NB: Weights can only be loaded from models of the same plugin as you intend " -"to train." -msgstr "" -"R|Cargue los pesos de un modelo preexistente en un modelo recién creado. " -"Para la mayoría de los modelos, esto cargará pesos del codificador del " -"modelo dado en el codificador del modelo recién creado. Algunos complementos " -"pueden tener opciones de configuración específicas que le permiten cargar " -"pesos de otras capas. Los pesos solo se cargarán al crear un nuevo modelo. " -"Esta opción se ignorará si está reanudando un modelo existente. En general, " -"también querrá 'congelar pesos' mientras el resto de su modelo se pone al " -"día con su codificador.\n" -"NB: Los pesos solo se pueden cargar desde modelos del mismo complemento que " -"desea entrenar." - -#: lib/cli/args.py:930 -msgid "" -"R|Select which trainer to use. Trainers can be configured from the Settings " -"menu or the config folder.\n" -"L|original: The original model created by /u/deepfakes.\n" -"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " -"for full dfaker method.\n" -"L|dfl-h128: 128px in/out model from deepfacelab\n" -"L|dfl-sae: Adaptable model from deepfacelab\n" -"L|dlight: A lightweight, high resolution DFaker variant.\n" -"L|iae: A model that uses intermediate layers to try to get better details\n" -"L|lightweight: A lightweight model for low-end cards. Don't expect great " -"results. Can train as low as 1.6GB with batch size 8.\n" -"L|realface: A high detail, dual density model based on DFaker, with " -"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " -"won't work so well. By andenixa et al. Very configurable.\n" -"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " -"unbalanced so B>A swaps won't work so well. Very configurable.\n" -"L|villain: 128px in/out model from villainguy. Very resource hungry (You " -"will require a GPU with a fair amount of VRAM). Good for details, but more " -"susceptible to color differences." -msgstr "" -"R|Seleccione el entrenador que desea utilizar. Los entrenadores se pueden " -"configurar desde el menú de configuración o la carpeta de configuración.\n" -"L|original: El modelo original creado por /u/deepfakes.\n" -"L|dfaker: Modelo de 64px in/128px out de dfaker. Habilitar 'warp-to-" -"landmarks' para el método completo de dfaker.\n" -"L|dfl-h128: modelo de 128px in/out de deepfacelab\n" -"L|dfl-sae: Modelo adaptable de deepfacelab\n" -"L|dlight: Una variante de DFaker ligera y de alta resolución.\n" -"L|iae: Un modelo que utiliza capas intermedias para tratar de obtener " -"mejores detalles.\n" -"L|lightweight: Un modelo ligero para tarjetas de gama baja. No esperes " -"grandes resultados. Puede entrenar hasta 1,6GB con tamaño de lote 8.\n" -"L|realface: Un modelo de alto detalle y doble densidad basado en DFaker, con " -"resolución de entrada y salida personalizable. Los autocodificadores están " -"desequilibrados, por lo que los intercambios B>A no funcionan tan bien. Por " -"andenixa et al. Muy configurable\n" -"L|Unbalanced: modelo de 128px de entrada/salida de andenixa. Los " -"autocodificadores están desequilibrados por lo que los intercambios B>A no " -"funcionarán tan bien. Muy configurable\n" -"L|villain: Modelo de 128px de entrada/salida de villainguy. Requiere muchos " -"recursos (se necesita una GPU con una buena cantidad de VRAM). Bueno para " -"los detalles, pero más susceptible a las diferencias de color." - -#: lib/cli/args.py:955 -msgid "" -"Output a summary of the model and exit. If a model folder is provided then a " -"summary of the saved model is displayed. Otherwise a summary of the model " -"that would be created by the chosen plugin and configuration settings is " -"displayed." -msgstr "" -"Genere un resumen del modelo y salga. Si se proporciona una carpeta de " -"modelo, se muestra un resumen del modelo guardado. De lo contrario, se " -"muestra un resumen del modelo que crearía el complemento elegido y los " -"ajustes de configuración." - -#: lib/cli/args.py:965 -msgid "" -"Freeze the weights of the model. Freezing weights means that some of the " -"parameters in the model will no longer continue to learn, but those that are " -"not frozen will continue to learn. For most models, this will freeze the " -"encoder, but some models may have configuration options for freezing other " -"layers." -msgstr "" -"Congele los pesos del modelo. Congelar pesos significa que algunos de los " -"parámetros del modelo ya no seguirán aprendiendo, pero los que no están " -"congelados seguirán aprendiendo. Para la mayoría de los modelos, esto " -"congelará el codificador, pero algunos modelos pueden tener opciones de " -"configuración para congelar otras capas." - -#: lib/cli/args.py:978 lib/cli/args.py:990 lib/cli/args.py:1001 -#: lib/cli/args.py:1087 -msgid "training" -msgstr "entrenamiento" - -#: lib/cli/args.py:979 -msgid "" -"Batch size. This is the number of images processed through the model for " -"each side per iteration. NB: As the model is fed 2 sides at a time, the " -"actual number of images within the model at any one time is double the " -"number that you set here. Larger batches require more GPU RAM." -msgstr "" -"Tamaño del lote. Este es el número de imágenes procesadas a través del " -"modelo para cada lado por iteración. Nota: Como el modelo se alimenta de 2 " -"lados a la vez, el número real de imágenes dentro del modelo en cualquier " -"momento es el doble del número que se establece aquí. Los lotes más grandes " -"requieren más RAM de la GPU." - -#: lib/cli/args.py:991 -msgid "" -"Length of training in iterations. This is only really used for automation. " -"There is no 'correct' number of iterations a model should be trained for. " -"You should stop training when you are happy with the previews. However, if " -"you want the model to stop automatically at a set number of iterations, you " -"can set that value here." -msgstr "" -"Duración del entrenamiento en iteraciones. Esto sólo se utiliza realmente " -"para la automatización. No hay un número 'correcto' de iteraciones para las " -"que deba entrenarse un modelo. Debe dejar de entrenar cuando esté satisfecho " -"con las previsiones. Sin embargo, si desea que el modelo se detenga " -"automáticamente en un número determinado de iteraciones, puede establecer " -"ese valor aquí." - -#: lib/cli/args.py:1002 -msgid "" -"Use the Tensorflow Mirrored Distrubution Strategy to train on multiple GPUs." -msgstr "" -"Utilice la estrategia de distribución en espejo de Tensorflow para entrenar " -"en múltiples GPUs." - -#: lib/cli/args.py:1012 lib/cli/args.py:1022 -msgid "Saving" -msgstr "Guardar" - -#: lib/cli/args.py:1013 -msgid "Sets the number of iterations between each model save." -msgstr "Establece el número de iteraciones entre cada guardado del modelo." - -#: lib/cli/args.py:1023 -msgid "" -"Sets the number of iterations before saving a backup snapshot of the model " -"in it's current state. Set to 0 for off." -msgstr "" -"Establece el número de iteraciones antes de guardar una copia de seguridad " -"del modelo en su estado actual. Establece 0 para que esté desactivado." - -#: lib/cli/args.py:1030 lib/cli/args.py:1041 lib/cli/args.py:1052 -msgid "timelapse" -msgstr "intervalo" - -#: lib/cli/args.py:1031 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. " -"This should be the input folder of 'A' faces that you would like to use for " -"creating the timelapse. You must also supply a --timelapse-output and a --" -"timelapse-input-B parameter." -msgstr "" -"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " -"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " -"Esta debe ser la carpeta de entrada de las caras \"A\" que desea utilizar " -"para crear el timelapse. También debe suministrar un parámetro --timelapse-" -"output y un parámetro --timelapse-input-B." - -#: lib/cli/args.py:1042 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. " -"This should be the input folder of 'B' faces that you would like to use for " -"creating the timelapse. You must also supply a --timelapse-output and a --" -"timelapse-input-A parameter." -msgstr "" -"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " -"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " -"Esta debe ser la carpeta de entrada de las caras \"B\" que desea utilizar " -"para crear el timelapse. También debe suministrar un parámetro --timelapse-" -"output y un parámetro --timelapse-input-A." - -#: lib/cli/args.py:1053 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. If " -"the input folders are supplied but no output folder, it will default to your " -"model folder /timelapse/" -msgstr "" -"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " -"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " -"Si se suministran las carpetas de entrada pero no la carpeta de salida, se " -"guardará por defecto en la carpeta del modelo /timelapse/" - -#: lib/cli/args.py:1065 lib/cli/args.py:1072 lib/cli/args.py:1079 -msgid "preview" -msgstr "previsualización" - -#: lib/cli/args.py:1066 -msgid "" -"Percentage amount to scale the preview by. 100%% is the model output size." -msgstr "" -"Cantidad porcentual para escalar la vista previa. 100%% es el tamaño de " -"salida del modelo." - -#: lib/cli/args.py:1073 -msgid "Show training preview output. in a separate window." -msgstr "" -"Mostrar la salida de la vista previa del entrenamiento. en una ventana " -"separada." - -#: lib/cli/args.py:1080 -msgid "" -"Writes the training result to a file. The image will be stored in the root " -"of your FaceSwap folder." -msgstr "" -"Escribe el resultado del entrenamiento en un archivo. La imagen se " -"almacenará en la raíz de su carpeta FaceSwap." - -#: lib/cli/args.py:1088 -msgid "" -"Disables TensorBoard logging. NB: Disabling logs means that you will not be " -"able to use the graph or analysis for this session in the GUI." -msgstr "" -"Desactiva el registro de TensorBoard. NB: Desactivar los registros significa " -"que no podrá utilizar el gráfico o el análisis de esta sesión en la GUI." - -#: lib/cli/args.py:1095 lib/cli/args.py:1104 lib/cli/args.py:1113 -#: lib/cli/args.py:1122 -msgid "augmentation" -msgstr "aumento" - -#: lib/cli/args.py:1096 -msgid "" -"Warps training faces to closely matched Landmarks from the opposite face-set " -"rather than randomly warping the face. This is the 'dfaker' way of doing " -"warping." -msgstr "" -"Deforma las caras de entrenamiento a puntos de referencia muy parecidos del " -"conjunto de caras opuestas en lugar de deformar la cara al azar. Esta es la " -"forma 'dfaker' de hacer la deformación." - -#: lib/cli/args.py:1105 -msgid "" -"To effectively learn, a random set of images are flipped horizontally. " -"Sometimes it is desirable for this not to occur. Generally this should be " -"left off except for during 'fit training'." -msgstr "" -"Para aprender de forma efectiva, se voltea horizontalmente un conjunto " -"aleatorio de imágenes. A veces es deseable que esto no ocurra. Por lo " -"general, esto debería dejarse sin efecto, excepto durante el 'entrenamiento " -"de ajuste'." - -#: lib/cli/args.py:1114 -msgid "" -"Color augmentation helps make the model less susceptible to color " -"differences between the A and B sets, at an increased training time cost. " -"Enable this option to disable color augmentation." -msgstr "" -"El aumento del color ayuda a que el modelo sea menos susceptible a las " -"diferencias de color entre los conjuntos A y B, con un mayor coste de tiempo " -"de entrenamiento. Activa esta opción para desactivar el aumento de color." - -#: lib/cli/args.py:1123 -msgid "" -"Warping is integral to training the Neural Network. This option should only " -"be enabled towards the very end of training to try to bring out more detail. " -"Think of it as 'fine-tuning'. Enabling this option from the beginning is " -"likely to kill a model and lead to terrible results." -msgstr "" -"La deformación es fundamental para el entrenamiento de la red neuronal. Esta " -"opción sólo debería activarse hacia el final del entrenamiento para tratar " -"de obtener más detalles. Piense en ello como un 'ajuste fino'. Si se activa " -"esta opción desde el principio, es probable que arruine el modelo y se " -"obtengan resultados terribles." - -#: lib/cli/args.py:1148 +#: lib/cli/args.py:319 msgid "Output to Shell console instead of GUI console" msgstr "Salida a la consola Shell en lugar de la consola GUI" - -#~ msgid "" -#~ "DEPRECATED - This option will be removed in a future update. Path to " -#~ "alignments file for training set A. Defaults to /alignments.json " -#~ "if not provided." -#~ msgstr "" -#~ "DEPRECIADO - Esta opción se eliminará en una futura actualización. Ruta " -#~ "al archivo de alineaciones para el conjunto de entrenamiento A. Por " -#~ "defecto es /alignments.json si no se proporciona." - -#~ msgid "" -#~ "DEPRECATED - This option will be removed in a future update. Path to " -#~ "alignments file for training set B. Defaults to /alignments.json " -#~ "if not provided." -#~ msgstr "" -#~ "DEPRECIADO - Esta opción se eliminará en una futura actualización. Ruta " -#~ "al archivo de alineaciones para el conjunto de entrenamiento B. Por " -#~ "defecto es /alignments.json si no se proporciona." diff --git a/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.mo b/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.mo new file mode 100644 index 0000000000..5ccd0c2f21 Binary files /dev/null and b/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.mo differ diff --git a/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.po b/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.po new file mode 100755 index 0000000000..4b2e299256 --- /dev/null +++ b/locales/es/LC_MESSAGES/lib.cli.args_extract_convert.po @@ -0,0 +1,720 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: faceswap.spanish\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-12 11:56+0100\n" +"PO-Revision-Date: 2024-04-12 12:02+0100\n" +"Last-Translator: \n" +"Language-Team: tokafondo\n" +"Language: es\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_extract_convert.py:46 lib/cli/args_extract_convert.py:56 +#: lib/cli/args_extract_convert.py:64 lib/cli/args_extract_convert.py:122 +#: lib/cli/args_extract_convert.py:483 lib/cli/args_extract_convert.py:492 +msgid "Data" +msgstr "Datos" + +#: lib/cli/args_extract_convert.py:48 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file. NB: This should be the source video/" +"frames NOT the source faces." +msgstr "" +"Directorio o vídeo de entrada. Un directorio que contenga los archivos de " +"imagen que desea procesar o la ruta a un archivo de vídeo. NB: Debe ser el " +"vídeo/los fotogramas de origen, NO las caras de origen." + +#: lib/cli/args_extract_convert.py:57 +msgid "Output directory. This is where the converted files will be saved." +msgstr "" +"Directorio de salida. Aquí es donde se guardarán los archivos convertidos." + +#: lib/cli/args_extract_convert.py:66 +msgid "" +"Optional path to an alignments file. Leave blank if the alignments file is " +"at the default location." +msgstr "" +"Ruta opcional a un archivo de alineaciones. Dejar en blanco si el archivo de " +"alineaciones está en la ubicación por defecto." + +#: lib/cli/args_extract_convert.py:97 +msgid "" +"Extract faces from image or video sources.\n" +"Extraction plugins can be configured in the 'Settings' Menu" +msgstr "" +"Extrae caras de fuentes de imagen o video.\n" +"Los plugins de extracción pueden ser configuradas en el menú de 'Ajustes'" + +#: lib/cli/args_extract_convert.py:124 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple videos and/or folders of images you wish to extract from. The faces " +"will be output to separate sub-folders in the output_dir." +msgstr "" +"Si se selecciona, input_dir debe ser una carpeta principal que contenga " +"varios videos y/o carpetas de imágenes de las que desea extraer. Las caras " +"se enviarán a subcarpetas separadas en output_dir." + +#: lib/cli/args_extract_convert.py:133 lib/cli/args_extract_convert.py:152 +#: lib/cli/args_extract_convert.py:167 lib/cli/args_extract_convert.py:206 +#: lib/cli/args_extract_convert.py:224 lib/cli/args_extract_convert.py:237 +#: lib/cli/args_extract_convert.py:247 lib/cli/args_extract_convert.py:257 +#: lib/cli/args_extract_convert.py:503 lib/cli/args_extract_convert.py:529 +#: lib/cli/args_extract_convert.py:568 +msgid "Plugins" +msgstr "Extensiones" + +#: lib/cli/args_extract_convert.py:135 +msgid "" +"R|Detector to use. Some of these have configurable settings in '/config/" +"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" +"L|cv2-dnn: A CPU only extractor which is the least reliable and least " +"resource intensive. Use this if not using a GPU and time is important.\n" +"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " +"than other GPU detectors but can often return more false positives.\n" +"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " +"fewer false positives than other GPU detectors, but is a lot more resource " +"intensive.\n" +"L|external: Import a face detection bounding box from a json file. " +"(configurable in Detect settings)" +msgstr "" +"R|Detector de caras a usar. Algunos tienen ajustes configurables en '/config/" +"extract.ini' o 'Ajustes > Configurar Extensiones de Extracción:\n" +"L|cv2-dnn: Extractor que usa sólo la CPU. Es el menos fiable y el que menos " +"recursos usa. Elegir este si necesita rapidez y no usar la GPU.\n" +"L|mtcnn: Buen detector. Rápido en la CPU y más rápido en la GPU. Usa menos " +"recursos que otros detectores basados en GPU, pero puede devolver más falsos " +"positivos.\n" +"L|s3fd: El mejor detector. Lento en la CPU, y más rápido en la GPU. Puede " +"detectar más caras y tiene menos falsos positivos que otros detectores " +"basados en GPU, pero uso muchos más recursos.\n" +"L|external: importe un cuadro de detección de detección de cara desde un " +"archivo JSON. (configurable en la configuración de detección)" + +#: lib/cli/args_extract_convert.py:154 +msgid "" +"R|Aligner to use.\n" +"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " +"but less accurate. Only use this if not using a GPU and time is important.\n" +"L|fan: Best aligner. Fast on GPU, slow on CPU.\n" +"L|external: Import 68 point 2D landmarks or an aligned bounding box from a " +"json file. (configurable in Align settings)" +msgstr "" +"R|Alineador a usar.\n" +"L|cv2-dnn: Detector que usa sólo la CPU. Más rápido, usa menos recursos, " +"pero es menos preciso. Elegir este si necesita rapidez y no usar la GPU.\n" +"L|fan: El mejor alineador. Rápido en la GPU, y lento en la CPU.\n" +"L|external: importar 68 puntos 2D Modos de referencia o un cuadro " +"delimitador alineado de un archivo JSON. (configurable en la configuración " +"alineada)" + +#: lib/cli/args_extract_convert.py:169 +msgid "" +"R|Additional Masker(s) to use. The masks generated here will all take up GPU " +"RAM. You can select none, one or multiple masks, but the extraction may take " +"longer the more you select. NB: The Extended and Components (landmark based) " +"masks are automatically generated on extraction.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU so will not use any additional VRAM.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"The auto generated masks are as follows:\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" +msgstr "" +"R|Enmascarador(es) adicional(es) a usar. Las máscaras generadas aquí usarán " +"todas RAM de la GPU. Puede seleccionar una, varias o ninguna máscaras, pero " +"la extracción tardará más cuanto más marque. Las máscaras Extended y " +"Components son siempre generadas durante la extracción.\n" +"L|bisenet-fp: Máscara relativamente ligera basada en NN que proporciona un " +"control más refinado sobre el área a enmascarar, incluido el enmascaramiento " +"completo de la cabeza (configurable en la configuración de la máscara).\n" +"L|custom: Una máscara ficticia que llena el área de la máscara con 1 o 0 " +"(configurable en la configuración). Esto solo es necesario si tiene la " +"intención de editar manualmente las máscaras personalizadas usted mismo en " +"la herramienta manual. Esta máscara no usa la GPU, por lo que no usará VRAM " +"adicional.\n" +"L|vgg-clear: Máscara diseñada para proporcionar una segmentación inteligente " +"de rostros principalmente frontales y libres de obstrucciones. Los rostros " +"de perfil y las obstrucciones pueden dar lugar a un rendimiento inferior.\n" +"L|vgg-obstructed: Máscara diseñada para proporcionar una segmentación " +"inteligente de rostros principalmente frontales. El modelo de la máscara ha " +"sido entrenado específicamente para reconocer algunas obstrucciones faciales " +"(manos y gafas). Los rostros de perfil pueden dar lugar a un rendimiento " +"inferior.\n" +"L|unet-dfl: Máscara diseñada para proporcionar una segmentación inteligente " +"de rostros principalmente frontales. El modelo de máscara ha sido entrenado " +"por los miembros de la comunidad y necesitará ser probado para una mayor " +"descripción. Los rostros de perfil pueden dar lugar a un rendimiento " +"inferior.\n" +"Las máscaras que siempre se generan son:\n" +"L|components: Máscara diseñada para proporcionar una segmentación facial " +"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " +"Se construye un casco convexo alrededor del exterior de los puntos de " +"referencia para crear una máscara.\n" +"L|extended: Máscara diseñada para proporcionar una segmentación facial " +"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " +"Se construye un casco convexo alrededor del exterior de los puntos de " +"referencia y la máscara se extiende hacia arriba en la frente.\n" +"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" + +#: lib/cli/args_extract_convert.py:208 +msgid "" +"R|Performing normalization can help the aligner better align faces with " +"difficult lighting conditions at an extraction speed cost. Different methods " +"will yield different results on different sets. NB: This does not impact the " +"output face, just the input to the aligner.\n" +"L|none: Don't perform normalization on the face.\n" +"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"L|hist: Equalize the histograms on the RGB channels.\n" +"L|mean: Normalize the face colors to the mean." +msgstr "" +"R|Realizar la normalización puede ayudar al alineador a alinear mejor las " +"caras con condiciones de iluminación difíciles a un coste de velocidad de " +"extracción. Diferentes métodos darán diferentes resultados en diferentes " +"conjuntos. NB: Esto no afecta a la cara de salida, sólo a la entrada del " +"alineador.\n" +"L|none: No realice la normalización en la cara.\n" +"L|clahe: Realice la ecualización adaptativa del histograma con contraste " +"limitado en el rostro.\n" +"L|hist: Iguala los histogramas de los canales RGB.\n" +"L|mean: Normalizar los colores de la cara a la media." + +#: lib/cli/args_extract_convert.py:226 +msgid "" +"The number of times to re-feed the detected face into the aligner. Each time " +"the face is re-fed into the aligner the bounding box is adjusted by a small " +"amount. The final landmarks are then averaged from each iteration. Helps to " +"remove 'micro-jitter' but at the cost of slower extraction speed. The more " +"times the face is re-fed into the aligner, the less micro-jitter should " +"occur but the longer extraction will take." +msgstr "" +"El número de veces que hay que volver a introducir la cara detectada en el " +"alineador. Cada vez que la cara se vuelve a introducir en el alineador, el " +"cuadro delimitador se ajusta en una pequeña cantidad. Los puntos de " +"referencia finales se promedian en cada iteración. Esto ayuda a eliminar el " +"'micro-jitter', pero a costa de una menor velocidad de extracción. Cuantas " +"más veces se vuelva a introducir la cara en el alineador, menos " +"microfluctuaciones se producirán, pero la extracción será más larga." + +#: lib/cli/args_extract_convert.py:239 +msgid "" +"Re-feed the initially found aligned face through the aligner. Can help " +"produce better alignments for faces that are rotated beyond 45 degrees in " +"the frame or are at extreme angles. Slows down extraction." +msgstr "" +"Vuelva a introducir la cara alineada encontrada inicialmente a través del " +"alineador. Puede ayudar a producir mejores alineaciones para las caras que " +"se giran más de 45 grados en el marco o se encuentran en ángulos extremos. " +"Ralentiza la extracción." + +#: lib/cli/args_extract_convert.py:249 +msgid "" +"If a face isn't found, rotate the images to try to find a face. Can find " +"more faces at the cost of extraction speed. Pass in a single number to use " +"increments of that size up to 360, or pass in a list of numbers to enumerate " +"exactly what angles to check." +msgstr "" +"Si no se encuentra una cara, gira las imágenes para intentar encontrar una " +"cara. Puede encontrar más caras a costa de la velocidad de extracción. Pase " +"un solo número para usar incrementos de ese tamaño hasta 360, o pase una " +"lista de números para enumerar exactamente qué ángulos comprobar." + +#: lib/cli/args_extract_convert.py:259 +msgid "" +"Obtain and store face identity encodings from VGGFace2. Slows down extract a " +"little, but will save time if using 'sort by face'" +msgstr "" +"Obtenga y almacene codificaciones de identidad facial de VGGFace2. Ralentiza " +"un poco la extracción, pero ahorrará tiempo si usa 'sort by face'" + +#: lib/cli/args_extract_convert.py:269 lib/cli/args_extract_convert.py:280 +#: lib/cli/args_extract_convert.py:293 lib/cli/args_extract_convert.py:307 +#: lib/cli/args_extract_convert.py:614 lib/cli/args_extract_convert.py:623 +#: lib/cli/args_extract_convert.py:638 lib/cli/args_extract_convert.py:651 +#: lib/cli/args_extract_convert.py:665 +msgid "Face Processing" +msgstr "Proceso de Caras" + +#: lib/cli/args_extract_convert.py:271 +msgid "" +"Filters out faces detected below this size. Length, in pixels across the " +"diagonal of the bounding box. Set to 0 for off" +msgstr "" +"Filtra las caras detectadas por debajo de este tamaño. Longitud, en píxeles " +"a lo largo de la diagonal del cuadro delimitador. Establecer a 0 para " +"desactivar" + +#: lib/cli/args_extract_convert.py:282 +msgid "" +"Optionally filter out people who you do not wish to extract by passing in " +"images of those people. Should be a small variety of images at different " +"angles and in different conditions. A folder containing the required images " +"or multiple image files, space separated, can be selected." +msgstr "" +"Opcionalmente, filtre a las personas que no desea extraer pasando imágenes " +"de esas personas. Debe ser una pequeña variedad de imágenes en diferentes " +"ángulos y en diferentes condiciones. Se puede seleccionar una carpeta que " +"contenga las imágenes requeridas o múltiples archivos de imágenes, separados " +"por espacios." + +#: lib/cli/args_extract_convert.py:295 +msgid "" +"Optionally select people you wish to extract by passing in images of that " +"person. Should be a small variety of images at different angles and in " +"different conditions A folder containing the required images or multiple " +"image files, space separated, can be selected." +msgstr "" +"Opcionalmente, seleccione las personas que desea extraer pasando imágenes de " +"esa persona. Debe haber una pequeña variedad de imágenes en diferentes " +"ángulos y en diferentes condiciones. Se puede seleccionar una carpeta que " +"contenga las imágenes requeridas o múltiples archivos de imágenes, separados " +"por espacios." + +#: lib/cli/args_extract_convert.py:309 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Higher values are stricter." +msgstr "" +"Para usar con los archivos nfilter/filter opcionales. Umbral para el " +"reconocimiento facial positivo. Los valores más altos son más estrictos." + +#: lib/cli/args_extract_convert.py:318 lib/cli/args_extract_convert.py:331 +#: lib/cli/args_extract_convert.py:344 lib/cli/args_extract_convert.py:356 +msgid "output" +msgstr "salida" + +#: lib/cli/args_extract_convert.py:320 +msgid "" +"The output size of extracted faces. Make sure that the model you intend to " +"train supports your required size. This will only need to be changed for hi-" +"res models." +msgstr "" +"El tamaño de salida de las caras extraídas. Asegúrese de que el modelo que " +"pretende entrenar admite el tamaño deseado. Esto sólo tendrá que ser " +"cambiado para los modelos de alta resolución." + +#: lib/cli/args_extract_convert.py:333 +msgid "" +"Extract every 'nth' frame. This option will skip frames when extracting " +"faces. For example a value of 1 will extract faces from every frame, a value " +"of 10 will extract faces from every 10th frame." +msgstr "" +"Extraer cada 'enésimo' fotograma. Esta opción omitirá los fotogramas al " +"extraer las caras. Por ejemplo, un valor de 1 extraerá las caras de cada " +"fotograma, un valor de 10 extraerá las caras de cada 10 fotogramas." + +#: lib/cli/args_extract_convert.py:346 +msgid "" +"Automatically save the alignments file after a set amount of frames. By " +"default the alignments file is only saved at the end of the extraction " +"process. NB: If extracting in 2 passes then the alignments file will only " +"start to be saved out during the second pass. WARNING: Don't interrupt the " +"script when writing the file because it might get corrupted. Set to 0 to " +"turn off" +msgstr "" +"Guardar automáticamente el archivo de alineaciones después de una cantidad " +"determinada de cuadros. Por defecto, el archivo de alineaciones sólo se " +"guarda al final del proceso de extracción. Nota: Si se extrae en 2 pases, el " +"archivo de alineaciones sólo se empezará a guardar durante el segundo pase. " +"ADVERTENCIA: No interrumpa el script al escribir el archivo porque podría " +"corromperse. Poner a 0 para desactivar" + +#: lib/cli/args_extract_convert.py:357 +msgid "Draw landmarks on the ouput faces for debugging purposes." +msgstr "" +"Dibujar puntos de referencia en las caras de salida para fines de depuración." + +#: lib/cli/args_extract_convert.py:363 lib/cli/args_extract_convert.py:373 +#: lib/cli/args_extract_convert.py:381 lib/cli/args_extract_convert.py:388 +#: lib/cli/args_extract_convert.py:678 lib/cli/args_extract_convert.py:691 +#: lib/cli/args_extract_convert.py:712 lib/cli/args_extract_convert.py:718 +msgid "settings" +msgstr "ajustes" + +#: lib/cli/args_extract_convert.py:365 +msgid "" +"Don't run extraction in parallel. Will run each part of the extraction " +"process separately (one after the other) rather than all at the same time. " +"Useful if VRAM is at a premium." +msgstr "" +"No ejecute la extracción en paralelo. Ejecutará cada parte del proceso de " +"extracción por separado (una tras otra) en lugar de hacerlo todo al mismo " +"tiempo. Útil si la VRAM es escasa." + +#: lib/cli/args_extract_convert.py:375 +msgid "" +"Skips frames that have already been extracted and exist in the alignments " +"file" +msgstr "" +"Omite los fotogramas que ya han sido extraídos y que existen en el archivo " +"de alineaciones" + +#: lib/cli/args_extract_convert.py:382 +msgid "Skip frames that already have detected faces in the alignments file" +msgstr "" +"Omitir los fotogramas que ya tienen caras detectadas en el archivo de " +"alineaciones" + +#: lib/cli/args_extract_convert.py:389 +msgid "Skip saving the detected faces to disk. Just create an alignments file" +msgstr "" +"No guardar las caras detectadas en el disco. Crear sólo un archivo de " +"alineaciones" + +#: lib/cli/args_extract_convert.py:463 +msgid "" +"Swap the original faces in a source video/images to your final faces.\n" +"Conversion plugins can be configured in the 'Settings' Menu" +msgstr "" +"Cambia las caras originales de un vídeo/imágenes de origen por las caras " +"finales.\n" +"Los plugins de conversión pueden ser configurados en el menú " +"\"Configuración\"" + +#: lib/cli/args_extract_convert.py:485 +msgid "" +"Only required if converting from images to video. Provide The original video " +"that the source frames were extracted from (for extracting the fps and " +"audio)." +msgstr "" +"Sólo es necesario si se convierte de imágenes a vídeo. Proporcione el vídeo " +"original del que se extrajeron los fotogramas de origen (para extraer los " +"fps y el audio)." + +#: lib/cli/args_extract_convert.py:494 +msgid "" +"Model directory. The directory containing the trained model you wish to use " +"for conversion." +msgstr "" +"Directorio del modelo. El directorio que contiene el modelo entrenado que " +"desea utilizar para la conversión." + +#: lib/cli/args_extract_convert.py:505 +msgid "" +"R|Performs color adjustment to the swapped face. Some of these options have " +"configurable settings in '/config/convert.ini' or 'Settings > Configure " +"Convert Plugins':\n" +"L|avg-color: Adjust the mean of each color channel in the swapped " +"reconstruction to equal the mean of the masked area in the original image.\n" +"L|color-transfer: Transfers the color distribution from the source to the " +"target image using the mean and standard deviations of the L*a*b* color " +"space.\n" +"L|manual-balance: Manually adjust the balance of the image in a variety of " +"color spaces. Best used with the Preview tool to set correct values.\n" +"L|match-hist: Adjust the histogram of each color channel in the swapped " +"reconstruction to equal the histogram of the masked area in the original " +"image.\n" +"L|seamless-clone: Use cv2's seamless clone function to remove extreme " +"gradients at the mask seam by smoothing colors. Generally does not give very " +"satisfactory results.\n" +"L|none: Don't perform color adjustment." +msgstr "" +"R|Realiza un ajuste de color a la cara intercambiada. Algunas de estas " +"opciones tienen ajustes configurables en '/config/convert.ini' o 'Ajustes > " +"Configurar Extensiones de Conversión':\n" +"L|avg-color: Ajuste la media de cada canal de color en la reconstrucción " +"intercambiada para igualar la media del área enmascarada en la imagen " +"original.\n" +"L|color-transfer: Transfiere la distribución del color de la imagen de " +"origen a la de destino utilizando la media y las desviaciones estándar del " +"espacio de color L*a*b*.\n" +"L|manual-balance: Ajuste manualmente el equilibrio de la imagen en una " +"variedad de espacios de color. Se utiliza mejor con la herramienta de vista " +"previa para establecer los valores correctos.\n" +"L|match-hist: Ajuste el histograma de cada canal de color en la " +"reconstrucción intercambiada para igualar el histograma del área enmascarada " +"en la imagen original.\n" +"L|seamless-clone: Utilice la función de clonación sin costuras de cv2 para " +"eliminar los gradientes extremos en la costura de la máscara, suavizando los " +"colores. Generalmente no da resultados muy satisfactorios.\n" +"L|none: No realice el ajuste de color." + +#: lib/cli/args_extract_convert.py:531 +msgid "" +"R|Masker to use. NB: The mask you require must exist within the alignments " +"file. You can add additional masks with the Mask Tool.\n" +"L|none: Don't use a mask.\n" +"L|bisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"L|bisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"L|custom_face: Custom user created, face centered mask.\n" +"L|custom_head: Custom user created, head centered mask.\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"L|predicted: If the 'Learn Mask' option was enabled during training, this " +"will use the mask that was created by the trained model." +msgstr "" +"R|Máscara a utilizar. NB: La máscara que necesita debe existir en el archivo " +"de alineaciones. Puede añadir máscaras adicionales con la herramienta de " +"máscaras.\n" +"L|none: No utilizar una máscara.\n" +"L|bisenet-fp-face: Máscara relativamente ligera basada en NN que proporciona " +"un control más refinado sobre el área a enmascarar (configurable en la " +"configuración de la máscara). Utilice esta versión de bisenet-fp si su " +"modelo está entrenado con centrado 'face' o 'legacy'.\n" +"L|bisenet-fp-head: Máscara relativamente ligera basada en NN que proporciona " +"un control más refinado sobre el área a enmascarar (configurable en la " +"configuración de la máscara). Utilice esta versión de bisenet-fp si su " +"modelo está entrenado con centrado de 'cabeza'.\n" +"L|custom_face: Máscara personalizada creada por el usuario y centrada en el " +"rostro..\n" +"L|custom_head: Máscara personalizada centrada en la cabeza creada por el " +"usuario.\n" +"L|components: Máscara diseñada para proporcionar una segmentación facial " +"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " +"Se construye un casco convexo alrededor del exterior de los puntos de " +"referencia para crear una máscara.\n" +"L|extended: Máscara diseñada para proporcionar una segmentación facial " +"basada en el posicionamiento de las ubicaciones de los puntos de referencia. " +"Se construye un casco convexo alrededor del exterior de los puntos de " +"referencia y la máscara se extiende hacia arriba en la frente.\n" +"L|vgg-clear: Máscara diseñada para proporcionar una segmentación inteligente " +"de rostros principalmente frontales y libres de obstrucciones. Los rostros " +"de perfil y las obstrucciones pueden dar lugar a un rendimiento inferior.\n" +"L|vgg-obstructed: Máscara diseñada para proporcionar una segmentación " +"inteligente de rostros principalmente frontales. El modelo de la máscara ha " +"sido entrenado específicamente para reconocer algunas obstrucciones faciales " +"(manos y gafas). Los rostros de perfil pueden dar lugar a un rendimiento " +"inferior.\n" +"L|unet-dfl: Máscara diseñada para proporcionar una segmentación inteligente " +"de rostros principalmente frontales. El modelo de máscara ha sido entrenado " +"por los miembros de la comunidad y necesitará ser probado para una mayor " +"descripción. Los rostros de perfil pueden dar lugar a un rendimiento " +"inferior.\n" +"L|predicted: Si la opción 'Learn Mask' se habilitó durante el entrenamiento, " +"esto usará la máscara que fue creada por el modelo entrenado." + +#: lib/cli/args_extract_convert.py:570 +msgid "" +"R|The plugin to use to output the converted images. The writers are " +"configurable in '/config/convert.ini' or 'Settings > Configure Convert " +"Plugins:'\n" +"L|ffmpeg: [video] Writes out the convert straight to video. When the input " +"is a series of images then the '-ref' (--reference-video) parameter must be " +"set.\n" +"L|gif: [animated image] Create an animated gif.\n" +"L|opencv: [images] The fastest image writer, but less options and formats " +"than other plugins.\n" +"L|patch: [images] Outputs the raw swapped face patch, along with the " +"transformation matrix required to re-insert the face back into the original " +"frame. Use this option if you wish to post-process and composite the final " +"face within external tools.\n" +"L|pillow: [images] Slower than opencv, but has more options and supports " +"more formats." +msgstr "" +"R|El plugin a utilizar para dar salida a las imágenes convertidas. Los " +"escritores son configurables en '/config/convert.ini' o 'Ajustes > " +"Configurar Extensiones de Conversión:'\n" +"L|ffmpeg: [video] Escribe la conversión directamente en vídeo. Cuando la " +"entrada es una serie de imágenes, el parámetro '-ref' (--reference-video) " +"debe ser establecido.\n" +"L|gif: [imagen animada] Crea un gif animado.\n" +"L|opencv: [images] El escritor de imágenes más rápido, pero con menos " +"opciones y formatos que otros plugins.\n" +"L|patch: [images] Genera el parche de cara intercambiado sin formato, junto " +"con la matriz de transformación necesaria para volver a insertar la cara en " +"el marco original.\n" +"L|pillow: [images] Más lento que opencv, pero tiene más opciones y soporta " +"más formatos." + +#: lib/cli/args_extract_convert.py:591 lib/cli/args_extract_convert.py:600 +#: lib/cli/args_extract_convert.py:703 +msgid "Frame Processing" +msgstr "Proceso de fotogramas" + +#: lib/cli/args_extract_convert.py:593 +#, python-format +msgid "" +"Scale the final output frames by this amount. 100%% will output the frames " +"at source dimensions. 50%% at half size 200%% at double size" +msgstr "" +"Escala los fotogramas finales de salida en esta cantidad. 100%% dará salida " +"a los fotogramas a las dimensiones de origen. 50%% a la mitad de tamaño. " +"200%% al doble de tamaño" + +#: lib/cli/args_extract_convert.py:602 +msgid "" +"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " +"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " +"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " +"converting from images, then the filenames must end with the frame-number!" +msgstr "" +"Rangos de fotogramas a los que aplicar la transferencia, por ejemplo, para " +"los fotogramas de 10 a 50 y de 90 a 100 utilice --frame-ranges 10-50 90-100. " +"Los fotogramas que queden fuera del rango seleccionado se descartarán a " +"menos que se seleccione '-k' (--keep-unchanged). Nota: Si está convirtiendo " +"imágenes, ¡los nombres de los archivos deben terminar con el número de " +"fotograma!" + +#: lib/cli/args_extract_convert.py:616 +msgid "" +"Scale the swapped face by this percentage. Positive values will enlarge the " +"face, Negative values will shrink the face." +msgstr "" +"Escale la cara intercambiada según este porcentaje. Los valores positivos " +"agrandarán la cara, los valores negativos la reducirán." + +#: lib/cli/args_extract_convert.py:625 +msgid "" +"If you have not cleansed your alignments file, then you can filter out faces " +"by defining a folder here that contains the faces extracted from your input " +"files/video. If this folder is defined, then only faces that exist within " +"your alignments file and also exist within the specified folder will be " +"converted. Leaving this blank will convert all faces that exist within the " +"alignments file." +msgstr "" +"Si no ha limpiado su archivo de alineaciones, puede filtrar las caras " +"definiendo aquí una carpeta que contenga las caras extraídas de sus archivos/" +"vídeos de entrada. Si se define esta carpeta, sólo se convertirán las caras " +"que existan en el archivo de alineaciones y también en la carpeta " +"especificada. Si se deja en blanco, se convertirán todas las caras que " +"existan en el archivo de alineaciones." + +#: lib/cli/args_extract_convert.py:640 +msgid "" +"Optionally filter out people who you do not wish to process by passing in an " +"image of that person. Should be a front portrait with a single person in the " +"image. Multiple images can be added space separated. NB: Using face filter " +"will significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"Opcionalmente, puede filtrar las personas que no desea procesar pasando una " +"imagen de esa persona. Debe ser un retrato frontal con una sola persona en " +"la imagen. Se pueden añadir varias imágenes separadas por espacios. NB: El " +"uso del filtro de caras disminuirá significativamente la velocidad de " +"extracción y no se puede garantizar su precisión." + +#: lib/cli/args_extract_convert.py:653 +msgid "" +"Optionally select people you wish to process by passing in an image of that " +"person. Should be a front portrait with a single person in the image. " +"Multiple images can be added space separated. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"Opcionalmente, seleccione las personas que desea procesar pasando una imagen " +"de esa persona. Debe ser un retrato frontal con una sola persona en la " +"imagen. Se pueden añadir varias imágenes separadas por espacios. NB: El uso " +"del filtro facial disminuirá significativamente la velocidad de extracción y " +"no se puede garantizar su precisión." + +#: lib/cli/args_extract_convert.py:667 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Lower values are stricter. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"Para usar con los archivos opcionales nfilter/filter. Umbral para el " +"reconocimiento positivo de caras. Los valores más bajos son más estrictos. " +"NB: El uso del filtro facial disminuirá significativamente la velocidad de " +"extracción y no se puede garantizar su precisión." + +#: lib/cli/args_extract_convert.py:680 +msgid "" +"The maximum number of parallel processes for performing conversion. " +"Converting images is system RAM heavy so it is possible to run out of memory " +"if you have a lot of processes and not enough RAM to accommodate them all. " +"Setting this to 0 will use the maximum available. No matter what you set " +"this to, it will never attempt to use more processes than are available on " +"your system. If singleprocess is enabled this setting will be ignored." +msgstr "" +"El número máximo de procesos paralelos para realizar la conversión. La " +"conversión de imágenes requiere mucha RAM del sistema, por lo que es posible " +"que se agote la memoria si tiene muchos procesos y no hay suficiente RAM " +"para acomodarlos a todos. Si se ajusta a 0, se utilizará el máximo " +"disponible. No importa lo que establezca, nunca intentará utilizar más " +"procesos que los disponibles en su sistema. Si 'singleprocess' está " +"habilitado, este ajuste será ignorado." + +#: lib/cli/args_extract_convert.py:693 +msgid "" +"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " +"alignments file for your destination video. However, if you wish you can " +"generate the alignments on-the-fly by enabling this option. This will use an " +"inferior extraction pipeline and will lead to substandard results. If an " +"alignments file is found, this option will be ignored." +msgstr "" +"Activar la conversión sobre la marcha. NO se recomienda. Debe generar un " +"archivo de alineación limpio para su vídeo de destino. Sin embargo, si lo " +"desea, puede generar las alineaciones sobre la marcha activando esta opción. " +"Esto utilizará una tubería de extracción inferior y conducirá a resultados " +"de baja calidad. Si se encuentra un archivo de alineaciones, esta opción " +"será ignorada." + +#: lib/cli/args_extract_convert.py:705 +msgid "" +"When used with --frame-ranges outputs the unchanged frames that are not " +"processed instead of discarding them." +msgstr "" +"Cuando se usa con --frame-ranges, la salida incluye los fotogramas no " +"procesados en vez de descartarlos." + +#: lib/cli/args_extract_convert.py:713 +msgid "Swap the model. Instead converting from of A -> B, converts B -> A" +msgstr "" +"Intercambiar el modelo. En vez de convertir de A a B, convierte de B a A" + +#: lib/cli/args_extract_convert.py:719 +msgid "Disable multiprocessing. Slower but less resource intensive." +msgstr "Desactiva el multiproceso. Es más lento, pero usa menos recursos." + +#~ msgid "" +#~ "[LEGACY] This only needs to be selected if a legacy model is being loaded " +#~ "or if there are multiple models in the model folder" +#~ msgstr "" +#~ "[LEGACY] Sólo es necesario seleccionar esta opción si se está cargando un " +#~ "modelo heredado si hay varios modelos en la carpeta de modelos" diff --git a/locales/es/LC_MESSAGES/lib.cli.args_train.mo b/locales/es/LC_MESSAGES/lib.cli.args_train.mo new file mode 100644 index 0000000000..5cb1754eb0 Binary files /dev/null and b/locales/es/LC_MESSAGES/lib.cli.args_train.mo differ diff --git a/locales/es/LC_MESSAGES/lib.cli.args_train.po b/locales/es/LC_MESSAGES/lib.cli.args_train.po new file mode 100755 index 0000000000..8d1d1d8270 --- /dev/null +++ b/locales/es/LC_MESSAGES/lib.cli.args_train.po @@ -0,0 +1,380 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: faceswap.spanish\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:04+0000\n" +"PO-Revision-Date: 2024-03-28 18:09+0000\n" +"Last-Translator: \n" +"Language-Team: tokafondo\n" +"Language: es\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_train.py:30 +msgid "" +"Train a model on extracted original (A) and swap (B) faces.\n" +"Training models can take a long time. Anything from 24hrs to over a week\n" +"Model plugins can be configured in the 'Settings' Menu" +msgstr "" +"Entrene un modelo con las caras originales (A) e intercambiadas (B) " +"extraídas.\n" +"El entrenamiento de los modelos puede llevar mucho tiempo. Desde 24 horas " +"hasta más de una semana.\n" +"Los plugins de los modelos pueden configurarse en el menú \"Ajustes\"" + +#: lib/cli/args_train.py:49 lib/cli/args_train.py:58 +msgid "faces" +msgstr "caras" + +#: lib/cli/args_train.py:51 +msgid "" +"Input directory. A directory containing training images for face A. This is " +"the original face, i.e. the face that you want to remove and replace with " +"face B." +msgstr "" +"Directorio de entrada. Un directorio que contiene imágenes de entrenamiento " +"para la cara A. Esta es la cara original, es decir, la cara que se quiere " +"eliminar y sustituir por la cara B." + +#: lib/cli/args_train.py:60 +msgid "" +"Input directory. A directory containing training images for face B. This is " +"the swap face, i.e. the face that you want to place onto the head of person " +"A." +msgstr "" +"Directorio de entrada. Un directorio que contiene imágenes de entrenamiento " +"para la cara B. Esta es la cara de intercambio, es decir, la cara que se " +"quiere colocar en la cabeza de la persona A." + +#: lib/cli/args_train.py:67 lib/cli/args_train.py:80 lib/cli/args_train.py:97 +#: lib/cli/args_train.py:123 lib/cli/args_train.py:133 +msgid "model" +msgstr "modelo" + +#: lib/cli/args_train.py:69 +msgid "" +"Model directory. This is where the training data will be stored. You should " +"always specify a new folder for new models. If starting a new model, select " +"either an empty folder, or a folder which does not exist (which will be " +"created). If continuing to train an existing model, specify the location of " +"the existing model." +msgstr "" +"Directorio del modelo. Aquí es donde se almacenarán los datos de " +"entrenamiento. Siempre debe especificar una nueva carpeta para los nuevos " +"modelos. Si se inicia un nuevo modelo, seleccione una carpeta vacía o una " +"carpeta que no exista (que se creará). Si continúa entrenando un modelo " +"existente, especifique la ubicación del modelo existente." + +#: lib/cli/args_train.py:82 +msgid "" +"R|Load the weights from a pre-existing model into a newly created model. For " +"most models this will load weights from the Encoder of the given model into " +"the encoder of the newly created model. Some plugins may have specific " +"configuration options allowing you to load weights from other layers. " +"Weights will only be loaded when creating a new model. This option will be " +"ignored if you are resuming an existing model. Generally you will also want " +"to 'freeze-weights' whilst the rest of your model catches up with your " +"Encoder.\n" +"NB: Weights can only be loaded from models of the same plugin as you intend " +"to train." +msgstr "" +"R|Cargue los pesos de un modelo preexistente en un modelo recién creado. " +"Para la mayoría de los modelos, esto cargará pesos del codificador del " +"modelo dado en el codificador del modelo recién creado. Algunos complementos " +"pueden tener opciones de configuración específicas que le permiten cargar " +"pesos de otras capas. Los pesos solo se cargarán al crear un nuevo modelo. " +"Esta opción se ignorará si está reanudando un modelo existente. En general, " +"también querrá 'congelar pesos' mientras el resto de su modelo se pone al " +"día con su codificador.\n" +"NB: Los pesos solo se pueden cargar desde modelos del mismo complemento que " +"desea entrenar." + +#: lib/cli/args_train.py:99 +msgid "" +"R|Select which trainer to use. Trainers can be configured from the Settings " +"menu or the config folder.\n" +"L|original: The original model created by /u/deepfakes.\n" +"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " +"for full dfaker method.\n" +"L|dfl-h128: 128px in/out model from deepfacelab\n" +"L|dfl-sae: Adaptable model from deepfacelab\n" +"L|dlight: A lightweight, high resolution DFaker variant.\n" +"L|iae: A model that uses intermediate layers to try to get better details\n" +"L|lightweight: A lightweight model for low-end cards. Don't expect great " +"results. Can train as low as 1.6GB with batch size 8.\n" +"L|realface: A high detail, dual density model based on DFaker, with " +"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " +"won't work so well. By andenixa et al. Very configurable.\n" +"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " +"unbalanced so B>A swaps won't work so well. Very configurable.\n" +"L|villain: 128px in/out model from villainguy. Very resource hungry (You " +"will require a GPU with a fair amount of VRAM). Good for details, but more " +"susceptible to color differences." +msgstr "" +"R|Seleccione el entrenador que desea utilizar. Los entrenadores se pueden " +"configurar desde el menú de configuración o la carpeta de configuración.\n" +"L|original: El modelo original creado por /u/deepfakes.\n" +"L|dfaker: Modelo de 64px in/128px out de dfaker. Habilitar 'warp-to-" +"landmarks' para el método completo de dfaker.\n" +"L|dfl-h128: modelo de 128px in/out de deepfacelab\n" +"L|dfl-sae: Modelo adaptable de deepfacelab\n" +"L|dlight: Una variante de DFaker ligera y de alta resolución.\n" +"L|iae: Un modelo que utiliza capas intermedias para tratar de obtener " +"mejores detalles.\n" +"L|lightweight: Un modelo ligero para tarjetas de gama baja. No esperes " +"grandes resultados. Puede entrenar hasta 1,6GB con tamaño de lote 8.\n" +"L|realface: Un modelo de alto detalle y doble densidad basado en DFaker, con " +"resolución de entrada y salida personalizable. Los autocodificadores están " +"desequilibrados, por lo que los intercambios B>A no funcionan tan bien. Por " +"andenixa et al. Muy configurable\n" +"L|Unbalanced: modelo de 128px de entrada/salida de andenixa. Los " +"autocodificadores están desequilibrados por lo que los intercambios B>A no " +"funcionarán tan bien. Muy configurable\n" +"L|villain: Modelo de 128px de entrada/salida de villainguy. Requiere muchos " +"recursos (se necesita una GPU con una buena cantidad de VRAM). Bueno para " +"los detalles, pero más susceptible a las diferencias de color." + +#: lib/cli/args_train.py:125 +msgid "" +"Output a summary of the model and exit. If a model folder is provided then a " +"summary of the saved model is displayed. Otherwise a summary of the model " +"that would be created by the chosen plugin and configuration settings is " +"displayed." +msgstr "" +"Genere un resumen del modelo y salga. Si se proporciona una carpeta de " +"modelo, se muestra un resumen del modelo guardado. De lo contrario, se " +"muestra un resumen del modelo que crearía el complemento elegido y los " +"ajustes de configuración." + +#: lib/cli/args_train.py:135 +msgid "" +"Freeze the weights of the model. Freezing weights means that some of the " +"parameters in the model will no longer continue to learn, but those that are " +"not frozen will continue to learn. For most models, this will freeze the " +"encoder, but some models may have configuration options for freezing other " +"layers." +msgstr "" +"Congele los pesos del modelo. Congelar pesos significa que algunos de los " +"parámetros del modelo ya no seguirán aprendiendo, pero los que no están " +"congelados seguirán aprendiendo. Para la mayoría de los modelos, esto " +"congelará el codificador, pero algunos modelos pueden tener opciones de " +"configuración para congelar otras capas." + +#: lib/cli/args_train.py:147 lib/cli/args_train.py:160 +#: lib/cli/args_train.py:175 lib/cli/args_train.py:191 +#: lib/cli/args_train.py:200 +msgid "training" +msgstr "entrenamiento" + +#: lib/cli/args_train.py:149 +msgid "" +"Batch size. This is the number of images processed through the model for " +"each side per iteration. NB: As the model is fed 2 sides at a time, the " +"actual number of images within the model at any one time is double the " +"number that you set here. Larger batches require more GPU RAM." +msgstr "" +"Tamaño del lote. Este es el número de imágenes procesadas a través del " +"modelo para cada lado por iteración. Nota: Como el modelo se alimenta de 2 " +"lados a la vez, el número real de imágenes dentro del modelo en cualquier " +"momento es el doble del número que se establece aquí. Los lotes más grandes " +"requieren más RAM de la GPU." + +#: lib/cli/args_train.py:162 +msgid "" +"Length of training in iterations. This is only really used for automation. " +"There is no 'correct' number of iterations a model should be trained for. " +"You should stop training when you are happy with the previews. However, if " +"you want the model to stop automatically at a set number of iterations, you " +"can set that value here." +msgstr "" +"Duración del entrenamiento en iteraciones. Esto sólo se utiliza realmente " +"para la automatización. No hay un número 'correcto' de iteraciones para las " +"que deba entrenarse un modelo. Debe dejar de entrenar cuando esté satisfecho " +"con las previsiones. Sin embargo, si desea que el modelo se detenga " +"automáticamente en un número determinado de iteraciones, puede establecer " +"ese valor aquí." + +#: lib/cli/args_train.py:177 +msgid "" +"R|Select the distribution stategy to use.\n" +"L|default: Use Tensorflow's default distribution strategy.\n" +"L|central-storage: Centralizes variables on the CPU whilst operations are " +"performed on 1 or more local GPUs. This can help save some VRAM at the cost " +"of some speed by not storing variables on the GPU. Note: Mixed-Precision is " +"not supported on multi-GPU setups.\n" +"L|mirrored: Supports synchronous distributed training across multiple local " +"GPUs. A copy of the model and all variables are loaded onto each GPU with " +"batches distributed to each GPU at each iteration." +msgstr "" +"562 / 5,000\n" +"Translation results\n" +"R|Seleccione la estrategia de distribución a utilizar.\n" +"L|default: utiliza la estrategia de distribución predeterminada de " +"Tensorflow.\n" +"L|central-storage: centraliza las variables en la CPU mientras que las " +"operaciones se realizan en 1 o más GPU locales. Esto puede ayudar a ahorrar " +"algo de VRAM a costa de cierta velocidad al no almacenar variables en la " +"GPU. Nota: Mixed-Precision no es compatible con configuraciones de múltiples " +"GPU.\n" +"L|mirrored: Admite el entrenamiento distribuido síncrono en varias GPU " +"locales. Se carga una copia del modelo y todas las variables en cada GPU con " +"lotes distribuidos a cada GPU en cada iteración." + +#: lib/cli/args_train.py:193 +msgid "" +"Disables TensorBoard logging. NB: Disabling logs means that you will not be " +"able to use the graph or analysis for this session in the GUI." +msgstr "" +"Desactiva el registro de TensorBoard. NB: Desactivar los registros significa " +"que no podrá utilizar el gráfico o el análisis de esta sesión en la GUI." + +#: lib/cli/args_train.py:202 +msgid "" +"Use the Learning Rate Finder to discover the optimal learning rate for " +"training. For new models, this will calculate the optimal learning rate for " +"the model. For existing models this will use the optimal learning rate that " +"was discovered when initializing the model. Setting this option will ignore " +"the manually configured learning rate (configurable in train settings)." +msgstr "" +"Utilice el Buscador de tasa de aprendizaje para descubrir la tasa de " +"aprendizaje óptima para la capacitación. Para modelos nuevos, esto calculará " +"la tasa de aprendizaje óptima para el modelo. Para los modelos existentes, " +"esto utilizará la tasa de aprendizaje óptima que se descubrió al inicializar " +"el modelo. Configurar esta opción ignorará la tasa de aprendizaje " +"configurada manualmente (configurable en la configuración del tren)." + +#: lib/cli/args_train.py:215 lib/cli/args_train.py:225 +msgid "Saving" +msgstr "Guardar" + +#: lib/cli/args_train.py:216 +msgid "Sets the number of iterations between each model save." +msgstr "Establece el número de iteraciones entre cada guardado del modelo." + +#: lib/cli/args_train.py:227 +msgid "" +"Sets the number of iterations before saving a backup snapshot of the model " +"in it's current state. Set to 0 for off." +msgstr "" +"Establece el número de iteraciones antes de guardar una copia de seguridad " +"del modelo en su estado actual. Establece 0 para que esté desactivado." + +#: lib/cli/args_train.py:234 lib/cli/args_train.py:246 +#: lib/cli/args_train.py:258 +msgid "timelapse" +msgstr "intervalo" + +#: lib/cli/args_train.py:236 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'A' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-B parameter." +msgstr "" +"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " +"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " +"Esta debe ser la carpeta de entrada de las caras \"A\" que desea utilizar " +"para crear el timelapse. También debe suministrar un parámetro --timelapse-" +"output y un parámetro --timelapse-input-B." + +#: lib/cli/args_train.py:248 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'B' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-A parameter." +msgstr "" +"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " +"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " +"Esta debe ser la carpeta de entrada de las caras \"B\" que desea utilizar " +"para crear el timelapse. También debe suministrar un parámetro --timelapse-" +"output y un parámetro --timelapse-input-A." + +#: lib/cli/args_train.py:260 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. If " +"the input folders are supplied but no output folder, it will default to your " +"model folder/timelapse/" +msgstr "" +"Opcional para crear un timelapse. Timelapse guardará una imagen de las caras " +"seleccionadas en la carpeta timelapse-output en cada iteración de guardado. " +"Si se suministran las carpetas de entrada pero no la carpeta de salida, se " +"guardará por defecto en la carpeta del modelo/timelapse/" + +#: lib/cli/args_train.py:269 lib/cli/args_train.py:276 +msgid "preview" +msgstr "previsualización" + +#: lib/cli/args_train.py:270 +msgid "Show training preview output. in a separate window." +msgstr "" +"Mostrar la salida de la vista previa del entrenamiento. en una ventana " +"separada." + +#: lib/cli/args_train.py:278 +msgid "" +"Writes the training result to a file. The image will be stored in the root " +"of your FaceSwap folder." +msgstr "" +"Escribe el resultado del entrenamiento en un archivo. La imagen se " +"almacenará en la raíz de su carpeta FaceSwap." + +#: lib/cli/args_train.py:285 lib/cli/args_train.py:295 +#: lib/cli/args_train.py:305 lib/cli/args_train.py:315 +msgid "augmentation" +msgstr "aumento" + +#: lib/cli/args_train.py:287 +msgid "" +"Warps training faces to closely matched Landmarks from the opposite face-set " +"rather than randomly warping the face. This is the 'dfaker' way of doing " +"warping." +msgstr "" +"Deforma las caras de entrenamiento a puntos de referencia muy parecidos del " +"conjunto de caras opuestas en lugar de deformar la cara al azar. Esta es la " +"forma 'dfaker' de hacer la deformación." + +#: lib/cli/args_train.py:297 +msgid "" +"To effectively learn, a random set of images are flipped horizontally. " +"Sometimes it is desirable for this not to occur. Generally this should be " +"left off except for during 'fit training'." +msgstr "" +"Para aprender de forma efectiva, se voltea horizontalmente un conjunto " +"aleatorio de imágenes. A veces es deseable que esto no ocurra. Por lo " +"general, esto debería dejarse sin efecto, excepto durante el 'entrenamiento " +"de ajuste'." + +#: lib/cli/args_train.py:307 +msgid "" +"Color augmentation helps make the model less susceptible to color " +"differences between the A and B sets, at an increased training time cost. " +"Enable this option to disable color augmentation." +msgstr "" +"El aumento del color ayuda a que el modelo sea menos susceptible a las " +"diferencias de color entre los conjuntos A y B, con un mayor coste de tiempo " +"de entrenamiento. Activa esta opción para desactivar el aumento de color." + +#: lib/cli/args_train.py:317 +msgid "" +"Warping is integral to training the Neural Network. This option should only " +"be enabled towards the very end of training to try to bring out more detail. " +"Think of it as 'fine-tuning'. Enabling this option from the beginning is " +"likely to kill a model and lead to terrible results." +msgstr "" +"La deformación es fundamental para el entrenamiento de la red neuronal. Esta " +"opción sólo debería activarse hacia el final del entrenamiento para tratar " +"de obtener más detalles. Piense en ello como un 'ajuste fino'. Si se activa " +"esta opción desde el principio, es probable que arruine el modelo y se " +"obtengan resultados terribles." diff --git a/locales/es/LC_MESSAGES/tools.alignments.cli.mo b/locales/es/LC_MESSAGES/tools.alignments.cli.mo index 1783184695..9499ca1e8a 100644 Binary files a/locales/es/LC_MESSAGES/tools.alignments.cli.mo and b/locales/es/LC_MESSAGES/tools.alignments.cli.mo differ diff --git a/locales/es/LC_MESSAGES/tools.alignments.cli.po b/locales/es/LC_MESSAGES/tools.alignments.cli.po index 0287766d44..ceb263f11a 100644 --- a/locales/es/LC_MESSAGES/tools.alignments.cli.po +++ b/locales/es/LC_MESSAGES/tools.alignments.cli.po @@ -6,26 +6,26 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" "Report-Msgid-Bugs-To: \n" -"POT-Creation-Date: 2022-05-24 12:38+0100\n" -"PO-Revision-Date: 2022-05-24 12:41+0100\n" +"POT-Creation-Date: 2024-04-19 11:28+0100\n" +"PO-Revision-Date: 2024-04-19 11:29+0100\n" "Last-Translator: \n" "Language-Team: tokafondo\n" "Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 3.0\n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" -#: tools/alignments/cli.py:15 +#: tools/alignments/cli.py:16 msgid "" "This command lets you perform various tasks pertaining to an alignments file." msgstr "" "Este comando le permite realizar varias tareas relacionadas con un archivo " "de alineación." -#: tools/alignments/cli.py:30 +#: tools/alignments/cli.py:31 msgid "" "Alignments tool\n" "This tool allows you to perform numerous actions on or using an alignments " @@ -36,57 +36,63 @@ msgstr "" "caras o una fuente de fotogramas, usando opcionalmente su correspondiente " "archivo de alineación." -#: tools/alignments/cli.py:41 -msgid " Must Pass in a frames folder/source video file (-fr)." +#: tools/alignments/cli.py:43 +msgid " Must Pass in a frames folder/source video file (-r)." msgstr "" -" Debe indicar una carpeta de fotogramas o archivo de vídeo de origen (-fr)." +" Debe indicar una carpeta de fotogramas o archivo de vídeo de origen (-r)." -#: tools/alignments/cli.py:42 -msgid " Must Pass in a faces folder (-fc)." -msgstr " Debe indicar una carpeta de caras (-fc)." +#: tools/alignments/cli.py:44 +msgid " Must Pass in a faces folder (-c)." +msgstr " Debe indicar una carpeta de caras (-c)." -#: tools/alignments/cli.py:43 +#: tools/alignments/cli.py:45 msgid "" -" Must Pass in either a frames folder/source video file OR afaces folder (-fr " -"or -fc)." +" Must Pass in either a frames folder/source video file OR a faces folder (-r " +"or -c)." msgstr "" " Debe indicar una carpeta de fotogramas o archivo de vídeo de origen, o una " -"carpeta de caras (-fr o -fc)." +"carpeta de caras (-r o -c)." -#: tools/alignments/cli.py:45 +#: tools/alignments/cli.py:47 msgid "" -" Must Pass in a frames folder/source video file AND a faces folder (-fr and -" -"fc)." +" Must Pass in a frames folder/source video file AND a faces folder (-r and -" +"c)." msgstr "" " Debe indicar una carpeta de fotogramas o archivo de vídeo de origen, y una " -"carpeta de caras (-fr y -fc)." +"carpeta de caras (-r y -c)." -#: tools/alignments/cli.py:47 +#: tools/alignments/cli.py:49 msgid " Use the output option (-o) to process results." msgstr " Usar la opción de salida (-o) para procesar los resultados." -#: tools/alignments/cli.py:55 tools/alignments/cli.py:94 +#: tools/alignments/cli.py:58 tools/alignments/cli.py:104 msgid "processing" msgstr "proceso" -#: tools/alignments/cli.py:57 +#: tools/alignments/cli.py:61 #, python-brace-format msgid "" "R|Choose which action you want to perform. NB: All actions require an " "alignments file (-a) to be passed in.\n" "L|'draw': Draw landmarks on frames in the selected folder/video. A subfolder " "will be created within the frames folder to hold the output.{0}\n" +"L|'export': Export the contents of an alignments file to a json file. Can be " +"used for editing alignment information in external tools and then re-" +"importing by using Faceswap's Extract 'Import' plugins. Note: masks and " +"identity vectors will not be included in the exported file, so will be re-" +"generated when the json file is imported back into Faceswap. All data is " +"exported with the origin (0, 0) at the top left of the canvas.\n" "L|'extract': Re-extract faces from the source frames/video based on " "alignment data. This is a lot quicker than re-detecting faces. Can pass in " "the '-een' (--extract-every-n) parameter to only extract every nth frame." "{1}\n" "L|'from-faces': Generate alignment file(s) from a folder of extracted faces. " "if the folder of faces comes from multiple sources, then multiple alignments " -"files will be created. NB: for faces which have been extracted folders of " -"source images, rather than a video, a single alignments file will be created " -"as there is no way for the process to know how many folders of images were " -"originally used. You do not need to provide an alignments file path to run " -"this job. {3}\n" +"files will be created. NB: for faces which have been extracted from folders " +"of source images, rather than a video, a single alignments file will be " +"created as there is no way for the process to know how many folders of " +"images were originally used. You do not need to provide an alignments file " +"path to run this job. {3}\n" "L|'missing-alignments': Identify frames that do not exist in the alignments " "file.{2}{0}\n" "L|'missing-frames': Identify frames in the alignments file that do not " @@ -110,6 +116,14 @@ msgstr "" "L|'draw': Dibuja puntos de referencia en los fotogramas de la carpeta o " "vídeo seleccionado. Se creará una subcarpeta dentro de la carpeta de " "fotogramas para guardar el resultado.{0}\n" +"L|'export': Exportar el contenido de un archivo de alineaciones a un archivo " +"JSON. Se puede utilizar para editar información de alineación en " +"herramientas externas y luego volver a importar mediante el uso de " +"complementos de 'import' de extracto de Faceswap. Nota: Las máscaras y los " +"vectores de identidad no se incluirán en el archivo exportado, por lo que se " +"volverán a generar cuando el archivo JSON se importe a FacesWap. Todos los " +"datos se exportan con el origen (0, 0) en la parte superior izquierda del " +"lienzo.\n" "L|'extract': Reextrae las caras de los fotogramas o vídeos de origen " "basándose en los datos de alineación. Esto es mucho más rápido que volver a " "detectar las caras. Se puede pasar el parámetro '-een' (--extract-every-n) " @@ -142,7 +156,7 @@ msgstr "" "L|'spatial': Realiza un filtrado espacial y temporal para suavizar las " "alineaciones (¡EXPERIMENTAL!)" -#: tools/alignments/cli.py:96 +#: tools/alignments/cli.py:107 msgid "" "R|How to output discovered items ('faces' and 'frames' only):\n" "L|'console': Print the list of frames to the screen. (DEFAULT)\n" @@ -158,37 +172,80 @@ msgstr "" "L|'move': Mueve los elementos descubiertos a una subcarpeta dentro del " "directorio de origen." -#: tools/alignments/cli.py:107 tools/alignments/cli.py:118 -#: tools/alignments/cli.py:125 +#: tools/alignments/cli.py:118 tools/alignments/cli.py:141 +#: tools/alignments/cli.py:148 msgid "data" msgstr "datos" -#: tools/alignments/cli.py:111 +#: tools/alignments/cli.py:125 msgid "" -"Full path to the alignments file to be processed. This is required for all " -"jobs except for 'from-faces' when the alignments file will be generated in " -"the specified faces folder." +"Full path to the alignments file to be processed. If you have input a " +"'frames_dir' and don't provide this option, the process will try to find the " +"alignments file at the default location. All jobs require an alignments file " +"with the exception of 'from-faces' when the alignments file will be " +"generated in the specified faces folder." msgstr "" -"Ruta completa del archivo de alineaciones a procesar. Esto es necesario para " -"todos los trabajos excepto para 'caras desde' cuando el archivo de " -"alineaciones se generará en la carpeta de caras especificada." +"Ruta completa al archivo de alineaciones a procesar. Si ingresó un " +"'frames_dir' y no proporciona esta opción, el proceso intentará encontrar el " +"archivo de alineaciones en la ubicación predeterminada. Todos los trabajos " +"requieren un archivo de alineaciones con la excepción de 'from-faces' cuando " +"el archivo de alineaciones se generará en la carpeta de caras especificada." -#: tools/alignments/cli.py:119 -msgid "Directory containing extracted faces." -msgstr "Directorio que contiene las caras extraídas." - -#: tools/alignments/cli.py:126 +#: tools/alignments/cli.py:142 msgid "Directory containing source frames that faces were extracted from." msgstr "" "Directorio que contiene los fotogramas de origen de los que se extrajeron " "las caras." -#: tools/alignments/cli.py:135 tools/alignments/cli.py:146 -#: tools/alignments/cli.py:156 +#: tools/alignments/cli.py:150 +msgid "" +"R|Run the aligmnents tool on multiple sources. The following jobs support " +"batch mode:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"If batch mode is selected then the other options should be set as follows:\n" +"L|alignments_file: For 'sort' and 'spatial' this should point to the parent " +"folder containing the alignments files to be processed. For all other jobs " +"this option is ignored, and the alignments files must exist at their default " +"location relative to the original frames folder/video.\n" +"L|faces_dir: For 'from-faces' this should be a parent folder, containing sub-" +"folders of extracted faces from which to generate alignments files. For " +"'extract' this should be a parent folder where sub-folders will be created " +"for each extraction to be run. For all other jobs this option is ignored.\n" +"L|frames_dir: For 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"and 'no-faces' this should be a parent folder containing video files or sub-" +"folders of images to perform the alignments job on. The alignments file " +"should exist at the default location. For all other jobs this option is " +"ignored." +msgstr "" +"R|Ejecute la herramienta de alineación en varias fuentes. Los siguientes " +"trabajos admiten el modo por lotes:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"Si se selecciona el modo por lotes, las otras opciones deben configurarse de " +"la siguiente manera:\n" +"L|alignments_file: para 'sort' y 'spatial', debe apuntar a la carpeta " +"principal que contiene los archivos de alineación que se van a procesar. " +"Para todos los demás trabajos, esta opción se ignora y los archivos de " +"alineaciones deben existir en su ubicación predeterminada en relación con la " +"carpeta/video de fotogramas originales.\n" +"L|faces_dir: para 'from-faces', esta debe ser una carpeta principal que " +"contenga subcarpetas de caras extraídas desde las cuales generar archivos de " +"alineación. Para 'extraer', esta debe ser una carpeta principal donde se " +"crearán subcarpetas para cada extracción que se ejecute. Para todos los " +"demás trabajos, esta opción se ignora.\n" +"L|frames_dir: para 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"y 'no-faces', esta debe ser una carpeta principal que contenga archivos de " +"video o subcarpetas de imágenes para realizar el trabajo de alineaciones en. " +"El archivo de alineaciones debe existir en la ubicación predeterminada. Para " +"todos los demás trabajos, esta opción se ignora." + +#: tools/alignments/cli.py:176 tools/alignments/cli.py:188 +#: tools/alignments/cli.py:198 msgid "extract" msgstr "extracción" -#: tools/alignments/cli.py:136 +#: tools/alignments/cli.py:178 msgid "" "[Extract only] Extract every 'nth' frame. This option will skip frames when " "extracting faces. For example a value of 1 will extract faces from every " @@ -199,11 +256,11 @@ msgstr "" "caras de cada fotograma, un valor de 10 extraerá las caras de cada 10 " "fotogramas." -#: tools/alignments/cli.py:147 +#: tools/alignments/cli.py:189 msgid "[Extract only] The output size of extracted faces." msgstr "[Sólo extracción] El tamaño de salida de las caras extraídas." -#: tools/alignments/cli.py:157 +#: tools/alignments/cli.py:200 msgid "" "[Extract only] Only extract faces that have been resized by this percent or " "more to meet the specified extract size (`-sz`, `--size`). Useful for " @@ -223,6 +280,9 @@ msgstr "" "desde 512 px o más. Una configuración de 200 solo extraerá las caras que se " "han reducido de 1024 px o más." +#~ msgid "Directory containing extracted faces." +#~ msgstr "Directorio que contiene las caras extraídas." + #~ msgid "Full path to the alignments file to be processed." #~ msgstr "Ruta completa del archivo de alineaciones a procesar." diff --git a/locales/es/LC_MESSAGES/tools.effmpeg.cli.mo b/locales/es/LC_MESSAGES/tools.effmpeg.cli.mo index e4212f5452..0b973d69f6 100644 Binary files a/locales/es/LC_MESSAGES/tools.effmpeg.cli.mo and b/locales/es/LC_MESSAGES/tools.effmpeg.cli.mo differ diff --git a/locales/es/LC_MESSAGES/tools.effmpeg.cli.po b/locales/es/LC_MESSAGES/tools.effmpeg.cli.po index dab3f62f5f..ea47680568 100644 --- a/locales/es/LC_MESSAGES/tools.effmpeg.cli.po +++ b/locales/es/LC_MESSAGES/tools.effmpeg.cli.po @@ -5,27 +5,28 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-02-19 16:39+0000\n" -"PO-Revision-Date: 2021-02-21 16:49+0000\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:50+0000\n" +"PO-Revision-Date: 2024-03-29 00:02+0000\n" +"Last-Translator: \n" "Language-Team: tokafondo\n" +"Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.2\n" -"Last-Translator: \n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" -"Language: es_ES\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" #: tools/effmpeg/cli.py:15 msgid "This command allows you to easily execute common ffmpeg tasks." msgstr "Este comando le permite ejecutar fácilmente tareas comunes de ffmpeg." -#: tools/effmpeg/cli.py:24 +#: tools/effmpeg/cli.py:52 msgid "A wrapper for ffmpeg for performing image <> video converting." msgstr "Un interfaz de ffmpeg para realizar la conversión de imagen <> vídeo." -#: tools/effmpeg/cli.py:51 +#: tools/effmpeg/cli.py:64 msgid "" "R|Choose which action you want ffmpeg ffmpeg to do.\n" "L|'extract': turns videos into images \n" @@ -45,17 +46,17 @@ msgstr "" "L|'mux-audio' añade audio de un vídeo a otro.\n" "L|'rescale' cambia el tamaño del vídeo.\n" "L|'rotate' rotar video\n" -"L|'slice' corta una parte del video en un archivo de video separado. " +"L|'slice' corta una parte del video en un archivo de video separado." -#: tools/effmpeg/cli.py:65 +#: tools/effmpeg/cli.py:78 msgid "Input file." msgstr "Archivo de entrada." -#: tools/effmpeg/cli.py:66 tools/effmpeg/cli.py:73 tools/effmpeg/cli.py:87 +#: tools/effmpeg/cli.py:79 tools/effmpeg/cli.py:86 tools/effmpeg/cli.py:100 msgid "data" msgstr "datos" -#: tools/effmpeg/cli.py:76 +#: tools/effmpeg/cli.py:89 msgid "" "Output file. If no output is specified then: if the output is meant to be a " "video then a video called 'out.mkv' will be created in the input directory; " @@ -71,18 +72,18 @@ msgstr "" "Nota: la extensión del archivo de salida elegida determinará la codificación " "del archivo." -#: tools/effmpeg/cli.py:89 +#: tools/effmpeg/cli.py:102 msgid "Path to reference video if 'input' was not a video." msgstr "" "Ruta de acceso al vídeo de referencia si se dio una carpeta con fotogramas " "en vez de un vídeo." -#: tools/effmpeg/cli.py:95 tools/effmpeg/cli.py:105 tools/effmpeg/cli.py:142 -#: tools/effmpeg/cli.py:171 +#: tools/effmpeg/cli.py:108 tools/effmpeg/cli.py:118 tools/effmpeg/cli.py:156 +#: tools/effmpeg/cli.py:185 msgid "output" msgstr "salida" -#: tools/effmpeg/cli.py:97 +#: tools/effmpeg/cli.py:110 msgid "" "Provide video fps. Can be an integer, float or fraction. Negative values " "will will make the program try to get the fps from the input or reference " @@ -92,7 +93,7 @@ msgstr "" "fracción. Los valores negativos harán que el programa intente obtener los " "fps de los vídeos de entrada o de referencia." -#: tools/effmpeg/cli.py:107 +#: tools/effmpeg/cli.py:120 msgid "" "Image format that extracted images should be saved as. '.bmp' will offer the " "fastest extraction speed, but will take the most storage space. '.png' will " @@ -103,11 +104,11 @@ msgstr "" "almacenamiento. '.png' será más lento pero ocupará menos espacio de " "almacenamiento." -#: tools/effmpeg/cli.py:114 tools/effmpeg/cli.py:123 tools/effmpeg/cli.py:132 +#: tools/effmpeg/cli.py:127 tools/effmpeg/cli.py:136 tools/effmpeg/cli.py:145 msgid "clip" msgstr "recorte" -#: tools/effmpeg/cli.py:116 +#: tools/effmpeg/cli.py:129 msgid "" "Enter the start time from which an action is to be applied. Default: " "00:00:00, in HH:MM:SS format. You can also enter the time with or without " @@ -117,7 +118,7 @@ msgstr "" "defecto: 00:00:00, en formato HH:MM:SS. También puede introducir la hora con " "o sin los dos puntos, por ejemplo, 00:0000 o 026010." -#: tools/effmpeg/cli.py:125 +#: tools/effmpeg/cli.py:138 msgid "" "Enter the end time to which an action is to be applied. If both an end time " "and duration are set, then the end time will be used and the duration will " @@ -127,7 +128,7 @@ msgstr "" "00:00:00, en formato HH:MM:SS. También puede introducir la hora con o sin " "los dos puntos, por ejemplo, 00:0000 o 026010." -#: tools/effmpeg/cli.py:134 +#: tools/effmpeg/cli.py:147 msgid "" "Enter the duration of the chosen action, for example if you enter 00:00:10 " "for slice, then the first 10 seconds after and including the start time will " @@ -138,7 +139,7 @@ msgstr "" "formato HH:MM:SS. También puede introducir la hora con o sin los dos puntos, " "por ejemplo, 00:0000 o 026010." -#: tools/effmpeg/cli.py:144 +#: tools/effmpeg/cli.py:158 msgid "" "Mux the audio from the reference video into the input video. This option is " "only used for the 'gen-vid' action. 'mux-audio' action has this turned on " @@ -148,11 +149,11 @@ msgstr "" "se utiliza para la acción 'gen-vid'. La acción 'mux-audio' la tiene activada " "implícitamente." -#: tools/effmpeg/cli.py:155 tools/effmpeg/cli.py:165 +#: tools/effmpeg/cli.py:169 tools/effmpeg/cli.py:179 msgid "rotate" msgstr "rotación" -#: tools/effmpeg/cli.py:157 +#: tools/effmpeg/cli.py:171 msgid "" "Transpose the video. If transpose is set, then degrees will be ignored. For " "cli you can enter either the number or the long command name, e.g. to use " @@ -163,23 +164,23 @@ msgstr "" "comando, por ejemplo, para usar (1, 90Clockwise) son válidas las opciones -" "tr 1 y -tr 90Clockwise" -#: tools/effmpeg/cli.py:166 +#: tools/effmpeg/cli.py:180 msgid "Rotate the video clockwise by the given number of degrees." msgstr "" "Gira el vídeo en el sentido de las agujas del reloj el número de grados " "indicado." -#: tools/effmpeg/cli.py:173 +#: tools/effmpeg/cli.py:187 msgid "Set the new resolution scale if the chosen action is 'rescale'." msgstr "" -"Establece la nueva escala de resolución si la acción elegida es \"reescalar" -"\"." +"Establece la nueva escala de resolución si la acción elegida es " +"\"reescalar\"." -#: tools/effmpeg/cli.py:178 tools/effmpeg/cli.py:186 +#: tools/effmpeg/cli.py:192 tools/effmpeg/cli.py:200 msgid "settings" msgstr "ajustes" -#: tools/effmpeg/cli.py:180 +#: tools/effmpeg/cli.py:194 msgid "" "Reduces output verbosity so that only serious errors are printed. If both " "quiet and verbose are set, verbose will override quiet." @@ -188,7 +189,7 @@ msgstr "" "errores graves. Si se establecen tanto 'quiet' como 'verbose', 'verbose' " "tendrá preferencia y anulará a 'quiet'." -#: tools/effmpeg/cli.py:188 +#: tools/effmpeg/cli.py:202 msgid "" "Increases output verbosity. If both quiet and verbose are set, verbose will " "override quiet." diff --git a/locales/es/LC_MESSAGES/tools.manual.mo b/locales/es/LC_MESSAGES/tools.manual.mo index a9c32b1e49..33cfdb8142 100644 Binary files a/locales/es/LC_MESSAGES/tools.manual.mo and b/locales/es/LC_MESSAGES/tools.manual.mo differ diff --git a/locales/es/LC_MESSAGES/tools.manual.po b/locales/es/LC_MESSAGES/tools.manual.po index 71b2b71d7a..0e03295c00 100644 --- a/locales/es/LC_MESSAGES/tools.manual.po +++ b/locales/es/LC_MESSAGES/tools.manual.po @@ -5,7 +5,8 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-06-08 19:24+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:55+0000\n" "PO-Revision-Date: \n" "Last-Translator: \n" "Language-Team: \n" @@ -13,9 +14,9 @@ msgstr "" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"X-Generator: Poedit 2.4.3\n" +"X-Generator: Poedit 3.4.2\n" -#: tools/manual\cli.py:13 +#: tools/manual/cli.py:13 msgid "" "This command lets you perform various actions on frames, faces and " "alignments files using visual tools." @@ -23,7 +24,7 @@ msgstr "" "Este comando le permite realizar varias acciones en los archivos de " "fotogramas, caras y alineaciones utilizando herramientas visuales." -#: tools/manual\cli.py:23 +#: tools/manual/cli.py:23 msgid "" "A tool to perform various actions on frames, faces and alignments files " "using visual tools" @@ -31,18 +32,18 @@ msgstr "" "Una herramienta que permite realizar diversas acciones en archivos de " "fotogramas, caras y alineaciones mediante herramientas visuales" -#: tools/manual\cli.py:35 tools/manual\cli.py:43 +#: tools/manual/cli.py:35 tools/manual/cli.py:44 msgid "data" msgstr "datos" -#: tools/manual\cli.py:37 +#: tools/manual/cli.py:38 msgid "" "Path to the alignments file for the input, if not at the default location" msgstr "" "Ruta del archivo de alineaciones para la entrada, si no está en la ubicación " "por defecto" -#: tools/manual\cli.py:44 +#: tools/manual/cli.py:46 msgid "" "Video file or directory containing source frames that faces were extracted " "from." @@ -50,11 +51,11 @@ msgstr "" "Archivo o directorio de vídeo que contiene los fotogramas de origen de los " "que se extrajeron las caras." -#: tools/manual\cli.py:51 tools/manual\cli.py:59 +#: tools/manual/cli.py:53 tools/manual/cli.py:62 msgid "options" msgstr "opciones" -#: tools/manual\cli.py:52 +#: tools/manual/cli.py:55 msgid "" "Force regeneration of the low resolution jpg thumbnails in the alignments " "file." @@ -62,7 +63,7 @@ msgstr "" "Forzar la regeneración de las miniaturas jpg de baja resolución en el " "archivo de alineaciones." -#: tools/manual\cli.py:60 +#: tools/manual/cli.py:64 msgid "" "The process attempts to speed up generation of thumbnails by extracting from " "the video in parallel threads. For some videos, this causes the caching " diff --git a/locales/es/LC_MESSAGES/tools.mask.cli.mo b/locales/es/LC_MESSAGES/tools.mask.cli.mo index efae7fa7c6..ad86f74687 100644 Binary files a/locales/es/LC_MESSAGES/tools.mask.cli.mo and b/locales/es/LC_MESSAGES/tools.mask.cli.mo differ diff --git a/locales/es/LC_MESSAGES/tools.mask.cli.po b/locales/es/LC_MESSAGES/tools.mask.cli.po index f09fe4cd72..1d70c85ba9 100644 --- a/locales/es/LC_MESSAGES/tools.mask.cli.po +++ b/locales/es/LC_MESSAGES/tools.mask.cli.po @@ -5,51 +5,62 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-05-17 18:17+0100\n" -"PO-Revision-Date: 2021-05-17 18:18+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-06-28 13:45+0100\n" +"PO-Revision-Date: 2024-06-28 13:47+0100\n" "Last-Translator: \n" "Language-Team: tokafondo\n" "Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.3\n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.4\n" #: tools/mask/cli.py:15 -msgid "This command lets you generate masks for existing alignments." +msgid "" +"This tool allows you to generate, import, export or preview masks for " +"existing alignments." msgstr "" -"Este comando permite generar máscaras para las alineaciones existentes." +"Esta herramienta le permite generar, importar, exportar o obtener una vista " +"previa de máscaras para alineaciones existentes.\n" +"Genere, importe, exporte o obtenga una vista previa de máscaras para " +"archivos de alineaciones existentes." -#: tools/mask/cli.py:24 +#: tools/mask/cli.py:25 msgid "" "Mask tool\n" -"Generate masks for existing alignments files." +"Generate, import, export or preview masks for existing alignments files." msgstr "" "Herramienta de máscara\n" -"Genera máscaras para los archivos de alineación existentes." +"Genere, importe, exporte o obtenga una vista previa de máscaras para " +"archivos de alineaciones existentes." -#: tools/mask/cli.py:32 tools/mask/cli.py:41 tools/mask/cli.py:51 +#: tools/mask/cli.py:35 tools/mask/cli.py:47 tools/mask/cli.py:58 +#: tools/mask/cli.py:69 msgid "data" msgstr "datos" -#: tools/mask/cli.py:35 +#: tools/mask/cli.py:39 msgid "" -"Full path to the alignments file to add the mask to. NB: if the mask already " -"exists in the alignments file it will be overwritten." +"Full path to the alignments file that contains the masks if not at the " +"default location. NB: If the input-type is faces and you wish to update the " +"corresponding alignments file, then you must provide a value here as the " +"location cannot be automatically detected." msgstr "" -"Ruta completa del archivo de alineaciones al que se añadirá la máscara. " -"Nota: si la máscara ya existe en el archivo de alineaciones, se " -"sobrescribirá." +"Ruta completa al archivo de alineaciones para agregar la máscara si no está " +"en la ubicación predeterminada. NB: si el tipo de entrada es caras y desea " +"actualizar el archivo de alineaciones correspondiente, debe proporcionar un " +"valor aquí ya que la ubicación no se puede detectar automáticamente." -#: tools/mask/cli.py:44 +#: tools/mask/cli.py:51 msgid "Directory containing extracted faces, source frames, or a video file." msgstr "" "Directorio que contiene las caras extraídas, los fotogramas de origen o un " "archivo de vídeo." -#: tools/mask/cli.py:53 +#: tools/mask/cli.py:61 msgid "" "R|Whether the `input` is a folder of faces or a folder frames/video\n" "L|faces: The input is a folder containing extracted faces.\n" @@ -59,11 +70,36 @@ msgstr "" "L|faces: La entrada es una carpeta que contiene caras extraídas.\n" "L|frames: La entrada es una carpeta que contiene fotogramas o es un vídeo" -#: tools/mask/cli.py:62 tools/mask/cli.py:90 +#: tools/mask/cli.py:71 +msgid "" +"R|Run the mask tool on multiple sources. If selected then the other options " +"should be set as follows:\n" +"L|input: A parent folder containing either all of the video files to be " +"processed, or containing sub-folders of frames/faces.\n" +"L|output-folder: If provided, then sub-folders will be created within the " +"given location to hold the previews for each input.\n" +"L|alignments: Alignments field will be ignored for batch processing. The " +"alignments files must exist at the default location (for frames). For batch " +"processing of masks with 'faces' as the input type, then only the PNG header " +"within the extracted faces will be updated." +msgstr "" +"R|Ejecute la herramienta de máscara en varias fuentes. Si se selecciona, las " +"otras opciones deben configurarse de la siguiente manera:\n" +"L|input: una carpeta principal que contiene todos los archivos de video que " +"se procesarán o que contiene subcarpetas de marcos/caras.\n" +"L|output-folder: si se proporciona, se crearán subcarpetas dentro de la " +"ubicación dada para contener las vistas previas de cada entrada.\n" +"L|alignments: el campo de alineaciones se ignorará para el procesamiento por " +"lotes. Los archivos de alineaciones deben existir en la ubicación " +"predeterminada (para marcos). Para el procesamiento por lotes de máscaras " +"con 'caras' como tipo de entrada, solo se actualizará el encabezado PNG " +"dentro de las caras extraídas." + +#: tools/mask/cli.py:87 tools/mask/cli.py:119 msgid "process" msgstr "proceso" -#: tools/mask/cli.py:63 +#: tools/mask/cli.py:89 msgid "" "R|Masker to use.\n" "L|bisenet-fp: Relatively lightweight NN based mask that provides more " @@ -72,6 +108,10 @@ msgid "" "L|components: Mask designed to provide facial segmentation based on the " "positioning of landmark locations. A convex hull is constructed around the " "exterior of the landmarks to create a mask.\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU.\n" "L|extended: Mask designed to provide facial segmentation based on the " "positioning of landmark locations. A convex hull is constructed around the " "exterior of the landmarks and the mask is extended upwards onto the " @@ -84,9 +124,8 @@ msgid "" "some facial obstructions (hands and eyeglasses). Profile faces may result in " "sub-par performance.\n" "L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " -"faces. The mask model has been trained by community members and will need " -"testing for further description. Profile faces may result in sub-par " -"performance." +"faces. The mask model has been trained by community members. Profile faces " +"may result in sub-par performance." msgstr "" "R|Máscara a utilizar.\n" "L|bisenet-fp: Máscara relativamente ligera basada en NN que proporciona un " @@ -96,6 +135,10 @@ msgstr "" "basada en la posición de los puntos de referencia. Se construye un casco " "convexo alrededor del exterior de los puntos de referencia para crear una " "máscara.\n" +"L|custom: Una máscara ficticia que llena el área de la máscara con 1 o 0 " +"(configurable en la configuración). Esto solo es necesario si tiene la " +"intención de editar manualmente las máscaras personalizadas usted mismo en " +"la herramienta manual. Esta máscara no utiliza la GPU.\n" "L|extended: Máscara diseñada para proporcionar una segmentación facial " "basada en el posicionamiento de las ubicaciones de los puntos de referencia. " "Se construye un casco convexo alrededor del exterior de los puntos de " @@ -114,32 +157,116 @@ msgstr "" "descripción. Los rostros de perfil pueden dar lugar a un rendimiento " "inferior." -#: tools/mask/cli.py:91 +#: tools/mask/cli.py:121 msgid "" -"R|Whether to update all masks in the alignments files, only those faces that " -"do not already have a mask of the given `mask type` or just to output the " -"masks to the `output` location.\n" -"L|all: Update the mask for all faces in the alignments file.\n" +"R|The Mask tool process to perform.\n" +"L|all: Update the mask for all faces in the alignments file for the selected " +"'masker'.\n" "L|missing: Create a mask for all faces in the alignments file where a mask " -"does not previously exist.\n" -"L|output: Don't update the masks, just output them for review in the given " -"output folder." -msgstr "" -"R|Si se actualizan todas las máscaras en los archivos de alineación, sólo " -"aquellas caras que no tienen ya una máscara del \"tipo de máscara\" dado o " -"sólo se envían las máscaras a la ubicación \"de salida\".\n" -"L|all: Actualiza la máscara de todas las caras del archivo de alineación.\n" -"L|missing: Crea una máscara para todas las caras del fichero de alineaciones " -"en las que no existe una máscara previamente.\n" -"L|output: No actualiza las máscaras, sólo las emite para su revisión en la " -"carpeta de salida dada." - -#: tools/mask/cli.py:104 tools/mask/cli.py:111 tools/mask/cli.py:124 -#: tools/mask/cli.py:137 tools/mask/cli.py:146 +"does not previously exist for the selected 'masker'.\n" +"L|output: Don't update the masks, just output the selected 'masker' for " +"review/editing in external tools to the given output folder.\n" +"L|import: Import masks that have been edited outside of faceswap into the " +"alignments file. Note: 'custom' must be the selected 'masker' and the masks " +"must be in the same format as the 'input-type' (frames or faces)" +msgstr "" +"R|Процесс инструмента «Маска», который необходимо выполнить.\n" +"L|all: обновить маску для всех лиц в файле выравниваний для выбранного " +"«masker».\n" +"L|missing: создать маску для всех граней в файле выравниваний, где маска " +"ранее не существовала для выбранного «masker».\n" +"L|output: не обновляйте маски, просто выведите выбранный «masker» для " +"просмотра/редактирования во внешних инструментах в данную выходную папку.\n" +"L|import: импортируйте маски, которые были отредактированы вне Facewap, в " +"файл выравниваний. Примечание. «custom» должен быть выбранным «masker», а " +"маски должны быть в том же формате, что и «input-type» (frames или faces)." + +#: tools/mask/cli.py:135 tools/mask/cli.py:154 tools/mask/cli.py:176 +msgid "import" +msgstr "importar" + +#: tools/mask/cli.py:137 +msgid "" +"R|Import only. The path to the folder that contains masks to be imported.\n" +"L|How the masks are provided is not important, but they will be stored, " +"internally, as 8-bit grayscale images.\n" +"L|If the input are images, then the masks must be named exactly the same as " +"input frames/faces (excluding the file extension).\n" +"L|If the input is a video file, then the filename of the masks is not " +"important but should contain the frame number at the end of the filename " +"(but before the file extension). The frame number can be separated from the " +"rest of the filename by any non-numeric character and can be padded by any " +"number of zeros. The frame number must correspond correctly to the frame " +"number in the original video (starting from frame 1)." +msgstr "" +"R|Sólo importar. La ruta a la carpeta que contiene las máscaras que se " +"importarán.\n" +"L|Cómo se proporcionan las máscaras no es importante, pero se almacenarán " +"internamente como imágenes en escala de grises de 8 bits.\n" +"L|Si la entrada son imágenes, entonces las máscaras deben tener el mismo " +"nombre que los cuadros/caras de entrada (excluyendo la extensión del " +"archivo).\n" +"L|Si la entrada es un archivo de vídeo, entonces el nombre del archivo de " +"las máscaras no es importante pero debe contener el número de fotograma al " +"final del nombre del archivo (pero antes de la extensión del archivo). El " +"número de fotograma se puede separar del resto del nombre del archivo " +"mediante cualquier carácter no numérico y se puede rellenar con cualquier " +"número de ceros. El número de fotograma debe corresponder correctamente al " +"número de fotograma del vídeo original (a partir del fotograma 1)." + +#: tools/mask/cli.py:156 +msgid "" +"R|Import/Output only. When importing masks, this is the centering to use. " +"For output this is only used for outputting custom imported masks, and " +"should correspond to the centering used when importing the mask. Note: For " +"any job other than 'import' and 'output' this option is ignored as mask " +"centering is handled internally.\n" +"L|face: Centers the mask on the center of the face, adjusting for pitch and " +"yaw. Outside of requirements for full head masking/training, this is likely " +"to be the best choice.\n" +"L|head: Centers the mask on the center of the head, adjusting for pitch and " +"yaw. Note: You should only select head centering if you intend to include " +"the full head (including hair) within the mask and are looking to train a " +"full head model.\n" +"L|legacy: The 'original' extraction technique. Centers the mask near the of " +"the nose with and crops closely to the face. Can result in the edges of the " +"mask appearing outside of the training area." +msgstr "" +"R|Solo importación/salida. Al importar máscaras, este es el centrado que se " +"debe utilizar. Para la salida, esto solo se utiliza para generar máscaras " +"importadas personalizadas y debe corresponder al centrado utilizado al " +"importar la máscara. Nota: Para cualquier trabajo que no sea \"importación\" " +"y \"salida\", esta opción se ignora ya que el centrado de la máscara se " +"maneja internamente.\n" +"L|cara: centra la máscara en el centro de la cara, ajustando el tono y la " +"orientación. Aparte de los requisitos para el entrenamiento/enmascaramiento " +"de cabeza completa, esta probablemente sea la mejor opción.\n" +"L|head: centra la máscara en el centro de la cabeza, ajustando el cabeceo y " +"la guiñada. Nota: Sólo debe seleccionar el centrado de la cabeza si desea " +"incluir la cabeza completa (incluido el cabello) dentro de la máscara y " +"desea entrenar un modelo de cabeza completa.\n" +"L|legacy: La técnica de extracción 'original'. Centra la máscara cerca de la " +"nariz y la recorta cerca de la cara. Puede provocar que los bordes de la " +"máscara aparezcan fuera del área de entrenamiento." + +#: tools/mask/cli.py:181 +msgid "" +"Import only. The size, in pixels to internally store the mask at.\n" +"The default is 128 which is fine for nearly all usecases. Larger sizes will " +"result in larger alignments files and longer processing." +msgstr "" +"Sólo importar. El tamaño, en píxeles, para almacenar internamente la " +"máscara.\n" +"El valor predeterminado es 128, que está bien para casi todos los casos de " +"uso. Los tamaños más grandes darán como resultado archivos de alineaciones " +"más grandes y un procesamiento más largo." + +#: tools/mask/cli.py:189 tools/mask/cli.py:197 tools/mask/cli.py:211 +#: tools/mask/cli.py:225 tools/mask/cli.py:235 msgid "output" msgstr "salida" -#: tools/mask/cli.py:105 +#: tools/mask/cli.py:191 msgid "" "Optional output location. If provided, a preview of the masks created will " "be output in the given folder." @@ -147,7 +274,7 @@ msgstr "" "Ubicación de salida opcional. Si se proporciona, se obtendrá una vista " "previa de las máscaras creadas en la carpeta indicada." -#: tools/mask/cli.py:115 +#: tools/mask/cli.py:202 msgid "" "Apply gaussian blur to the mask output. Has the effect of smoothing the " "edges of the mask giving less of a hard edge. the size is in pixels. This " @@ -160,7 +287,7 @@ msgstr "" "redondeará al siguiente número impar. NB: Sólo afecta a la vista previa de " "salida. Si se ajusta a 0, se desactiva" -#: tools/mask/cli.py:128 +#: tools/mask/cli.py:216 msgid "" "Helps reduce 'blotchiness' on some masks by making light shades white and " "dark shades black. Higher values will impact more of the mask. NB: Only " @@ -171,7 +298,7 @@ msgstr "" "más a la máscara. NB: Sólo afecta a la vista previa de salida. Si se ajusta " "a 0, se desactiva" -#: tools/mask/cli.py:138 +#: tools/mask/cli.py:227 msgid "" "R|How to format the output when processing is set to 'output'.\n" "L|combined: The image contains the face/frame, face mask and masked face.\n" @@ -186,7 +313,7 @@ msgstr "" "enmascarada.\n" "L|mask: Sólo emite la máscara como una imagen de un solo canal." -#: tools/mask/cli.py:147 +#: tools/mask/cli.py:237 msgid "" "R|Whether to output the whole frame or only the face box when using output " "processing. Only has an effect when using frames as input." @@ -194,3 +321,31 @@ msgstr "" "R|Marcar esta opción dará como salida el fotograma completo, en vez de sólo " "el cuadro de la cara cuando se utiliza el procesamiento de salida. Sólo " "tiene efecto cuando se utilizan cuadros como entrada." + +#~ msgid "" +#~ "R|Whether to update all masks in the alignments files, only those faces " +#~ "that do not already have a mask of the given `mask type` or just to " +#~ "output the masks to the `output` location.\n" +#~ "L|all: Update the mask for all faces in the alignments file.\n" +#~ "L|missing: Create a mask for all faces in the alignments file where a " +#~ "mask does not previously exist.\n" +#~ "L|output: Don't update the masks, just output them for review in the " +#~ "given output folder." +#~ msgstr "" +#~ "R|Si se actualizan todas las máscaras en los archivos de alineación, sólo " +#~ "aquellas caras que no tienen ya una máscara del \"tipo de máscara\" dado " +#~ "o sólo se envían las máscaras a la ubicación \"de salida\".\n" +#~ "L|all: Actualiza la máscara de todas las caras del archivo de " +#~ "alineación.\n" +#~ "L|missing: Crea una máscara para todas las caras del fichero de " +#~ "alineaciones en las que no existe una máscara previamente.\n" +#~ "L|output: No actualiza las máscaras, sólo las emite para su revisión en " +#~ "la carpeta de salida dada." + +#~ msgid "" +#~ "Full path to the alignments file to add the mask to. NB: if the mask " +#~ "already exists in the alignments file it will be overwritten." +#~ msgstr "" +#~ "Ruta completa del archivo de alineaciones al que se añadirá la máscara. " +#~ "Nota: si la máscara ya existe en el archivo de alineaciones, se " +#~ "sobrescribirá." diff --git a/locales/es/LC_MESSAGES/tools.mo b/locales/es/LC_MESSAGES/tools.mo deleted file mode 100644 index afc32e4c19..0000000000 Binary files a/locales/es/LC_MESSAGES/tools.mo and /dev/null differ diff --git a/locales/es/LC_MESSAGES/tools.model.cli.mo b/locales/es/LC_MESSAGES/tools.model.cli.mo new file mode 100644 index 0000000000..55dd5dba0e Binary files /dev/null and b/locales/es/LC_MESSAGES/tools.model.cli.mo differ diff --git a/locales/es/LC_MESSAGES/tools.model.cli.po b/locales/es/LC_MESSAGES/tools.model.cli.po new file mode 100644 index 0000000000..56079517ca --- /dev/null +++ b/locales/es/LC_MESSAGES/tools.model.cli.po @@ -0,0 +1,90 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:51+0000\n" +"PO-Revision-Date: 2024-03-29 00:00+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: es\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=2; plural=(n != 1);\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/model/cli.py:13 +msgid "This tool lets you perform actions on saved Faceswap models." +msgstr "" +"Esta herramienta le permite realizar acciones en modelos Faceswap guardados." + +#: tools/model/cli.py:22 +msgid "A tool for performing actions on Faceswap trained model files" +msgstr "" +"Una herramienta para realizar acciones en archivos de modelos entrenados " +"Faceswap" + +#: tools/model/cli.py:34 +msgid "" +"Model directory. A directory containing the model you wish to perform an " +"action on." +msgstr "" +"Directorio de modelo. Un directorio que contiene el modelo en el que desea " +"realizar una acción." + +#: tools/model/cli.py:43 +msgid "" +"R|Choose which action you want to perform.\n" +"L|'inference' - Create an inference only copy of the model. Strips any " +"layers from the model which are only required for training. NB: This is for " +"exporting the model for use in external applications. Inference generated " +"models cannot be used within Faceswap. See the 'format' option for " +"specifying the model output format.\n" +"L|'nan-scan' - Scan the model file for NaNs or Infs (invalid data).\n" +"L|'restore' - Restore a model from backup." +msgstr "" +"R|Elige qué acción quieres realizar.\n" +"L|'inference': crea una copia del modelo solo de inferencia. Elimina las " +"capas del modelo que solo se requieren para el entrenamiento. NB: Esto es " +"para exportar el modelo para su uso en aplicaciones externas. Los modelos " +"generados por inferencia no se pueden usar en Faceswap. Consulte la opción " +"'formato' para especificar el formato de salida del modelo.\n" +"L|'nan-scan': escanea el archivo del modelo en busca de NaN o Inf (datos no " +"válidos).\n" +"L|'restore': restaura un modelo desde una copia de seguridad." + +#: tools/model/cli.py:57 tools/model/cli.py:69 +msgid "inference" +msgstr "inferencia" + +#: tools/model/cli.py:59 +msgid "" +"R|The format to save the model as. Note: Only used for 'inference' job.\n" +"L|'h5' - Standard Keras H5 format. Does not store any custom layer " +"information. Layers will need to be loaded from Faceswap to use.\n" +"L|'saved-model' - Tensorflow's Saved Model format. Contains all information " +"required to load the model outside of Faceswap." +msgstr "" +"R|El formato para guardar el modelo. Nota: Solo se usa para el trabajo de " +"'inference'.\n" +"L|'h5' - Formato estándar de Keras H5. No almacena ninguna información de " +"capa personalizada. Las capas deberán cargarse desde Faceswap para usar.\n" +"L|'saved-model': formato de modelo guardado de Tensorflow. Contiene toda la " +"información necesaria para cargar el modelo fuera de Faceswap." + +#: tools/model/cli.py:71 +#, fuzzy +#| msgid "" +#| "Only used for 'inference' job. Generate the inference model for B -> A " +#| "instead of A -> B." +msgid "" +"Only used for 'inference' job. Generate the inference model for B -> A " +"instead of A -> B." +msgstr "" +"Solo se usa para el trabajo de 'inference'. Genere el modelo de inferencia " +"para B -> A en lugar de A -> B." diff --git a/locales/es/LC_MESSAGES/tools.po b/locales/es/LC_MESSAGES/tools.po deleted file mode 100644 index ac20cfe6e9..0000000000 --- a/locales/es/LC_MESSAGES/tools.po +++ /dev/null @@ -1,27 +0,0 @@ -# SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION -# FIRST AUTHOR , YEAR. -# -msgid "" -msgstr "" -"Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-02-18 23:49-0000\n" -"PO-Revision-Date: 2021-02-19 18:00+0000\n" -"Language-Team: tokafondo\n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=UTF-8\n" -"Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.3\n" -"Last-Translator: \n" -"Plural-Forms: nplurals=2; plural=(n != 1);\n" -"Language: es_ES\n" - -#: tools.py:46 -msgid "" -"Please backup your data and/or test the tool you want to use with a smaller " -"data set to make sure you understand how it works." -msgstr "" -"Por favor, haga una copia de seguridad de sus datos, y pruebe la " -"herramienta que quiere utilizar con un conjunto de datos más pequeño para " -"asegurarse de que entiende cómo funciona." diff --git a/locales/es/LC_MESSAGES/tools.preview.mo b/locales/es/LC_MESSAGES/tools.preview.mo index ccf4c668d1..955c957645 100644 Binary files a/locales/es/LC_MESSAGES/tools.preview.mo and b/locales/es/LC_MESSAGES/tools.preview.mo differ diff --git a/locales/es/LC_MESSAGES/tools.preview.po b/locales/es/LC_MESSAGES/tools.preview.po index 8b17fc3ab9..f9cfb9218a 100644 --- a/locales/es/LC_MESSAGES/tools.preview.po +++ b/locales/es/LC_MESSAGES/tools.preview.po @@ -5,25 +5,26 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-02-18 23:09-0000\n" -"PO-Revision-Date: 2021-03-19 14:28+0000\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:00+0000\n" +"Last-Translator: \n" "Language-Team: tokafondo\n" +"Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.3\n" -"Last-Translator: \n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" -"Language: es_ES\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" -#: tools/preview/cli.py:14 +#: tools/preview/cli.py:15 msgid "This command allows you to preview swaps to tweak convert settings." msgstr "" "Este comando permite previsualizar los intercambios para ajustar la " "configuración de la conversión." -#: tools/preview/cli.py:23 +#: tools/preview/cli.py:30 msgid "" "Preview tool\n" "Allows you to configure your convert settings with a live preview" @@ -31,11 +32,11 @@ msgstr "" "Herramienta de vista previa\n" "Permite configurar los ajustes de conversión con una vista previa en directo" -#: tools/preview/cli.py:33 tools/preview/cli.py:42 tools/preview/cli.py:49 +#: tools/preview/cli.py:47 tools/preview/cli.py:57 tools/preview/cli.py:65 msgid "data" msgstr "datos" -#: tools/preview/cli.py:35 +#: tools/preview/cli.py:50 msgid "" "Input directory or video. Either a directory containing the image files you " "wish to process or path to a video file." @@ -43,14 +44,14 @@ msgstr "" "Directorio o vídeo de entrada. Un directorio que contenga los archivos de " "imagen que desea procesar o la ruta a un archivo de vídeo." -#: tools/preview/cli.py:44 +#: tools/preview/cli.py:60 msgid "" "Path to the alignments file for the input, if not at the default location" msgstr "" -"Ruta del archivo de alineaciones para la entrada, si no está en la " -"ubicación por defecto" +"Ruta del archivo de alineaciones para la entrada, si no está en la ubicación " +"por defecto" -#: tools/preview/cli.py:51 +#: tools/preview/cli.py:68 msgid "" "Model directory. A directory containing the trained model you wish to " "process." @@ -58,31 +59,35 @@ msgstr "" "Directorio del modelo. Un directorio que contiene el modelo entrenado que " "desea procesar." -#: tools/preview/cli.py:58 +#: tools/preview/cli.py:74 msgid "Swap the model. Instead of A -> B, swap B -> A" -msgstr "" -"Intercambiar el modelo. En lugar de convertir A en B, convierte B en A" +msgstr "Intercambiar el modelo. En lugar de convertir A en B, convierte B en A" -#: tools/preview\preview.py:1303 +#: tools/preview/control_panels.py:510 msgid "Save full config" msgstr "Guardar la configuración completa" -#: tools/preview\preview.py:1306 +#: tools/preview/control_panels.py:513 msgid "Reset full config to default values" msgstr "Restablecer la configuración completa a los valores por defecto" -#: tools/preview\preview.py:1309 +#: tools/preview/control_panels.py:516 msgid "Reset full config to saved values" msgstr "Restablecer la configuración completa a los valores guardados" -#: tools/preview\preview.py:1453 -msgid "Save {} config" -msgstr "Guardar la configuración de {}" +#: tools/preview/control_panels.py:667 +#, python-brace-format +msgid "Save {title} config" +msgstr "Guardar la configuración de {title}" -#: tools/preview\preview.py:1456 -msgid "Reset {} config to default values" -msgstr "Restablecer la configuración completa de {} a los valores por defecto" +#: tools/preview/control_panels.py:670 +#, python-brace-format +msgid "Reset {title} config to default values" +msgstr "" +"Restablecer la configuración completa de {title} a los valores por defecto" -#: tools/preview\preview.py:1459 -msgid "Reset {} config to saved values" -msgstr "Restablecer la configuración completa de {} a los valores guardados" +#: tools/preview/control_panels.py:673 +#, python-brace-format +msgid "Reset {title} config to saved values" +msgstr "" +"Restablecer la configuración completa de {title} a los valores guardados" diff --git a/locales/es/LC_MESSAGES/tools.restore.cli.mo b/locales/es/LC_MESSAGES/tools.restore.cli.mo deleted file mode 100644 index 40170fef76..0000000000 Binary files a/locales/es/LC_MESSAGES/tools.restore.cli.mo and /dev/null differ diff --git a/locales/es/LC_MESSAGES/tools.restore.cli.po b/locales/es/LC_MESSAGES/tools.restore.cli.po deleted file mode 100644 index e9a34b92fe..0000000000 --- a/locales/es/LC_MESSAGES/tools.restore.cli.po +++ /dev/null @@ -1,36 +0,0 @@ -# SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION -# FIRST AUTHOR , YEAR. -# -msgid "" -msgstr "" -"Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-02-18 23:06-0000\n" -"PO-Revision-Date: 2021-02-19 18:05+0000\n" -"Language-Team: tokafondo\n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=UTF-8\n" -"Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.3\n" -"Last-Translator: \n" -"Plural-Forms: nplurals=2; plural=(n != 1);\n" -"Language: es_ES\n" - -#: tools/restore/cli.py:13 -msgid "This command lets you restore models from backup." -msgstr "Este comando permite restaurar modelos desde una copia de seguridad." - -#: tools/restore/cli.py:22 -msgid "A tool for restoring models from backup (.bk) files" -msgstr "" -"Una herramienta para restaurar modelos a partir de archivos de copia de " -"seguridad (.bk)" - -#: tools/restore/cli.py:33 -msgid "" -"Model directory. A directory containing the model you wish to restore from " -"backup." -msgstr "" -"Directorio del modelo. Un directorio que contiene el modelo que desea " -"restaurar desde la copia de seguridad." diff --git a/locales/es/LC_MESSAGES/tools.sort.cli.mo b/locales/es/LC_MESSAGES/tools.sort.cli.mo index 504ee6da24..1eea2cf248 100644 Binary files a/locales/es/LC_MESSAGES/tools.sort.cli.mo and b/locales/es/LC_MESSAGES/tools.sort.cli.mo differ diff --git a/locales/es/LC_MESSAGES/tools.sort.cli.po b/locales/es/LC_MESSAGES/tools.sort.cli.po index 2a585be3ed..0914d6011c 100644 --- a/locales/es/LC_MESSAGES/tools.sort.cli.po +++ b/locales/es/LC_MESSAGES/tools.sort.cli.po @@ -5,205 +5,379 @@ msgid "" msgstr "" "Project-Id-Version: faceswap.spanish\n" -"POT-Creation-Date: 2021-08-07 12:34+0100\n" -"PO-Revision-Date: 2021-08-07 12:38+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:03+0000\n" "Last-Translator: \n" "Language-Team: tokafondo\n" "Language: es_ES\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.3\n" "Plural-Forms: nplurals=2; plural=(n != 1);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" -#: tools/sort/cli.py:14 +#: tools/sort/cli.py:15 msgid "This command lets you sort images using various methods." msgstr "" "Este comando le permite ordenar las imágenes utilizando varios métodos." -#: tools/sort/cli.py:23 +#: tools/sort/cli.py:21 +msgid "" +" Adjust the '-t' ('--threshold') parameter to control the strength of " +"grouping." +msgstr "" +" Ajuste el parámetro '-t' ('--threshold') para controlar la fuerza de la " +"agrupación." + +#: tools/sort/cli.py:22 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the percentage of color pixels " +"that appear in the image." +msgstr "" +" Ajuste el parámetro '-b' ('--bins') para controlar el número de " +"contenedores para agrupar. Cada imagen se asigna a un contenedor por el " +"porcentaje de píxeles de color que aparecen en la imagen." + +#: tools/sort/cli.py:25 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the number of degrees the face " +"is orientated from center." +msgstr "" +" Ajuste el parámetro '-b' ('--bins') para controlar el número de " +"contenedores para agrupar. Cada imagen se asigna a un contenedor por el " +"número de grados que la cara está orientada desde el centro." + +#: tools/sort/cli.py:28 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. The minimum and maximum values are taken for the chosen sort " +"metric. The bins are then populated with the results from the group sorting." +msgstr "" +" Ajuste el parámetro '-b' ('--bins') para controlar el número de " +"contenedores para agrupar. Los valores mínimo y máximo se toman para la " +"métrica de clasificación elegida. Luego, los contenedores se llenan con los " +"resultados de la clasificación de grupos." + +#: tools/sort/cli.py:32 +msgid "faces by blurriness." +msgstr "rostros por desenfoque." + +#: tools/sort/cli.py:33 +msgid "faces by fft filtered blurriness." +msgstr "caras por borrosidad filtrada fft." + +#: tools/sort/cli.py:34 +msgid "" +"faces by the estimated distance of the alignments from an 'average' face. " +"This can be useful for eliminating misaligned faces. Sorts from most like an " +"average face to least like an average face." +msgstr "" +"caras por la distancia estimada de las alineaciones desde una cara " +"'promedio'. Esto puede ser útil para eliminar caras desalineadas. Ordena de " +"más parecido a un rostro promedio a menos parecido a un rostro promedio." + +#: tools/sort/cli.py:37 +msgid "" +"faces using VGG Face2 by face similarity. This uses a pairwise clustering " +"algorithm to check the distances between 512 features on every face in your " +"set and order them appropriately." +msgstr "" +"caras usando VGG Face2 por similitud de caras. Esto utiliza un algoritmo de " +"agrupamiento por pares para verificar las distancias entre 512 " +"características en cada cara de su conjunto y ordenarlas apropiadamente." + +#: tools/sort/cli.py:40 +msgid "faces by their landmarks." +msgstr "caras por sus puntos de referencia." + +#: tools/sort/cli.py:41 +msgid "Like 'face-cnn' but sorts by dissimilarity." +msgstr "Como 'face-cnn' pero ordenada por la similitud." + +#: tools/sort/cli.py:42 +msgid "faces by Yaw (rotation left to right)." +msgstr "caras por guiñada (rotación de izquierda a derecha)." + +#: tools/sort/cli.py:43 +msgid "faces by Pitch (rotation up and down)." +msgstr "caras por Pitch (rotación arriba y abajo)." + +#: tools/sort/cli.py:44 +msgid "" +"faces by Roll (rotation). Aligned faces should have a roll value close to " +"zero. The further the Roll value from zero the higher liklihood the face is " +"misaligned." +msgstr "" +"caras por Roll (rotación). Las caras alineadas deben tener un valor de " +"balanceo cercano a cero. Cuanto más lejos esté el valor de Roll de cero, " +"mayor será la probabilidad de que la cara esté desalineada." + +#: tools/sort/cli.py:46 +msgid "faces by their color histogram." +msgstr "caras por su histograma de color." + +#: tools/sort/cli.py:47 +msgid "Like 'hist' but sorts by dissimilarity." +msgstr "Como 'hist' pero ordenada por la disimilitud." + +#: tools/sort/cli.py:48 +msgid "" +"images by the average intensity of the converted grayscale color channel." +msgstr "" +"imágenes por la intensidad media del canal de color en escala de grises " +"convertido." + +#: tools/sort/cli.py:49 +msgid "" +"images by their number of black pixels. Useful when faces are near borders " +"and a large part of the image is black." +msgstr "" +"imágenes por su número de píxeles negros. Útil cuando las caras están cerca " +"de los bordes y una gran parte de la imagen es negra." + +#: tools/sort/cli.py:51 +msgid "" +"images by the average intensity of the converted Y color channel. Bright " +"lighting and oversaturated images will be ranked first." +msgstr "" +"imágenes por la intensidad media del canal de color Y convertido. La " +"iluminación brillante y las imágenes sobresaturadas se clasificarán en " +"primer lugar." + +#: tools/sort/cli.py:53 +msgid "" +"images by the average intensity of the converted Cg color channel. Green " +"images will be ranked first and red images will be last." +msgstr "" +"imágenes por la intensidad media del canal de color Cg convertido. Las " +"imágenes verdes se clasificarán primero y las imágenes rojas serán las " +"últimas." + +#: tools/sort/cli.py:55 +msgid "" +"images by the average intensity of the converted Co color channel. Orange " +"images will be ranked first and blue images will be last." +msgstr "" +"imágenes por la intensidad media del canal de color Co convertido. Las " +"imágenes naranjas se clasificarán en primer lugar y las imágenes azules en " +"último lugar." + +#: tools/sort/cli.py:57 +msgid "" +"images by their size in the original frame. Faces further from the camera " +"and from lower resolution sources will be sorted first, whilst faces closer " +"to the camera and from higher resolution sources will be sorted last." +msgstr "" +"imágenes por su tamaño en el marco original. Las caras más alejadas de la " +"cámara y de fuentes de menor resolución se ordenarán primero, mientras que " +"las caras más cercanas a la cámara y de fuentes de mayor resolución se " +"ordenarán en último lugar." + +#: tools/sort/cli.py:81 msgid "Sort faces using a number of different techniques" msgstr "Clasificar los rostros mediante diferentes técnicas" -#: tools/sort/cli.py:33 tools/sort/cli.py:40 +#: tools/sort/cli.py:91 tools/sort/cli.py:98 tools/sort/cli.py:110 +#: tools/sort/cli.py:150 msgid "data" msgstr "datos" -#: tools/sort/cli.py:34 +#: tools/sort/cli.py:92 msgid "Input directory of aligned faces." msgstr "Directorio de entrada de caras alineadas." -#: tools/sort/cli.py:41 -msgid "Output directory for sorted aligned faces." -msgstr "Directorio de salida para las caras alineadas ordenadas." +#: tools/sort/cli.py:100 +msgid "" +"Output directory for sorted aligned faces. If not provided and 'keep' is " +"selected then a new folder called 'sorted' will be created within the input " +"folder to house the output. If not provided and 'keep' is not selected then " +"the images will be sorted in-place, overwriting the original contents of the " +"'input_dir'" +msgstr "" +"Directorio de salida para caras alineadas ordenadas. Si no se proporciona y " +"se selecciona 'keep', se creará una nueva carpeta llamada 'sorted' dentro de " +"la carpeta de entrada para albergar la salida. Si no se proporciona y no se " +"selecciona 'keep', las imágenes se ordenarán en el lugar, sobrescribiendo el " +"contenido original de 'input_dir'" + +#: tools/sort/cli.py:112 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple folders of faces you wish to sort. The faces will be output to " +"separate sub-folders in the output_dir" +msgstr "" +"R|Si se selecciona, input_dir debe ser una carpeta principal que contenga " +"varias carpetas de caras que desea ordenar. Las caras se enviarán a " +"subcarpetas separadas en output_dir" -#: tools/sort/cli.py:50 tools/sort/cli.py:99 +#: tools/sort/cli.py:121 msgid "sort settings" msgstr "ajustes de ordenación" -#: tools/sort/cli.py:52 -msgid "" -"R|Sort by method. Choose how images are sorted. \n" -"L|'blur': Sort faces by blurriness.\n" -"L|'blur-fft': Sort faces by fft filtered blurriness.\n" -"L|'distance' Sort faces by the estimated distance of the alignments from an " -"'average' face. This can be useful for eliminating misaligned faces.\n" -"L|'face': Use VGG Face to sort by face similarity. This uses a pairwise " -"clustering algorithm to check the distances between 512 features on every " -"face in your set and order them appropriately.\n" -"L|'face-cnn': Sort faces by their landmarks. You can adjust the threshold " -"with the '-t' (--ref_threshold) option.\n" -"L|'face-cnn-dissim': Like 'face-cnn' but sorts by dissimilarity.\n" -"L|'face-yaw': Sort faces by Yaw (rotation left to right).\n" -"L|'hist': Sort faces by their color histogram. You can adjust the threshold " -"with the '-t' (--ref_threshold) option.\n" -"L|'hist-dissim': Like 'hist' but sorts by dissimilarity.\n" -"L|'color-gray': Sort images by the average intensity of the converted " -"grayscale color channel.\n" -"L|'color-luma': Sort images by the average intensity of the converted Y " -"color channel. Bright lighting and oversaturated images will be ranked " -"first.\n" -"L|'color-green': Sort images by the average intensity of the converted Cg " -"color channel. Green images will be ranked first and red images will be " -"last.\n" -"L|'color-orange': Sort images by the average intensity of the converted Co " -"color channel. Orange images will be ranked first and blue images will be " -"last.\n" -"L|'size': Sort images by their size in the original frame. Faces closer to " -"the camera and from higher resolution sources will be sorted first, whilst " -"faces further from the camera and from lower resolution sources will be " -"sorted last.\n" -"L|'black-pixels': Sort images by their number of black pixels. Useful when " -"faces are near borders and a large part of the image is black.\n" -"Default: face" -msgstr "" -"R|Método de ordenación. Elige cómo se ordenan las imágenes. \n" -"L|'blur': Ordena las caras por desenfoque.\n" -"L|'blur-fft': Ordena las caras por fft filtrado desenfoque.\n" -"L|'distance' Ordene las caras por la distancia estimada de las alineaciones " -"desde una cara \"promedio\". Esto puede resultar útil para eliminar caras " -"desalineadas.\n" -"L|'face': Utiliza VGG Face para ordenar por similitud de caras. Esto utiliza " -"un algoritmo de agrupación por pares para comprobar las distancias entre 512 " -"características en cada cara en su conjunto y ordenarlos adecuadamente.\n" -"L|'face-cnn': Ordena las caras por sus puntos de referencia. Puedes ajustar " -"el umbral con la opción '-t' (--ref_threshold).\n" -"L|'face-cnn-dissim': Como 'face-cnn' pero ordena por disimilitud.\n" -"L|'face-yaw': Ordena las caras por Yaw (rotación de izquierda a derecha).\n" -"L|'hist': Ordena las caras por su histograma de color. Puedes ajustar el " -"umbral con la opción '-t' (--ref_threshold).\n" -"L|'hist-dissim': Como 'hist' pero ordena por disimilitud.\n" -"L|'color-gray': Ordena las imágenes por la intensidad media del canal de " -"color previa conversión a escala de grises convertido.\n" -"L|'color-luma': Ordena las imágenes por la intensidad media del canal de " -"color Y. Las imágenes muy brillantes y sobresaturadas se clasificarán " -"primero.\n" -"L|'color-green': Ordena las imágenes por la intensidad media del canal de " -"color Cg. Las imágenes verdes serán clasificadas primero y las rojas serán " -"las últimas.\n" -"L|'color-orange': Ordena las imágenes por la intensidad media del canal de " -"color Co. Las imágenes naranjas serán clasificadas primero y las azules " -"serán las últimas.\n" -"L|'size': Ordena las imágenes por su tamaño en el marco original. Los " -"rostros más cercanos a la cámara y de fuentes de mayor resolución se " -"ordenarán primero, mientras que los rostros más alejados de la cámara y de " -"fuentes de menor resolución se ordenarán en último lugar.\n" -"\vL|'black-pixels': Ordene las imágenes por su número de píxeles negros. " -"Útil cuando los rostros están cerca de los bordes y una gran parte de la " -"imagen es negra .\n" -"Por defecto: face" - -#: tools/sort/cli.py:88 tools/sort/cli.py:115 tools/sort/cli.py:127 -#: tools/sort/cli.py:138 -msgid "output" -msgstr "salida" - -#: tools/sort/cli.py:89 -msgid "" -"Keeps the original files in the input directory. Be careful when using this " -"with rename grouping and no specified output directory as this would keep " -"the original and renamed files in the same directory." -msgstr "" -"Mantiene los archivos originales en el directorio de entrada. Tenga cuidado " -"al usar esto con la agrupación de renombre y sin especificar el directorio " -"de salida, ya que esto mantendría los archivos originales y renombrados en " -"el mismo directorio." - -#: tools/sort/cli.py:101 -msgid "" -"Float value. Minimum threshold to use for grouping comparison with 'face-" -"cnn' and 'hist' methods. The lower the value the more discriminating the " -"grouping is. Leaving -1.0 will allow the program set the default value " -"automatically. For face-cnn 7.2 should be enough, with 4 being very " -"discriminating. For hist 0.3 should be enough, with 0.2 being very " -"discriminating. Be careful setting a value that's too low in a directory " -"with many images, as this could result in a lot of directories being " -"created. Defaults: face-cnn 7.2, hist 0.3" -msgstr "" -"Valor flotante. Umbral mínimo a utilizar para la comparación de agrupaciones " -"con los métodos 'face-cnn' e 'hist'. Cuanto más bajo sea el valor, más " -"discriminante será la agrupación. Si se deja -1.0, el programa establecerá " -"el valor por defecto automáticamente. Para 'face-cnn' 7.2 debería ser " -"suficiente, siendo 4 muy discriminante. Para 'hist' 0.3 debería ser " -"suficiente, siendo 0,2 muy discriminante. Tenga cuidado al establecer un " -"valor demasiado bajo en un directorio con muchas imágenes, ya que esto " -"podría resultar en la creación de muchos directorios. Por defecto: 'face-" -"cnn' = 7.2, 'hist' = 0.3" - -#: tools/sort/cli.py:116 -msgid "" -"R|Default: rename.\n" -"L|'folders': files are sorted using the -s/--sort-by method, then they are " -"organized into folders using the -g/--group-by grouping method.\n" -"L|'rename': files are sorted using the -s/--sort-by then they are renamed." -msgstr "" -"R|Por defecto: renombrar.\n" -"L|'folders': los archivos se ordenan utilizando el método -s/--sort-by, y " -"luego se organizan en carpetas utilizando el método de agrupación -g/--group-" -"by.\n" -"L|'rename': los archivos se ordenan utilizando el método -s/--sort-by y " -"luego se renombran." - -#: tools/sort/cli.py:129 -msgid "" -"Group by method. When -fp/--final-processing by folders choose the how the " -"images are grouped after sorting. Default: hist" -msgstr "" -"Método de agrupamiento. Elija la forma de agrupar las imágenes, en el caso " -"de hacerlo por carpetas, después de la clasificación. Por defecto: hist" - -#: tools/sort/cli.py:140 -msgid "" -"Integer value. Number of folders that will be used to group by blur, face-" -"yaw and black-pixels. For blur folder 0 will be the least blurry, while the " -"last folder will be the blurriest. For face-yaw the number of bins is by how " -"much 180 degrees is divided. So if you use 18, then each folder will be a 10 " -"degree increment. Folder 0 will contain faces looking the most to the left " -"whereas the last folder will contain the faces looking the most to the " -"right. If the number of images doesn't divide evenly into the number of " -"bins, the remaining images get put in the last bin. For black-pixels it " -"represents the divider of the percentage of black pixels. For 10, first " -"folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, " -"etc. Default value: 5" -msgstr "" -"Valor entero. Número de carpetas que se utilizarán al agrupar por 'blur' y " -"'face-yaw'. Para 'blur' la carpeta 0 será la menos borrosa, mientras que la " -"última carpeta será la más borrosa. Para 'face-yaw' el número de carpetas es " -"por cuanto se dividen los 180 grados. Así que si usas 18, entonces cada " -"carpeta será un incremento de 10 grados. La carpeta 0 contendrá las caras " -"que miren más a la izquierda, mientras que la última carpeta contendrá las " -"caras que miren más a la derecha. Si el número de imágenes no se divide " -"uniformemente en el número de carpetas, las imágenes restantes se colocan en " -"la última carpeta. Para píxeles negros, representa el divisor del porcentaje " -"de píxeles negros. Para 10, la primera carpeta tendrá las caras con 0 a 10%% " -"de píxeles negros, la segunda de 11 a 20%%, etc. Valor por defecto: 5" - -#: tools/sort/cli.py:154 tools/sort/cli.py:164 +#: tools/sort/cli.py:124 +msgid "" +"R|Choose how images are sorted. Selecting a sort method gives the images a " +"new filename based on the order the image appears within the given method.\n" +"L|'none': Don't sort the images. When a 'group-by' method is selected, " +"selecting 'none' means that the files will be moved/copied into their " +"respective bins, but the files will keep their original filenames. Selecting " +"'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|Elige cómo se ordenan las imágenes. Al seleccionar un método de " +"clasificación, las imágenes reciben un nuevo nombre de archivo basado en el " +"orden en que aparece la imagen dentro del método dado.\n" +"L|'none': No ordenar las imágenes. Cuando se selecciona un método de " +"'agrupar por', seleccionar 'none' significa que los archivos se moverán/" +"copiarán en sus contenedores respectivos, pero los archivos mantendrán sus " +"nombres de archivo originales. Seleccionar 'none' para 'sort-by' y 'group-" +"by' no hará nada" + +#: tools/sort/cli.py:136 tools/sort/cli.py:164 tools/sort/cli.py:184 +msgid "group settings" +msgstr "ajustes de grupo" + +#: tools/sort/cli.py:139 +#, fuzzy +#| msgid "" +#| "R|Selecting a group by method will move/copy files into numbered bins " +#| "based on the selected method.\n" +#| "L|'none': Don't bin the images. Folders will be sorted by the selected " +#| "'sort-by' but will not be binned, instead they will be sorted into a " +#| "single folder. Selecting 'none' for both 'sort-by' and 'group-by' will " +#| "do nothing" +msgid "" +"R|Selecting a group by method will move/copy files into numbered bins based " +"on the selected method.\n" +"L|'none': Don't bin the images. Folders will be sorted by the selected 'sort-" +"by' but will not be binned, instead they will be sorted into a single " +"folder. Selecting 'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|Al seleccionar un grupo por método, los archivos se moverán/copiarán en " +"contenedores numerados según el método seleccionado.\n" +"L|'none': No agrupar las imágenes. Las carpetas se ordenarán por el 'sort-" +"by' seleccionado, pero no se agruparán, sino que se ordenarán en una sola " +"carpeta. Seleccionar 'none' para 'sort-by' y 'group-by' no hará nada" + +#: tools/sort/cli.py:152 +msgid "" +"Whether to keep the original files in their original location. Choosing a " +"'sort-by' method means that the files have to be renamed. Selecting 'keep' " +"means that the original files will be kept, and the renamed files will be " +"created in the specified output folder. Unselecting keep means that the " +"original files will be moved and renamed based on the selected sort/group " +"criteria." +msgstr "" +"Ya sea para mantener los archivos originales en su ubicación original. " +"Elegir un método de 'sort-by' significa que los archivos tienen que ser " +"renombrados. Seleccionar 'keep' significa que los archivos originales se " +"mantendrán y los archivos renombrados se crearán en la carpeta de salida " +"especificada. Deseleccionar 'keep' significa que los archivos originales se " +"moverán y cambiarán de nombre en función de los criterios de clasificación/" +"grupo seleccionados." + +#: tools/sort/cli.py:167 +msgid "" +"R|Float value. Minimum threshold to use for grouping comparison with 'face-" +"cnn' 'hist' and 'face' methods.\n" +"The lower the value the more discriminating the grouping is. Leaving -1.0 " +"will allow Faceswap to choose the default value.\n" +"L|For 'face-cnn' 7.2 should be enough, with 4 being very discriminating. \n" +"L|For 'hist' 0.3 should be enough, with 0.2 being very discriminating. \n" +"L|For 'face' between 0.1 (more bins) to 0.5 (fewer bins) should be about " +"right.\n" +"Be careful setting a value that's too extrene in a directory with many " +"images, as this could result in a lot of folders being created. Defaults: " +"face-cnn 7.2, hist 0.3, face 0.25" +msgstr "" +"R|Valor flotante. Umbral mínimo a usar para agrupar la comparación con los " +"métodos 'face-cnn' 'hist' y 'face'.\n" +"Cuanto más bajo es el valor, más discriminatoria es la agrupación. Dejar " +"-1.0 permitirá que Faceswap elija el valor predeterminado.\n" +"L|Para 'face-cnn' 7.2 debería ser suficiente, siendo 4 muy discriminatorio.\n" +"L|Para 'hist' 0.3 debería ser suficiente, siendo 0.2 muy discriminatorio.\n" +"L|Para 'face', entre 0,1 (más contenedores) y 0,4 (pocos contenedores) " +"debería ser correcto.\n" +"Tenga cuidado al establecer un valor que sea demasiado extremo en un " +"directorio con muchas imágenes, ya que esto podría resultar en la creación " +"de muchas carpetas. Valores predeterminados: face-cnn 7.2, hist 0.3, face " +"0.25" + +#: tools/sort/cli.py:187 +#, fuzzy, python-format +#| msgid "" +#| "R|Integer value. Used to control the number of bins created for grouping " +#| "by: any 'blur' methods, 'color' methods or 'face metric' methods " +#| "('distance', 'size') and 'orientation; methods ('yaw', 'pitch'). For any " +#| "other grouping methods see the '-t' ('--threshold') option.\n" +#| "L|For 'face metric' methods the bins are filled, according the the " +#| "distribution of faces between the minimum and maximum chosen metric.\n" +#| "L|For 'color' methods the number of bins represents the divider of the " +#| "percentage of colored pixels. Eg. For a bin number of '5': The first " +#| "folder will have the faces with 0%% to 20%% colored pixels, second 21%% " +#| "to 40%%, etc. Any empty bins will be deleted, so you may end up with " +#| "fewer bins than selected.\n" +#| "L|For 'blur' methods folder 0 will be the least blurry, while the last " +#| "folder will be the blurriest.\n" +#| "L|For 'orientation' methods the number of bins is dictated by how much " +#| "180 degrees is divided. Eg. If 18 is selected, then each folder will be a " +#| "10 degree increment. Folder 0 will contain faces looking the most to the " +#| "left/down whereas the last folder will contain the faces looking the most " +#| "to the right/up. NB: Some bins may be empty if faces do not fit the " +#| "criteria.\n" +#| "Default value: 5" +msgid "" +"R|Integer value. Used to control the number of bins created for grouping by: " +"any 'blur' methods, 'color' methods or 'face metric' methods ('distance', " +"'size') and 'orientation; methods ('yaw', 'pitch'). For any other grouping " +"methods see the '-t' ('--threshold') option.\n" +"L|For 'face metric' methods the bins are filled, according the the " +"distribution of faces between the minimum and maximum chosen metric.\n" +"L|For 'color' methods the number of bins represents the divider of the " +"percentage of colored pixels. Eg. For a bin number of '5': The first folder " +"will have the faces with 0%% to 20%% colored pixels, second 21%% to 40%%, " +"etc. Any empty bins will be deleted, so you may end up with fewer bins than " +"selected.\n" +"L|For 'blur' methods folder 0 will be the least blurry, while the last " +"folder will be the blurriest.\n" +"L|For 'orientation' methods the number of bins is dictated by how much 180 " +"degrees is divided. Eg. If 18 is selected, then each folder will be a 10 " +"degree increment. Folder 0 will contain faces looking the most to the left/" +"down whereas the last folder will contain the faces looking the most to the " +"right/up. NB: Some bins may be empty if faces do not fit the criteria. \n" +"Default value: 5" +msgstr "" +"R|Valor entero. Se utiliza para controlar el número de contenedores creados " +"para agrupar por: cualquier método de 'blur', método de 'color' o método de " +"'face metric' ('distance', 'size') y 'orientación; métodos ('yaw', 'pitch'). " +"Para cualquier otro método de agrupación, consulte la opción '-t' ('--" +"threshold').\n" +"L|Para los métodos de 'face metric', los contenedores se llenan de acuerdo " +"con la distribución de caras entre la métrica mínima y máxima elegida.\n" +"L|Para los métodos de 'color', el número de contenedores representa el " +"divisor del porcentaje de píxeles coloreados. P.ej. Para un número de " +"contenedor de '5': la primera carpeta tendrá las caras con 0%% a 20%% " +"píxeles de color, la segunda 21%% a 40%%, etc. Se eliminarán todos los " +"contenedores vacíos, por lo que puede terminar con menos contenedores que " +"los seleccionados.\n" +"L|Para los métodos 'blur', la carpeta 0 será la menos borrosa, mientras que " +"la última carpeta será la más borrosa.\n" +"L|Para los métodos de 'orientation', el número de contenedores está dictado " +"por cuánto se dividen 180 grados. P.ej. Si se selecciona 18, cada carpeta " +"tendrá un incremento de 10 grados. La carpeta 0 contendrá las caras que " +"miran más hacia la izquierda/abajo, mientras que la última carpeta contendrá " +"las caras que miran más hacia la derecha/arriba. NB: algunos contenedores " +"pueden estar vacíos si las caras no se ajustan a los criterios.\n" +"Valor predeterminado: 5" + +#: tools/sort/cli.py:207 tools/sort/cli.py:217 msgid "settings" msgstr "ajustes" -#: tools/sort/cli.py:156 +#: tools/sort/cli.py:210 msgid "" "Logs file renaming changes if grouping by renaming, or it logs the file " "copying/movement if grouping by folders. If no log file is specified with " @@ -215,7 +389,7 @@ msgstr "" "se especifica ningún archivo de registro con '--log-file', se creará un " "archivo 'sort_log.json' en el directorio de entrada." -#: tools/sort/cli.py:167 +#: tools/sort/cli.py:221 msgid "" "Specify a log file to use for saving the renaming or grouping information. " "If specified extension isn't 'json' or 'yaml', then json will be used as the " @@ -225,3 +399,153 @@ msgstr "" "información de renombrado o agrupación. Si la extensión especificada no es " "'json' o 'yaml', se utilizará json como serializador, con el nombre de " "archivo suministrado. Por defecto: sort_log.json" + +#~ msgid " option is deprecated. Use 'yaw'" +#~ msgstr " la opción está en desuso. Usa 'yaw'" + +#~ msgid " option is deprecated. Use 'color-black'" +#~ msgstr " la opción está en desuso. Usa 'color-black'" + +#~ msgid "output" +#~ msgstr "salida" + +#~ msgid "" +#~ "Deprecated and no longer used. The final processing will be dictated by " +#~ "the sort/group by methods and whether 'keep_original' is selected." +#~ msgstr "" +#~ "En desuso y ya no se usa. El procesamiento final será dictado por los " +#~ "métodos de ordenación/agrupación y si se selecciona 'keepl'." + +#~ msgid "Output directory for sorted aligned faces." +#~ msgstr "Directorio de salida para las caras alineadas ordenadas." + +#~ msgid "" +#~ "R|Sort by method. Choose how images are sorted. \n" +#~ "L|'blur': Sort faces by blurriness.\n" +#~ "L|'blur-fft': Sort faces by fft filtered blurriness.\n" +#~ "L|'distance' Sort faces by the estimated distance of the alignments from " +#~ "an 'average' face. This can be useful for eliminating misaligned faces.\n" +#~ "L|'face': Use VGG Face to sort by face similarity. This uses a pairwise " +#~ "clustering algorithm to check the distances between 512 features on every " +#~ "face in your set and order them appropriately.\n" +#~ "L|'face-cnn': Sort faces by their landmarks. You can adjust the threshold " +#~ "with the '-t' (--ref_threshold) option.\n" +#~ "L|'face-cnn-dissim': Like 'face-cnn' but sorts by dissimilarity.\n" +#~ "L|'face-yaw': Sort faces by Yaw (rotation left to right).\n" +#~ "L|'hist': Sort faces by their color histogram. You can adjust the " +#~ "threshold with the '-t' (--ref_threshold) option.\n" +#~ "L|'hist-dissim': Like 'hist' but sorts by dissimilarity.\n" +#~ "L|'color-gray': Sort images by the average intensity of the converted " +#~ "grayscale color channel.\n" +#~ "L|'color-luma': Sort images by the average intensity of the converted Y " +#~ "color channel. Bright lighting and oversaturated images will be ranked " +#~ "first.\n" +#~ "L|'color-green': Sort images by the average intensity of the converted Cg " +#~ "color channel. Green images will be ranked first and red images will be " +#~ "last.\n" +#~ "L|'color-orange': Sort images by the average intensity of the converted " +#~ "Co color channel. Orange images will be ranked first and blue images will " +#~ "be last.\n" +#~ "L|'size': Sort images by their size in the original frame. Faces closer " +#~ "to the camera and from higher resolution sources will be sorted first, " +#~ "whilst faces further from the camera and from lower resolution sources " +#~ "will be sorted last.\n" +#~ "L|'black-pixels': Sort images by their number of black pixels. Useful " +#~ "when faces are near borders and a large part of the image is black.\n" +#~ "Default: face" +#~ msgstr "" +#~ "R|Método de ordenación. Elige cómo se ordenan las imágenes. \n" +#~ "L|'blur': Ordena las caras por desenfoque.\n" +#~ "L|'blur-fft': Ordena las caras por fft filtrado desenfoque.\n" +#~ "L|'distance' Ordene las caras por la distancia estimada de las " +#~ "alineaciones desde una cara \"promedio\". Esto puede resultar útil para " +#~ "eliminar caras desalineadas.\n" +#~ "L|'face': Utiliza VGG Face para ordenar por similitud de caras. Esto " +#~ "utiliza un algoritmo de agrupación por pares para comprobar las " +#~ "distancias entre 512 características en cada cara en su conjunto y " +#~ "ordenarlos adecuadamente.\n" +#~ "L|'face-cnn': Ordena las caras por sus puntos de referencia. Puedes " +#~ "ajustar el umbral con la opción '-t' (--ref_threshold).\n" +#~ "L|'face-cnn-dissim': Como 'face-cnn' pero ordena por disimilitud.\n" +#~ "L|'face-yaw': Ordena las caras por Yaw (rotación de izquierda a " +#~ "derecha).\n" +#~ "L|'hist': Ordena las caras por su histograma de color. Puedes ajustar el " +#~ "umbral con la opción '-t' (--ref_threshold).\n" +#~ "L|'hist-dissim': Como 'hist' pero ordena por disimilitud.\n" +#~ "L|'color-gray': Ordena las imágenes por la intensidad media del canal de " +#~ "color previa conversión a escala de grises convertido.\n" +#~ "L|'color-luma': Ordena las imágenes por la intensidad media del canal de " +#~ "color Y. Las imágenes muy brillantes y sobresaturadas se clasificarán " +#~ "primero.\n" +#~ "L|'color-green': Ordena las imágenes por la intensidad media del canal de " +#~ "color Cg. Las imágenes verdes serán clasificadas primero y las rojas " +#~ "serán las últimas.\n" +#~ "L|'color-orange': Ordena las imágenes por la intensidad media del canal " +#~ "de color Co. Las imágenes naranjas serán clasificadas primero y las " +#~ "azules serán las últimas.\n" +#~ "L|'size': Ordena las imágenes por su tamaño en el marco original. Los " +#~ "rostros más cercanos a la cámara y de fuentes de mayor resolución se " +#~ "ordenarán primero, mientras que los rostros más alejados de la cámara y " +#~ "de fuentes de menor resolución se ordenarán en último lugar.\n" +#~ "\vL|'black-pixels': Ordene las imágenes por su número de píxeles negros. " +#~ "Útil cuando los rostros están cerca de los bordes y una gran parte de la " +#~ "imagen es negra .\n" +#~ "Por defecto: face" + +#~ msgid "" +#~ "Keeps the original files in the input directory. Be careful when using " +#~ "this with rename grouping and no specified output directory as this would " +#~ "keep the original and renamed files in the same directory." +#~ msgstr "" +#~ "Mantiene los archivos originales en el directorio de entrada. Tenga " +#~ "cuidado al usar esto con la agrupación de renombre y sin especificar el " +#~ "directorio de salida, ya que esto mantendría los archivos originales y " +#~ "renombrados en el mismo directorio." + +#~ msgid "" +#~ "R|Default: rename.\n" +#~ "L|'folders': files are sorted using the -s/--sort-by method, then they " +#~ "are organized into folders using the -g/--group-by grouping method.\n" +#~ "L|'rename': files are sorted using the -s/--sort-by then they are renamed." +#~ msgstr "" +#~ "R|Por defecto: renombrar.\n" +#~ "L|'folders': los archivos se ordenan utilizando el método -s/--sort-by, y " +#~ "luego se organizan en carpetas utilizando el método de agrupación -g/--" +#~ "group-by.\n" +#~ "L|'rename': los archivos se ordenan utilizando el método -s/--sort-by y " +#~ "luego se renombran." + +#~ msgid "" +#~ "Group by method. When -fp/--final-processing by folders choose the how " +#~ "the images are grouped after sorting. Default: hist" +#~ msgstr "" +#~ "Método de agrupamiento. Elija la forma de agrupar las imágenes, en el " +#~ "caso de hacerlo por carpetas, después de la clasificación. Por defecto: " +#~ "hist" + +#, python-format +#~ msgid "" +#~ "Integer value. Number of folders that will be used to group by blur, face-" +#~ "yaw and black-pixels. For blur folder 0 will be the least blurry, while " +#~ "the last folder will be the blurriest. For face-yaw the number of bins is " +#~ "by how much 180 degrees is divided. So if you use 18, then each folder " +#~ "will be a 10 degree increment. Folder 0 will contain faces looking the " +#~ "most to the left whereas the last folder will contain the faces looking " +#~ "the most to the right. If the number of images doesn't divide evenly into " +#~ "the number of bins, the remaining images get put in the last bin. For " +#~ "black-pixels it represents the divider of the percentage of black pixels. " +#~ "For 10, first folder will have the faces with 0 to 10%% black pixels, " +#~ "second 11 to 20%%, etc. Default value: 5" +#~ msgstr "" +#~ "Valor entero. Número de carpetas que se utilizarán al agrupar por 'blur' " +#~ "y 'face-yaw'. Para 'blur' la carpeta 0 será la menos borrosa, mientras " +#~ "que la última carpeta será la más borrosa. Para 'face-yaw' el número de " +#~ "carpetas es por cuanto se dividen los 180 grados. Así que si usas 18, " +#~ "entonces cada carpeta será un incremento de 10 grados. La carpeta 0 " +#~ "contendrá las caras que miren más a la izquierda, mientras que la última " +#~ "carpeta contendrá las caras que miren más a la derecha. Si el número de " +#~ "imágenes no se divide uniformemente en el número de carpetas, las " +#~ "imágenes restantes se colocan en la última carpeta. Para píxeles negros, " +#~ "representa el divisor del porcentaje de píxeles negros. Para 10, la " +#~ "primera carpeta tendrá las caras con 0 a 10%% de píxeles negros, la " +#~ "segunda de 11 a 20%%, etc. Valor por defecto: 5" diff --git a/locales/gui.menu.pot b/locales/gui.menu.pot new file mode 100644 index 0000000000..a20a799f11 --- /dev/null +++ b/locales/gui.menu.pot @@ -0,0 +1,154 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-07 13:54+0100\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: ./lib/gui/menu.py:37 +msgid "faceswap.dev - Guides and Forum" +msgstr "" + +#: ./lib/gui/menu.py:38 +msgid "Patreon - Support this project" +msgstr "" + +#: ./lib/gui/menu.py:39 +msgid "Discord - The FaceSwap Discord server" +msgstr "" + +#: ./lib/gui/menu.py:40 +msgid "Github - Our Source Code" +msgstr "" + +#: ./lib/gui/menu.py:60 +msgid "File" +msgstr "" + +#: ./lib/gui/menu.py:61 +msgid "Settings" +msgstr "" + +#: ./lib/gui/menu.py:62 +msgid "Help" +msgstr "" + +#: ./lib/gui/menu.py:85 +msgid "Configure Settings..." +msgstr "" + +#: ./lib/gui/menu.py:116 +msgid "New Project..." +msgstr "" + +#: ./lib/gui/menu.py:121 +msgid "Open Project..." +msgstr "" + +#: ./lib/gui/menu.py:126 +msgid "Save Project" +msgstr "" + +#: ./lib/gui/menu.py:131 +msgid "Save Project as..." +msgstr "" + +#: ./lib/gui/menu.py:136 +msgid "Reload Project from Disk" +msgstr "" + +#: ./lib/gui/menu.py:141 +msgid "Close Project" +msgstr "" + +#: ./lib/gui/menu.py:147 +msgid "Open Task..." +msgstr "" + +#: ./lib/gui/menu.py:154 +msgid "Open recent" +msgstr "" + +#: ./lib/gui/menu.py:156 +msgid "Quit" +msgstr "" + +#: ./lib/gui/menu.py:211 +msgid "{} Task" +msgstr "" + +#: ./lib/gui/menu.py:223 +msgid "Clear recent files" +msgstr "" + +#: ./lib/gui/menu.py:391 +msgid "Check for updates..." +msgstr "" + +#: ./lib/gui/menu.py:394 +msgid "Update Faceswap..." +msgstr "" + +#: ./lib/gui/menu.py:398 +msgid "Switch Branch" +msgstr "" + +#: ./lib/gui/menu.py:401 +msgid "Resources" +msgstr "" + +#: ./lib/gui/menu.py:404 +msgid "Output System Information" +msgstr "" + +#: ./lib/gui/menu.py:589 +msgid "currently selected Task" +msgstr "" + +#: ./lib/gui/menu.py:589 +msgid "Project" +msgstr "" + +#: ./lib/gui/menu.py:591 +msgid "Reload {} from disk" +msgstr "" + +#: ./lib/gui/menu.py:593 +msgid "Create a new {}..." +msgstr "" + +#: ./lib/gui/menu.py:595 +msgid "Reset {} to default" +msgstr "" + +#: ./lib/gui/menu.py:597 +msgid "Save {}" +msgstr "" + +#: ./lib/gui/menu.py:599 +msgid "Save {} as..." +msgstr "" + +#: ./lib/gui/menu.py:603 +msgid " from a task or project file" +msgstr "" + +#: ./lib/gui/menu.py:604 +msgid "Load {}..." +msgstr "" + +#: ./lib/gui/menu.py:659 +msgid "Configure {} settings..." +msgstr "" diff --git a/locales/gui.tooltips.pot b/locales/gui.tooltips.pot index be16034bd2..f6973d6152 100644 --- a/locales/gui.tooltips.pot +++ b/locales/gui.tooltips.pot @@ -119,62 +119,6 @@ msgstr "" msgid "Enable or disable {} display" msgstr "" -#: ./lib/gui/menu.py:32 -msgid "faceswap.dev - Guides and Forum" -msgstr "" - -#: ./lib/gui/menu.py:33 -msgid "Patreon - Support this project" -msgstr "" - -#: ./lib/gui/menu.py:34 -msgid "Discord - The FaceSwap Discord server" -msgstr "" - -#: ./lib/gui/menu.py:35 -msgid "Github - Our Source Code" -msgstr "" - -#: ./lib/gui/menu.py:527 -msgid "Configure {} settings..." -msgstr "" - -#: ./lib/gui/menu.py:535 -msgid "Project" -msgstr "" - -#: ./lib/gui/menu.py:535 -msgid "currently selected Task" -msgstr "" - -#: ./lib/gui/menu.py:537 -msgid "Reload {} from disk" -msgstr "" - -#: ./lib/gui/menu.py:539 -msgid "Create a new {}..." -msgstr "" - -#: ./lib/gui/menu.py:541 -msgid "Reset {} to default" -msgstr "" - -#: ./lib/gui/menu.py:543 -msgid "Save {}" -msgstr "" - -#: ./lib/gui/menu.py:545 -msgid "Save {} as..." -msgstr "" - -#: ./lib/gui/menu.py:549 -msgid " from a task or project file" -msgstr "" - -#: ./lib/gui/menu.py:550 -msgid "Load {}..." -msgstr "" - #: ./lib/gui/popup_configure.py:209 msgid "Close without saving" msgstr "" diff --git a/locales/kr/LC_MESSAGES/faceswap.mo b/locales/kr/LC_MESSAGES/faceswap.mo new file mode 100644 index 0000000000..4613eb7345 Binary files /dev/null and b/locales/kr/LC_MESSAGES/faceswap.mo differ diff --git a/locales/kr/LC_MESSAGES/faceswap.po b/locales/kr/LC_MESSAGES/faceswap.po new file mode 100644 index 0000000000..c4829dac52 --- /dev/null +++ b/locales/kr/LC_MESSAGES/faceswap.po @@ -0,0 +1,34 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"POT-Creation-Date: 2021-02-18 23:48-0000\n" +"PO-Revision-Date: 2022-11-24 12:21+0900\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.2\n" + +#: faceswap.py:43 +msgid "Extract the faces from pictures or a video" +msgstr "그림들 또는 비디오에서 얼굴을 추출합니다" + +#: faceswap.py:44 +msgid "Train a model for the two faces A and B" +msgstr "얼굴들 A와 B에 대한 모델을 훈련시킵니다" + +#: faceswap.py:47 +msgid "Convert source pictures or video to a new one with the face swapped" +msgstr "원본 이미지 또는 비디오를 얼굴이 뒤바뀐 새로운 이미지 또는 영상으로 변환합니다" + +#: faceswap.py:48 +msgid "Launch the Faceswap Graphical User Interface" +msgstr "Faceswap GUI를 실행합니다" diff --git a/locales/kr/LC_MESSAGES/gui.menu.mo b/locales/kr/LC_MESSAGES/gui.menu.mo new file mode 100644 index 0000000000..2bab76f5b0 Binary files /dev/null and b/locales/kr/LC_MESSAGES/gui.menu.mo differ diff --git a/locales/kr/LC_MESSAGES/gui.menu.po b/locales/kr/LC_MESSAGES/gui.menu.po new file mode 100644 index 0000000000..b20b5dca54 --- /dev/null +++ b/locales/kr/LC_MESSAGES/gui.menu.po @@ -0,0 +1,155 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-07 13:54+0100\n" +"PO-Revision-Date: 2023-06-07 14:11+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/gui/menu.py:37 +msgid "faceswap.dev - Guides and Forum" +msgstr "faceswap.dev - Guides and Forum" + +#: lib/gui/menu.py:38 +msgid "Patreon - Support this project" +msgstr "Patreon - Support this project" + +#: lib/gui/menu.py:39 +msgid "Discord - The FaceSwap Discord server" +msgstr "Discord - The FaceSwap Discord server" + +#: lib/gui/menu.py:40 +msgid "Github - Our Source Code" +msgstr "Github - Our Source Code" + +#: lib/gui/menu.py:60 +msgid "File" +msgstr "" + +#: lib/gui/menu.py:61 +msgid "Settings" +msgstr "" + +#: lib/gui/menu.py:62 +msgid "Help" +msgstr "" + +#: lib/gui/menu.py:85 +msgid "Configure Settings..." +msgstr "" + +#: lib/gui/menu.py:116 +msgid "New Project..." +msgstr "" + +#: lib/gui/menu.py:121 +msgid "Open Project..." +msgstr "" + +#: lib/gui/menu.py:126 +msgid "Save Project" +msgstr "" + +#: lib/gui/menu.py:131 +msgid "Save Project as..." +msgstr "" + +#: lib/gui/menu.py:136 +msgid "Reload Project from Disk" +msgstr "" + +#: lib/gui/menu.py:141 +msgid "Close Project" +msgstr "" + +#: lib/gui/menu.py:147 +msgid "Open Task..." +msgstr "" + +#: lib/gui/menu.py:154 +msgid "Open recent" +msgstr "" + +#: lib/gui/menu.py:156 +msgid "Quit" +msgstr "" + +#: lib/gui/menu.py:211 +msgid "{} Task" +msgstr "" + +#: lib/gui/menu.py:223 +msgid "Clear recent files" +msgstr "" + +#: lib/gui/menu.py:391 +msgid "Check for updates..." +msgstr "" + +#: lib/gui/menu.py:394 +msgid "Update Faceswap..." +msgstr "" + +#: lib/gui/menu.py:398 +msgid "Switch Branch" +msgstr "" + +#: lib/gui/menu.py:401 +msgid "Resources" +msgstr "" + +#: lib/gui/menu.py:404 +msgid "Output System Information" +msgstr "" + +#: lib/gui/menu.py:589 +msgid "currently selected Task" +msgstr "현재 선택된 작업" + +#: lib/gui/menu.py:589 +msgid "Project" +msgstr "프로젝트" + +#: lib/gui/menu.py:591 +msgid "Reload {} from disk" +msgstr "디스크에서 {}를 다시 가져옵니다" + +#: lib/gui/menu.py:593 +msgid "Create a new {}..." +msgstr "새로운 {}를 만들기." + +#: lib/gui/menu.py:595 +msgid "Reset {} to default" +msgstr "{} 기본으로 재설정" + +#: lib/gui/menu.py:597 +msgid "Save {}" +msgstr "{} 저장" + +#: lib/gui/menu.py:599 +msgid "Save {} as..." +msgstr "{}를 다른 이름으로 저장." + +#: lib/gui/menu.py:603 +msgid " from a task or project file" +msgstr " 작업 또는 프로젝트 파일에서" + +#: lib/gui/menu.py:604 +msgid "Load {}..." +msgstr "{} 가져오기." + +#: lib/gui/menu.py:659 +msgid "Configure {} settings..." +msgstr "{} 세팅 설정하기." diff --git a/locales/kr/LC_MESSAGES/gui.tooltips.mo b/locales/kr/LC_MESSAGES/gui.tooltips.mo new file mode 100644 index 0000000000..bce4cb2148 Binary files /dev/null and b/locales/kr/LC_MESSAGES/gui.tooltips.mo differ diff --git a/locales/kr/LC_MESSAGES/gui.tooltips.po b/locales/kr/LC_MESSAGES/gui.tooltips.po new file mode 100644 index 0000000000..16c1631fcc --- /dev/null +++ b/locales/kr/LC_MESSAGES/gui.tooltips.po @@ -0,0 +1,205 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"POT-Creation-Date: 2021-03-22 18:37+0000\n" +"PO-Revision-Date: 2023-06-07 14:13+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/gui/command.py:184 +msgid "Output command line options to the console" +msgstr "명령어 옵션들을 콘솔에 출력" + +#: lib/gui/command.py:195 +msgid "Run the {} script" +msgstr "{} 스크립트 실행" + +#: lib/gui/control_helper.py:1234 +msgid "Select a folder..." +msgstr "폴더 선택." + +#: lib/gui/control_helper.py:1235 lib/gui/control_helper.py:1236 +msgid "Select a file..." +msgstr "파일 선택." + +#: lib/gui/control_helper.py:1237 +msgid "Select a folder of images..." +msgstr "이미지들의 폴더 선택." + +#: lib/gui/control_helper.py:1238 +msgid "Select a video..." +msgstr "비디오 선택." + +#: lib/gui/control_helper.py:1239 +msgid "Select a model folder..." +msgstr "모델 폴더 선택하기." + +#: lib/gui/control_helper.py:1240 +msgid "Select one or more files..." +msgstr "하나 이상의 파일들 선택." + +#: lib/gui/control_helper.py:1241 +msgid "Select a file or folder..." +msgstr "파일 또는 폴더 선택." + +#: lib/gui/control_helper.py:1242 +msgid "Select a save location..." +msgstr "저장 위치 선택." + +#: lib/gui/display.py:71 +msgid "Summary statistics for each training session" +msgstr "각 훈련 세션들에 대한 통계 요약" + +#: lib/gui/display.py:113 +msgid "Preview updates every 5 seconds" +msgstr "5초마다 미리보기를 업데이트하기" + +#: lib/gui/display.py:122 +msgid "Graph showing Loss vs Iterations" +msgstr "반복에 따른 손실율 그래프" + +#: lib/gui/display.py:125 +msgid "Training preview. Updated on every save iteration" +msgstr "훈련 미리보기. 매 저장된 반복마다 업데이트됩니다" + +#: lib/gui/display_analysis.py:342 +msgid "Load/Refresh stats for the currently training session" +msgstr "현재 훈련 세션에 대한 통계 가져오기/새로고침" + +#: lib/gui/display_analysis.py:344 +msgid "Clear currently displayed session stats" +msgstr "현재 보여지는 세션 통계 지우기" + +#: lib/gui/display_analysis.py:346 +msgid "Save session stats to csv" +msgstr "세션 통계 csv로 저장하기" + +#: lib/gui/display_analysis.py:348 +msgid "Load saved session stats" +msgstr "저장된 세션 통계 가져오기" + +#: lib/gui/display_command.py:94 +msgid "Preview updates at every model save. Click to refresh now." +msgstr "" +"모델을 저장할 때마다 미리보기를 업데이트합니다. 지금 새로고침하기 위해 누르세" +"요." + +#: lib/gui/display_command.py:261 +msgid "Graph updates at every model save. Click to refresh now." +msgstr "" +"모델을 저장할 때마다 그래프를 업데이트합니다. 지금 새로고침하기 위해 누르세" +"요." + +#: lib/gui/display_command.py:275 +msgid "Display the raw loss data" +msgstr "원시 손실 데이터 보이기" + +#: lib/gui/display_command.py:287 +msgid "Display the smoothed loss data" +msgstr "매끄러운 손실 데이터 보이기" + +#: lib/gui/display_command.py:294 +msgid "Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing." +msgstr "" +"매끄러움 정도를 설정합니다. 0이면 매끄러움이 없고, 0.99이면 최대로 매끄러워집" +"니다." + +#: lib/gui/display_command.py:324 +msgid "Set the number of iterations to display. 0 displays the full session." +msgstr "" +"화면에 보여질 반복 횟수를 설정합니다. 0 displays는 모든 세션에서 보여줍니다." + +#: lib/gui/display_page.py:238 +msgid "Save {}(s) to file" +msgstr "{}(s)를 파일에 저장합니다" + +#: lib/gui/display_page.py:250 +msgid "Enable or disable {} display" +msgstr "{} display를 활성화 또는 비활성화" + +#: lib/gui/popup_configure.py:209 +msgid "Close without saving" +msgstr "저장하지 않고 닫기" + +#: lib/gui/popup_configure.py:210 +msgid "Save this page's config" +msgstr "이 페이지의 설정을 저장" + +#: lib/gui/popup_configure.py:211 +msgid "Reset this page's config to default values" +msgstr "이 페이지의 설정을 기본값으로 재설정" + +#: lib/gui/popup_configure.py:213 +msgid "Save all settings for the currently selected config" +msgstr "현재 선택된 모든 설정을 저장" + +#: lib/gui/popup_configure.py:216 +msgid "Reset all settings for the currently selected config to default values" +msgstr "현재 선택된 모든 설정을 기본값으로 재설정" + +#: lib/gui/popup_configure.py:538 +msgid "Select a plugin to configure:" +msgstr "구성할 플러그인 선택:" + +#: lib/gui/popup_session.py:191 +msgid "Display {}" +msgstr "{} 보이기" + +#: lib/gui/popup_session.py:342 +msgid "Refresh graph" +msgstr "그래프 새로고침" + +#: lib/gui/popup_session.py:344 +msgid "Save display data to csv" +msgstr "디스플레이 데이터를 csv로 저장" + +#: lib/gui/popup_session.py:346 +msgid "Number of data points to sample for rolling average" +msgstr "샘플의 이동평균 데이터 포인트 개수" + +#: lib/gui/popup_session.py:348 +msgid "Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing" +msgstr "" +"매끄러움 정도를 설정합니다. 0이면 매끄러움이 없고, 0.99이면 최대로 매끄러워집" +"니다" + +#: lib/gui/popup_session.py:350 +msgid "" +"Flatten data points that fall more than 1 standard deviation from the mean " +"to the mean value." +msgstr "평균에서 값까지 1 표준 편차보다 더 멀리 떨어진 데이터들 펴기." + +#: lib/gui/popup_session.py:353 +msgid "Display rolling average of the data" +msgstr "데이터의 이동평균 보이기" + +#: lib/gui/popup_session.py:355 +msgid "Smooth the data" +msgstr "데이터 매끄럽게 하기" + +#: lib/gui/popup_session.py:357 +msgid "Display raw data" +msgstr "원시 데이터 보이기" + +#: lib/gui/popup_session.py:359 +msgid "Display polynormal data trend" +msgstr "다항 데이터 트렌드 보이기" + +#: lib/gui/popup_session.py:361 +msgid "Set the data to display" +msgstr "데이터를 display에 설정하기" + +#: lib/gui/popup_session.py:363 +msgid "Change y-axis scale" +msgstr "변경합니다 y축의 범위를" diff --git a/locales/kr/LC_MESSAGES/lib.cli.args.mo b/locales/kr/LC_MESSAGES/lib.cli.args.mo new file mode 100644 index 0000000000..bc4b9fce78 Binary files /dev/null and b/locales/kr/LC_MESSAGES/lib.cli.args.mo differ diff --git a/locales/kr/LC_MESSAGES/lib.cli.args.po b/locales/kr/LC_MESSAGES/lib.cli.args.po new file mode 100644 index 0000000000..ec473df1f1 --- /dev/null +++ b/locales/kr/LC_MESSAGES/lib.cli.args.po @@ -0,0 +1,57 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:06+0000\n" +"PO-Revision-Date: 2024-03-28 18:17+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args.py:188 lib/cli/args.py:199 lib/cli/args.py:208 +#: lib/cli/args.py:219 +msgid "Global Options" +msgstr "전역 옵션들" + +#: lib/cli/args.py:190 +msgid "" +"R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond " +"to any GPU(s) that you do not wish to be made available to Faceswap. " +"Selecting all GPUs here will force Faceswap into CPU mode.\n" +"L|{}" +msgstr "" +"R|Faceswap에서 사용되는 GPUs를 제외합니다. Faceswap에서 사용되게 하고 싶지 " +"않은 GPU(s)에 해당하는 번호를 선택하세요. 모든 GPUs를 선택하면 Faceswap으로 " +"하여금 CPU mode를 강제로 사용하게 합니다.\n" +"L|{}" + +#: lib/cli/args.py:201 +msgid "" +"Optionally overide the saved config with the path to a custom config file." +msgstr "선택적으로 저장된 설정을 경로와 함께 개인 설정 파일에 덮어씌웁니다." + +#: lib/cli/args.py:210 +msgid "" +"Log level. Stick with INFO or VERBOSE unless you need to file an error " +"report. Be careful with TRACE as it will generate a lot of data" +msgstr "" +"로그 레벨. 오류 리포트가 필요하지 않다면 INFO와 VERBOSE를 사용하세요. 단, 굉" +"장히 많은 데이터를 생성할 수 있는 TRACE는 조심하세요" + +#: lib/cli/args.py:220 +msgid "Path to store the logfile. Leave blank to store in the faceswap folder" +msgstr "로그파일을 저장할 경로. faceswap 폴더에 저장하고 싶으면 비워두세요" + +#: lib/cli/args.py:319 +msgid "Output to Shell console instead of GUI console" +msgstr "결과를 GUI 콘솔이 아닌 쉘 콘솔에 출력합니다" diff --git a/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.mo b/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.mo new file mode 100644 index 0000000000..1f0c43722c Binary files /dev/null and b/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.mo differ diff --git a/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.po b/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.po new file mode 100644 index 0000000000..a504625624 --- /dev/null +++ b/locales/kr/LC_MESSAGES/lib.cli.args_extract_convert.po @@ -0,0 +1,656 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-12 11:56+0100\n" +"PO-Revision-Date: 2024-04-12 12:00+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_extract_convert.py:46 lib/cli/args_extract_convert.py:56 +#: lib/cli/args_extract_convert.py:64 lib/cli/args_extract_convert.py:122 +#: lib/cli/args_extract_convert.py:483 lib/cli/args_extract_convert.py:492 +msgid "Data" +msgstr "데이터" + +#: lib/cli/args_extract_convert.py:48 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file. NB: This should be the source video/" +"frames NOT the source faces." +msgstr "" +"폴더나 비디오를 입력하세요. 당신이 사용하고 싶은 이미지 파일들을 가진 폴더 또" +"는 비디오 파일의 경로여야 합니다. NB: 이 폴더는 원본 비디오여야 합니다." + +#: lib/cli/args_extract_convert.py:57 +msgid "Output directory. This is where the converted files will be saved." +msgstr "출력 폴더. 변환된 파일들이 저장될 곳입니다." + +#: lib/cli/args_extract_convert.py:66 +msgid "" +"Optional path to an alignments file. Leave blank if the alignments file is " +"at the default location." +msgstr "" +"(선택적) alignments 파일의 경로. 비워두면 alignments 파일이 기본 위치에 저장" +"됩니다." + +#: lib/cli/args_extract_convert.py:97 +msgid "" +"Extract faces from image or video sources.\n" +"Extraction plugins can be configured in the 'Settings' Menu" +msgstr "" +"얼굴들을 이미지 또는 비디오에서 추출합니다.\n" +"추출 플러그인은 '설정' 메뉴에서 설정할 수 있습니다" + +#: lib/cli/args_extract_convert.py:124 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple videos and/or folders of images you wish to extract from. The faces " +"will be output to separate sub-folders in the output_dir." +msgstr "" +"R|만약 선택된다면 input_dir은 당신이 추출하고자 하는 여러개의 비디오 그리고/" +"또는 이미지들을 가진 부모 폴더가 되야 합니다. 얼굴들은 output_dir에 분리된 하" +"위 폴더에 저장됩니다." + +#: lib/cli/args_extract_convert.py:133 lib/cli/args_extract_convert.py:152 +#: lib/cli/args_extract_convert.py:167 lib/cli/args_extract_convert.py:206 +#: lib/cli/args_extract_convert.py:224 lib/cli/args_extract_convert.py:237 +#: lib/cli/args_extract_convert.py:247 lib/cli/args_extract_convert.py:257 +#: lib/cli/args_extract_convert.py:503 lib/cli/args_extract_convert.py:529 +#: lib/cli/args_extract_convert.py:568 +msgid "Plugins" +msgstr "플러그인들" + +#: lib/cli/args_extract_convert.py:135 +msgid "" +"R|Detector to use. Some of these have configurable settings in '/config/" +"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" +"L|cv2-dnn: A CPU only extractor which is the least reliable and least " +"resource intensive. Use this if not using a GPU and time is important.\n" +"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " +"than other GPU detectors but can often return more false positives.\n" +"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " +"fewer false positives than other GPU detectors, but is a lot more resource " +"intensive.\n" +"L|external: Import a face detection bounding box from a json file. " +"(configurable in Detect settings)" +msgstr "" +"R|사용할 감지기. 몇몇 감지기들은 '/config/extract.ini' 또는 '설정 > 추출 플러" +"그인 설정'에서 설정이 가능합니다:\n" +"L|cv2-dnn: 가장 믿을 수 없고 가장 자원을 덜 사용하며 CPU만을 사용하는 추출기" +"입니다. 만약 GPU를 사용하지 않고 시간이 중요하다면 사용하세요.\n" +"L|mtcnn: 좋은 감지기. CPU에서도 빠르고 GPU에서도 빠릅니다. 다른 GPU 감지기들" +"보다 더 적은 자원을 사용하지만 가끔 더 많은 false positives를 돌려줄 수 있습" +"니다.\n" +"L|s3fd: 가장 좋은 감지기. CPU에선 느리고 GPU에선 빠릅니다. 다른 GPU 감지기들" +"보다 더 많은 얼굴들을 감지할 수 있고 과 더 적은 false positives를 돌려주지만 " +"자원을 굉장히 많이 사용합니다.\n" +"L|external: JSON 파일에서 얼굴 감지 경계 박스를 가져옵니다. (설정 감지에서 구" +"성 가능)" + +#: lib/cli/args_extract_convert.py:154 +msgid "" +"R|Aligner to use.\n" +"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " +"but less accurate. Only use this if not using a GPU and time is important.\n" +"L|fan: Best aligner. Fast on GPU, slow on CPU.\n" +"L|external: Import 68 point 2D landmarks or an aligned bounding box from a " +"json file. (configurable in Align settings)" +msgstr "" +"R|사용할 Aligner.\n" +"L|cv2-dnn: CPU만을 사용하는 특징점 감지기. 빠르고 자원을 덜 사용하지만 부정확" +"합니다. GPU를 사용하지 않고 시간이 중요할 때에만 사용하세요.\n" +"L|fan: 가장 좋은 aligner. GPU에선 빠르고 CPU에선 느립니다.\n" +"L|external: JSON 파일에서 68 포인트 2D 랜드 마크 또는 정렬 된 경계 상자를 가" +"져옵니다. (정렬 설정에서 구성 가능)" + +#: lib/cli/args_extract_convert.py:169 +msgid "" +"R|Additional Masker(s) to use. The masks generated here will all take up GPU " +"RAM. You can select none, one or multiple masks, but the extraction may take " +"longer the more you select. NB: The Extended and Components (landmark based) " +"masks are automatically generated on extraction.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU so will not use any additional VRAM.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"The auto generated masks are as follows:\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" +msgstr "" +"R|사용할 추가 Mask입니다. 여기서 생성된 마스크는 모두 GPU RAM을 차지합니다. " +"마스크를 0개, 1개 또는 여러 개 선택할 수 있지만 더 많이 선택할수록 추출에 시" +"간이 더 걸릴 수 있습니다. NB: 확장 및 구성 요소(특징점 기반) 마스크는 추출 " +"시 자동으로 생성됩니다.\n" +"L|bisnet-fp: 전체 헤드 마스킹(마스크 설정에서 구성 가능)을 포함하여 마스킹할 " +"영역에 대한 보다 정교한 제어를 제공하는 비교적 가벼운 NN 기반 마스크입니다.\n" +"L|custom: 마스크 영역을 모든 1 또는 0으로 채우는 dummy 마스크입니다(설정에서 " +"구성 가능). 수동 도구에서 사용자 정의 마스크를 직접 수동으로 편집하려는 경우" +"에만 필요합니다. 이 마스크는 GPU를 사용하지 않으므로 추가 VRAM을 사용하지 않" +"습니다.\n" +"L|vgg-clear: 대부분의 정면에 장애물이 없는 스마트한 분할을 제공하도록 설계된 " +"마스크입니다. 프로필 얼굴들 및 장애물들로 인해 성능이 저하될 수 있습니다.\n" +"L|vgg-obstructed: 대부분의 정면 얼굴을 스마트하게 분할할 수 있도록 설계된 마" +"스크입니다. 마스크 모델은 일부 안면 장애물(손과 안경)을 인식하도록 특별히 훈" +"련되었습니다. 프로필 얼굴들은 평균 이하의 성능을 초래할 수 있습니다.\n" +"L|unet-dfl: 대부분 정면 얼굴을 스마트하게 분할하도록 설계된 마스크. 마스크 모" +"델은 커뮤니티 구성원들에 의해 훈련되었으며 추가 설명을 위해 테스트가 필요하" +"다. 프로필 얼굴들은 평균 이하의 성능을 초래할 수 있습니다.\n" +"자동 생성 마스크는 다음과 같습니다.\n" +"L|components: 특징점 위치의 위치를 기반으로 얼굴 분할을 제공하도록 설계된 마" +"스크입니다. 특징점의 외부에는 마스크를 만들기 위해 convex hull가 형성되어 있" +"습니다.\n" +"L|extended: 특징점 위치의 위치를 기반으로 얼굴 분할을 제공하도록 설계된 마스" +"크입니다. 특징점의 외부에는 convex hull가 형성되어 있으며, 마스크는 이마 위" +"로 뻗어 있습ㄴ다.\n" +"(예: '-M unet-dfl vgg-clear', '--masker vgg-obstructed')" + +#: lib/cli/args_extract_convert.py:208 +msgid "" +"R|Performing normalization can help the aligner better align faces with " +"difficult lighting conditions at an extraction speed cost. Different methods " +"will yield different results on different sets. NB: This does not impact the " +"output face, just the input to the aligner.\n" +"L|none: Don't perform normalization on the face.\n" +"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"L|hist: Equalize the histograms on the RGB channels.\n" +"L|mean: Normalize the face colors to the mean." +msgstr "" +"R|정규화를 수행하면 aligner가 추출 속도 비용으로 어려운 조명 조건의 얼굴을 " +"더 잘 정렬할 수 있습니다. 방법이 다르면 세트마다 결과가 다릅니다. NB: 출력 얼" +"굴에는 영향을 주지 않으며 aligner에 대한 입력에만 영향을 줍니다.\n" +"L|none: 얼굴에 정규화를 수행하지 마십시오.\n" +"L|clahe: 얼굴에 Contrast Limited Adaptive Histogram Equalization를 수행합니" +"다.\n" +"L|hist: RGB 채널의 히스토그램을 동일하게 합니다.\n" +"L|mean: 얼굴 색상을 평균으로 정규화합니다." + +#: lib/cli/args_extract_convert.py:226 +msgid "" +"The number of times to re-feed the detected face into the aligner. Each time " +"the face is re-fed into the aligner the bounding box is adjusted by a small " +"amount. The final landmarks are then averaged from each iteration. Helps to " +"remove 'micro-jitter' but at the cost of slower extraction speed. The more " +"times the face is re-fed into the aligner, the less micro-jitter should " +"occur but the longer extraction will take." +msgstr "" +"검출된 얼굴을 aligner에 다시 공급하는 횟수입니다. 얼굴이 aligner에 다시 공급" +"될 때마다 경계 상자가 소량 조정됩니다. 그런 다음 각 반복에서 최종 특징점의 평" +"균을 구한다. 'micro-jitter'를 제거하는 데 도움이 되지만 추출 속도가 느려집니" +"다. 얼굴이 aligner에 다시 공급되는 횟수가 많을수록 micro-jitter 적게 발생하지" +"만 추출에 더 오랜 시간이 걸립니다." + +#: lib/cli/args_extract_convert.py:239 +msgid "" +"Re-feed the initially found aligned face through the aligner. Can help " +"produce better alignments for faces that are rotated beyond 45 degrees in " +"the frame or are at extreme angles. Slows down extraction." +msgstr "" +"_aligner를 통해 처음 발견된 정렬된 얼굴을 재공급합니다. 프레임에서 45도 이상 " +"회전하거나 극단적인 각도에 있는 얼굴을 더 잘 정렬할 수 있습니다. 추출 속도가 " +"느려집니다." + +#: lib/cli/args_extract_convert.py:249 +msgid "" +"If a face isn't found, rotate the images to try to find a face. Can find " +"more faces at the cost of extraction speed. Pass in a single number to use " +"increments of that size up to 360, or pass in a list of numbers to enumerate " +"exactly what angles to check." +msgstr "" +"얼굴이 발견되지 않으면 이미지를 회전하여 얼굴을 찾습니다. 추출 속도를 희생하" +"면서 더 많은 얼굴을 찾을 수 있습니다. 단일 숫자를 입력하여 해당 크기의 증분" +"을 360까지 사용하거나 숫자 목록을 입력하여 확인할 각도를 정확하게 열거합니다." + +#: lib/cli/args_extract_convert.py:259 +msgid "" +"Obtain and store face identity encodings from VGGFace2. Slows down extract a " +"little, but will save time if using 'sort by face'" +msgstr "" +"VGGFace2에서 얼굴 식별 인코딩을 가져와 저장합니다. 추출 속도를 약간 늦추지만 " +"'얼굴별로 정렬'을 사용하면 시간을 절약할 수 있습니다." + +#: lib/cli/args_extract_convert.py:269 lib/cli/args_extract_convert.py:280 +#: lib/cli/args_extract_convert.py:293 lib/cli/args_extract_convert.py:307 +#: lib/cli/args_extract_convert.py:614 lib/cli/args_extract_convert.py:623 +#: lib/cli/args_extract_convert.py:638 lib/cli/args_extract_convert.py:651 +#: lib/cli/args_extract_convert.py:665 +msgid "Face Processing" +msgstr "얼굴 처리" + +#: lib/cli/args_extract_convert.py:271 +msgid "" +"Filters out faces detected below this size. Length, in pixels across the " +"diagonal of the bounding box. Set to 0 for off" +msgstr "" +"이 크기 미만으로 탐지된 얼굴을 필터링합니다. 길이, 경계 상자의 대각선에 걸친 " +"픽셀 단위입니다. 0으로 설정하면 꺼집니다" + +#: lib/cli/args_extract_convert.py:282 +msgid "" +"Optionally filter out people who you do not wish to extract by passing in " +"images of those people. Should be a small variety of images at different " +"angles and in different conditions. A folder containing the required images " +"or multiple image files, space separated, can be selected." +msgstr "" +"선택적으로 추출하지 않을 사람의 이미지들을 전달하여 그 사람들을 제외합니다. " +"각도와 조건이 다른 작은 다양한 이미지여야 합니다. 추출되지 않는데 필요한 이미" +"지들 또는 공백으로 구분된 여러 이미지 파일이 들어 있는 폴더를 선택할 수 있습" +"니다." + +#: lib/cli/args_extract_convert.py:295 +msgid "" +"Optionally select people you wish to extract by passing in images of that " +"person. Should be a small variety of images at different angles and in " +"different conditions A folder containing the required images or multiple " +"image files, space separated, can be selected." +msgstr "" +"선택적으로 추출하고 싶은 사람의 이미지를 전달하여 그 사람을 선택합니다. 각도" +"와 조건이 다른 작은 다양한 이미지여야 합니다. 추출할 때 필요한 이미지들 또는 " +"공백으로 구분된 여러 이미지 파일이 들어 있는 폴더를 선택할 수 있습니다." + +#: lib/cli/args_extract_convert.py:309 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Higher values are stricter." +msgstr "" +"옵션인 nfilter/filter 파일과 함께 사용합니다. 긍정적인 얼굴 인식을 위한 임계" +"값. 값이 높을수록 엄격합니다." + +#: lib/cli/args_extract_convert.py:318 lib/cli/args_extract_convert.py:331 +#: lib/cli/args_extract_convert.py:344 lib/cli/args_extract_convert.py:356 +msgid "output" +msgstr "출력" + +#: lib/cli/args_extract_convert.py:320 +msgid "" +"The output size of extracted faces. Make sure that the model you intend to " +"train supports your required size. This will only need to be changed for hi-" +"res models." +msgstr "" +"추출된 얼굴의 출력 크기입니다. 훈련하려는 모델이 필요한 크기를 지원하는지 꼭 " +"확인하세요. 이것은 고해상도 모델에 대해서만 변경하면 됩니다." + +#: lib/cli/args_extract_convert.py:333 +msgid "" +"Extract every 'nth' frame. This option will skip frames when extracting " +"faces. For example a value of 1 will extract faces from every frame, a value " +"of 10 will extract faces from every 10th frame." +msgstr "" +"모든 'n번째' 프레임을 추출합니다. 이 옵션은 얼굴을 추출할 때 건너뛸 프레임을 " +"설정합니다. 예를 들어, 값이 1이면 모든 프레임에서 얼굴이 추출되고, 값이 10이" +"면 모든 10번째 프레임에서 얼굴이 추출됩니다." + +#: lib/cli/args_extract_convert.py:346 +msgid "" +"Automatically save the alignments file after a set amount of frames. By " +"default the alignments file is only saved at the end of the extraction " +"process. NB: If extracting in 2 passes then the alignments file will only " +"start to be saved out during the second pass. WARNING: Don't interrupt the " +"script when writing the file because it might get corrupted. Set to 0 to " +"turn off" +msgstr "" +"프레임 수가 설정된 후 alignments 파일을 자동으로 저장합니다. 기본적으로 " +"alignments 파일은 추출 프로세스가 끝날 때만 저장됩니다. NB: 2번째 추출에서 성" +"공하면 두 번째 추출 중에만 alignments 파일이 저장되기 시작합니다. 경고: 파일" +"을 쓸 때 스크립트가 손상될 수 있으므로 스크립트를 중단하지 마십시오. 해제하려" +"면 0으로 설정" + +#: lib/cli/args_extract_convert.py:357 +msgid "Draw landmarks on the ouput faces for debugging purposes." +msgstr "디버깅을 위해 출력 얼굴에 특징점을 그립니다." + +#: lib/cli/args_extract_convert.py:363 lib/cli/args_extract_convert.py:373 +#: lib/cli/args_extract_convert.py:381 lib/cli/args_extract_convert.py:388 +#: lib/cli/args_extract_convert.py:678 lib/cli/args_extract_convert.py:691 +#: lib/cli/args_extract_convert.py:712 lib/cli/args_extract_convert.py:718 +msgid "settings" +msgstr "설정" + +#: lib/cli/args_extract_convert.py:365 +msgid "" +"Don't run extraction in parallel. Will run each part of the extraction " +"process separately (one after the other) rather than all at the same time. " +"Useful if VRAM is at a premium." +msgstr "" +"추출을 병렬로 실행하지 마십시오. 추출 프로세스의 각 부분을 동시에 모두 실행하" +"는 것이 아니라 개별적으로(하나씩) 실행합니다. VRAM이 프리미엄인 경우 유용합니" +"다." + +#: lib/cli/args_extract_convert.py:375 +msgid "" +"Skips frames that have already been extracted and exist in the alignments " +"file" +msgstr "이미 추출되었거나 alignments 파일에 존재하는 프레임들을 스킵합니다" + +#: lib/cli/args_extract_convert.py:382 +msgid "Skip frames that already have detected faces in the alignments file" +msgstr "이미 얼굴을 탐지하여 alignments 파일에 존재하는 프레임들을 스킵합니다" + +#: lib/cli/args_extract_convert.py:389 +msgid "Skip saving the detected faces to disk. Just create an alignments file" +msgstr "" +"탐지된 얼굴을 디스크에 저장하지 않습니다. 그저 alignments 파일을 만듭니다" + +#: lib/cli/args_extract_convert.py:463 +msgid "" +"Swap the original faces in a source video/images to your final faces.\n" +"Conversion plugins can be configured in the 'Settings' Menu" +msgstr "" +"원본 비디오/이미지의 원래 얼굴을 최종 얼굴으로 바꿉니다.\n" +"변환 플러그인은 '설정' 메뉴에서 구성할 수 있습니다" + +#: lib/cli/args_extract_convert.py:485 +msgid "" +"Only required if converting from images to video. Provide The original video " +"that the source frames were extracted from (for extracting the fps and " +"audio)." +msgstr "" +"이미지에서 비디오로 변환하는 경우에만 필요합니다. 소스 프레임이 추출된 원본 " +"비디오(fps 및 오디오 추출용)를 입력하세요." + +#: lib/cli/args_extract_convert.py:494 +msgid "" +"Model directory. The directory containing the trained model you wish to use " +"for conversion." +msgstr "" +"모델 폴더. 당신이 변환에 사용하고자 하는 훈련된 모델을 가진 폴더입니다." + +#: lib/cli/args_extract_convert.py:505 +msgid "" +"R|Performs color adjustment to the swapped face. Some of these options have " +"configurable settings in '/config/convert.ini' or 'Settings > Configure " +"Convert Plugins':\n" +"L|avg-color: Adjust the mean of each color channel in the swapped " +"reconstruction to equal the mean of the masked area in the original image.\n" +"L|color-transfer: Transfers the color distribution from the source to the " +"target image using the mean and standard deviations of the L*a*b* color " +"space.\n" +"L|manual-balance: Manually adjust the balance of the image in a variety of " +"color spaces. Best used with the Preview tool to set correct values.\n" +"L|match-hist: Adjust the histogram of each color channel in the swapped " +"reconstruction to equal the histogram of the masked area in the original " +"image.\n" +"L|seamless-clone: Use cv2's seamless clone function to remove extreme " +"gradients at the mask seam by smoothing colors. Generally does not give very " +"satisfactory results.\n" +"L|none: Don't perform color adjustment." +msgstr "" +"R|스왑된 얼굴의 색상 조정을 수행합니다. 이러한 옵션 중 일부에는 '/config/" +"convert.ini' 또는 '설정 > 변환 플러그인 구성'에서 구성 가능한 설정이 있습니" +"다.\n" +"L|avg-color: 스왑된 재구성에서 각 색상 채널의 평균이 원본 영상에서 마스킹된 " +"영역의 평균과 동일하도록 조정합니다.\n" +"L|color-transfer: L*a*b* 색 공간의 평균 및 표준 편차를 사용하여 소스에서 대" +"상 이미지로 색 분포를 전송합니다.\n" +"L|manual-balance: 다양한 색 공간에서 이미지의 밸런스를 수동으로 조정합니다. " +"올바른 값을 설정하려면 미리 보기 도구와 함께 사용하는 것이 좋습니다.\n" +"L|match-hist: 스왑된 재구성에서 각 색상 채널의 히스토그램을 조정하여 원래 영" +"상에서 마스킹된 영역의 히스토그램과 동일하게 만듭니다.\n" +"L|seamless-clone: cv2의 원활한 복제 기능을 사용하여 색상을 평활화하여 마스크 " +"심에서 극단적인 gradients을 제거합니다. 일반적으로 매우 만족스러운 결과를 제" +"공하지 않습니다.\n" +"L|none: 색상 조정을 수행하지 않습니다." + +#: lib/cli/args_extract_convert.py:531 +msgid "" +"R|Masker to use. NB: The mask you require must exist within the alignments " +"file. You can add additional masks with the Mask Tool.\n" +"L|none: Don't use a mask.\n" +"L|bisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"L|bisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"L|custom_face: Custom user created, face centered mask.\n" +"L|custom_head: Custom user created, head centered mask.\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"L|predicted: If the 'Learn Mask' option was enabled during training, this " +"will use the mask that was created by the trained model." +msgstr "" +"R|사용할 마스크. NB: 필요한 마스크는 alignments 파일 내에 있어야 합니다. 마스" +"크 도구를 사용하여 마스크를 추가할 수 있습니다.\n" +"L|none: 마스크 쓰지 마세요.\n" +"L|bisnet-fp_face: 마스크할 영역을 보다 정교하게 제어할 수 있는 비교적 가벼운 " +"NN 기반 마스크입니다(마스크 설정에서 구성 가능). 모델이 '얼굴' 또는 '레거시' " +"중심으로 훈련된 경우 이 버전의 bisnet-fp를 사용하십시오.\n" +"L|bisnet-fp_head: 마스크할 영역을 보다 정교하게 제어할 수 있는 비교적 가벼운 " +"NN 기반 마스크입니다(마스크 설정에서 구성 가능). 모델이 '헤드' 중심으로 훈련" +"된 경우 이 버전의 bisnet-fp를 사용하십시오.\n" +"L|custom_face: 사용자 지정 사용자가 생성한 얼굴 중심 마스크입니다.\n" +"L|custom_head: 사용자 지정 사용자가 생성한 머리 중심 마스크입니다.\n" +"L|components: 특징점 위치의 배치를 기반으로 얼굴 분할을 제공하도록 설계된 마" +"스크입니다. 특징점의 외부에는 마스크를 만들기 위해 convex hull가 형성되어 있" +"습니다.\n" +"L|extended: 특징점 위치의 배치를 기반으로 얼굴 분할을 제공하도록 설계된 마스" +"크입니다. 지형지물의 외부에는 convex hull가 형성되어 있으며, 마스크는 이마 위" +"로 뻗어 있습니다.\n" +"L|vgg-clear: 대부분의 정면에 장애물이 없는 스마트한 분할을 제공하도록 설계된 " +"마스크입니다. 옆 얼굴 및 장애물로 인해 성능이 저하될 수 있습니다.\n" +"L|vgg-obstructed: 대부분의 정면 얼굴을 스마트하게 분할할 수 있도록 설계된 마" +"스크입니다. 마스크 모델은 일부 안면 장애물(손과 안경)을 인식하도록 특별히 훈" +"련되었습니다. 옆 얼굴은 평균 이하의 성능을 초래할 수 있습니다.\n" +"L|unet-dfl: 대부분 정면 얼굴을 스마트하게 분할하도록 설계된 마스크. 마스크 모" +"델은 커뮤니티 구성원들에 의해 훈련되었으며 추가 설명을 위해 테스트가 필요하" +"다. 옆 얼굴은 평균 이하의 성능을 초래할 수 있습니다.\n" +"L|predicted: 교육 중에 'Learn Mask(마스크 학습)' 옵션이 활성화된 경우에는 교" +"육을 받은 모델이 만든 마스크가 사용됩니다." + +#: lib/cli/args_extract_convert.py:570 +msgid "" +"R|The plugin to use to output the converted images. The writers are " +"configurable in '/config/convert.ini' or 'Settings > Configure Convert " +"Plugins:'\n" +"L|ffmpeg: [video] Writes out the convert straight to video. When the input " +"is a series of images then the '-ref' (--reference-video) parameter must be " +"set.\n" +"L|gif: [animated image] Create an animated gif.\n" +"L|opencv: [images] The fastest image writer, but less options and formats " +"than other plugins.\n" +"L|patch: [images] Outputs the raw swapped face patch, along with the " +"transformation matrix required to re-insert the face back into the original " +"frame. Use this option if you wish to post-process and composite the final " +"face within external tools.\n" +"L|pillow: [images] Slower than opencv, but has more options and supports " +"more formats." +msgstr "" +"R|변환된 이미지를 출력하는 데 사용할 플러그인입니다. 기록 장치는 '/config/" +"convert.ini' 또는 '설정 > 변환 플러그인 구성:'에서 구성할 수 있습니다.\n" +"L|ffmpeg: [video] 변환된 결과를 바로 video로 씁니다. 입력이 영상 시리즈인 경" +"우 '-ref'(--reference-video) 파라미터를 설정해야 합니다.\n" +"L|gif : [애니메이션 이미지] 애니메이션 gif를 만듭니다.\n" +"L|opencv: [이미지] 가장 빠른 이미지 작성기이지만 다른 플러그인에 비해 옵션과 " +"형식이 적습니다.\n" +"L|patch: [이미지] 원래 프레임에 얼굴을 다시 삽입하는 데 필요한 변환 행렬과 함" +"께 원시 교체된 얼굴 패치를 출력합니다.\n" +"L|pillow: [images] opencv보다 느리지만 더 많은 옵션이 있고 더 많은 형식을 지" +"원합니다." + +#: lib/cli/args_extract_convert.py:591 lib/cli/args_extract_convert.py:600 +#: lib/cli/args_extract_convert.py:703 +msgid "Frame Processing" +msgstr "프레임 처리" + +#: lib/cli/args_extract_convert.py:593 +#, python-format +msgid "" +"Scale the final output frames by this amount. 100%% will output the frames " +"at source dimensions. 50%% at half size 200%% at double size" +msgstr "" +"최종 출력 프레임의 크기를 이 양만큼 조정합니다. 100%%는 원본의 차원에서 프레" +"임을 출력합니다. 50%%는 절반 크기에서, 200%%는 두 배 크기에서" + +#: lib/cli/args_extract_convert.py:602 +msgid "" +"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " +"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " +"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " +"converting from images, then the filenames must end with the frame-number!" +msgstr "" +"예를 들어 전송을 적용할 프레임 범위 프레임 10 - 50 및 90 - 100의 경우 --" +"frame-ranges 10-50 90-100을 사용합니다. '-k'(--keep-unchanged)를 선택하지 않" +"으면 선택한 범위를 벗어나는 프레임이 삭제됩니다. NB: 이미지에서 변환하는 경" +"우 파일 이름은 프레임 번호로 끝나야 합니다!" + +#: lib/cli/args_extract_convert.py:616 +msgid "" +"Scale the swapped face by this percentage. Positive values will enlarge the " +"face, Negative values will shrink the face." +msgstr "" +"이 백분율로 교체된 면의 크기를 조정합니다. 양수 값은 얼굴을 확대하고, 음수 값" +"은 얼굴을 축소합니다." + +#: lib/cli/args_extract_convert.py:625 +msgid "" +"If you have not cleansed your alignments file, then you can filter out faces " +"by defining a folder here that contains the faces extracted from your input " +"files/video. If this folder is defined, then only faces that exist within " +"your alignments file and also exist within the specified folder will be " +"converted. Leaving this blank will convert all faces that exist within the " +"alignments file." +msgstr "" +"만약 alignments 파일을 지우지 않은 경우 입력 파일/비디오에서 추출된 얼굴이 포" +"함된 폴더를 정의하여 얼굴을 걸러낼 수 있습니다. 이 폴더가 정의된 경우 " +"alignments 파일 내에 존재하거나 지정된 폴더 내에 존재하는 얼굴만 변환됩니다. " +"이 항목을 공백으로 두면 alignments 파일 내에 있는 모든 얼굴이 변환됩니다." + +#: lib/cli/args_extract_convert.py:640 +msgid "" +"Optionally filter out people who you do not wish to process by passing in an " +"image of that person. Should be a front portrait with a single person in the " +"image. Multiple images can be added space separated. NB: Using face filter " +"will significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"선택적으로 처리하고 싶지 않은 사람의 이미지를 전달하여 그 사람을 걸러낼 수 있" +"습니다. 이미지는 한 사람의 정면 모습이여야 합니다. 여러 이미지를 공백으로 구" +"분하여 추가할 수 있습니다. 주의: 얼굴 필터를 사용하면 추출 속도가 현저히 감소" +"하므로 정확성을 보장할 수 없습니다." + +#: lib/cli/args_extract_convert.py:653 +msgid "" +"Optionally select people you wish to process by passing in an image of that " +"person. Should be a front portrait with a single person in the image. " +"Multiple images can be added space separated. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"선택적으로 해당 사용자의 이미지를 전달하여 처리할 사용자를 선택합니다. 이미지" +"에 한 사람이 있는 정면 초상화여야 합니다. 여러 이미지를 공백으로 구분하여 추" +"가할 수 있습니다. 주의: 얼굴 필터를 사용하면 추출 속도가 현저히 감소하므로 정" +"확성을 보장할 수 없습니다." + +#: lib/cli/args_extract_convert.py:667 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Lower values are stricter. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"옵션인 nfilter/filter 파일을 함께 사용합니다. 긍정적인 얼굴 인식을 위한 임계" +"값. 낮은 값이 더 엄격합니다. 주의: 얼굴 필터를 사용하면 추출 속도가 현저히 감" +"소하므로 정확성을 보장할 수 없습니다." + +#: lib/cli/args_extract_convert.py:680 +msgid "" +"The maximum number of parallel processes for performing conversion. " +"Converting images is system RAM heavy so it is possible to run out of memory " +"if you have a lot of processes and not enough RAM to accommodate them all. " +"Setting this to 0 will use the maximum available. No matter what you set " +"this to, it will never attempt to use more processes than are available on " +"your system. If singleprocess is enabled this setting will be ignored." +msgstr "" +"변환을 수행하기 위한 최대 병렬 프로세스 수입니다. 이미지 변환은 시스템 RAM에 " +"부담이 크기 때문에 프로세스가 많고 모든 프로세스를 수용할 RAM이 충분하지 않" +"은 경우 메모리가 부족할 수 있습니다. 이것을 0으로 설정하면 사용 가능한 최대값" +"을 사용합니다. 얼마를 설정하든 시스템에서 사용 가능한 것보다 더 많은 프로세스" +"를 사용하려고 시도하지 않습니다. 단일 프로세스가 활성화된 경우 이 설정은 무시" +"됩니다." + +#: lib/cli/args_extract_convert.py:693 +msgid "" +"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " +"alignments file for your destination video. However, if you wish you can " +"generate the alignments on-the-fly by enabling this option. This will use an " +"inferior extraction pipeline and will lead to substandard results. If an " +"alignments file is found, this option will be ignored." +msgstr "" +"실시간 변환을 활성화합니다. 권장하지 않습니다. 당신은 변환 비디오에 대한 깨끗" +"한 alignments 파일을 생성해야 합니다. 그러나 원하는 경우 이 옵션을 활성화하" +"여 즉시 alignments 파일을 생성할 수 있습니다. 이것은 안좋은 추출 과정을 사용" +"하고 표준 이하의 결과로 이어질 것입니다. alignments 파일이 발견되면 이 옵션" +"은 무시됩니다." + +#: lib/cli/args_extract_convert.py:705 +msgid "" +"When used with --frame-ranges outputs the unchanged frames that are not " +"processed instead of discarding them." +msgstr "" +"사용시 --frame-ranges 인자를 사용하면 변경되지 않은 프레임을 버리지 않은 결과" +"가 출력됩니다." + +#: lib/cli/args_extract_convert.py:713 +msgid "Swap the model. Instead converting from of A -> B, converts B -> A" +msgstr "모델을 바꿉니다. A -> B에서 변환하는 대신 B -> A로 변환" + +#: lib/cli/args_extract_convert.py:719 +msgid "Disable multiprocessing. Slower but less resource intensive." +msgstr "멀티프로세싱을 쓰지 않습니다. 느리지만 자원을 덜 소모합니다." + +#~ msgid "" +#~ "[LEGACY] This only needs to be selected if a legacy model is being loaded " +#~ "or if there are multiple models in the model folder" +#~ msgstr "" +#~ "[LEGACY] 이것은 레거시 모델을 로드 중이거나 모델 폴더에 여러 모델이 있는 " +#~ "경우에만 선택되어야 합니다" diff --git a/locales/kr/LC_MESSAGES/lib.cli.args_train.mo b/locales/kr/LC_MESSAGES/lib.cli.args_train.mo new file mode 100644 index 0000000000..1e5985f335 Binary files /dev/null and b/locales/kr/LC_MESSAGES/lib.cli.args_train.mo differ diff --git a/locales/kr/LC_MESSAGES/lib.cli.args_train.po b/locales/kr/LC_MESSAGES/lib.cli.args_train.po new file mode 100644 index 0000000000..c90082345d --- /dev/null +++ b/locales/kr/LC_MESSAGES/lib.cli.args_train.po @@ -0,0 +1,350 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:04+0000\n" +"PO-Revision-Date: 2024-03-28 18:16+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_train.py:30 +msgid "" +"Train a model on extracted original (A) and swap (B) faces.\n" +"Training models can take a long time. Anything from 24hrs to over a week\n" +"Model plugins can be configured in the 'Settings' Menu" +msgstr "" +"추출된 원래(A) 얼굴과 스왑(B) 얼굴에 대한 모델을 훈련합니다.\n" +"모델을 훈련하는 데 시간이 오래 걸릴 수 있습니다. 24시간에서 일주일 이상의 시" +"간이 필요합니다.\n" +"모델 플러그인은 '설정' 메뉴에서 구성할 수 있습니다" + +#: lib/cli/args_train.py:49 lib/cli/args_train.py:58 +msgid "faces" +msgstr "얼굴들" + +#: lib/cli/args_train.py:51 +msgid "" +"Input directory. A directory containing training images for face A. This is " +"the original face, i.e. the face that you want to remove and replace with " +"face B." +msgstr "" +"입력 디렉토리. 얼굴 A에 대한 훈련 이미지가 포함된 디렉토리입니다. 이것은 원" +"래 얼굴, 즉 제거하고 B 얼굴로 대체하려는 얼굴입니다." + +#: lib/cli/args_train.py:60 +msgid "" +"Input directory. A directory containing training images for face B. This is " +"the swap face, i.e. the face that you want to place onto the head of person " +"A." +msgstr "" +"입력 디렉터리. 얼굴 B에 대한 훈련 이미지를 포함하는 디렉토리. 이것은 대체 얼" +"굴, 즉 사람 A의 얼굴 앞에 배치하려는 얼굴이다." + +#: lib/cli/args_train.py:67 lib/cli/args_train.py:80 lib/cli/args_train.py:97 +#: lib/cli/args_train.py:123 lib/cli/args_train.py:133 +msgid "model" +msgstr "모델" + +#: lib/cli/args_train.py:69 +msgid "" +"Model directory. This is where the training data will be stored. You should " +"always specify a new folder for new models. If starting a new model, select " +"either an empty folder, or a folder which does not exist (which will be " +"created). If continuing to train an existing model, specify the location of " +"the existing model." +msgstr "" +"모델 디렉토리. 여기에 훈련 데이터가 저장됩니다. 새 모델의 경우 항상 새 폴더" +"를 지정해야 합니다. 새 모델을 시작할 경우 빈 폴더 또는 존재하지 않는 폴더(생" +"성될 폴더)를 선택합니다. 기존 모델을 계속 학습하는 경우 기존 모델의 위치를 지" +"정합니다." + +#: lib/cli/args_train.py:82 +msgid "" +"R|Load the weights from a pre-existing model into a newly created model. For " +"most models this will load weights from the Encoder of the given model into " +"the encoder of the newly created model. Some plugins may have specific " +"configuration options allowing you to load weights from other layers. " +"Weights will only be loaded when creating a new model. This option will be " +"ignored if you are resuming an existing model. Generally you will also want " +"to 'freeze-weights' whilst the rest of your model catches up with your " +"Encoder.\n" +"NB: Weights can only be loaded from models of the same plugin as you intend " +"to train." +msgstr "" +"R|기존 모델의 가중치를 새로 생성된 모델로 로드합니다. 대부분의 모델에서는 주" +"어진 모델의 인코더에서 새로 생성된 모델의 인코더로 가중치를 로드합니다. 일부 " +"플러그인에는 다른 층에서 가중치를 로드할 수 있는 특정 구성 옵션이 있을 수 있" +"습니다. 가중치는 새 모델을 생성할 때만 로드됩니다. 기존 모델을 재개하는 경우 " +"이 옵션은 무시됩니다. 일반적으로 나머지 모델이 인코더를 따라잡는 동안에도 '가" +"중치 동결'이 필요합니다.\n" +"주의: 가중치는 훈련하려는 플러그인 모델에서만 로드할 수 있습니다." + +#: lib/cli/args_train.py:99 +msgid "" +"R|Select which trainer to use. Trainers can be configured from the Settings " +"menu or the config folder.\n" +"L|original: The original model created by /u/deepfakes.\n" +"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " +"for full dfaker method.\n" +"L|dfl-h128: 128px in/out model from deepfacelab\n" +"L|dfl-sae: Adaptable model from deepfacelab\n" +"L|dlight: A lightweight, high resolution DFaker variant.\n" +"L|iae: A model that uses intermediate layers to try to get better details\n" +"L|lightweight: A lightweight model for low-end cards. Don't expect great " +"results. Can train as low as 1.6GB with batch size 8.\n" +"L|realface: A high detail, dual density model based on DFaker, with " +"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " +"won't work so well. By andenixa et al. Very configurable.\n" +"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " +"unbalanced so B>A swaps won't work so well. Very configurable.\n" +"L|villain: 128px in/out model from villainguy. Very resource hungry (You " +"will require a GPU with a fair amount of VRAM). Good for details, but more " +"susceptible to color differences." +msgstr "" +"R|사용할 훈련 모델을 선택합니다. 훈련 모델은 설정 메뉴 또는 구성 폴더에서 구" +"성할 수 있습니다.\n" +"L|original: /u/deepfakes로 만든 원래 모델입니다.\n" +"L|dfaker: 64px in/128px out 모델 from dfaker. Full dfaker 메서드에 대해 '특징" +"점으로 변환'를 활성화합니다.\n" +"L|dfl-h128: Deepfake lab의 128px in/out 모델\n" +"L|dfl-sae: Deepface Lab의 적응형 모델\n" +"L|dlight: 경량, 고해상도 DFaker 변형입니다.\n" +"L|iae: 중간 층들을 사용하여 더 나은 세부 정보를 얻기 위해 노력하는 모델.\n" +"L|lightweight: 저가형 카드용 경량 모델. 좋은 결과를 기대하지 마세요. 최대한 " +"낮게 잡아서 배치 사이즈 8에 1.6GB까지 훈련이 가능합니다.\n" +"L|realface: DFaker를 기반으로 한 높은 디테일의 이중 밀도 모델로, 사용자 정의 " +"가능한 입/출력 해상도를 제공합니다. 오토인코더가 불균형하여 B>A 스왑이 잘 작" +"동하지 않습니다. Andenixa 등에 의해. 매우 구성 가능합니다.\n" +"L|unbalanced: andenixa의 128px in/out 모델. 오토인코더가 불균형하여 B>A 스왑" +"이 잘 작동하지 않습니다. 매우 구성 가능합니다.\n" +"L|villain : villainguy의 128px in/out 모델. 리소스가 매우 부족합니다( 상당한 " +"양의 VRAM이 있는 GPU가 필요합니다). 세부 사항에는 좋지만 색상 차이에 더 취약" +"합니다." + +#: lib/cli/args_train.py:125 +msgid "" +"Output a summary of the model and exit. If a model folder is provided then a " +"summary of the saved model is displayed. Otherwise a summary of the model " +"that would be created by the chosen plugin and configuration settings is " +"displayed." +msgstr "" +"모델 요약을 출력하고 종료합니다. 모델 폴더가 제공되면 저장된 모델의 요약이 표" +"시됩니다. 그렇지 않으면 선택한 플러그인 및 구성 설정에 의해 생성되는 모델 요" +"약이 표시됩니다." + +#: lib/cli/args_train.py:135 +msgid "" +"Freeze the weights of the model. Freezing weights means that some of the " +"parameters in the model will no longer continue to learn, but those that are " +"not frozen will continue to learn. For most models, this will freeze the " +"encoder, but some models may have configuration options for freezing other " +"layers." +msgstr "" +"모델의 가중치를 동결합니다. 가중치를 고정하면 모델의 일부 매개변수가 더 이상 " +"학습되지 않지만 고정되지 않은 매개변수는 계속 학습됩니다. 대부분의 모델에서 " +"이렇게 하면 인코더가 고정되지만 일부 모델에는 다른 레이어를 고정하기 위한 구" +"성 옵션이 있을 수 있습니다." + +#: lib/cli/args_train.py:147 lib/cli/args_train.py:160 +#: lib/cli/args_train.py:175 lib/cli/args_train.py:191 +#: lib/cli/args_train.py:200 +msgid "training" +msgstr "훈련" + +#: lib/cli/args_train.py:149 +msgid "" +"Batch size. This is the number of images processed through the model for " +"each side per iteration. NB: As the model is fed 2 sides at a time, the " +"actual number of images within the model at any one time is double the " +"number that you set here. Larger batches require more GPU RAM." +msgstr "" +"배치 크기. 반복당 각 측면에 대해 모델을 통해 처리되는 이미지 수입니다. NB: " +"한 번에 모델에게 2개의 측면이 공급되므로 한 번에 모델 내의 실제 이미지 수는 " +"여기에서 설정한 수의 두 배입니다. 더 큰 배치에는 더 많은 GPU RAM이 필요합니" +"다." + +#: lib/cli/args_train.py:162 +msgid "" +"Length of training in iterations. This is only really used for automation. " +"There is no 'correct' number of iterations a model should be trained for. " +"You should stop training when you are happy with the previews. However, if " +"you want the model to stop automatically at a set number of iterations, you " +"can set that value here." +msgstr "" +"반복에서 훈련 길이. 이것은 실제로 자동화에만 사용됩니다. 모델을 훈련해야 하" +"는 '올바른' 반복 횟수는 없습니다. 미리 보기에 만족하면 훈련을 중단해야 합니" +"다. 그러나 설정된 반복 횟수에서 모델이 자동으로 중지되도록 하려면 여기에서 해" +"당 값을 설정할 수 있습니다." + +#: lib/cli/args_train.py:177 +msgid "" +"R|Select the distribution stategy to use.\n" +"L|default: Use Tensorflow's default distribution strategy.\n" +"L|central-storage: Centralizes variables on the CPU whilst operations are " +"performed on 1 or more local GPUs. This can help save some VRAM at the cost " +"of some speed by not storing variables on the GPU. Note: Mixed-Precision is " +"not supported on multi-GPU setups.\n" +"L|mirrored: Supports synchronous distributed training across multiple local " +"GPUs. A copy of the model and all variables are loaded onto each GPU with " +"batches distributed to each GPU at each iteration." +msgstr "" +"R|사용할 배포 상태를 선택합니다.\n" +"L|default: Tensorflow의 기본 배포 전략을 사용합니다.\n" +"L|central-storage: 작업이 1개 이상의 로컬 GPU에서 수행되는 동안 CPU의 변수를 " +"중앙 집중화합니다. 이렇게 하면 GPU에 변수를 저장하지 않음으로써 약간의 속도" +"를 희생하여 일부 VRAM을 절약할 수 있습니다. 참고: 다중 정밀도는 다중 GPU 설정" +"에서 지원되지 않습니다.\n" +"L|mirrored: 여러 로컬 GPU에서 동기화 분산 훈련을 지원합니다. 모델의 복사본과 " +"모든 변수는 각 반복에서 각 GPU에 배포된 배치들와 함께 각 GPU에 로드됩니다." + +#: lib/cli/args_train.py:193 +msgid "" +"Disables TensorBoard logging. NB: Disabling logs means that you will not be " +"able to use the graph or analysis for this session in the GUI." +msgstr "" +"텐서보드 로깅을 비활성화합니다. 주의: 로그를 비활성화하면 GUI에서 이 세션에 " +"대한 그래프 또는 분석을 사용할 수 없습니다." + +#: lib/cli/args_train.py:202 +msgid "" +"Use the Learning Rate Finder to discover the optimal learning rate for " +"training. For new models, this will calculate the optimal learning rate for " +"the model. For existing models this will use the optimal learning rate that " +"was discovered when initializing the model. Setting this option will ignore " +"the manually configured learning rate (configurable in train settings)." +msgstr "" +"학습률 찾기를 사용하여 훈련을 위한 최적의 학습률을 찾아보세요. 새 모델의 경" +"우 모델에 대한 최적의 학습률을 계산합니다. 기존 모델의 경우 모델을 초기화할 " +"때 발견된 최적의 학습률을 사용합니다. 이 옵션을 설정하면 수동으로 구성된 학습" +"률(기차 설정에서 구성 가능)이 무시됩니다." + +#: lib/cli/args_train.py:215 lib/cli/args_train.py:225 +msgid "Saving" +msgstr "저장" + +#: lib/cli/args_train.py:216 +msgid "Sets the number of iterations between each model save." +msgstr "각 모델 저장 사이의 반복 횟수를 설정합니다." + +#: lib/cli/args_train.py:227 +msgid "" +"Sets the number of iterations before saving a backup snapshot of the model " +"in it's current state. Set to 0 for off." +msgstr "" +"현재 상태에서 모델의 백업 스냅샷을 저장하기 전에 반복할 횟수를 설정합니다. 0" +"으로 설정하면 꺼집니다." + +#: lib/cli/args_train.py:234 lib/cli/args_train.py:246 +#: lib/cli/args_train.py:258 +msgid "timelapse" +msgstr "타임랩스" + +#: lib/cli/args_train.py:236 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'A' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-B parameter." +msgstr "" +"타임랩스를 만드는 옵션입니다. Timelapse(시간 경과)는 저장을 반복할 때마다 선" +"택한 얼굴의 이미지를 Timelapse-output(시간 경과 출력) 폴더에 저장합니다. 타임" +"랩스를 만드는 데 사용할 'A' 얼굴의 입력 폴더여야 합니다. 또한 사용자는 --" +"timelapse-output 및 --timelapse-input-B 매개 변수를 제공해야 합니다." + +#: lib/cli/args_train.py:248 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'B' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-A parameter." +msgstr "" +"타임 랩스를 만드는 데 선택적입니다. Timelapse(시간 경과)는 저장을 반복할 때마" +"다 선택한 얼굴의 이미지를 Timelapse-output(시간 경과 출력) 폴더에 저장합니" +"다. 타임 랩스를 만드는 데 사용할 'B' 얼굴의 입력 폴더여야 합니다. 또한 사용자" +"는 --timelapse-output 및 --timelapse-input-A 매개 변수를 제공해야 합니다." + +#: lib/cli/args_train.py:260 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. If " +"the input folders are supplied but no output folder, it will default to your " +"model folder/timelapse/" +msgstr "" +"타임랩스를 만드는 데 선택적입니다. Timelapse(시간 경과)는 저장을 반복할 때마" +"다 선택한 얼굴의 이미지를 Timelapse-output(시간 경과 출력) 폴더에 저장합니" +"다. 입력 폴더가 제공되었지만 출력 폴더가 없는 경우 모델 폴더에/timelapse/로 " +"기본 설정됩니다" + +#: lib/cli/args_train.py:269 lib/cli/args_train.py:276 +msgid "preview" +msgstr "미리보기" + +#: lib/cli/args_train.py:270 +msgid "Show training preview output. in a separate window." +msgstr "훈련 미리보기 결과를 각기 다른 창에서 보여줍니다." + +#: lib/cli/args_train.py:278 +msgid "" +"Writes the training result to a file. The image will be stored in the root " +"of your FaceSwap folder." +msgstr "" +"훈련 결과를 파일에 씁니다. 이미지는 Faceswap 폴더의 최상위 폴더에 저장됩니다." + +#: lib/cli/args_train.py:285 lib/cli/args_train.py:295 +#: lib/cli/args_train.py:305 lib/cli/args_train.py:315 +msgid "augmentation" +msgstr "보정" + +#: lib/cli/args_train.py:287 +msgid "" +"Warps training faces to closely matched Landmarks from the opposite face-set " +"rather than randomly warping the face. This is the 'dfaker' way of doing " +"warping." +msgstr "" +"무작위로 얼굴을 변환하지 않고 반대쪽 얼굴 세트에서 특징점과 밀접하게 일치하도" +"록 훈련 얼굴을 변환해줍니다. 이것은 변환하는 'dfaker' 방식이다." + +#: lib/cli/args_train.py:297 +msgid "" +"To effectively learn, a random set of images are flipped horizontally. " +"Sometimes it is desirable for this not to occur. Generally this should be " +"left off except for during 'fit training'." +msgstr "" +"효과적으로 학습하기 위해 임의의 이미지 세트를 수평으로 뒤집습니다. 때때로 이" +"런 일이 일어나지 않는 것이 바람직합니다. 일반적으로 'fit training' 중을 제외" +"하고는 이 작업을 중단해야 합니다." + +#: lib/cli/args_train.py:307 +msgid "" +"Color augmentation helps make the model less susceptible to color " +"differences between the A and B sets, at an increased training time cost. " +"Enable this option to disable color augmentation." +msgstr "" +"색상 보정은 모델이 A와 B 세트 사이의 색상 차이에 덜 민감하게 만드는 데 도움" +"이 되며, 훈련 시간 비용이 증가합니다. 색상 보저를 사용하지 않으려면 이 옵션" +"을 사용합니다." + +#: lib/cli/args_train.py:317 +msgid "" +"Warping is integral to training the Neural Network. This option should only " +"be enabled towards the very end of training to try to bring out more detail. " +"Think of it as 'fine-tuning'. Enabling this option from the beginning is " +"likely to kill a model and lead to terrible results." +msgstr "" +"변환은 신경망을 훈련하는 데 필수적입니다. 이 옵션은 보다 세부적인 것들을 뽑아" +"내위하여 훈련 막바지까지 활성화하여야 합니다. 이것은 '미세 조정'이라고 생각하" +"면 됩니다. 처음부터 이 옵션을 활성화하면 모델이 죽을 수있고 끔찍한 결과를 초" +"래할 수 있습니다." diff --git a/locales/kr/LC_MESSAGES/tools.alignments.cli.mo b/locales/kr/LC_MESSAGES/tools.alignments.cli.mo new file mode 100644 index 0000000000..0c4f74d355 Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.alignments.cli.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.alignments.cli.po b/locales/kr/LC_MESSAGES/tools.alignments.cli.po new file mode 100644 index 0000000000..1bedaa0a22 --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.alignments.cli.po @@ -0,0 +1,254 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-19 11:28+0100\n" +"PO-Revision-Date: 2024-04-19 11:30+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/alignments/cli.py:16 +msgid "" +"This command lets you perform various tasks pertaining to an alignments file." +msgstr "" +"이 명령을 사용하여 alignments 파일과 관련된 다양한 작ㅇ를 수행할 수 있습니다." + +#: tools/alignments/cli.py:31 +msgid "" +"Alignments tool\n" +"This tool allows you to perform numerous actions on or using an alignments " +"file against its corresponding faceset/frame source." +msgstr "" +"_alignments 도구\n" +"이 도구를 사용하면 해당 얼굴 세트/프레임 원본에 해당하는 alignments 파일을 사" +"용하거나 여러 작업을 수행할 수 있습니다." + +#: tools/alignments/cli.py:43 +msgid " Must Pass in a frames folder/source video file (-r)." +msgstr "" +" 프레임들이 저장된 폴더나 원본 비디오 파일을 무조건 전달해야 합니다 (-r)." + +#: tools/alignments/cli.py:44 +msgid " Must Pass in a faces folder (-c)." +msgstr " 얼굴 폴더를 무조건 전달해야 합니다 (-c)." + +#: tools/alignments/cli.py:45 +msgid "" +" Must Pass in either a frames folder/source video file OR a faces folder (-r " +"or -c)." +msgstr "" +" 프레임 폴더나 원본 비디오 파일 또는 얼굴 폴더중 하나를 무조건 전달해야 합니" +"다 (-r and -c)." + +#: tools/alignments/cli.py:47 +msgid "" +" Must Pass in a frames folder/source video file AND a faces folder (-r and -" +"c)." +msgstr "" +" 프레임 폴더나 원본 비디오 파일 그리고 얼굴 폴더를 무조건 전달해야 합니다 (-" +"r and -c)." + +#: tools/alignments/cli.py:49 +msgid " Use the output option (-o) to process results." +msgstr " 결과를 진행하려면 (-o) 출력 옵션을 사용하세요." + +#: tools/alignments/cli.py:58 tools/alignments/cli.py:104 +msgid "processing" +msgstr "처리" + +#: tools/alignments/cli.py:61 +#, python-brace-format +msgid "" +"R|Choose which action you want to perform. NB: All actions require an " +"alignments file (-a) to be passed in.\n" +"L|'draw': Draw landmarks on frames in the selected folder/video. A subfolder " +"will be created within the frames folder to hold the output.{0}\n" +"L|'export': Export the contents of an alignments file to a json file. Can be " +"used for editing alignment information in external tools and then re-" +"importing by using Faceswap's Extract 'Import' plugins. Note: masks and " +"identity vectors will not be included in the exported file, so will be re-" +"generated when the json file is imported back into Faceswap. All data is " +"exported with the origin (0, 0) at the top left of the canvas.\n" +"L|'extract': Re-extract faces from the source frames/video based on " +"alignment data. This is a lot quicker than re-detecting faces. Can pass in " +"the '-een' (--extract-every-n) parameter to only extract every nth frame." +"{1}\n" +"L|'from-faces': Generate alignment file(s) from a folder of extracted faces. " +"if the folder of faces comes from multiple sources, then multiple alignments " +"files will be created. NB: for faces which have been extracted from folders " +"of source images, rather than a video, a single alignments file will be " +"created as there is no way for the process to know how many folders of " +"images were originally used. You do not need to provide an alignments file " +"path to run this job. {3}\n" +"L|'missing-alignments': Identify frames that do not exist in the alignments " +"file.{2}{0}\n" +"L|'missing-frames': Identify frames in the alignments file that do not " +"appear within the frames folder/video.{2}{0}\n" +"L|'multi-faces': Identify where multiple faces exist within the alignments " +"file.{2}{4}\n" +"L|'no-faces': Identify frames that exist within the alignment file but no " +"faces were detected.{2}{0}\n" +"L|'remove-faces': Remove deleted faces from an alignments file. The original " +"alignments file will be backed up.{3}\n" +"L|'rename' - Rename faces to correspond with their parent frame and position " +"index in the alignments file (i.e. how they are named after running extract)." +"{3}\n" +"L|'sort': Re-index the alignments from left to right. For alignments with " +"multiple faces this will ensure that the left-most face is at index 0.\n" +"L|'spatial': Perform spatial and temporal filtering to smooth alignments " +"(EXPERIMENTAL!)" +msgstr "" +"R|실행할 작업을 선택합니다. 주의: 모든 작업을 수행하려면 alignments 파일(-a)" +"을 전달해야 합니다.\n" +"L|'draw': 선택한 폴더/비디오의 프레임에 특징점을 그립니다. 출력을 저장할 하" +"위 폴더가 프레임 폴더 내에 생성됩니다.{0}\n" +"L|'export': 정렬 파일의 내용을 JSON 파일로 내보내십시오. 외부 도구에서 정렬 " +"정보를 편집 한 다음 FaceSwap의 추출물 'Import'플러그인을 사용하여 다시 인상하" +"는 데 사용할 수 있습니다. 참고 : 마스크 및 ID 벡터는 내보내기 파일에 포함되" +"지 않으므로 JSON 파일이 다시 FaceSwap으로 가져 오면 다시 생성됩니다. 모든 데" +"이터는 캔버스의 왼쪽 상단에있는 원점 (0, 0)으로 내 보냅니다.\n" +"L|'extract': alignments 데이터를 기반으로 소스 프레임/비디오에서 얼굴을 재추" +"출합니다. 이것은 얼굴을 재감지하는 것보다 훨씬 더 빠릅니다. '-een'(--extract-" +"every-n) 매개 변수를 전달하여 모든 n번째 프레임을 추출할 수 있습니다.{1}\n" +"L|'from-faces': 추출된 얼굴 폴더에서 alignments 파일을 생성합니다. 폴더 내의 " +"얼굴들을 여러 소스에서 가져온 경우 여러 alignments 파일이 생성됩니다. 참고: " +"비디오가 아닌 원본 이미지의 폴더를 추출한 얼굴의 경우, 원래 사용된 이미지의 " +"폴더 수를 알 수 없으므로 단일 alignments 파일이 생성됩니다. 이 작업을 실행하" +"기 위해 alignments 파일 경로를 제공할 필요는 없습니다. {3}\n" +"L|'missing-alignments': alignments 파일에 없는 프레임을 식별합니다.{2}{0}\n" +"L|'missing-frames': alignments 파일에서 [프레임 폴더/비디오] 내에 나타나지 않" +"는 프레임을 식별합니다.{2}{0}\n" +"L|'multi-faces': alignments 파일 내에서 여러 얼굴이 있는 위치를 식별합니다." +"{2}{4}\n" +"L|'no faces': alignments 파일 내에 있지만 얼굴이 탐지되지 않은 프레임을 식별" +"합니다.{2}{0}\n" +"L|'removes-faces': alignments 파일에서 삭제된 얼굴을 제거합니다. 원래 " +"alignments 파일은 백업됩니다.{3}\n" +"L|'rename' : alignments 파일의 상위 프레임 및 위치 색인에 해당하도록 얼굴 이" +"름을 바꿉니다(즉, 추출을 실행한 후에 얼굴 이름을 짓는 방법).{3}\n" +"L|'sort': alignments을 왼쪽에서 오른쪽으로 다시 인덱싱합니다. 얼굴이 여러 개" +"인 alignments의 경우 맨 왼쪽 얼굴이 색인 0에 있습니다.\n" +"L| 'spatial': 공간 및 시간 필터링을 수행하여 alignments를 원활하게 수행합니다" +"(실험적!)." + +#: tools/alignments/cli.py:107 +msgid "" +"R|How to output discovered items ('faces' and 'frames' only):\n" +"L|'console': Print the list of frames to the screen. (DEFAULT)\n" +"L|'file': Output the list of frames to a text file (stored within the source " +"directory).\n" +"L|'move': Move the discovered items to a sub-folder within the source " +"directory." +msgstr "" +"R|검색된 항목을 출력하는 방법('얼굴' 및 '프레임'만 해당):\n" +"L|'console': 프레임 목록을 화면에 인쇄합니다. (기본값)\n" +"L|'파일': 프레임 목록을 텍스트 파일(소스 디렉토리에 저장)로 출력합니다.\n" +"L|'이동': 검색된 항목을 원본 디렉토리 내의 하위 폴더로 이동합니다." + +#: tools/alignments/cli.py:118 tools/alignments/cli.py:141 +#: tools/alignments/cli.py:148 +msgid "data" +msgstr "데이터" + +#: tools/alignments/cli.py:125 +msgid "" +"Full path to the alignments file to be processed. If you have input a " +"'frames_dir' and don't provide this option, the process will try to find the " +"alignments file at the default location. All jobs require an alignments file " +"with the exception of 'from-faces' when the alignments file will be " +"generated in the specified faces folder." +msgstr "" +"처리할 alignments 파일의 전체 경로입니다. 'frames_dir'을 입력했는데 이 옵션" +"을 제공하지 않으면 프로세스는 기본 위치에서 alignments 파일을 찾으려고 합니" +"다. 지정된 얼굴 폴더에 alignments 파일이 생성될 때 모든 작업은 'from-" +"faces'를 제외한 alignments 파일이 필요로 합니다." + +#: tools/alignments/cli.py:142 +msgid "Directory containing source frames that faces were extracted from." +msgstr "얼굴 추출의 소스로 쓰인 원본 프레임이 저장된 디렉토리." + +#: tools/alignments/cli.py:150 +msgid "" +"R|Run the aligmnents tool on multiple sources. The following jobs support " +"batch mode:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"If batch mode is selected then the other options should be set as follows:\n" +"L|alignments_file: For 'sort' and 'spatial' this should point to the parent " +"folder containing the alignments files to be processed. For all other jobs " +"this option is ignored, and the alignments files must exist at their default " +"location relative to the original frames folder/video.\n" +"L|faces_dir: For 'from-faces' this should be a parent folder, containing sub-" +"folders of extracted faces from which to generate alignments files. For " +"'extract' this should be a parent folder where sub-folders will be created " +"for each extraction to be run. For all other jobs this option is ignored.\n" +"L|frames_dir: For 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"and 'no-faces' this should be a parent folder containing video files or sub-" +"folders of images to perform the alignments job on. The alignments file " +"should exist at the default location. For all other jobs this option is " +"ignored." +msgstr "" +"R|여러 소스에서 정렬 도구를 실행합니다. 다음 작업은 배치 모드를 지원합니다.\n" +"L|그리기, 추출, 얼굴부터, 정렬 누락, 프레임 누락, 얼굴 없음, 정렬, 공간.\n" +"배치 모드를 선택한 경우 다른 옵션을 다음과 같이 설정해야 합니다.\n" +"L|alignments_file: 'sort'및 'spatial'의 경우 처리할 정렬 파일이 포함된 상위 " +"폴더를 가리켜야 합니다. 다른 모든 작업의 경우 이 옵션은 무시되며 정렬 파일은 " +"원본 프레임 폴더/비디오에 상대적인 기본 위치에 있어야 합니다.\n" +"L|faces_dir: 'from-faces'의 경우 정렬 파일을 생성할 추출된 면의 하위 폴더를 " +"포함하는 상위 폴더여야 합니다. '추출'의 경우 실행할 각 추출에 대해 하위 폴더" +"가 생성되는 상위 폴더여야 합니다. 다른 모든 작업의 경우 이 옵션은 무시됩니" +"다.\n" +"L|frames_dir: 'draw', 'extract', 'missing-alignments', 'missing-frames' 및 " +"'no-faces'의 경우 비디오 파일이 포함된 상위 폴더 또는 정렬 작업을 수행할 이미" +"지의 하위 폴더여야 합니다. 에. 정렬 파일은 기본 위치에 있어야 합니다. 다른 모" +"든 작업의 경우 이 옵션은 무시됩니다." + +#: tools/alignments/cli.py:176 tools/alignments/cli.py:188 +#: tools/alignments/cli.py:198 +msgid "extract" +msgstr "추출" + +#: tools/alignments/cli.py:178 +msgid "" +"[Extract only] Extract every 'nth' frame. This option will skip frames when " +"extracting faces. For example a value of 1 will extract faces from every " +"frame, a value of 10 will extract faces from every 10th frame." +msgstr "" +"[Extract only] 모든 'n번째' 프레임을 추출합니다. 이 옵션은 얼굴을 추출할 때 " +"프레임을 건너뜁니다. 예를 들어, 값이 1이면 모든 프레임에서 얼굴이 추출되고, " +"값이 10이면 모든 10번째 프레임에서 얼굴이 추출됩니다." + +#: tools/alignments/cli.py:189 +msgid "[Extract only] The output size of extracted faces." +msgstr "[Extract only] 추출된 얼굴들의 결과 크기입니다." + +#: tools/alignments/cli.py:200 +msgid "" +"[Extract only] Only extract faces that have been resized by this percent or " +"more to meet the specified extract size (`-sz`, `--size`). Useful for " +"excluding low-res images from a training set. Set to 0 to extract all faces. " +"Eg: For an extract size of 512px, A setting of 50 will only include faces " +"that have been resized from 256px or above. Setting to 100 will only extract " +"faces that have been resized from 512px or above. A setting of 200 will only " +"extract faces that have been downscaled from 1024px or above." +msgstr "" +"[Extract only] 지정된 추출 크기('-sz', '--size')를 맞추기 위하여 크기가 이 비" +"율 이상 resize된 얼굴들만 추출합니다. 훈련 세트에서 저해상도 이미지를 제외하" +"는 데 유용합니다. 모든 얼굴을 추출하려면 0으로 설정합니다. 예: 추출 크기가 " +"512px인 경우, 50으로 설정하면 크기가 256px 이상인 면만 포함됩니다. 100으로 설" +"정하면 512px 이상에서 크기가 조정된 얼굴만 추출됩니다. 200으로 설정하면 " +"1024px 이상에서 축소된 얼굴만 추출됩니다." + +#~ msgid "Directory containing extracted faces." +#~ msgstr "추출된 얼굴들이 저장된 디렉토리." diff --git a/locales/kr/LC_MESSAGES/tools.effmpeg.cli.mo b/locales/kr/LC_MESSAGES/tools.effmpeg.cli.mo new file mode 100644 index 0000000000..f6f563913f Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.effmpeg.cli.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.effmpeg.cli.po b/locales/kr/LC_MESSAGES/tools.effmpeg.cli.po new file mode 100644 index 0000000000..58b106c585 --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.effmpeg.cli.po @@ -0,0 +1,185 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:50+0000\n" +"PO-Revision-Date: 2024-03-29 00:05+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/effmpeg/cli.py:15 +msgid "This command allows you to easily execute common ffmpeg tasks." +msgstr "" +"이 명령어는 사용자에게 일반 ffmpeg 작업을 쉽게 실행할 수 있도록 해줍니다." + +#: tools/effmpeg/cli.py:52 +msgid "A wrapper for ffmpeg for performing image <> video converting." +msgstr "이미지 <> 비디오 변환을 수행하기 위한 ffmpeg용 wrapper입니다." + +#: tools/effmpeg/cli.py:64 +msgid "" +"R|Choose which action you want ffmpeg ffmpeg to do.\n" +"L|'extract': turns videos into images \n" +"L|'gen-vid': turns images into videos \n" +"L|'get-fps' returns the chosen video's fps.\n" +"L|'get-info' returns information about a video.\n" +"L|'mux-audio' add audio from one video to another.\n" +"L|'rescale' resize video.\n" +"L|'rotate' rotate video.\n" +"L|'slice' cuts a portion of the video into a separate video file." +msgstr "" +"R|ffmpeg ffmpeg에서 수행할 작업을 선택합니다.\n" +"L|'extraction': 비디오를 이미지로 바꿉니다.\n" +"L|'gen-vid': 이미지를 비디오로 바꿉니다.\n" +"L|'get-fps'는 선택한 비디오의 fps를 반환합니다.\n" +"L|'get-info'는 동영상에 대한 정보를 반환합니다.\n" +"L|'mux-audio'는 한 비디오에서 다른 비디오로 오디오를 추가합니다.\n" +"L|'rescale' 크기 조정 비디오.\n" +"L|'rotate' 비디오 회전.\n" +"L| 'slice'는 동영상의 일부를 별도의 동영상 파일로 잘라냅니다." + +#: tools/effmpeg/cli.py:78 +msgid "Input file." +msgstr "입력 파일." + +#: tools/effmpeg/cli.py:79 tools/effmpeg/cli.py:86 tools/effmpeg/cli.py:100 +msgid "data" +msgstr "데이터" + +#: tools/effmpeg/cli.py:89 +msgid "" +"Output file. If no output is specified then: if the output is meant to be a " +"video then a video called 'out.mkv' will be created in the input directory; " +"if the output is meant to be a directory then a directory called 'out' will " +"be created inside the input directory. Note: the chosen output file " +"extension will determine the file encoding." +msgstr "" +"출력 파일. 출력이 지정되지 않은 경우: 출력이 비디오여야 한다면 입력 디렉토리" +"에 'out.mkv'라는 비디오가 생성됩니다. 출력이 디렉토리여야 한다면 입력 디렉토" +"리 내에 'out'이라는 디렉터리가 생성됩니다. 참고: 선택한 출력 파일 확장자가 파" +"일 인코딩을 결정합니다." + +#: tools/effmpeg/cli.py:102 +msgid "Path to reference video if 'input' was not a video." +msgstr "만약 input이 비디오가 아닐 경우 참고 비디으의 경로." + +#: tools/effmpeg/cli.py:108 tools/effmpeg/cli.py:118 tools/effmpeg/cli.py:156 +#: tools/effmpeg/cli.py:185 +msgid "output" +msgstr "출력" + +#: tools/effmpeg/cli.py:110 +msgid "" +"Provide video fps. Can be an integer, float or fraction. Negative values " +"will will make the program try to get the fps from the input or reference " +"videos." +msgstr "" +"비디오 fps를 제공합니다. 정수, 부동 또는 분수가 될 수 있습니다. 음수 값을 지" +"정하면 프로그램이 입력 또는 참조 비디오에서 fps를 가져오려고 합니다." + +#: tools/effmpeg/cli.py:120 +msgid "" +"Image format that extracted images should be saved as. '.bmp' will offer the " +"fastest extraction speed, but will take the most storage space. '.png' will " +"be slower but will take less storage." +msgstr "" +"추출된 이미지의 확장자는 '.bmp'로 저장되어야 합니다. '.bmp'는 가장 빠른 추출 " +"속도를 제공하지만 가장 많은 저장 공간을 차지합니다. '.png'은 속도는 더 느리지" +"만 저장 공간은 더 적게 차지합니다." + +#: tools/effmpeg/cli.py:127 tools/effmpeg/cli.py:136 tools/effmpeg/cli.py:145 +msgid "clip" +msgstr "클립" + +#: tools/effmpeg/cli.py:129 +msgid "" +"Enter the start time from which an action is to be applied. Default: " +"00:00:00, in HH:MM:SS format. You can also enter the time with or without " +"the colons, e.g. 00:0000 or 026010." +msgstr "" +"작업을 적용할 시작 시간을 입력합니다. 기본값: 00:00:00, HH:MM:SS 형식입니다. " +"콜론을 포함하거나 포함하지 않은 시간(예: 00:0000 또는 026010)을 입력할 수도 " +"있습니다." + +#: tools/effmpeg/cli.py:138 +msgid "" +"Enter the end time to which an action is to be applied. If both an end time " +"and duration are set, then the end time will be used and the duration will " +"be ignored. Default: 00:00:00, in HH:MM:SS." +msgstr "" +"적용된 작업의 종료 시간을 입력합니다. 종료 시간과 기간이 모두 설정된 경우 종" +"료 시간이 사용되고 기간이 무시됩니다. 기본값: 00:00:00, HH:MM:SS." + +#: tools/effmpeg/cli.py:147 +msgid "" +"Enter the duration of the chosen action, for example if you enter 00:00:10 " +"for slice, then the first 10 seconds after and including the start time will " +"be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. You can " +"also enter the time with or without the colons, e.g. 00:0000 or 026010." +msgstr "" +"선택한 작업의 지속 시간을 입력합니다. 예를 들어 슬라이스에 00:00:10을 입력하" +"면 시작 시간 이후의 첫 10초가 새 비디오로 잘라집니다. 기본값: 00:00:00, HH:" +"MM:SS 형식입니다. 콜론을 포함하거나 포함하지 않은 시간(예: 00:0000 또는 " +"026010)을 입력할 수도 있습니다." + +#: tools/effmpeg/cli.py:158 +msgid "" +"Mux the audio from the reference video into the input video. This option is " +"only used for the 'gen-vid' action. 'mux-audio' action has this turned on " +"implicitly." +msgstr "" +"참조 비디오의 오디오를 입력 비디오에 병합합니다. 이 옵션은 'gen-vid' 작업에" +"만 사용됩니다. 'mux-timeout' 작업은 이 작업을 암시적으로 활성화했습니다." + +#: tools/effmpeg/cli.py:169 tools/effmpeg/cli.py:179 +msgid "rotate" +msgstr "회전" + +#: tools/effmpeg/cli.py:171 +msgid "" +"Transpose the video. If transpose is set, then degrees will be ignored. For " +"cli you can enter either the number or the long command name, e.g. to use " +"(1, 90Clockwise) -tr 1 or -tr 90Clockwise" +msgstr "" +"비디오를 전치합니다. 전치를 설정하면 각도가 무시됩니다. cli의 경우 숫자 또는 " +"긴 명령 이름을 입력할 수 있습니다(예: (1, 90Clockwise) (-tr 1 또는 -tr " +"90Clockwise)" + +#: tools/effmpeg/cli.py:180 +msgid "Rotate the video clockwise by the given number of degrees." +msgstr "비디오를 주어진 입력 각도에 따라 시계방향으로 회전합니다." + +#: tools/effmpeg/cli.py:187 +msgid "Set the new resolution scale if the chosen action is 'rescale'." +msgstr "선택한 작업이 'rescale'이라면 새로운 해상도 크기를 설정합니다." + +#: tools/effmpeg/cli.py:192 tools/effmpeg/cli.py:200 +msgid "settings" +msgstr "설정" + +#: tools/effmpeg/cli.py:194 +msgid "" +"Reduces output verbosity so that only serious errors are printed. If both " +"quiet and verbose are set, verbose will override quiet." +msgstr "" +"출력 상세도를 줄여 심각한 오류만 출력합니다. quiet와 verbose가 모두 설정된 경" +"우 verbose가 quiet를 재정의합니다." + +#: tools/effmpeg/cli.py:202 +msgid "" +"Increases output verbosity. If both quiet and verbose are set, verbose will " +"override quiet." +msgstr "" +"출력 상세도를 높입니다. quiet와 verbose가 모두 설정된 경우 verbose가 quiet를 " +"재정의합니다." diff --git a/locales/kr/LC_MESSAGES/tools.manual.mo b/locales/kr/LC_MESSAGES/tools.manual.mo new file mode 100644 index 0000000000..2a801da9a2 Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.manual.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.manual.po b/locales/kr/LC_MESSAGES/tools.manual.po new file mode 100644 index 0000000000..0fbcc99572 --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.manual.po @@ -0,0 +1,283 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:55+0000\n" +"PO-Revision-Date: 2024-03-29 00:05+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/manual/cli.py:13 +msgid "" +"This command lets you perform various actions on frames, faces and " +"alignments files using visual tools." +msgstr "" +"이 명령어는 visual 도구들을 사용하여 프레임, 얼굴, alignments 파일들에 대한 " +"다양한 작업을 수행할 수 있도록 해줍니다." + +#: tools/manual/cli.py:23 +msgid "" +"A tool to perform various actions on frames, faces and alignments files " +"using visual tools" +msgstr "" +"프레임, 얼굴, alignments 파일들에 대한 다양한 작업을 수행할 수 있도록 해주는 " +"도구" + +#: tools/manual/cli.py:35 tools/manual/cli.py:44 +msgid "data" +msgstr "데이터" + +#: tools/manual/cli.py:38 +msgid "" +"Path to the alignments file for the input, if not at the default location" +msgstr "" +"입력에 대한 alignments 파일의 경로, 만약 설정되지 않았다면 기본 경로입니다" + +#: tools/manual/cli.py:46 +msgid "" +"Video file or directory containing source frames that faces were extracted " +"from." +msgstr "얼굴이 추출된 소스 프레임을 가지고 있는 비디오 파일 또는 디렉토리." + +#: tools/manual/cli.py:53 tools/manual/cli.py:62 +msgid "options" +msgstr "설정" + +#: tools/manual/cli.py:55 +msgid "" +"Force regeneration of the low resolution jpg thumbnails in the alignments " +"file." +msgstr "_alignments 파일에서 저해상도 jpg 미리 보기를 강제로 재생성합니다." + +#: tools/manual/cli.py:64 +msgid "" +"The process attempts to speed up generation of thumbnails by extracting from " +"the video in parallel threads. For some videos, this causes the caching " +"process to hang. If this happens, then set this option to generate the " +"thumbnails in a slower, but more stable single thread." +msgstr "" +"프로세스는 병렬 스레드에서 비디오를 추출하여 썸네일 생성 속도를 높이려고 시도" +"합니다. 일부 비디오의 경우 캐싱 프로세스가 중단될 수 있습니다. 이런 경우 이 " +"옵션을 설정하여 더 느리지만 안정적인 단일 스레드에서 썸네일를 생성하십시오." + +#: tools/manual\faceviewer\frame.py:163 +msgid "Display the landmarks mesh" +msgstr "특징점 망 보이기" + +#: tools/manual\faceviewer\frame.py:164 +msgid "Display the mask" +msgstr "마스크 보이기" + +#: tools/manual\frameviewer\editor\_base.py:628 +#: tools/manual\frameviewer\editor\landmarks.py:44 +#: tools/manual\frameviewer\editor\mask.py:75 +msgid "Magnify/Demagnify the View" +msgstr "보기를 확대/축소 합니다" + +#: tools/manual\frameviewer\editor\bounding_box.py:33 +#: tools/manual\frameviewer\editor\extract_box.py:32 +msgid "Delete Face" +msgstr "얼굴 삭제" + +#: tools/manual\frameviewer\editor\bounding_box.py:36 +msgid "" +"Bounding Box Editor\n" +"Edit the bounding box being fed into the aligner to recalculate the " +"landmarks.\n" +"\n" +" - Grab the corner anchors to resize the bounding box.\n" +" - Click and drag the bounding box to relocate.\n" +" - Click in empty space to create a new bounding box.\n" +" - Right click a bounding box to delete a face." +msgstr "" +"경계 상자 편집기\n" +"aligner 에 공급되는 경계 상자를 편집하여 특징점을 다시 계산합니다.\n" +"\n" +"- corner anchors를 사용하여 경계 상자의 크기를 재조정합니다.\n" +"- 경계 상자를 클릭하고 끌어서 재배치합니다.\n" +"- 빈 공간을 클릭하여 새 경계 상자를 만듭니다.\n" +"- 경계 상자를 마우스 오른쪽 단추로 클릭하여 얼굴을 삭제합니다." + +#: tools/manual\frameviewer\editor\bounding_box.py:70 +msgid "" +"Aligner to use. FAN will obtain better alignments, but cv2-dnn can be useful " +"if FAN cannot get decent alignments and you want to set a base to edit from." +msgstr "" +"사용할 aligner. FAN은 더 나은 alignments을 얻을 수 있지만, 만약 FAN이 적절한 " +"alignments을 얻을 수 없고 편집을 시작할 기준점을 설정하려는 경우 cv2-dnn이 유" +"용할 수 있습니다." + +#: tools/manual\frameviewer\editor\bounding_box.py:83 +msgid "" +"Normalization method to use for feeding faces to the aligner. This can help " +"the aligner better align faces with difficult lighting conditions. Different " +"methods will yield different results on different sets. NB: This does not " +"impact the output face, just the input to the aligner.\n" +"\tnone: Don't perform normalization on the face.\n" +"\tclahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"\thist: Equalize the histograms on the RGB channels.\n" +"\tmean: Normalize the face colors to the mean." +msgstr "" +"_aligner에 얼굴을 공급하는 데 사용할 정규화 방법입니다. 이렇게 하면 aligner" +"가 어려운 조명 조건에서 얼굴을 더 잘 정렬할 수 있습니다. 방법이 다르면 세트마" +"다 결과가 다릅니다. NB: 출력 얼굴에는 영향을 주지 않으며 aligner에게 주는 입" +"력에만 영향을 줍니다.\n" +"\tnone: 얼굴에 정규화를 수행하지 않습니다.\n" +"\tclahe: 얼굴에 Contrast Limited Adaptive Histogram Equalization를 수행합니" +"다.\n" +"\thist: RGB 채널의 히스토그램을 균등화합니다.\n" +"\tmean: 얼굴 색상을 평균으로 정규화합니다." + +#: tools/manual\frameviewer\editor\extract_box.py:35 +msgid "" +"Extract Box Editor\n" +"Move the extract box that has been generated by the aligner. Click and " +"drag:\n" +"\n" +" - Inside the bounding box to relocate the landmarks.\n" +" - The corner anchors to resize the landmarks.\n" +" - Outside of the corners to rotate the landmarks." +msgstr "" +"Box Editor 추출\n" +"aligner에서 생성한 추출 box를 이동합니다. click & drag:\n" +"\n" +"- bouding box 내부에서 특징점을 재배치.\n" +"- 특징점들의 크기를 조정하는 corner anchors.\n" +"- 모서리를 벗어나 특징점을 회전합니다." + +#: tools/manual\frameviewer\editor\landmarks.py:27 +msgid "" +"Landmark Point Editor\n" +"Edit the individual landmark points.\n" +"\n" +" - Click and drag individual points to relocate.\n" +" - Draw a box to select multiple points to relocate." +msgstr "" +"특징점 편집기\n" +"개별 특징점들을 편집합니다.\n" +"\n" +" - 개별 특징점들을 클릭 & 드래그 하여 재배치합니다.\n" +" - 재배치할 여러개의 점들을 박스를 그려서 선택합니다." + +#: tools/manual\frameviewer\editor\mask.py:33 +msgid "" +"Mask Editor\n" +"Edit the mask.\n" +" - NB: For Landmark based masks (e.g. components/extended) it is better to " +"make sure the landmarks are correct rather than editing the mask directly. " +"Any change to the landmarks after editing the mask will override your manual " +"edits." +msgstr "" +"마스크 편집기\n" +"마스크를 편집합니다.\n" +"- 주의: 특징점 기반 마스크(예: 구성 요소/확장)의 경우 마스크를 직접 편집하기" +"보다는 특징점이 올바른지 확인하는 것이 좋습니다. 마스크를 편집한 후 특징점들 " +"변경하면 변경된 특징점들이 수동으로 편집한 마스크에 덮어 씌워집니다." + +#: tools/manual\frameviewer\editor\mask.py:77 +msgid "Draw Tool" +msgstr "그리기 도구" + +#: tools/manual\frameviewer\editor\mask.py:78 +msgid "Erase Tool" +msgstr "지우개 도구" + +#: tools/manual\frameviewer\editor\mask.py:97 +msgid "Select which mask to edit" +msgstr "편집할 마스크를 선택" + +#: tools/manual\frameviewer\editor\mask.py:104 +msgid "Set the brush size. ([ - decrease, ] - increase)" +msgstr "붓 크기 설정. ([ - decrease, ] - increase)" + +#: tools/manual\frameviewer\editor\mask.py:111 +msgid "Select the brush cursor color." +msgstr "붓 커서 색깔 선택." + +#: tools/manual\frameviewer\frame.py:78 +msgid "Play/Pause (SPACE)" +msgstr "재생/멈춤 (스페이스 바)" + +#: tools/manual\frameviewer\frame.py:79 +msgid "Go to First Frame (HOME)" +msgstr "첫 번째 프레임으로 이동 (HOME)" + +#: tools/manual\frameviewer\frame.py:80 +msgid "Go to Previous Frame (Z)" +msgstr "이전 프레임으로 이동 (Z)" + +#: tools/manual\frameviewer\frame.py:81 +msgid "Go to Next Frame (X)" +msgstr "다음 프레임으로 이동 (X)" + +#: tools/manual\frameviewer\frame.py:82 +msgid "Go to Last Frame (END)" +msgstr "마지막 프레임으로 이동 (END)" + +#: tools/manual\frameviewer\frame.py:83 +msgid "Extract the faces to a folder... (Ctrl+E)" +msgstr "폴더에 얼굴 추출... (Ctrl+E)" + +#: tools/manual\frameviewer\frame.py:84 +msgid "Save the Alignments file (Ctrl+S)" +msgstr "_Alignments file 저장 (Ctrl + S" + +#: tools/manual\frameviewer\frame.py:85 +msgid "Filter Frames to only those Containing the Selected Item (F)" +msgstr "오로지 선택된 아이템들을 가지고 있는 필터 프레임 (F)" + +#: tools/manual\frameviewer\frame.py:86 +msgid "" +"Set the distance from an 'average face' to be considered misaligned. Higher " +"distances are more restrictive" +msgstr "" +"'평균 얼굴'로부터의 거리를 잘못 정렬된 것으로 간주하도록 설정. 먼 거리에서 조" +"금 더 제한적입니다" + +#: tools/manual\frameviewer\frame.py:391 +msgid "View alignments" +msgstr "보기 정렬" + +#: tools/manual\frameviewer\frame.py:392 +msgid "Bounding box editor" +msgstr "경계 상자 편집기" + +#: tools/manual\frameviewer\frame.py:393 +msgid "Location editor" +msgstr "위치 편집기" + +#: tools/manual\frameviewer\frame.py:394 +msgid "Mask editor" +msgstr "마스크 편집기" + +#: tools/manual\frameviewer\frame.py:395 +msgid "Landmark point editor" +msgstr "특징점 편집기" + +#: tools/manual\frameviewer\frame.py:470 +msgid "Next" +msgstr "다음" + +#: tools/manual\frameviewer\frame.py:470 +msgid "Previous" +msgstr "이전" + +#: tools/manual\frameviewer\frame.py:481 +msgid "Revert to saved Alignments ({})" +msgstr "저장된 Alignments로 돌아가기 ({})" + +#: tools/manual\frameviewer\frame.py:487 +msgid "Copy {} Alignments ({})" +msgstr "{} Alignments를 복사 ({})" diff --git a/locales/kr/LC_MESSAGES/tools.mask.cli.mo b/locales/kr/LC_MESSAGES/tools.mask.cli.mo new file mode 100644 index 0000000000..9de146e0bc Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.mask.cli.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.mask.cli.po b/locales/kr/LC_MESSAGES/tools.mask.cli.po new file mode 100644 index 0000000000..94080f7faf --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.mask.cli.po @@ -0,0 +1,318 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-06-28 13:45+0100\n" +"PO-Revision-Date: 2024-06-28 13:48+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.4\n" + +#: tools/mask/cli.py:15 +msgid "" +"This tool allows you to generate, import, export or preview masks for " +"existing alignments." +msgstr "" +"이 도구를 사용하면 기존 정렬에 대한 마스크를 생성, 가져오기, 내보내기 또는 미" +"리 볼 수 있습니다." + +#: tools/mask/cli.py:25 +msgid "" +"Mask tool\n" +"Generate, import, export or preview masks for existing alignments files." +msgstr "" +"마스크 도구\n" +"기존 alignments 파일에 대한 마스크를 생성, 가져오기, 내보내기 또는 미리 봅니" +"다." + +#: tools/mask/cli.py:35 tools/mask/cli.py:47 tools/mask/cli.py:58 +#: tools/mask/cli.py:69 +msgid "data" +msgstr "데이터" + +#: tools/mask/cli.py:39 +msgid "" +"Full path to the alignments file that contains the masks if not at the " +"default location. NB: If the input-type is faces and you wish to update the " +"corresponding alignments file, then you must provide a value here as the " +"location cannot be automatically detected." +msgstr "" +"기본 위치가 아닌 경우 마스크를 추가할 정렬 파일의 전체 경로입니다. NB: 입력 " +"유형이 얼굴이고 해당 정렬 파일을 업데이트하려는 경우 위치를 자동으로 감지할 " +"수 없으므로 여기에 값을 제공해야 합니다." + +#: tools/mask/cli.py:51 +msgid "Directory containing extracted faces, source frames, or a video file." +msgstr "추출된 얼굴들, 원본 프레임들, 또는 비디오 파일이 존재하는 디렉토리." + +#: tools/mask/cli.py:61 +msgid "" +"R|Whether the `input` is a folder of faces or a folder frames/video\n" +"L|faces: The input is a folder containing extracted faces.\n" +"L|frames: The input is a folder containing frames or is a video" +msgstr "" +"R|'입력'이 얼굴의 폴더인지 아니면 폴더 프레임/비디오인지\n" +"L|faces: 입력은 추출된 얼굴을 포함된 폴더입니다.\n" +"L|frames: 입력이 프레임을 포함된 폴더이거나 비디오입니다" + +#: tools/mask/cli.py:71 +msgid "" +"R|Run the mask tool on multiple sources. If selected then the other options " +"should be set as follows:\n" +"L|input: A parent folder containing either all of the video files to be " +"processed, or containing sub-folders of frames/faces.\n" +"L|output-folder: If provided, then sub-folders will be created within the " +"given location to hold the previews for each input.\n" +"L|alignments: Alignments field will be ignored for batch processing. The " +"alignments files must exist at the default location (for frames). For batch " +"processing of masks with 'faces' as the input type, then only the PNG header " +"within the extracted faces will be updated." +msgstr "" +"R|여러 소스에서 마스크 도구를 실행합니다. 선택한 경우 다른 옵션을 다음과 같" +"이 설정해야 합니다.\n" +"L|input: 처리할 모든 비디오 파일을 포함하거나 프레임/얼굴의 하위 폴더를 포함" +"하는 상위 폴더입니다.\n" +"L|output-folder: 제공된 경우 각 입력에 대한 미리 보기를 보관하기 위해 지정된 " +"위치 내에 하위 폴더가 생성됩니다.\n" +"L|alignments: 일괄 처리에서는 정렬 필드가 무시됩니다. 정렬 파일은 기본 위치" +"(프레임용)에 있어야 합니다. 입력 유형이 '얼굴'인 마스크를 일괄 처리하는 경우 " +"추출된 얼굴 내의 PNG 헤더만 업데이트됩니다." + +#: tools/mask/cli.py:87 tools/mask/cli.py:119 +msgid "process" +msgstr "진행" + +#: tools/mask/cli.py:89 +msgid "" +"R|Masker to use.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members. Profile faces " +"may result in sub-par performance." +msgstr "" +"R|사용할 마스크.\n" +"L|bisnet-fp: 전체 얼굴 마스킹(마스크 설정에서 구성 가능)을 포함하여 마스킹할 " +"영역에 대한 보다 정교한 제어를 제공하는 비교적 가벼운 NN 기반 마스크입니다.\n" +"L|components: 특징점 위치를 기반으로 얼굴 분할을 제공하도록 설계된 마스크입니" +"다. 특징점의 외부에는 마스크를 만들기 위해 convex hull이가 형성되어 있습니" +"다.\n" +"L|custom: 마스크 영역을 모든 1 또는 0으로 채우는 더미 마스크입니다(설정에서 " +"구성 가능). 수동 도구에서 사용자 정의 마스크를 직접 수동으로 편집하려는 경우" +"에만 필요합니다. 이 마스크는 GPU를 사용하지 않습니다.\n" +"L|extended: 특징점 위치를 기반으로 얼굴 분할을 제공하도록 설계된 마스크입니" +"다. 지형지물의 외부에는 convex hull이 형성되어 있으며, 마스크는 이마 위로 뻗" +"어 있습니다.\n" +"L|vgg-clear: 대부분의 정면에 장애물이 없는 스마트한 분할을 제공하도록 설계된 " +"마스크입니다. 프로필 면 및 장애물로 인해 성능이 저하될 수 있습니다.\n" +"L|vgg-obstructed: 대부분의 정면 얼굴을 스마트하게 분할할 수 있도록 설계된 마" +"스크입니다. 마스크 모델은 일부 안면 장애물(손과 안경)을 인식하도록 특별히 훈" +"련되었습니다. 옆 얼굴은 평균 이하의 성능을 초래할 수 있습니다.\n" +"L|unet-dfl: 대부분 정면 얼굴을 스마트하게 분할하도록 설계된 마스크. 마스크 모" +"델은 커뮤니티 구성원들에 의해 훈련되었으며 추가 설명을 위해 테스트가 필요합니" +"다. 옆 얼굴은 평균 이하의 성능을 초래할 수 있습니다." + +#: tools/mask/cli.py:121 +msgid "" +"R|The Mask tool process to perform.\n" +"L|all: Update the mask for all faces in the alignments file for the selected " +"'masker'.\n" +"L|missing: Create a mask for all faces in the alignments file where a mask " +"does not previously exist for the selected 'masker'.\n" +"L|output: Don't update the masks, just output the selected 'masker' for " +"review/editing in external tools to the given output folder.\n" +"L|import: Import masks that have been edited outside of faceswap into the " +"alignments file. Note: 'custom' must be the selected 'masker' and the masks " +"must be in the same format as the 'input-type' (frames or faces)" +msgstr "" +"R|수행할 마스크 도구 프로세스입니다.\n" +"L|all: 선택한 'masker'에 대한 정렬 파일의 모든 면에 대한 마스크를 업데이트합" +"니다.\n" +"L|missing: 선택한 'masker'에 대해 이전에 마스크가 존재하지 않았던 정렬 파일" +"의 모든 면에 대한 마스크를 생성합니다.\n" +"L|output: 마스크를 업데이트하지 않고 외부 도구에서 검토/편집하기 위해 선택한 " +"'masker'를 지정된 출력 폴더로 출력합니다.\n" +"L|import: Faceswap 외부에서 편집된 마스크를 정렬 파일로 가져옵니다. 참고: " +"'custom'은 선택된 'masker'여야 하며 마스크는 'input-type'(frames 또는 faces)" +"과 동일한 형식이어야 합니다." + +#: tools/mask/cli.py:135 tools/mask/cli.py:154 tools/mask/cli.py:176 +msgid "import" +msgstr "수입" + +#: tools/mask/cli.py:137 +msgid "" +"R|Import only. The path to the folder that contains masks to be imported.\n" +"L|How the masks are provided is not important, but they will be stored, " +"internally, as 8-bit grayscale images.\n" +"L|If the input are images, then the masks must be named exactly the same as " +"input frames/faces (excluding the file extension).\n" +"L|If the input is a video file, then the filename of the masks is not " +"important but should contain the frame number at the end of the filename " +"(but before the file extension). The frame number can be separated from the " +"rest of the filename by any non-numeric character and can be padded by any " +"number of zeros. The frame number must correspond correctly to the frame " +"number in the original video (starting from frame 1)." +msgstr "" +"R|가져오기만 가능합니다. 가져올 마스크가 포함된 폴더의 경로입니다.\n" +"L|마스크 제공 방법은 중요하지 않지만 내부적으로 8비트 회색조 이미지로 저장됩" +"니다.\n" +"L|입력이 이미지인 경우 마스크 이름은 입력 프레임/얼굴과 정확히 동일하게 지정" +"되어야 합니다(파일 확장자 제외).\n" +"L|입력이 비디오 파일인 경우 마스크의 파일 이름은 중요하지 않지만 파일 이름 끝" +"에(파일 확장자 앞에) 프레임 번호가 포함되어야 합니다. 프레임 번호는 숫자가 아" +"닌 문자로 파일 이름의 나머지 부분과 구분될 수 있으며 임의 개수의 0으로 채워" +"질 수 있습니다. 프레임 번호는 원본 비디오의 프레임 번호(프레임 1부터 시작)와 " +"정확하게 일치해야 합니다." + +#: tools/mask/cli.py:156 +msgid "" +"R|Import/Output only. When importing masks, this is the centering to use. " +"For output this is only used for outputting custom imported masks, and " +"should correspond to the centering used when importing the mask. Note: For " +"any job other than 'import' and 'output' this option is ignored as mask " +"centering is handled internally.\n" +"L|face: Centers the mask on the center of the face, adjusting for pitch and " +"yaw. Outside of requirements for full head masking/training, this is likely " +"to be the best choice.\n" +"L|head: Centers the mask on the center of the head, adjusting for pitch and " +"yaw. Note: You should only select head centering if you intend to include " +"the full head (including hair) within the mask and are looking to train a " +"full head model.\n" +"L|legacy: The 'original' extraction technique. Centers the mask near the of " +"the nose with and crops closely to the face. Can result in the edges of the " +"mask appearing outside of the training area." +msgstr "" +"R|Import/Output only. 마스크를 가져올 때, 이것은 사용할 중앙 정렬입니다. 출력" +"의 경우, 이것은 사용자 지정 가져온 마스크를 출력하는 데만 사용되며, 마스크를 " +"가져올 때 사용된 중앙 정렬과 일치해야 합니다. 참고: 'import' 및 'output' 이외" +"의 모든 작업의 ​​경우 마스크 중앙 정렬이 내부적으로 처리되므로 이 옵션은 무시됩" +"니다.\n" +"L|면: 피치와 요를 조정하여 마스크를 얼굴 중앙에 배치합니다. 머리 전체 마스킹/" +"훈련에 대한 요구 사항을 제외하면 이것이 최선의 선택일 가능성이 높습니다.\n" +"L|head: 마스크를 머리 중앙에 배치하여 피치와 요를 조정합니다. 참고: 마스크 내" +"에 머리 전체(머리카락 포함)를 포함하고 머리 전체 모델을 훈련시키려는 경우 머" +"리 중심 맞추기만 선택해야 합니다.\n" +"L|레거시: '원래' 추출 기술입니다. 마스크를 코 근처 중앙에 배치하고 얼굴에 가" +"깝게 자릅니다. 마스크 가장자리가 훈련 영역 외부에 나타날 수 있습니다." + +#: tools/mask/cli.py:181 +msgid "" +"Import only. The size, in pixels to internally store the mask at.\n" +"The default is 128 which is fine for nearly all usecases. Larger sizes will " +"result in larger alignments files and longer processing." +msgstr "" +"가져오기만. 마스크를 내부적으로 저장할 크기(픽셀)입니다.\n" +"기본값은 128이며 거의 모든 사용 사례에 적합합니다. 크기가 클수록 정렬 파일도 " +"커지고 처리 시간도 길어집니다." + +#: tools/mask/cli.py:189 tools/mask/cli.py:197 tools/mask/cli.py:211 +#: tools/mask/cli.py:225 tools/mask/cli.py:235 +msgid "output" +msgstr "출력" + +#: tools/mask/cli.py:191 +msgid "" +"Optional output location. If provided, a preview of the masks created will " +"be output in the given folder." +msgstr "" +"선택적 출력 위치. 만약 값이 제공된다면 생성된 마스크 미리 보기가 주어진 폴더" +"에 출력됩니다." + +#: tools/mask/cli.py:202 +msgid "" +"Apply gaussian blur to the mask output. Has the effect of smoothing the " +"edges of the mask giving less of a hard edge. the size is in pixels. This " +"value should be odd, if an even number is passed in then it will be rounded " +"to the next odd number. NB: Only effects the output preview. Set to 0 for off" +msgstr "" +"마스크 출력에 gaussian blur를 적용합니다. 마스크의 가장자리를 매끄럽게 하여 " +"단단한 가장자리를 덜 제공하는 효과가 있습니다. 크기는 픽셀 단위입니다. 이 값" +"은 홀수여야 하며 짝수가 전달되면 다음 홀수로 반올림됩니다. NB: 출력 미리 보기" +"에만 영향을 줍니다. 0으로 설정하면 꺼집니다" + +#: tools/mask/cli.py:216 +msgid "" +"Helps reduce 'blotchiness' on some masks by making light shades white and " +"dark shades black. Higher values will impact more of the mask. NB: Only " +"effects the output preview. Set to 0 for off" +msgstr "" +"밝은 색조를 흰색으로, 어두운 색조를 검은색으로 만들어 일부 마스크의 '흐림'을 " +"줄이는 데 도움이 됩니다. 값이 클수록 마스크에 더 많은 영향을 미칩니다. NB: 출" +"력 미리 보기에만 영향을 줍니다. 0으로 설정하면 꺼집니다" + +#: tools/mask/cli.py:227 +msgid "" +"R|How to format the output when processing is set to 'output'.\n" +"L|combined: The image contains the face/frame, face mask and masked face.\n" +"L|masked: Output the face/frame as rgba image with the face masked.\n" +"L|mask: Only output the mask as a single channel image." +msgstr "" +"R|처리가 'output'으로 설정되어 있을 때 출력을 구성하는 방법.\n" +"L|combined: 이미지에는 얼굴/프레임, 얼굴 마스크 및 마스크된 얼굴이 포함됩니" +"다.\n" +"L|masked: 마스크된 얼굴/프레임을 Rgba 이미지로 출력합니다.\n" +"L|mask: 마스크를 단일 채널 이미지로만 출력합니다." + +#: tools/mask/cli.py:237 +msgid "" +"R|Whether to output the whole frame or only the face box when using output " +"processing. Only has an effect when using frames as input." +msgstr "" +"R|출력 처리를 사용할 때 전체 프레임을 출력할지 또는 페이스 박스만 출력할지 여" +"부. 프레임을 입력으로 사용할 때만 효과가 있습니다." + +#~ msgid "" +#~ "R|Whether to update all masks in the alignments files, only those faces " +#~ "that do not already have a mask of the given `mask type` or just to " +#~ "output the masks to the `output` location.\n" +#~ "L|all: Update the mask for all faces in the alignments file.\n" +#~ "L|missing: Create a mask for all faces in the alignments file where a " +#~ "mask does not previously exist.\n" +#~ "L|output: Don't update the masks, just output them for review in the " +#~ "given output folder." +#~ msgstr "" +#~ "R|alignments 파일의 모든 마스크를 업데이트할지, 지정된 '마스크 유형'의 마" +#~ "스크가 아직 없는 페이스만 업데이트할지, 아니면 단순히 '출력' 위치로 마스크" +#~ "를 출력할지 여부.\n" +#~ "L|all: alignments 파일의 모든 얼굴에 대한 마스크를 업데이트합니다.\n" +#~ "L|missing: 마스크가 없었던 alignments 파일의 모든 얼굴에 대한 마스크를 만" +#~ "듭니다.\n" +#~ "L|output: 마스크를 업데이트하지 말고 지정된 출력 폴더에서 검토할 수 있도" +#~ "록 출력하십시오." + +#~ msgid "" +#~ "Full path to the alignments file to add the mask to. NB: if the mask " +#~ "already exists in the alignments file it will be overwritten." +#~ msgstr "" +#~ "마스크를 추가할 alignments 파일의 전체 경로입니다. 주의: alignments 파일" +#~ "에 마스크가 이미 있으면 alignments 파일이 덮어 씌워집니다." diff --git a/locales/kr/LC_MESSAGES/tools.model.cli.mo b/locales/kr/LC_MESSAGES/tools.model.cli.mo new file mode 100644 index 0000000000..9ccdfde6dc Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.model.cli.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.model.cli.po b/locales/kr/LC_MESSAGES/tools.model.cli.po new file mode 100644 index 0000000000..524943a5c0 --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.model.cli.po @@ -0,0 +1,82 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:51+0000\n" +"PO-Revision-Date: 2024-03-29 00:05+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/model/cli.py:13 +msgid "This tool lets you perform actions on saved Faceswap models." +msgstr "" +"이 도구를 사용하여 저장된 Faceswap 모델에서 작업을 수행할 수 있습니다." + +#: tools/model/cli.py:22 +msgid "A tool for performing actions on Faceswap trained model files" +msgstr "_Faceswap 훈련을 받은 모델 파일에서 작업을 수행하기 위한 도구" + +#: tools/model/cli.py:34 +msgid "" +"Model directory. A directory containing the model you wish to perform an " +"action on." +msgstr "모델 디렉토리. 작업을 수행할 모델이 들어 있는 디렉토리입니다." + +#: tools/model/cli.py:43 +msgid "" +"R|Choose which action you want to perform.\n" +"L|'inference' - Create an inference only copy of the model. Strips any " +"layers from the model which are only required for training. NB: This is for " +"exporting the model for use in external applications. Inference generated " +"models cannot be used within Faceswap. See the 'format' option for " +"specifying the model output format.\n" +"L|'nan-scan' - Scan the model file for NaNs or Infs (invalid data).\n" +"L|'restore' - Restore a model from backup." +msgstr "" +"R|실행할 작업을 선택합니다.\n" +"L|'inference' - 모델의 추론 전용 사본을 만듭니다. 모델에서 훈련에만 필요한 " +"모든 레이어를 제거합니다. NB: 이것은 외부 응용 프로그램에서 사용하기 위해 모" +"델을 내보내기 위한 것입니다. 추론 생성 모델은 Faceswap 내에서 사용할 수 없습" +"니다. 모델 출력 형식을 지정하려면 'format' 옵션을 참조하십시오.\n" +"L|'nan-scan' - 모델 파일에서 NaN 또는 Infs(잘못된 데이터)를 검색합니다.\n" +"L|'restore' - 백업에서 모델을 복원합니다." + +#: tools/model/cli.py:57 tools/model/cli.py:69 +msgid "inference" +msgstr "추론" + +#: tools/model/cli.py:59 +msgid "" +"R|The format to save the model as. Note: Only used for 'inference' job.\n" +"L|'h5' - Standard Keras H5 format. Does not store any custom layer " +"information. Layers will need to be loaded from Faceswap to use.\n" +"L|'saved-model' - Tensorflow's Saved Model format. Contains all information " +"required to load the model outside of Faceswap." +msgstr "" +"R|모델을 저장할 형식입니다. 참고: '추론' 작업에만 사용됩니다.\n" +"L|'h5' - 표준 케라스 H5 형식. 사용자 지정 레이어 정보를 저장하지 않습니다. " +"사용하려면 Faceswap에서 레이어를 로드해야 합니다.\n" +"L| 'saved-model' - 텐서플로의 저장된 모델 형식. Faceswap 외부에서 모델을 로" +"드하는 데 필요한 모든 정보를 포함합니다." + +#: tools/model/cli.py:71 +#, fuzzy +#| msgid "" +#| "Only used for 'inference' job. Generate the inference model for B -> A " +#| "instead of A -> B." +msgid "" +"Only used for 'inference' job. Generate the inference model for B -> A " +"instead of A -> B." +msgstr "" +"'추론' 작업에만 쓰입니다. A -> B 대신 B -> A에 대한 추론 모델을 생성합니다." diff --git a/locales/kr/LC_MESSAGES/tools.preview.mo b/locales/kr/LC_MESSAGES/tools.preview.mo new file mode 100644 index 0000000000..13d6841ba1 Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.preview.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.preview.po b/locales/kr/LC_MESSAGES/tools.preview.po new file mode 100644 index 0000000000..03e8b6631f --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.preview.po @@ -0,0 +1,87 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:04+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/preview/cli.py:15 +msgid "This command allows you to preview swaps to tweak convert settings." +msgstr "" +"이 명령어는 변환 설정을 변경하기 위한 변환 미리보기를 가능하게 해줍니다." + +#: tools/preview/cli.py:30 +msgid "" +"Preview tool\n" +"Allows you to configure your convert settings with a live preview" +msgstr "" +"미리보기 도구\n" +"라이브로 미리보기를 보면서 변환 설정을 구성할 수 있도록 해줍니다" + +#: tools/preview/cli.py:47 tools/preview/cli.py:57 tools/preview/cli.py:65 +msgid "data" +msgstr "데이터" + +#: tools/preview/cli.py:50 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file." +msgstr "" +"입력 디렉토리 또는 비디오. 처리할 이미지 파일이 들어 있는 디렉토리 또는 비디" +"오 파일의 경로입니다." + +#: tools/preview/cli.py:60 +msgid "" +"Path to the alignments file for the input, if not at the default location" +msgstr "입력 alignments 파일의 경로, 만약 제공되지 않는다면 기본 위치" + +#: tools/preview/cli.py:68 +msgid "" +"Model directory. A directory containing the trained model you wish to " +"process." +msgstr "" +"모델 디렉토리. 사용자가 처리하고 싶어하는 훈련된 모델이 있는 디렉토리." + +#: tools/preview/cli.py:74 +msgid "Swap the model. Instead of A -> B, swap B -> A" +msgstr "모델을 스왑함. A -> B 대신, B -> A로 스왑함" + +#: tools/preview/control_panels.py:510 +msgid "Save full config" +msgstr "전체 설정을 저장" + +#: tools/preview/control_panels.py:513 +msgid "Reset full config to default values" +msgstr "전체 설정을 기본 값으로 초기화" + +#: tools/preview/control_panels.py:516 +msgid "Reset full config to saved values" +msgstr "전체 설정을 저장된 값으로 초기화" + +#: tools/preview/control_panels.py:667 +#, python-brace-format +msgid "Save {title} config" +msgstr "{title} 설정 저장" + +#: tools/preview/control_panels.py:670 +#, python-brace-format +msgid "Reset {title} config to default values" +msgstr "{title} 설정을 기본 값으로 초기화" + +#: tools/preview/control_panels.py:673 +#, python-brace-format +msgid "Reset {title} config to saved values" +msgstr "{title} 설정을 저장된 값으로 초기화" diff --git a/locales/kr/LC_MESSAGES/tools.sort.cli.mo b/locales/kr/LC_MESSAGES/tools.sort.cli.mo new file mode 100644 index 0000000000..39509c675d Binary files /dev/null and b/locales/kr/LC_MESSAGES/tools.sort.cli.mo differ diff --git a/locales/kr/LC_MESSAGES/tools.sort.cli.po b/locales/kr/LC_MESSAGES/tools.sort.cli.po new file mode 100644 index 0000000000..19d99c1628 --- /dev/null +++ b/locales/kr/LC_MESSAGES/tools.sort.cli.po @@ -0,0 +1,388 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:04+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ko_KR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=1; plural=0;\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/sort/cli.py:15 +msgid "This command lets you sort images using various methods." +msgstr "이 명령어는 다양한 메소드를 이용하여 이미지를 정렬해줍니다." + +#: tools/sort/cli.py:21 +msgid "" +" Adjust the '-t' ('--threshold') parameter to control the strength of " +"grouping." +msgstr " 그룹화의 강도를 제어하기 위해 '-t' ('--threshold') 인자를 조정하세요." + +#: tools/sort/cli.py:22 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the percentage of color pixels " +"that appear in the image." +msgstr "" +" '-b'('--bins') 매개 변수를 조정하여 그룹화할 bins의 수를 제어합니다. 각 이미" +"지는 이미지에 나타나는 색상 픽셀의 백분율에 따라 bin에 할당됩니다." + +#: tools/sort/cli.py:25 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the number of degrees the face " +"is orientated from center." +msgstr "" +" '-b'('--bins') 매개 변수를 조정하여 그룹화할 bins의 수를 제어합니다. 각 이미" +"지는 얼굴이 이미지 중심에서 떨어진 각도에 따라 bin에 할당됩니다." + +#: tools/sort/cli.py:28 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. The minimum and maximum values are taken for the chosen sort " +"metric. The bins are then populated with the results from the group sorting." +msgstr "" +" '-b'('--bins') 매개 변수를 조정하여 그룹화할 bins의 수를 제어합니다. 선택한 " +"정렬 방법에 대해 최소값과 최대값이 사용됩니다. 그런 다음 bins가 그룹 정렬의 " +"결과로 채워집니다." + +#: tools/sort/cli.py:32 +msgid "faces by blurriness." +msgstr "흐릿한 얼굴." + +#: tools/sort/cli.py:33 +msgid "faces by fft filtered blurriness." +msgstr "fft 필터링된 흐릿한 얼굴." + +#: tools/sort/cli.py:34 +msgid "" +"faces by the estimated distance of the alignments from an 'average' face. " +"This can be useful for eliminating misaligned faces. Sorts from most like an " +"average face to least like an average face." +msgstr "" +"'평균' 얼굴에서 alignments의 추정 거리를 기준으로 하는 얼굴. 이는 잘못 정렬" +"된 얼굴을 제거하는 데 유용할 수 있습니다. 가장 평균 얼굴에서 가장 덜 평균 얼" +"굴순으로 정렬합니다." + +#: tools/sort/cli.py:37 +msgid "" +"faces using VGG Face2 by face similarity. This uses a pairwise clustering " +"algorithm to check the distances between 512 features on every face in your " +"set and order them appropriately." +msgstr "" +"얼굴 유사성에 따라 VGG Face2를 사용하는 얼굴. 이 알고리즘은 쌍별 클러스터링 " +"알고리즘을 사용하여 세트의 모든 얼굴에서 512개의 특징 사이의 거리를 확인하고 " +"적절하게 정렬합니다." + +#: tools/sort/cli.py:40 +msgid "faces by their landmarks." +msgstr "특징점이 있는 얼굴." + +#: tools/sort/cli.py:41 +msgid "Like 'face-cnn' but sorts by dissimilarity." +msgstr "'face-cnn'과 비슷하지만 비유사성에 따라 정렬된." + +#: tools/sort/cli.py:42 +msgid "faces by Yaw (rotation left to right)." +msgstr "yaw (왼쪽에서 오른쪽으로 회전)에 의한 얼굴." + +#: tools/sort/cli.py:43 +msgid "faces by Pitch (rotation up and down)." +msgstr "pitch (위에서 아래로 회전)에 의한 얼굴." + +#: tools/sort/cli.py:44 +msgid "" +"faces by Roll (rotation). Aligned faces should have a roll value close to " +"zero. The further the Roll value from zero the higher liklihood the face is " +"misaligned." +msgstr "" +"이동 (회전)에 의한 얼굴. 정렬된 얼굴들은 0에 가까운 이동 값을 가져야 한다. 이" +"동 값이 0에서 멀수록 얼굴들이 잘못 정렬되었을 가능성이 높습니다." + +#: tools/sort/cli.py:46 +msgid "faces by their color histogram." +msgstr "색상 히스토그램에 의한 얼굴." + +#: tools/sort/cli.py:47 +msgid "Like 'hist' but sorts by dissimilarity." +msgstr "'hist' 같지만 비유사성에 따라 정렬된." + +#: tools/sort/cli.py:48 +msgid "" +"images by the average intensity of the converted grayscale color channel." +msgstr "변환된 회색 계열 색상 채널의 평균 강도에 따른 이미지." + +#: tools/sort/cli.py:49 +msgid "" +"images by their number of black pixels. Useful when faces are near borders " +"and a large part of the image is black." +msgstr "" +"검은색 픽셀의 개수에 따른 이미지들. 얼굴이 테두리 근처에 있고 이미지의 대부분" +"이 검은색일 때 유용합니다." + +#: tools/sort/cli.py:51 +msgid "" +"images by the average intensity of the converted Y color channel. Bright " +"lighting and oversaturated images will be ranked first." +msgstr "" +"변환된 Y 색상 채널의 평균 강도를 기준으로 한 이미지. 밝은 조명과 과포화 이미" +"지가 1위를 차지할 것이다." + +#: tools/sort/cli.py:53 +msgid "" +"images by the average intensity of the converted Cg color channel. Green " +"images will be ranked first and red images will be last." +msgstr "" +"변환된 Cg 컬러 채널의 평균 강도를 기준으로 한 이미지. 녹색 이미지가 먼저 순위" +"가 매겨지고 빨간색 이미지가 마지막 순위가 됩니다." + +#: tools/sort/cli.py:55 +msgid "" +"images by the average intensity of the converted Co color channel. Orange " +"images will be ranked first and blue images will be last." +msgstr "" +"변환된 Co 색상 채널의 평균 강도를 기준으로 한 이미지. 주황색 이미지가 먼저 순" +"위가 매겨지고 파란색 이미지가 마지막 순위가 됩니다." + +#: tools/sort/cli.py:57 +msgid "" +"images by their size in the original frame. Faces further from the camera " +"and from lower resolution sources will be sorted first, whilst faces closer " +"to the camera and from higher resolution sources will be sorted last." +msgstr "" +"이미지를 원래 프레임의 크기별로 표시합니다. 카메라에서 더 멀리 떨어져 있고 저" +"해상도 원본에서 온 얼굴이 먼저 정렬되고, 카메라에 더 가까이 있고 고해상도 원" +"본에서 온 얼굴이 마지막으로 정렬됩니다." + +#: tools/sort/cli.py:81 +msgid "Sort faces using a number of different techniques" +msgstr "얼굴을 정렬하는데 사용되는 서로 다른 기술들의 개수" + +#: tools/sort/cli.py:91 tools/sort/cli.py:98 tools/sort/cli.py:110 +#: tools/sort/cli.py:150 +msgid "data" +msgstr "데이터" + +#: tools/sort/cli.py:92 +msgid "Input directory of aligned faces." +msgstr "정렬된 얼굴들의 입력 디렉토리." + +#: tools/sort/cli.py:100 +msgid "" +"Output directory for sorted aligned faces. If not provided and 'keep' is " +"selected then a new folder called 'sorted' will be created within the input " +"folder to house the output. If not provided and 'keep' is not selected then " +"the images will be sorted in-place, overwriting the original contents of the " +"'input_dir'" +msgstr "" +"정렬된 aligned 얼굴의 출력 디렉토리입니다. 제공되지 않은 상태에서 'keep'을 선" +"택하면 출력을 저장하기 위해 입력 폴더 내에 'sorted'라는 새 폴더가 생성됩니" +"다. 제공되지 않고 'keep'을 선택하지 않으면 이미지가 제자리에 정렬되어 " +"'input_dir'의 원래 내용을 덮어씁니다." + +#: tools/sort/cli.py:112 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple folders of faces you wish to sort. The faces will be output to " +"separate sub-folders in the output_dir" +msgstr "" +"R|선택되면 input_dir는 정렬할 여러 개의 얼굴 폴더를 포함하는 상위 폴더여야 합" +"니다. 얼굴은 output_dir의 별도 하위 폴더로 출력됩니다" + +#: tools/sort/cli.py:121 +msgid "sort settings" +msgstr "정렬 설정" + +#: tools/sort/cli.py:124 +msgid "" +"R|Choose how images are sorted. Selecting a sort method gives the images a " +"new filename based on the order the image appears within the given method.\n" +"L|'none': Don't sort the images. When a 'group-by' method is selected, " +"selecting 'none' means that the files will be moved/copied into their " +"respective bins, but the files will keep their original filenames. Selecting " +"'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|이미지 정렬 방법을 선택합니다. 정렬 방법을 선택하면 이미지가 주어진 방법 내" +"에 나타나는 순서에 따라 이미지에 새 파일 이름이 지정됩니다.\n" +"L|'none': 이미지를 정렬하지 않습니다. 'group-by' 메서드를 선택한 경우 " +"'none'을 선택하면 파일이 각 bin으로 이동/복사되지만 파일은 원래 파일 이름을 " +"유지합니다. 'sort-by' 및 'group-by' 모두에 대해 'none'을 선택해도 아무 효과" +"가 없습니다" + +#: tools/sort/cli.py:136 tools/sort/cli.py:164 tools/sort/cli.py:184 +msgid "group settings" +msgstr "그룹 설정" + +#: tools/sort/cli.py:139 +#, fuzzy +#| msgid "" +#| "R|Selecting a group by method will move/copy files into numbered bins " +#| "based on the selected method.\n" +#| "L|'none': Don't bin the images. Folders will be sorted by the selected " +#| "'sort-by' but will not be binned, instead they will be sorted into a " +#| "single folder. Selecting 'none' for both 'sort-by' and 'group-by' will " +#| "do nothing" +msgid "" +"R|Selecting a group by method will move/copy files into numbered bins based " +"on the selected method.\n" +"L|'none': Don't bin the images. Folders will be sorted by the selected 'sort-" +"by' but will not be binned, instead they will be sorted into a single " +"folder. Selecting 'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|방법별로 그룹을 선택하면 선택한 방법에 따라 파일이 번호가 매겨진 빈으로 이" +"동/복사됩니다.\n" +"L|'none': 이미지를 버리지 않습니다. 폴더는 선택한 '정렬 기준'에 따라 정렬되지" +"만 버려지진 않고 단일 폴더로 정렬됩니다. 'sort-by' 및 'group-by' 모두에 대해 " +"'none'을 선택해도 아무 효과가 없습니다" + +#: tools/sort/cli.py:152 +msgid "" +"Whether to keep the original files in their original location. Choosing a " +"'sort-by' method means that the files have to be renamed. Selecting 'keep' " +"means that the original files will be kept, and the renamed files will be " +"created in the specified output folder. Unselecting keep means that the " +"original files will be moved and renamed based on the selected sort/group " +"criteria." +msgstr "" +"원본 파일을 원래 위치에 유지할지 여부입니다. '정렬 기준' 방법을 선택하면 파" +"일 이름을 변경해야 합니다. 'keep'을 선택하면 원래 파일이 유지되고 이름이 변경" +"된 파일이 지정된 출력 폴더에 생성됩니다. keep을 선택취소하면 선택한 정렬/그" +"룹 기준에 따라 원래 파일이 이동되고 이름이 변경됩니다." + +#: tools/sort/cli.py:167 +msgid "" +"R|Float value. Minimum threshold to use for grouping comparison with 'face-" +"cnn' 'hist' and 'face' methods.\n" +"The lower the value the more discriminating the grouping is. Leaving -1.0 " +"will allow Faceswap to choose the default value.\n" +"L|For 'face-cnn' 7.2 should be enough, with 4 being very discriminating. \n" +"L|For 'hist' 0.3 should be enough, with 0.2 being very discriminating. \n" +"L|For 'face' between 0.1 (more bins) to 0.5 (fewer bins) should be about " +"right.\n" +"Be careful setting a value that's too extrene in a directory with many " +"images, as this could result in a lot of folders being created. Defaults: " +"face-cnn 7.2, hist 0.3, face 0.25" +msgstr "" +"R|float 값. 'face-cnn', 'hist' 및 'face' 메서드와의 그룹 비교에 사용할 최소 " +"임계값입니다.\n" +"값이 낮을수록 그룹을 더 잘 구별할 수 있습니다. -1.0을 그대로 두면 Faceswap에" +"서 기본값을 선택할 수 있습니다.\n" +"L|'face-cnn'의 경우 7.2이면 충분하며, 4는 매우 많이 구별된다. \n" +"L|'hist'의 경우 0.3이면 충분하며, 0.2는 매우 많이 구별된다. \n" +"L|0.1(더 많은 빈)에서 0.5(더 적은 빈) 사이의 '얼굴'의 경우는 거의 오른쪽이어" +"야 합니다.\n" +"이미지가 많은 디렉터리에서 너무 극단적인 값을 설정하면 폴더가 많이 생성될 수 " +"있으므로 주의하십시오. 기본값: face-cnn 7.2, hist 0.3, face 0.25" + +#: tools/sort/cli.py:187 +#, fuzzy, python-format +#| msgid "" +#| "R|Integer value. Used to control the number of bins created for grouping " +#| "by: any 'blur' methods, 'color' methods or 'face metric' methods " +#| "('distance', 'size') and 'orientation; methods ('yaw', 'pitch'). For any " +#| "other grouping methods see the '-t' ('--threshold') option.\n" +#| "L|For 'face metric' methods the bins are filled, according the the " +#| "distribution of faces between the minimum and maximum chosen metric.\n" +#| "L|For 'color' methods the number of bins represents the divider of the " +#| "percentage of colored pixels. Eg. For a bin number of '5': The first " +#| "folder will have the faces with 0%% to 20%% colored pixels, second 21%% " +#| "to 40%%, etc. Any empty bins will be deleted, so you may end up with " +#| "fewer bins than selected.\n" +#| "L|For 'blur' methods folder 0 will be the least blurry, while the last " +#| "folder will be the blurriest.\n" +#| "L|For 'orientation' methods the number of bins is dictated by how much " +#| "180 degrees is divided. Eg. If 18 is selected, then each folder will be a " +#| "10 degree increment. Folder 0 will contain faces looking the most to the " +#| "left/down whereas the last folder will contain the faces looking the most " +#| "to the right/up. NB: Some bins may be empty if faces do not fit the " +#| "criteria.\n" +#| "Default value: 5" +msgid "" +"R|Integer value. Used to control the number of bins created for grouping by: " +"any 'blur' methods, 'color' methods or 'face metric' methods ('distance', " +"'size') and 'orientation; methods ('yaw', 'pitch'). For any other grouping " +"methods see the '-t' ('--threshold') option.\n" +"L|For 'face metric' methods the bins are filled, according the the " +"distribution of faces between the minimum and maximum chosen metric.\n" +"L|For 'color' methods the number of bins represents the divider of the " +"percentage of colored pixels. Eg. For a bin number of '5': The first folder " +"will have the faces with 0%% to 20%% colored pixels, second 21%% to 40%%, " +"etc. Any empty bins will be deleted, so you may end up with fewer bins than " +"selected.\n" +"L|For 'blur' methods folder 0 will be the least blurry, while the last " +"folder will be the blurriest.\n" +"L|For 'orientation' methods the number of bins is dictated by how much 180 " +"degrees is divided. Eg. If 18 is selected, then each folder will be a 10 " +"degree increment. Folder 0 will contain faces looking the most to the left/" +"down whereas the last folder will contain the faces looking the most to the " +"right/up. NB: Some bins may be empty if faces do not fit the criteria. \n" +"Default value: 5" +msgstr "" +"R| 정수 값. 그룹화를 위해 생성된 bins의 수를 제어하는 데 사용됩니다. 임의의 " +"'blur' 방법, 'color' 방법 또는 'face metric' 방법('거리', '크기'), " +"'orientation' 방법('yaw', 'pitch'). 다른 그룹화 방법은 '-t'('--임계값') 옵션" +"을 참조하십시오.\n" +"L|'face metric' 방법의 경우 선택한 최소 메트릭과 최대 메트릭 사이의 얼굴 분포" +"에 따라 bins가 채워집니다.\n" +"L|'color' 방법의 경우 bins의 수는 색상 픽셀의 백분율을 나눈 값을 나타냅니다. " +"예: bin 번호가 '5'인 경우: 첫 번째 폴더는 0%%에서 20%%의 색상 픽셀을 가진 얼" +"굴을 가질 것이고, 두 번째는 21%%에서 40%% 등을 가질 것이다. 텅 빈 bins는 삭제" +"되므로 선택한 bins보다 더 적은 bins을 가질 수 있습니다.\n" +"L|'blur' 메서드의 경우 폴더 0이 가장 흐림이 적으며 마지막 폴더가 가장 흐림이 " +"많습니다.\n" +"L|'orientation' 방법의 경우 bins의 수는 180도를 얼마나 나누느냐에 따라 결정됩" +"니다. 예: 18을 선택하면 각 폴더가 10도씩 증가합니다. 폴더 0은 왼쪽/아래쪽 얼" +"굴을 가장 많이 포함하는 반면, 마지막 폴더는 오른쪽/위 얼굴을 가장 많이 포함합" +"니다. 주의: 얼굴이 기준에 맞지 않으면 일부 bins가 비어 있을 수 있습니다.\n" +"기본값: 5" + +#: tools/sort/cli.py:207 tools/sort/cli.py:217 +msgid "settings" +msgstr "설정" + +#: tools/sort/cli.py:210 +msgid "" +"Logs file renaming changes if grouping by renaming, or it logs the file " +"copying/movement if grouping by folders. If no log file is specified with " +"'--log-file', then a 'sort_log.json' file will be created in the input " +"directory." +msgstr "" +"만약 renaming별로 그룹화하면 로그 파일에서 renaming이 변경됩니다. 또는 폴더별" +"로 그룹화하는 경우 파일 복사/이동을 기록합니다. '--log-file'로 로그 파일을 지" +"정하지 않으면 'sort_log.json' 파일이 입력 디렉토리에 생성됩니다." + +#: tools/sort/cli.py:221 +msgid "" +"Specify a log file to use for saving the renaming or grouping information. " +"If specified extension isn't 'json' or 'yaml', then json will be used as the " +"serializer, with the supplied filename. Default: sort_log.json" +msgstr "" +"_renaming 또는 grouping 정보를 저장하는 데 사용할 로그 파일을 지정합니다. 지" +"정된 확장자가 'json' 또는 'yaml'이 아니면 json이 제공된 파일 이름과 함께 직렬" +"화기로 사용됩니다. 기본값: sort_log.json" + +#~ msgid " option is deprecated. Use 'yaw'" +#~ msgstr " 이 옵션은 더 이상 사용되지 않습니다. 'yaw'를 사용하세요" + +#~ msgid " option is deprecated. Use 'color-black'" +#~ msgstr " 이 옵션은 더 이상 사용되지 않습니다. 'color-black'을 사용하세요" + +#~ msgid "output" +#~ msgstr "출력" + +#~ msgid "" +#~ "Deprecated and no longer used. The final processing will be dictated by " +#~ "the sort/group by methods and whether 'keep_original' is selected." +#~ msgstr "" +#~ "폐기되었고 더 이상 사용되지 않습니다. 최종 처리는 sort/group-by 메서드와 " +#~ "'keep_original'이 선택되었는지 여부에 의해 결정됩니다." diff --git a/locales/lib.cli.args.pot b/locales/lib.cli.args.pot index f788ab79c6..03a77c5a74 100644 --- a/locales/lib.cli.args.pot +++ b/locales/lib.cli.args.pot @@ -1,414 +1,50 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-05-17 18:04+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:06+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - -#: lib/cli/args.py:177 lib/cli/args.py:187 lib/cli/args.py:195 -#: lib/cli/args.py:205 +#: lib/cli/args.py:188 lib/cli/args.py:199 lib/cli/args.py:208 +#: lib/cli/args.py:219 msgid "Global Options" msgstr "" -#: lib/cli/args.py:178 +#: lib/cli/args.py:190 msgid "" -"R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond to any GPU(s) that you do not wish to be made available to Faceswap. Selecting all GPUs here will force Faceswap into CPU mode.\n" +"R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond " +"to any GPU(s) that you do not wish to be made available to Faceswap. " +"Selecting all GPUs here will force Faceswap into CPU mode.\n" "L|{}" msgstr "" -#: lib/cli/args.py:188 -msgid "Optionally overide the saved config with the path to a custom config file." -msgstr "" - -#: lib/cli/args.py:196 -msgid "Log level. Stick with INFO or VERBOSE unless you need to file an error report. Be careful with TRACE as it will generate a lot of data" -msgstr "" - -#: lib/cli/args.py:206 -msgid "Path to store the logfile. Leave blank to store in the faceswap folder" -msgstr "" - -#: lib/cli/args.py:299 lib/cli/args.py:308 lib/cli/args.py:316 -#: lib/cli/args.py:630 lib/cli/args.py:639 -msgid "Data" -msgstr "" - -#: lib/cli/args.py:300 -msgid "Input directory or video. Either a directory containing the image files you wish to process or path to a video file. NB: This should be the source video/frames NOT the source faces." -msgstr "" - -#: lib/cli/args.py:309 -msgid "Output directory. This is where the converted files will be saved." -msgstr "" - -#: lib/cli/args.py:317 -msgid "Optional path to an alignments file. Leave blank if the alignments file is at the default location." -msgstr "" - -#: lib/cli/args.py:340 -msgid "" -"Extract faces from image or video sources.\n" -"Extraction plugins can be configured in the 'Settings' Menu" -msgstr "" - -#: lib/cli/args.py:365 lib/cli/args.py:381 lib/cli/args.py:393 -#: lib/cli/args.py:428 lib/cli/args.py:446 lib/cli/args.py:458 -#: lib/cli/args.py:649 lib/cli/args.py:676 lib/cli/args.py:712 -msgid "Plugins" -msgstr "" - -#: lib/cli/args.py:366 -msgid "" -"R|Detector to use. Some of these have configurable settings in '/config/extract.ini' or 'Settings > Configure Extract 'Plugins':\n" -"L|cv2-dnn: A CPU only extractor which is the least reliable and least resource intensive. Use this if not using a GPU and time is important.\n" -"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources than other GPU detectors but can often return more false positives.\n" -"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and fewer false positives than other GPU detectors, but is a lot more resource intensive." -msgstr "" - -#: lib/cli/args.py:382 -msgid "" -"R|Aligner to use.\n" -"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, but less accurate. Only use this if not using a GPU and time is important.\n" -"L|fan: Best aligner. Fast on GPU, slow on CPU." -msgstr "" - -#: lib/cli/args.py:394 -msgid "" -"R|Additional Masker(s) to use. The masks generated here will all take up GPU RAM. You can select none, one or multiple masks, but the extraction may take longer the more you select. NB: The Extended and Components (landmark based) masks are automatically generated on extraction.\n" -"L|bisenet-fp: Relatively lightweight NN based mask that provides more refined control over the area to be masked including full head masking (configurable in mask settings).\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal faces clear of obstructions. Profile faces and obstructions may result in sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been specifically trained to recognize some facial obstructions (hands and eyeglasses). Profile faces may result in sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been trained by community members and will need testing for further description. Profile faces may result in sub-par performance.\n" -"The auto generated masks are as follows:\n" -"L|components: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks and the mask is extended upwards onto the forehead.\n" -"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" -msgstr "" - -#: lib/cli/args.py:429 -msgid "" -"R|Performing normalization can help the aligner better align faces with difficult lighting conditions at an extraction speed cost. Different methods will yield different results on different sets. NB: This does not impact the output face, just the input to the aligner.\n" -"L|none: Don't perform normalization on the face.\n" -"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the face.\n" -"L|hist: Equalize the histograms on the RGB channels.\n" -"L|mean: Normalize the face colors to the mean." -msgstr "" - -#: lib/cli/args.py:447 -msgid "The number of times to re-feed the detected face into the aligner. Each time the face is re-fed into the aligner the bounding box is adjusted by a small amount. The final landmarks are then averaged from each iteration. Helps to remove 'micro-jitter' but at the cost of slower extraction speed. The more times the face is re-fed into the aligner, the less micro-jitter should occur but the longer extraction will take." -msgstr "" - -#: lib/cli/args.py:459 -msgid "If a face isn't found, rotate the images to try to find a face. Can find more faces at the cost of extraction speed. Pass in a single number to use increments of that size up to 360, or pass in a list of numbers to enumerate exactly what angles to check." -msgstr "" - -#: lib/cli/args.py:471 lib/cli/args.py:481 lib/cli/args.py:494 -#: lib/cli/args.py:508 lib/cli/args.py:749 lib/cli/args.py:763 -#: lib/cli/args.py:776 lib/cli/args.py:790 -msgid "Face Processing" -msgstr "" - -#: lib/cli/args.py:472 -msgid "Filters out faces detected below this size. Length, in pixels across the diagonal of the bounding box. Set to 0 for off" -msgstr "" - -#: lib/cli/args.py:482 lib/cli/args.py:764 -msgid "Optionally filter out people who you do not wish to process by passing in an image of that person. Should be a front portrait with a single person in the image. Multiple images can be added space separated. NB: Using face filter will significantly decrease extraction speed and its accuracy cannot be guaranteed." -msgstr "" - -#: lib/cli/args.py:495 lib/cli/args.py:777 -msgid "Optionally select people you wish to process by passing in an image of that person. Should be a front portrait with a single person in the image. Multiple images can be added space separated. NB: Using face filter will significantly decrease extraction speed and its accuracy cannot be guaranteed." -msgstr "" - -#: lib/cli/args.py:509 lib/cli/args.py:791 -msgid "For use with the optional nfilter/filter files. Threshold for positive face recognition. Lower values are stricter. NB: Using face filter will significantly decrease extraction speed and its accuracy cannot be guaranteed." -msgstr "" - -#: lib/cli/args.py:520 lib/cli/args.py:532 lib/cli/args.py:544 -#: lib/cli/args.py:556 -msgid "output" -msgstr "" - -#: lib/cli/args.py:521 -msgid "The output size of extracted faces. Make sure that the model you intend to train supports your required size. This will only need to be changed for hi-res models." -msgstr "" - -#: lib/cli/args.py:533 -msgid "Extract every 'nth' frame. This option will skip frames when extracting faces. For example a value of 1 will extract faces from every frame, a value of 10 will extract faces from every 10th frame." -msgstr "" - -#: lib/cli/args.py:545 -msgid "Automatically save the alignments file after a set amount of frames. By default the alignments file is only saved at the end of the extraction process. NB: If extracting in 2 passes then the alignments file will only start to be saved out during the second pass. WARNING: Don't interrupt the script when writing the file because it might get corrupted. Set to 0 to turn off" -msgstr "" - -#: lib/cli/args.py:557 -msgid "Draw landmarks on the ouput faces for debugging purposes." -msgstr "" - -#: lib/cli/args.py:563 lib/cli/args.py:572 lib/cli/args.py:580 -#: lib/cli/args.py:587 lib/cli/args.py:803 lib/cli/args.py:814 -#: lib/cli/args.py:822 lib/cli/args.py:841 lib/cli/args.py:847 -msgid "settings" -msgstr "" - -#: lib/cli/args.py:564 -msgid "Don't run extraction in parallel. Will run each part of the extraction process separately (one after the other) rather than all at the smae time. Useful if VRAM is at a premium." -msgstr "" - -#: lib/cli/args.py:573 -msgid "Skips frames that have already been extracted and exist in the alignments file" -msgstr "" - -#: lib/cli/args.py:581 -msgid "Skip frames that already have detected faces in the alignments file" -msgstr "" - -#: lib/cli/args.py:588 -msgid "Skip saving the detected faces to disk. Just create an alignments file" -msgstr "" - -#: lib/cli/args.py:610 -msgid "" -"Swap the original faces in a source video/images to your final faces.\n" -"Conversion plugins can be configured in the 'Settings' Menu" -msgstr "" - -#: lib/cli/args.py:631 -msgid "Only required if converting from images to video. Provide The original video that the source frames were extracted from (for extracting the fps and audio)." -msgstr "" - -#: lib/cli/args.py:640 -msgid "Model directory. The directory containing the trained model you wish to use for conversion." -msgstr "" - -#: lib/cli/args.py:650 -msgid "" -"R|Performs color adjustment to the swapped face. Some of these options have configurable settings in '/config/convert.ini' or 'Settings > Configure Convert Plugins':\n" -"L|avg-color: Adjust the mean of each color channel in the swapped reconstruction to equal the mean of the masked area in the original image.\n" -"L|color-transfer: Transfers the color distribution from the source to the target image using the mean and standard deviations of the L*a*b* color space.\n" -"L|manual-balance: Manually adjust the balance of the image in a variety of color spaces. Best used with the Preview tool to set correct values.\n" -"L|match-hist: Adjust the histogram of each color channel in the swapped reconstruction to equal the histogram of the masked area in the original image.\n" -"L|seamless-clone: Use cv2's seamless clone function to remove extreme gradients at the mask seam by smoothing colors. Generally does not give very satisfactory results.\n" -"L|none: Don't perform color adjustment." -msgstr "" - -#: lib/cli/args.py:677 -msgid "" -"R|Masker to use. NB: The mask you require must exist within the alignments file. You can add additional masks with the Mask Tool.\n" -"L|none: Don't use a mask.\n" -"L|bisenet-fp-face: Relatively lightweight NN based mask that provides more refined control over the area to be masked (configurable in mask settings). Use this version of bisenet-fp if your model is trained with 'face' or 'legacy' centering.\n" -"L|bisenet-fp-head: Relatively lightweight NN based mask that provides more refined control over the area to be masked (configurable in mask settings). Use this version of bisenet-fp if your model is trained with 'head' centering.\n" -"L|components: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks and the mask is extended upwards onto the forehead.\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal faces clear of obstructions. Profile faces and obstructions may result in sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been specifically trained to recognize some facial obstructions (hands and eyeglasses). Profile faces may result in sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been trained by community members and will need testing for further description. Profile faces may result in sub-par performance.\n" -"L|predicted: If the 'Learn Mask' option was enabled during training, this will use the mask that was created by the trained model." -msgstr "" - -#: lib/cli/args.py:713 -msgid "" -"R|The plugin to use to output the converted images. The writers are configurable in '/config/convert.ini' or 'Settings > Configure Convert Plugins:'\n" -"L|ffmpeg: [video] Writes out the convert straight to video. When the input is a series of images then the '-ref' (--reference-video) parameter must be set.\n" -"L|gif: [animated image] Create an animated gif.\n" -"L|opencv: [images] The fastest image writer, but less options and formats than other plugins.\n" -"L|pillow: [images] Slower than opencv, but has more options and supports more formats." -msgstr "" - -#: lib/cli/args.py:732 lib/cli/args.py:739 lib/cli/args.py:833 -msgid "Frame Processing" -msgstr "" - -#: lib/cli/args.py:733 -msgid "Scale the final output frames by this amount. 100%% will output the frames at source dimensions. 50%% at half size 200%% at double size" -msgstr "" - -#: lib/cli/args.py:740 -msgid "Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use --frame-ranges 10-50 90-100. Frames falling outside of the selected range will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are converting from images, then the filenames must end with the frame-number!" -msgstr "" - -#: lib/cli/args.py:750 -msgid "If you have not cleansed your alignments file, then you can filter out faces by defining a folder here that contains the faces extracted from your input files/video. If this folder is defined, then only faces that exist within your alignments file and also exist within the specified folder will be converted. Leaving this blank will convert all faces that exist within the alignments file." -msgstr "" - -#: lib/cli/args.py:804 -msgid "The maximum number of parallel processes for performing conversion. Converting images is system RAM heavy so it is possible to run out of memory if you have a lot of processes and not enough RAM to accommodate them all. Setting this to 0 will use the maximum available. No matter what you set this to, it will never attempt to use more processes than are available on your system. If singleprocess is enabled this setting will be ignored." -msgstr "" - -#: lib/cli/args.py:815 -msgid "[LEGACY] This only needs to be selected if a legacy model is being loaded or if there are multiple models in the model folder" -msgstr "" - -#: lib/cli/args.py:823 -msgid "Enable On-The-Fly Conversion. NOT recommended. You should generate a clean alignments file for your destination video. However, if you wish you can generate the alignments on-the-fly by enabling this option. This will use an inferior extraction pipeline and will lead to substandard results. If an alignments file is found, this option will be ignored." -msgstr "" - -#: lib/cli/args.py:834 -msgid "When used with --frame-ranges outputs the unchanged frames that are not processed instead of discarding them." -msgstr "" - -#: lib/cli/args.py:842 -msgid "Swap the model. Instead converting from of A -> B, converts B -> A" -msgstr "" - -#: lib/cli/args.py:848 -msgid "Disable multiprocessing. Slower but less resource intensive." -msgstr "" - -#: lib/cli/args.py:864 -msgid "" -"Train a model on extracted original (A) and swap (B) faces.\n" -"Training models can take a long time. Anything from 24hrs to over a week\n" -"Model plugins can be configured in the 'Settings' Menu" -msgstr "" - -#: lib/cli/args.py:883 lib/cli/args.py:892 -msgid "faces" -msgstr "" - -#: lib/cli/args.py:884 -msgid "Input directory. A directory containing training images for face A. This is the original face, i.e. the face that you want to remove and replace with face B." -msgstr "" - -#: lib/cli/args.py:893 -msgid "Input directory. A directory containing training images for face B. This is the swap face, i.e. the face that you want to place onto the head of person A." -msgstr "" - -#: lib/cli/args.py:901 lib/cli/args.py:913 lib/cli/args.py:929 -#: lib/cli/args.py:954 lib/cli/args.py:964 -msgid "model" -msgstr "" - -#: lib/cli/args.py:902 -msgid "Model directory. This is where the training data will be stored. You should always specify a new folder for new models. If starting a new model, select either an empty folder, or a folder which does not exist (which will be created). If continuing to train an existing model, specify the location of the existing model." -msgstr "" - -#: lib/cli/args.py:914 +#: lib/cli/args.py:201 msgid "" -"R|Load the weights from a pre-existing model into a newly created model. For most models this will load weights from the Encoder of the given model into the encoder of the newly created model. Some plugins may have specific configuration options allowing you to load weights from other layers. Weights will only be loaded when creating a new model. This option will be ignored if you are resuming an existing model. Generally you will also want to 'freeze-weights' whilst the rest of your model catches up with your Encoder.\n" -"NB: Weights can only be loaded from models of the same plugin as you intend to train." +"Optionally overide the saved config with the path to a custom config file." msgstr "" -#: lib/cli/args.py:930 +#: lib/cli/args.py:210 msgid "" -"R|Select which trainer to use. Trainers can be configured from the Settings menu or the config folder.\n" -"L|original: The original model created by /u/deepfakes.\n" -"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' for full dfaker method.\n" -"L|dfl-h128: 128px in/out model from deepfacelab\n" -"L|dfl-sae: Adaptable model from deepfacelab\n" -"L|dlight: A lightweight, high resolution DFaker variant.\n" -"L|iae: A model that uses intermediate layers to try to get better details\n" -"L|lightweight: A lightweight model for low-end cards. Don't expect great results. Can train as low as 1.6GB with batch size 8.\n" -"L|realface: A high detail, dual density model based on DFaker, with customizable in/out resolution. The autoencoders are unbalanced so B>A swaps won't work so well. By andenixa et al. Very configurable.\n" -"L|unbalanced: 128px in/out model from andenixa. The autoencoders are unbalanced so B>A swaps won't work so well. Very configurable.\n" -"L|villain: 128px in/out model from villainguy. Very resource hungry (You will require a GPU with a fair amount of VRAM). Good for details, but more susceptible to color differences." -msgstr "" - -#: lib/cli/args.py:955 -msgid "Output a summary of the model and exit. If a model folder is provided then a summary of the saved model is displayed. Otherwise a summary of the model that would be created by the chosen plugin and configuration settings is displayed." -msgstr "" - -#: lib/cli/args.py:965 -msgid "Freeze the weights of the model. Freezing weights means that some of the parameters in the model will no longer continue to learn, but those that are not frozen will continue to learn. For most models, this will freeze the encoder, but some models may have configuration options for freezing other layers." -msgstr "" - -#: lib/cli/args.py:978 lib/cli/args.py:990 lib/cli/args.py:1001 -#: lib/cli/args.py:1087 -msgid "training" +"Log level. Stick with INFO or VERBOSE unless you need to file an error " +"report. Be careful with TRACE as it will generate a lot of data" msgstr "" -#: lib/cli/args.py:979 -msgid "Batch size. This is the number of images processed through the model for each side per iteration. NB: As the model is fed 2 sides at a time, the actual number of images within the model at any one time is double the number that you set here. Larger batches require more GPU RAM." -msgstr "" - -#: lib/cli/args.py:991 -msgid "Length of training in iterations. This is only really used for automation. There is no 'correct' number of iterations a model should be trained for. You should stop training when you are happy with the previews. However, if you want the model to stop automatically at a set number of iterations, you can set that value here." -msgstr "" - -#: lib/cli/args.py:1002 -msgid "Use the Tensorflow Mirrored Distrubution Strategy to train on multiple GPUs." -msgstr "" - -#: lib/cli/args.py:1012 lib/cli/args.py:1022 -msgid "Saving" -msgstr "" - -#: lib/cli/args.py:1013 -msgid "Sets the number of iterations between each model save." -msgstr "" - -#: lib/cli/args.py:1023 -msgid "Sets the number of iterations before saving a backup snapshot of the model in it's current state. Set to 0 for off." -msgstr "" - -#: lib/cli/args.py:1030 lib/cli/args.py:1041 lib/cli/args.py:1052 -msgid "timelapse" -msgstr "" - -#: lib/cli/args.py:1031 -msgid "Optional for creating a timelapse. Timelapse will save an image of your selected faces into the timelapse-output folder at every save iteration. This should be the input folder of 'A' faces that you would like to use for creating the timelapse. You must also supply a --timelapse-output and a --timelapse-input-B parameter." -msgstr "" - -#: lib/cli/args.py:1042 -msgid "Optional for creating a timelapse. Timelapse will save an image of your selected faces into the timelapse-output folder at every save iteration. This should be the input folder of 'B' faces that you would like to use for creating the timelapse. You must also supply a --timelapse-output and a --timelapse-input-A parameter." -msgstr "" - -#: lib/cli/args.py:1053 -msgid "Optional for creating a timelapse. Timelapse will save an image of your selected faces into the timelapse-output folder at every save iteration. If the input folders are supplied but no output folder, it will default to your model folder /timelapse/" -msgstr "" - -#: lib/cli/args.py:1065 lib/cli/args.py:1072 lib/cli/args.py:1079 -msgid "preview" -msgstr "" - -#: lib/cli/args.py:1066 -msgid "Percentage amount to scale the preview by. 100%% is the model output size." -msgstr "" - -#: lib/cli/args.py:1073 -msgid "Show training preview output. in a separate window." -msgstr "" - -#: lib/cli/args.py:1080 -msgid "Writes the training result to a file. The image will be stored in the root of your FaceSwap folder." -msgstr "" - -#: lib/cli/args.py:1088 -msgid "Disables TensorBoard logging. NB: Disabling logs means that you will not be able to use the graph or analysis for this session in the GUI." -msgstr "" - -#: lib/cli/args.py:1095 lib/cli/args.py:1104 lib/cli/args.py:1113 -#: lib/cli/args.py:1122 -msgid "augmentation" -msgstr "" - -#: lib/cli/args.py:1096 -msgid "Warps training faces to closely matched Landmarks from the opposite face-set rather than randomly warping the face. This is the 'dfaker' way of doing warping." -msgstr "" - -#: lib/cli/args.py:1105 -msgid "To effectively learn, a random set of images are flipped horizontally. Sometimes it is desirable for this not to occur. Generally this should be left off except for during 'fit training'." -msgstr "" - -#: lib/cli/args.py:1114 -msgid "Color augmentation helps make the model less susceptible to color differences between the A and B sets, at an increased training time cost. Enable this option to disable color augmentation." -msgstr "" - -#: lib/cli/args.py:1123 -msgid "Warping is integral to training the Neural Network. This option should only be enabled towards the very end of training to try to bring out more detail. Think of it as 'fine-tuning'. Enabling this option from the beginning is likely to kill a model and lead to terrible results." +#: lib/cli/args.py:220 +msgid "Path to store the logfile. Leave blank to store in the faceswap folder" msgstr "" -#: lib/cli/args.py:1148 +#: lib/cli/args.py:319 msgid "Output to Shell console instead of GUI console" msgstr "" - diff --git a/locales/lib.cli.args_extract_convert.pot b/locales/lib.cli.args_extract_convert.pot new file mode 100644 index 0000000000..d650f81997 --- /dev/null +++ b/locales/lib.cli.args_extract_convert.pot @@ -0,0 +1,455 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-12 11:56+0100\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: lib/cli/args_extract_convert.py:46 lib/cli/args_extract_convert.py:56 +#: lib/cli/args_extract_convert.py:64 lib/cli/args_extract_convert.py:122 +#: lib/cli/args_extract_convert.py:483 lib/cli/args_extract_convert.py:492 +msgid "Data" +msgstr "" + +#: lib/cli/args_extract_convert.py:48 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file. NB: This should be the source video/" +"frames NOT the source faces." +msgstr "" + +#: lib/cli/args_extract_convert.py:57 +msgid "Output directory. This is where the converted files will be saved." +msgstr "" + +#: lib/cli/args_extract_convert.py:66 +msgid "" +"Optional path to an alignments file. Leave blank if the alignments file is " +"at the default location." +msgstr "" + +#: lib/cli/args_extract_convert.py:97 +msgid "" +"Extract faces from image or video sources.\n" +"Extraction plugins can be configured in the 'Settings' Menu" +msgstr "" + +#: lib/cli/args_extract_convert.py:124 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple videos and/or folders of images you wish to extract from. The faces " +"will be output to separate sub-folders in the output_dir." +msgstr "" + +#: lib/cli/args_extract_convert.py:133 lib/cli/args_extract_convert.py:152 +#: lib/cli/args_extract_convert.py:167 lib/cli/args_extract_convert.py:206 +#: lib/cli/args_extract_convert.py:224 lib/cli/args_extract_convert.py:237 +#: lib/cli/args_extract_convert.py:247 lib/cli/args_extract_convert.py:257 +#: lib/cli/args_extract_convert.py:503 lib/cli/args_extract_convert.py:529 +#: lib/cli/args_extract_convert.py:568 +msgid "Plugins" +msgstr "" + +#: lib/cli/args_extract_convert.py:135 +msgid "" +"R|Detector to use. Some of these have configurable settings in '/config/" +"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" +"L|cv2-dnn: A CPU only extractor which is the least reliable and least " +"resource intensive. Use this if not using a GPU and time is important.\n" +"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " +"than other GPU detectors but can often return more false positives.\n" +"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " +"fewer false positives than other GPU detectors, but is a lot more resource " +"intensive.\n" +"L|external: Import a face detection bounding box from a json file. " +"(configurable in Detect settings)" +msgstr "" + +#: lib/cli/args_extract_convert.py:154 +msgid "" +"R|Aligner to use.\n" +"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " +"but less accurate. Only use this if not using a GPU and time is important.\n" +"L|fan: Best aligner. Fast on GPU, slow on CPU.\n" +"L|external: Import 68 point 2D landmarks or an aligned bounding box from a " +"json file. (configurable in Align settings)" +msgstr "" + +#: lib/cli/args_extract_convert.py:169 +msgid "" +"R|Additional Masker(s) to use. The masks generated here will all take up GPU " +"RAM. You can select none, one or multiple masks, but the extraction may take " +"longer the more you select. NB: The Extended and Components (landmark based) " +"masks are automatically generated on extraction.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU so will not use any additional VRAM.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"The auto generated masks are as follows:\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" +msgstr "" + +#: lib/cli/args_extract_convert.py:208 +msgid "" +"R|Performing normalization can help the aligner better align faces with " +"difficult lighting conditions at an extraction speed cost. Different methods " +"will yield different results on different sets. NB: This does not impact the " +"output face, just the input to the aligner.\n" +"L|none: Don't perform normalization on the face.\n" +"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"L|hist: Equalize the histograms on the RGB channels.\n" +"L|mean: Normalize the face colors to the mean." +msgstr "" + +#: lib/cli/args_extract_convert.py:226 +msgid "" +"The number of times to re-feed the detected face into the aligner. Each time " +"the face is re-fed into the aligner the bounding box is adjusted by a small " +"amount. The final landmarks are then averaged from each iteration. Helps to " +"remove 'micro-jitter' but at the cost of slower extraction speed. The more " +"times the face is re-fed into the aligner, the less micro-jitter should " +"occur but the longer extraction will take." +msgstr "" + +#: lib/cli/args_extract_convert.py:239 +msgid "" +"Re-feed the initially found aligned face through the aligner. Can help " +"produce better alignments for faces that are rotated beyond 45 degrees in " +"the frame or are at extreme angles. Slows down extraction." +msgstr "" + +#: lib/cli/args_extract_convert.py:249 +msgid "" +"If a face isn't found, rotate the images to try to find a face. Can find " +"more faces at the cost of extraction speed. Pass in a single number to use " +"increments of that size up to 360, or pass in a list of numbers to enumerate " +"exactly what angles to check." +msgstr "" + +#: lib/cli/args_extract_convert.py:259 +msgid "" +"Obtain and store face identity encodings from VGGFace2. Slows down extract a " +"little, but will save time if using 'sort by face'" +msgstr "" + +#: lib/cli/args_extract_convert.py:269 lib/cli/args_extract_convert.py:280 +#: lib/cli/args_extract_convert.py:293 lib/cli/args_extract_convert.py:307 +#: lib/cli/args_extract_convert.py:614 lib/cli/args_extract_convert.py:623 +#: lib/cli/args_extract_convert.py:638 lib/cli/args_extract_convert.py:651 +#: lib/cli/args_extract_convert.py:665 +msgid "Face Processing" +msgstr "" + +#: lib/cli/args_extract_convert.py:271 +msgid "" +"Filters out faces detected below this size. Length, in pixels across the " +"diagonal of the bounding box. Set to 0 for off" +msgstr "" + +#: lib/cli/args_extract_convert.py:282 +msgid "" +"Optionally filter out people who you do not wish to extract by passing in " +"images of those people. Should be a small variety of images at different " +"angles and in different conditions. A folder containing the required images " +"or multiple image files, space separated, can be selected." +msgstr "" + +#: lib/cli/args_extract_convert.py:295 +msgid "" +"Optionally select people you wish to extract by passing in images of that " +"person. Should be a small variety of images at different angles and in " +"different conditions A folder containing the required images or multiple " +"image files, space separated, can be selected." +msgstr "" + +#: lib/cli/args_extract_convert.py:309 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Higher values are stricter." +msgstr "" + +#: lib/cli/args_extract_convert.py:318 lib/cli/args_extract_convert.py:331 +#: lib/cli/args_extract_convert.py:344 lib/cli/args_extract_convert.py:356 +msgid "output" +msgstr "" + +#: lib/cli/args_extract_convert.py:320 +msgid "" +"The output size of extracted faces. Make sure that the model you intend to " +"train supports your required size. This will only need to be changed for hi-" +"res models." +msgstr "" + +#: lib/cli/args_extract_convert.py:333 +msgid "" +"Extract every 'nth' frame. This option will skip frames when extracting " +"faces. For example a value of 1 will extract faces from every frame, a value " +"of 10 will extract faces from every 10th frame." +msgstr "" + +#: lib/cli/args_extract_convert.py:346 +msgid "" +"Automatically save the alignments file after a set amount of frames. By " +"default the alignments file is only saved at the end of the extraction " +"process. NB: If extracting in 2 passes then the alignments file will only " +"start to be saved out during the second pass. WARNING: Don't interrupt the " +"script when writing the file because it might get corrupted. Set to 0 to " +"turn off" +msgstr "" + +#: lib/cli/args_extract_convert.py:357 +msgid "Draw landmarks on the ouput faces for debugging purposes." +msgstr "" + +#: lib/cli/args_extract_convert.py:363 lib/cli/args_extract_convert.py:373 +#: lib/cli/args_extract_convert.py:381 lib/cli/args_extract_convert.py:388 +#: lib/cli/args_extract_convert.py:678 lib/cli/args_extract_convert.py:691 +#: lib/cli/args_extract_convert.py:712 lib/cli/args_extract_convert.py:718 +msgid "settings" +msgstr "" + +#: lib/cli/args_extract_convert.py:365 +msgid "" +"Don't run extraction in parallel. Will run each part of the extraction " +"process separately (one after the other) rather than all at the same time. " +"Useful if VRAM is at a premium." +msgstr "" + +#: lib/cli/args_extract_convert.py:375 +msgid "" +"Skips frames that have already been extracted and exist in the alignments " +"file" +msgstr "" + +#: lib/cli/args_extract_convert.py:382 +msgid "Skip frames that already have detected faces in the alignments file" +msgstr "" + +#: lib/cli/args_extract_convert.py:389 +msgid "Skip saving the detected faces to disk. Just create an alignments file" +msgstr "" + +#: lib/cli/args_extract_convert.py:463 +msgid "" +"Swap the original faces in a source video/images to your final faces.\n" +"Conversion plugins can be configured in the 'Settings' Menu" +msgstr "" + +#: lib/cli/args_extract_convert.py:485 +msgid "" +"Only required if converting from images to video. Provide The original video " +"that the source frames were extracted from (for extracting the fps and " +"audio)." +msgstr "" + +#: lib/cli/args_extract_convert.py:494 +msgid "" +"Model directory. The directory containing the trained model you wish to use " +"for conversion." +msgstr "" + +#: lib/cli/args_extract_convert.py:505 +msgid "" +"R|Performs color adjustment to the swapped face. Some of these options have " +"configurable settings in '/config/convert.ini' or 'Settings > Configure " +"Convert Plugins':\n" +"L|avg-color: Adjust the mean of each color channel in the swapped " +"reconstruction to equal the mean of the masked area in the original image.\n" +"L|color-transfer: Transfers the color distribution from the source to the " +"target image using the mean and standard deviations of the L*a*b* color " +"space.\n" +"L|manual-balance: Manually adjust the balance of the image in a variety of " +"color spaces. Best used with the Preview tool to set correct values.\n" +"L|match-hist: Adjust the histogram of each color channel in the swapped " +"reconstruction to equal the histogram of the masked area in the original " +"image.\n" +"L|seamless-clone: Use cv2's seamless clone function to remove extreme " +"gradients at the mask seam by smoothing colors. Generally does not give very " +"satisfactory results.\n" +"L|none: Don't perform color adjustment." +msgstr "" + +#: lib/cli/args_extract_convert.py:531 +msgid "" +"R|Masker to use. NB: The mask you require must exist within the alignments " +"file. You can add additional masks with the Mask Tool.\n" +"L|none: Don't use a mask.\n" +"L|bisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"L|bisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"L|custom_face: Custom user created, face centered mask.\n" +"L|custom_head: Custom user created, head centered mask.\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"L|predicted: If the 'Learn Mask' option was enabled during training, this " +"will use the mask that was created by the trained model." +msgstr "" + +#: lib/cli/args_extract_convert.py:570 +msgid "" +"R|The plugin to use to output the converted images. The writers are " +"configurable in '/config/convert.ini' or 'Settings > Configure Convert " +"Plugins:'\n" +"L|ffmpeg: [video] Writes out the convert straight to video. When the input " +"is a series of images then the '-ref' (--reference-video) parameter must be " +"set.\n" +"L|gif: [animated image] Create an animated gif.\n" +"L|opencv: [images] The fastest image writer, but less options and formats " +"than other plugins.\n" +"L|patch: [images] Outputs the raw swapped face patch, along with the " +"transformation matrix required to re-insert the face back into the original " +"frame. Use this option if you wish to post-process and composite the final " +"face within external tools.\n" +"L|pillow: [images] Slower than opencv, but has more options and supports " +"more formats." +msgstr "" + +#: lib/cli/args_extract_convert.py:591 lib/cli/args_extract_convert.py:600 +#: lib/cli/args_extract_convert.py:703 +msgid "Frame Processing" +msgstr "" + +#: lib/cli/args_extract_convert.py:593 +#, python-format +msgid "" +"Scale the final output frames by this amount. 100%% will output the frames " +"at source dimensions. 50%% at half size 200%% at double size" +msgstr "" + +#: lib/cli/args_extract_convert.py:602 +msgid "" +"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " +"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " +"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " +"converting from images, then the filenames must end with the frame-number!" +msgstr "" + +#: lib/cli/args_extract_convert.py:616 +msgid "" +"Scale the swapped face by this percentage. Positive values will enlarge the " +"face, Negative values will shrink the face." +msgstr "" + +#: lib/cli/args_extract_convert.py:625 +msgid "" +"If you have not cleansed your alignments file, then you can filter out faces " +"by defining a folder here that contains the faces extracted from your input " +"files/video. If this folder is defined, then only faces that exist within " +"your alignments file and also exist within the specified folder will be " +"converted. Leaving this blank will convert all faces that exist within the " +"alignments file." +msgstr "" + +#: lib/cli/args_extract_convert.py:640 +msgid "" +"Optionally filter out people who you do not wish to process by passing in an " +"image of that person. Should be a front portrait with a single person in the " +"image. Multiple images can be added space separated. NB: Using face filter " +"will significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" + +#: lib/cli/args_extract_convert.py:653 +msgid "" +"Optionally select people you wish to process by passing in an image of that " +"person. Should be a front portrait with a single person in the image. " +"Multiple images can be added space separated. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" + +#: lib/cli/args_extract_convert.py:667 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Lower values are stricter. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" + +#: lib/cli/args_extract_convert.py:680 +msgid "" +"The maximum number of parallel processes for performing conversion. " +"Converting images is system RAM heavy so it is possible to run out of memory " +"if you have a lot of processes and not enough RAM to accommodate them all. " +"Setting this to 0 will use the maximum available. No matter what you set " +"this to, it will never attempt to use more processes than are available on " +"your system. If singleprocess is enabled this setting will be ignored." +msgstr "" + +#: lib/cli/args_extract_convert.py:693 +msgid "" +"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " +"alignments file for your destination video. However, if you wish you can " +"generate the alignments on-the-fly by enabling this option. This will use an " +"inferior extraction pipeline and will lead to substandard results. If an " +"alignments file is found, this option will be ignored." +msgstr "" + +#: lib/cli/args_extract_convert.py:705 +msgid "" +"When used with --frame-ranges outputs the unchanged frames that are not " +"processed instead of discarding them." +msgstr "" + +#: lib/cli/args_extract_convert.py:713 +msgid "Swap the model. Instead converting from of A -> B, converts B -> A" +msgstr "" + +#: lib/cli/args_extract_convert.py:719 +msgid "Disable multiprocessing. Slower but less resource intensive." +msgstr "" diff --git a/locales/lib.cli.args_train.pot b/locales/lib.cli.args_train.pot new file mode 100644 index 0000000000..9902e629a5 --- /dev/null +++ b/locales/lib.cli.args_train.pot @@ -0,0 +1,255 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:04+0000\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: lib/cli/args_train.py:30 +msgid "" +"Train a model on extracted original (A) and swap (B) faces.\n" +"Training models can take a long time. Anything from 24hrs to over a week\n" +"Model plugins can be configured in the 'Settings' Menu" +msgstr "" + +#: lib/cli/args_train.py:49 lib/cli/args_train.py:58 +msgid "faces" +msgstr "" + +#: lib/cli/args_train.py:51 +msgid "" +"Input directory. A directory containing training images for face A. This is " +"the original face, i.e. the face that you want to remove and replace with " +"face B." +msgstr "" + +#: lib/cli/args_train.py:60 +msgid "" +"Input directory. A directory containing training images for face B. This is " +"the swap face, i.e. the face that you want to place onto the head of person " +"A." +msgstr "" + +#: lib/cli/args_train.py:67 lib/cli/args_train.py:80 lib/cli/args_train.py:97 +#: lib/cli/args_train.py:123 lib/cli/args_train.py:133 +msgid "model" +msgstr "" + +#: lib/cli/args_train.py:69 +msgid "" +"Model directory. This is where the training data will be stored. You should " +"always specify a new folder for new models. If starting a new model, select " +"either an empty folder, or a folder which does not exist (which will be " +"created). If continuing to train an existing model, specify the location of " +"the existing model." +msgstr "" + +#: lib/cli/args_train.py:82 +msgid "" +"R|Load the weights from a pre-existing model into a newly created model. For " +"most models this will load weights from the Encoder of the given model into " +"the encoder of the newly created model. Some plugins may have specific " +"configuration options allowing you to load weights from other layers. " +"Weights will only be loaded when creating a new model. This option will be " +"ignored if you are resuming an existing model. Generally you will also want " +"to 'freeze-weights' whilst the rest of your model catches up with your " +"Encoder.\n" +"NB: Weights can only be loaded from models of the same plugin as you intend " +"to train." +msgstr "" + +#: lib/cli/args_train.py:99 +msgid "" +"R|Select which trainer to use. Trainers can be configured from the Settings " +"menu or the config folder.\n" +"L|original: The original model created by /u/deepfakes.\n" +"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " +"for full dfaker method.\n" +"L|dfl-h128: 128px in/out model from deepfacelab\n" +"L|dfl-sae: Adaptable model from deepfacelab\n" +"L|dlight: A lightweight, high resolution DFaker variant.\n" +"L|iae: A model that uses intermediate layers to try to get better details\n" +"L|lightweight: A lightweight model for low-end cards. Don't expect great " +"results. Can train as low as 1.6GB with batch size 8.\n" +"L|realface: A high detail, dual density model based on DFaker, with " +"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " +"won't work so well. By andenixa et al. Very configurable.\n" +"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " +"unbalanced so B>A swaps won't work so well. Very configurable.\n" +"L|villain: 128px in/out model from villainguy. Very resource hungry (You " +"will require a GPU with a fair amount of VRAM). Good for details, but more " +"susceptible to color differences." +msgstr "" + +#: lib/cli/args_train.py:125 +msgid "" +"Output a summary of the model and exit. If a model folder is provided then a " +"summary of the saved model is displayed. Otherwise a summary of the model " +"that would be created by the chosen plugin and configuration settings is " +"displayed." +msgstr "" + +#: lib/cli/args_train.py:135 +msgid "" +"Freeze the weights of the model. Freezing weights means that some of the " +"parameters in the model will no longer continue to learn, but those that are " +"not frozen will continue to learn. For most models, this will freeze the " +"encoder, but some models may have configuration options for freezing other " +"layers." +msgstr "" + +#: lib/cli/args_train.py:147 lib/cli/args_train.py:160 +#: lib/cli/args_train.py:175 lib/cli/args_train.py:191 +#: lib/cli/args_train.py:200 +msgid "training" +msgstr "" + +#: lib/cli/args_train.py:149 +msgid "" +"Batch size. This is the number of images processed through the model for " +"each side per iteration. NB: As the model is fed 2 sides at a time, the " +"actual number of images within the model at any one time is double the " +"number that you set here. Larger batches require more GPU RAM." +msgstr "" + +#: lib/cli/args_train.py:162 +msgid "" +"Length of training in iterations. This is only really used for automation. " +"There is no 'correct' number of iterations a model should be trained for. " +"You should stop training when you are happy with the previews. However, if " +"you want the model to stop automatically at a set number of iterations, you " +"can set that value here." +msgstr "" + +#: lib/cli/args_train.py:177 +msgid "" +"R|Select the distribution stategy to use.\n" +"L|default: Use Tensorflow's default distribution strategy.\n" +"L|central-storage: Centralizes variables on the CPU whilst operations are " +"performed on 1 or more local GPUs. This can help save some VRAM at the cost " +"of some speed by not storing variables on the GPU. Note: Mixed-Precision is " +"not supported on multi-GPU setups.\n" +"L|mirrored: Supports synchronous distributed training across multiple local " +"GPUs. A copy of the model and all variables are loaded onto each GPU with " +"batches distributed to each GPU at each iteration." +msgstr "" + +#: lib/cli/args_train.py:193 +msgid "" +"Disables TensorBoard logging. NB: Disabling logs means that you will not be " +"able to use the graph or analysis for this session in the GUI." +msgstr "" + +#: lib/cli/args_train.py:202 +msgid "" +"Use the Learning Rate Finder to discover the optimal learning rate for " +"training. For new models, this will calculate the optimal learning rate for " +"the model. For existing models this will use the optimal learning rate that " +"was discovered when initializing the model. Setting this option will ignore " +"the manually configured learning rate (configurable in train settings)." +msgstr "" + +#: lib/cli/args_train.py:215 lib/cli/args_train.py:225 +msgid "Saving" +msgstr "" + +#: lib/cli/args_train.py:216 +msgid "Sets the number of iterations between each model save." +msgstr "" + +#: lib/cli/args_train.py:227 +msgid "" +"Sets the number of iterations before saving a backup snapshot of the model " +"in it's current state. Set to 0 for off." +msgstr "" + +#: lib/cli/args_train.py:234 lib/cli/args_train.py:246 +#: lib/cli/args_train.py:258 +msgid "timelapse" +msgstr "" + +#: lib/cli/args_train.py:236 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'A' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-B parameter." +msgstr "" + +#: lib/cli/args_train.py:248 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'B' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-A parameter." +msgstr "" + +#: lib/cli/args_train.py:260 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. If " +"the input folders are supplied but no output folder, it will default to your " +"model folder/timelapse/" +msgstr "" + +#: lib/cli/args_train.py:269 lib/cli/args_train.py:276 +msgid "preview" +msgstr "" + +#: lib/cli/args_train.py:270 +msgid "Show training preview output. in a separate window." +msgstr "" + +#: lib/cli/args_train.py:278 +msgid "" +"Writes the training result to a file. The image will be stored in the root " +"of your FaceSwap folder." +msgstr "" + +#: lib/cli/args_train.py:285 lib/cli/args_train.py:295 +#: lib/cli/args_train.py:305 lib/cli/args_train.py:315 +msgid "augmentation" +msgstr "" + +#: lib/cli/args_train.py:287 +msgid "" +"Warps training faces to closely matched Landmarks from the opposite face-set " +"rather than randomly warping the face. This is the 'dfaker' way of doing " +"warping." +msgstr "" + +#: lib/cli/args_train.py:297 +msgid "" +"To effectively learn, a random set of images are flipped horizontally. " +"Sometimes it is desirable for this not to occur. Generally this should be " +"left off except for during 'fit training'." +msgstr "" + +#: lib/cli/args_train.py:307 +msgid "" +"Color augmentation helps make the model less susceptible to color " +"differences between the A and B sets, at an increased training time cost. " +"Enable this option to disable color augmentation." +msgstr "" + +#: lib/cli/args_train.py:317 +msgid "" +"Warping is integral to training the Neural Network. This option should only " +"be enabled towards the very end of training to try to bring out more detail. " +"Think of it as 'fine-tuning'. Enabling this option from the beginning is " +"likely to kill a model and lead to terrible results." +msgstr "" diff --git a/locales/lib.config.pot b/locales/lib.config.pot new file mode 100644 index 0000000000..b1144f809c --- /dev/null +++ b/locales/lib.config.pot @@ -0,0 +1,61 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-11 23:28+0100\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: lib/config.py:393 +msgid "" +"\n" +"This option can be updated for existing models.\n" +msgstr "" + +#: lib/config.py:395 +msgid "" +"\n" +"If selecting multiple options then each option should be separated by a " +"space or a comma (e.g. item1, item2, item3)\n" +msgstr "" + +#: lib/config.py:398 +msgid "" +"\n" +"Choose from: {}" +msgstr "" + +#: lib/config.py:400 +msgid "" +"\n" +"Choose from: True, False" +msgstr "" + +#: lib/config.py:404 +msgid "" +"\n" +"Select an integer between {} and {}" +msgstr "" + +#: lib/config.py:408 +msgid "" +"\n" +"Select a decimal number between {} and {}" +msgstr "" + +#: lib/config.py:409 +msgid "" +"\n" +"[Default: {}]" +msgstr "" diff --git a/locales/plugins.extract._config.pot b/locales/plugins.extract._config.pot new file mode 100644 index 0000000000..a65b8fda60 --- /dev/null +++ b/locales/plugins.extract._config.pot @@ -0,0 +1,117 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-08 16:43+0100\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: plugins/extract/_config.py:32 +msgid "Options that apply to all extraction plugins" +msgstr "" + +#: plugins/extract/_config.py:38 +msgid "settings" +msgstr "" + +#: plugins/extract/_config.py:39 +msgid "" +"Enable the Tensorflow GPU `allow_growth` configuration " +"option. This option prevents Tensorflow from allocating all of the GPU VRAM " +"at launch but can lead to higher VRAM fragmentation and slower performance. " +"Should only be enabled if you are having problems running extraction." +msgstr "" + +#: plugins/extract/_config.py:50 plugins/extract/_config.py:64 +#: plugins/extract/_config.py:78 plugins/extract/_config.py:89 +#: plugins/extract/_config.py:99 plugins/extract/_config.py:108 +#: plugins/extract/_config.py:119 +msgid "filters" +msgstr "" + +#: plugins/extract/_config.py:51 +msgid "" +"Filters out faces below this size. This is a multiplier of the minimum " +"dimension of the frame (i.e. 1280x720 = 720). If the original face extract " +"box is smaller than the minimum dimension times this multiplier, it is " +"considered a false positive and discarded. Faces which are found to be " +"unusually smaller than the frame tend to be misaligned images, except in " +"extreme long-shots. These can be usually be safely discarded." +msgstr "" + +#: plugins/extract/_config.py:65 +msgid "" +"Filters out faces above this size. This is a multiplier of the minimum " +"dimension of the frame (i.e. 1280x720 = 720). If the original face extract " +"box is larger than the minimum dimension times this multiplier, it is " +"considered a false positive and discarded. Faces which are found to be " +"unusually larger than the frame tend to be misaligned images except in " +"extreme close-ups. These can be usually be safely discarded." +msgstr "" + +#: plugins/extract/_config.py:79 +msgid "" +"Filters out faces who's landmarks are above this distance from an 'average' " +"face. Values above 15 tend to be fairly safe. Values above 10 will remove " +"more false positives, but may also filter out some faces at extreme angles." +msgstr "" + +#: plugins/extract/_config.py:90 +msgid "" +"Filters out faces who's calculated roll is greater than zero +/- this value " +"in degrees. Aligned faces should have a roll value close to zero. Values " +"that are a significant distance from 0 degrees tend to be misaligned images. " +"These can usually be safely disgarded." +msgstr "" + +#: plugins/extract/_config.py:100 +msgid "" +"Filters out faces where the lowest point of the aligned face's eye or " +"eyebrow is lower than the highest point of the aligned face's mouth. Any " +"faces where this occurs are misaligned and can be safely disgarded." +msgstr "" + +#: plugins/extract/_config.py:109 +msgid "" +"If enabled, and 're-feed' has been selected for extraction, then interim " +"alignments will be filtered prior to averaging the final landmarks. This can " +"help improve the final alignments by removing any obvious misaligns from the " +"interim results, and may also help pick up difficult alignments. If " +"disabled, then all re-feed results will be averaged." +msgstr "" + +#: plugins/extract/_config.py:120 +msgid "" +"If enabled, saves any filtered out images into a sub-folder during the " +"extraction process. If disabled, filtered faces are deleted. Note: The faces " +"will always be filtered out of the alignments file, regardless of whether " +"you keep the faces or not." +msgstr "" + +#: plugins/extract/_config.py:129 plugins/extract/_config.py:138 +msgid "re-align" +msgstr "" + +#: plugins/extract/_config.py:130 +msgid "" +"If enabled, and 're-align' has been selected for extraction, then all re-" +"feed iterations are re-aligned. If disabled, then only the final averaged " +"output from re-feed will be re-aligned." +msgstr "" + +#: plugins/extract/_config.py:139 +msgid "" +"If enabled, and 're-align' has been selected for extraction, then any " +"alignments which would be filtered out will not be re-aligned." +msgstr "" diff --git a/locales/plugins.train._config.pot b/locales/plugins.train._config.pot new file mode 100644 index 0000000000..38e94f3152 --- /dev/null +++ b/locales/plugins.train._config.pot @@ -0,0 +1,614 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-26 17:37+0000\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: plugins/train/_config.py:17 +msgid "" +"\n" +"NB: Unless specifically stated, values changed here will only take effect " +"when creating a new model." +msgstr "" + +#: plugins/train/_config.py:22 +msgid "" +"Focal Frequency Loss. Analyzes the frequency spectrum of the images rather " +"than the images themselves. This loss function can be used on its own, but " +"the original paper found increased benefits when using it as a complementary " +"loss to another spacial loss function (e.g. MSE). Ref: Focal Frequency Loss " +"for Image Reconstruction and Synthesis https://arxiv.org/pdf/2012.12821.pdf " +"NB: This loss does not currently work on AMD cards." +msgstr "" + +#: plugins/train/_config.py:29 +msgid "" +"Nvidia FLIP. A perceptual loss measure that approximates the difference " +"perceived by humans as they alternate quickly (or flip) between two images. " +"Used on its own and this loss function creates a distinct grid on the " +"output. However it can be helpful when used as a complimentary loss " +"function. Ref: FLIP: A Difference Evaluator for Alternating Images: https://" +"research.nvidia.com/sites/default/files/node/3260/FLIP_Paper.pdf" +msgstr "" + +#: plugins/train/_config.py:36 +msgid "" +"Gradient Magnitude Similarity Deviation seeks to match the global standard " +"deviation of the pixel to pixel differences between two images. Similar in " +"approach to SSIM. Ref: Gradient Magnitude Similarity Deviation: An Highly " +"Efficient Perceptual Image Quality Index https://arxiv.org/ftp/arxiv/" +"papers/1308/1308.3052.pdf" +msgstr "" + +#: plugins/train/_config.py:41 +msgid "" +"The L_inf norm will reduce the largest individual pixel error in an image. " +"As each largest error is minimized sequentially, the overall error is " +"improved. This loss will be extremely focused on outliers." +msgstr "" + +#: plugins/train/_config.py:45 +msgid "" +"Laplacian Pyramid Loss. Attempts to improve results by focussing on edges " +"using Laplacian Pyramids. As this loss function gives priority to edges over " +"other low-frequency information, like color, it should not be used on its " +"own. The original implementation uses this loss as a complimentary function " +"to MSE. Ref: Optimizing the Latent Space of Generative Networks https://" +"arxiv.org/abs/1707.05776" +msgstr "" + +#: plugins/train/_config.py:52 +msgid "" +"LPIPS is a perceptual loss that uses the feature outputs of other pretrained " +"models as a loss metric. Be aware that this loss function will use more " +"VRAM. Used on its own and this loss will create a distinct moire pattern on " +"the output, however it can be helpful as a complimentary loss function. The " +"output of this function is strong, so depending on your chosen primary loss " +"function, you are unlikely going to want to set the weight above about 25%. " +"Ref: The Unreasonable Effectiveness of Deep Features as a Perceptual Metric " +"http://arxiv.org/abs/1801.03924\n" +"This variant uses the AlexNet backbone. A fairly light and old model which " +"performed best in the paper's original implementation.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" + +#: plugins/train/_config.py:62 +msgid "" +"Same as lpips_alex, but using the SqueezeNet backbone. A more lightweight " +"version of AlexNet.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" + +#: plugins/train/_config.py:65 +msgid "" +"Same as lpips_alex, but using the VGG16 backbone. A more heavyweight model.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" + +#: plugins/train/_config.py:68 +msgid "" +"log(cosh(x)) acts similar to MSE for small errors and to MAE for large " +"errors. Like MSE, it is very stable and prevents overshoots when errors are " +"near zero. Like MAE, it is robust to outliers." +msgstr "" + +#: plugins/train/_config.py:72 +msgid "" +"Mean absolute error will guide reconstructions of each pixel towards its " +"median value in the training dataset. Robust to outliers but as a median, it " +"can potentially ignore some infrequent image types in the dataset." +msgstr "" + +#: plugins/train/_config.py:76 +msgid "" +"Mean squared error will guide reconstructions of each pixel towards its " +"average value in the training dataset. As an avg, it will be susceptible to " +"outliers and typically produces slightly blurrier results. Ref: Multi-Scale " +"Structural Similarity for Image Quality Assessment https://www.cns.nyu.edu/" +"pub/eero/wang03b.pdf" +msgstr "" + +#: plugins/train/_config.py:81 +msgid "" +"Multiscale Structural Similarity Index Metric is similar to SSIM except that " +"it performs the calculations along multiple scales of the input image." +msgstr "" + +#: plugins/train/_config.py:84 +msgid "" +"Smooth_L1 is a modification of the MAE loss to correct two of its " +"disadvantages. This loss has improved stability and guidance for small " +"errors. Ref: A General and Adaptive Robust Loss Function https://arxiv.org/" +"pdf/1701.03077.pdf" +msgstr "" + +#: plugins/train/_config.py:88 +msgid "" +"Structural Similarity Index Metric is a perception-based loss that considers " +"changes in texture, luminance, contrast, and local spatial statistics of an " +"image. Potentially delivers more realistic looking images. Ref: Image " +"Quality Assessment: From Error Visibility to Structural Similarity http://" +"www.cns.nyu.edu/pub/eero/wang03-reprint.pdf" +msgstr "" + +#: plugins/train/_config.py:93 +msgid "" +"Instead of minimizing the difference between the absolute value of each " +"pixel in two reference images, compute the pixel to pixel spatial difference " +"in each image and then minimize that difference between two images. Allows " +"for large color shifts, but maintains the structure of the image." +msgstr "" + +#: plugins/train/_config.py:97 +msgid "Do not use an additional loss function." +msgstr "" + +#: plugins/train/_config.py:117 +msgid "Options that apply to all models" +msgstr "" + +#: plugins/train/_config.py:126 plugins/train/_config.py:150 +msgid "face" +msgstr "" + +#: plugins/train/_config.py:128 +msgid "" +"How to center the training image. The extracted images are centered on the " +"middle of the skull based on the face's estimated pose. A subsection of " +"these images are used for training. The centering used dictates how this " +"subsection will be cropped from the aligned images.\n" +"\tface: Centers the training image on the center of the face, adjusting for " +"pitch and yaw.\n" +"\thead: Centers the training image on the center of the head, adjusting for " +"pitch and yaw. NB: You should only select head centering if you intend to " +"include the full head (including hair) in the final swap. This may give " +"mixed results. Additionally, it is only worth choosing head centering if you " +"are training with a mask that includes the hair (e.g. BiSeNet-FP-Head).\n" +"\tlegacy: The 'original' extraction technique. Centers the training image " +"near the tip of the nose with no adjustment. Can result in the edges of the " +"face appearing outside of the training area." +msgstr "" + +#: plugins/train/_config.py:152 +msgid "" +"How much of the extracted image to train on. A lower coverage will limit the " +"model's scope to a zoomed-in central area while higher amounts can include " +"the entire face. A trade-off exists between lower amounts given more detail " +"versus higher amounts avoiding noticeable swap transitions. For 'Face' " +"centering you will want to leave this above 75%. For Head centering you will " +"most likely want to set this to 100%. Sensible values for 'Legacy' centering " +"are:\n" +"\t62.5% spans from eyebrow to eyebrow.\n" +"\t75.0% spans from temple to temple.\n" +"\t87.5% spans from ear to ear.\n" +"\t100.0% is a mugshot." +msgstr "" + +#: plugins/train/_config.py:168 plugins/train/_config.py:179 +msgid "initialization" +msgstr "" + +#: plugins/train/_config.py:170 +msgid "" +"Use ICNR to tile the default initializer in a repeating pattern. This " +"strategy is designed for pairing with sub-pixel / pixel shuffler to reduce " +"the 'checkerboard effect' in image reconstruction. \n" +"\t https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf" +msgstr "" + +#: plugins/train/_config.py:181 +msgid "" +"Use Convolution Aware Initialization for convolutional layers. This can help " +"eradicate the vanishing and exploding gradient problem as well as lead to " +"higher accuracy, lower loss and faster convergence.\n" +"NB:\n" +"\t This can use more VRAM when creating a new model so you may want to lower " +"the batch size for the first run. The batch size can be raised again when " +"reloading the model. \n" +"\t Multi-GPU is not supported for this option, so you should start the model " +"on a single GPU. Once training has started, you can stop training, enable " +"multi-GPU and resume.\n" +"\t Building the model will likely take several minutes as the calculations " +"for this initialization technique are expensive. This will only impact " +"starting a new model." +msgstr "" + +#: plugins/train/_config.py:198 plugins/train/_config.py:223 +#: plugins/train/_config.py:238 plugins/train/_config.py:256 +#: plugins/train/_config.py:337 +msgid "optimizer" +msgstr "" + +#: plugins/train/_config.py:202 +msgid "" +"The optimizer to use.\n" +"\t adabelief - Adapting Stepsizes by the Belief in Observed Gradients. An " +"optimizer with the aim to converge faster, generalize better and remain more " +"stable. (https://arxiv.org/abs/2010.07468). NB: Epsilon for AdaBelief needs " +"to be set to a smaller value than other Optimizers. Generally setting the " +"'Epsilon Exponent' to around '-16' should work.\n" +"\t adam - Adaptive Moment Optimization. A stochastic gradient descent method " +"that is based on adaptive estimation of first-order and second-order " +"moments.\n" +"\t nadam - Adaptive Moment Optimization with Nesterov Momentum. Much like " +"Adam but uses a different formula for calculating momentum.\n" +"\t rms-prop - Root Mean Square Propagation. Maintains a moving (discounted) " +"average of the square of the gradients. Divides the gradient by the root of " +"this average." +msgstr "" + +#: plugins/train/_config.py:225 +msgid "" +"Learning rate - how fast your network will learn (how large are the " +"modifications to the model weights after one batch of training). Values that " +"are too large might result in model crashes and the inability of the model " +"to find the best solution. Values that are too small might be unable to " +"escape from dead-ends and find the best global minimum." +msgstr "" + +#: plugins/train/_config.py:240 +msgid "" +"The epsilon adds a small constant to weight updates to attempt to avoid " +"'divide by zero' errors. Unless you are using the AdaBelief Optimizer, then " +"Generally this option should be left at default value, For AdaBelief, " +"setting this to around '-16' should work.\n" +"In all instances if you are getting 'NaN' loss values, and have been unable " +"to resolve the issue any other way (for example, increasing batch size, or " +"lowering learning rate), then raising the epsilon can lead to a more stable " +"model. It may, however, come at the cost of slower training and a less " +"accurate final result.\n" +"NB: The value given here is the 'exponent' to the epsilon. For example, " +"choosing '-7' will set the epsilon to 1e-7. Choosing '-3' will set the " +"epsilon to 0.001 (1e-3)." +msgstr "" + +#: plugins/train/_config.py:262 +msgid "" +"When to save the Optimizer Weights. Saving the optimizer weights is not " +"necessary and will increase the model file size 3x (and by extension the " +"amount of time it takes to save the model). However, it can be useful to " +"save these weights if you want to guarantee that a resumed model carries off " +"exactly from where it left off, rather than spending a few hundred " +"iterations catching up.\n" +"\t never - Don't save optimizer weights.\n" +"\t always - Save the optimizer weights at every save iteration. Model saving " +"will take longer, due to the increased file size, but you will always have " +"the last saved optimizer state in your model file.\n" +"\t exit - Only save the optimizer weights when explicitly terminating a " +"model. This can be when the model is actively stopped or when the target " +"iterations are met. Note: If the training session ends because of another " +"reason (e.g. power outage, Out of Memory Error, NaN detected) then the " +"optimizer weights will NOT be saved." +msgstr "" + +#: plugins/train/_config.py:285 plugins/train/_config.py:297 +#: plugins/train/_config.py:314 +msgid "Learning Rate Finder" +msgstr "" + +#: plugins/train/_config.py:287 +msgid "" +"The number of iterations to process to find the optimal learning rate. " +"Higher values will take longer, but will be more accurate." +msgstr "" + +#: plugins/train/_config.py:299 +msgid "" +"The operation mode for the learning rate finder. Only applicable to new " +"models. For existing models this will always default to 'set'.\n" +"\tset - Train with the discovered optimal learning rate.\n" +"\tgraph_and_set - Output a graph in the training folder showing the " +"discovered learning rates and train with the optimal learning rate.\n" +"\tgraph_and_exit - Output a graph in the training folder with the discovered " +"learning rates and exit." +msgstr "" + +#: plugins/train/_config.py:316 +msgid "" +"How aggressively to set the Learning Rate. More aggressive can learn faster, " +"but is more likely to lead to exploding gradients.\n" +"\tdefault - The default optimal learning rate. A safe choice for nearly all " +"use cases.\n" +"\taggressive - Set's a higher learning rate than the default. May learn " +"faster but with a higher chance of exploding gradients.\n" +"\textreme - The highest optimal learning rate. A much higher risk of " +"exploding gradients." +msgstr "" + +#: plugins/train/_config.py:330 +msgid "" +"Apply AutoClipping to the gradients. AutoClip analyzes the gradient weights " +"and adjusts the normalization value dynamically to fit the data. Can help " +"prevent NaNs and improve model optimization at the expense of VRAM. Ref: " +"AutoClip: Adaptive Gradient Clipping for Source Separation Networks https://" +"arxiv.org/abs/2007.14469" +msgstr "" + +#: plugins/train/_config.py:343 plugins/train/_config.py:355 +#: plugins/train/_config.py:369 plugins/train/_config.py:386 +msgid "network" +msgstr "" + +#: plugins/train/_config.py:345 +msgid "" +"Use reflection padding rather than zero padding with convolutions. Each " +"convolution must pad the image boundaries to maintain the proper sizing. " +"More complex padding schemes can reduce artifacts at the border of the " +"image.\n" +"\t http://www-cs.engr.ccny.cuny.edu/~wolberg/cs470/hw/hw2_pad.txt" +msgstr "" + +#: plugins/train/_config.py:358 +msgid "" +"Enable the Tensorflow GPU 'allow_growth' configuration option. This option " +"prevents Tensorflow from allocating all of the GPU VRAM at launch but can " +"lead to higher VRAM fragmentation and slower performance. Should only be " +"enabled if you are receiving errors regarding 'cuDNN fails to initialize' " +"when commencing training." +msgstr "" + +#: plugins/train/_config.py:371 +msgid "" +"NVIDIA GPUs can run operations in float16 faster than in float32. Mixed " +"precision allows you to use a mix of float16 with float32, to get the " +"performance benefits from float16 and the numeric stability benefits from " +"float32.\n" +"\n" +"This is untested on DirectML backend, but will run on most Nvidia models. it " +"will only speed up training on more recent GPUs. Those with compute " +"capability 7.0 or higher will see the greatest performance benefit from " +"mixed precision because they have Tensor Cores. Older GPUs offer no math " +"performance benefit for using mixed precision, however memory and bandwidth " +"savings can enable some speedups. Generally RTX GPUs and later will offer " +"the most benefit." +msgstr "" + +#: plugins/train/_config.py:388 +msgid "" +"If a 'NaN' is generated in the model, this means that the model has " +"corrupted and the model is likely to start deteriorating from this point on. " +"Enabling NaN protection will stop training immediately in the event of a " +"NaN. The last save will not contain the NaN, so you may still be able to " +"rescue your model." +msgstr "" + +#: plugins/train/_config.py:401 +msgid "convert" +msgstr "" + +#: plugins/train/_config.py:403 +msgid "" +"[GPU Only]. The number of faces to feed through the model at once when " +"running the Convert process.\n" +"\n" +"NB: Increasing this figure is unlikely to improve convert speed, however, if " +"you are getting Out of Memory errors, then you may want to reduce the batch " +"size." +msgstr "" + +#: plugins/train/_config.py:422 +msgid "" +"Loss configuration options\n" +"Loss is the mechanism by which a Neural Network judges how well it thinks " +"that it is recreating a face." +msgstr "" + +#: plugins/train/_config.py:429 plugins/train/_config.py:441 +#: plugins/train/_config.py:454 plugins/train/_config.py:474 +#: plugins/train/_config.py:486 plugins/train/_config.py:506 +#: plugins/train/_config.py:518 plugins/train/_config.py:538 +#: plugins/train/_config.py:554 plugins/train/_config.py:570 +#: plugins/train/_config.py:587 +msgid "loss" +msgstr "" + +#: plugins/train/_config.py:433 +msgid "The loss function to use." +msgstr "" + +#: plugins/train/_config.py:445 +msgid "" +"The second loss function to use. If using a structural based loss (such as " +"SSIM, MS-SSIM or GMSD) it is common to add an L1 regularization(MAE) or L2 " +"regularization (MSE) function. You can adjust the weighting of this loss " +"function with the loss_weight_2 option." +msgstr "" + +#: plugins/train/_config.py:460 +msgid "" +"The amount of weight to apply to the second loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the second loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the second loss function will be reduced by " +"a quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the second loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the second loss function altogether." +msgstr "" + +#: plugins/train/_config.py:478 +msgid "" +"The third loss function to use. You can adjust the weighting of this loss " +"function with the loss_weight_3 option." +msgstr "" + +#: plugins/train/_config.py:492 +msgid "" +"The amount of weight to apply to the third loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the third loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the third loss function will be reduced by a " +"quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the third loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the third loss function altogether." +msgstr "" + +#: plugins/train/_config.py:510 +msgid "" +"The fourth loss function to use. You can adjust the weighting of this loss " +"function with the loss_weight_3 option." +msgstr "" + +#: plugins/train/_config.py:524 +msgid "" +"The amount of weight to apply to the fourth loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the fourth loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the fourth loss function will be reduced by " +"a quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the fourth loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the fourth loss function altogether." +msgstr "" + +#: plugins/train/_config.py:543 +msgid "" +"The loss function to use when learning a mask.\n" +"\t MAE - Mean absolute error will guide reconstructions of each pixel " +"towards its median value in the training dataset. Robust to outliers but as " +"a median, it can potentially ignore some infrequent image types in the " +"dataset.\n" +"\t MSE - Mean squared error will guide reconstructions of each pixel towards " +"its average value in the training dataset. As an average, it will be " +"susceptible to outliers and typically produces slightly blurrier results." +msgstr "" + +#: plugins/train/_config.py:560 +msgid "" +"The amount of priority to give to the eyes.\n" +"\n" +"The value given here is as a multiplier of the main loss score. For " +"example:\n" +"\t 1 - The eyes will receive the same priority as the rest of the face. \n" +"\t 10 - The eyes will be given a score 10 times higher than the rest of the " +"face.\n" +"\n" +"NB: Penalized Mask Loss must be enable to use this option." +msgstr "" + +#: plugins/train/_config.py:576 +msgid "" +"The amount of priority to give to the mouth.\n" +"\n" +"The value given here is as a multiplier of the main loss score. For " +"Example:\n" +"\t 1 - The mouth will receive the same priority as the rest of the face. \n" +"\t 10 - The mouth will be given a score 10 times higher than the rest of the " +"face.\n" +"\n" +"NB: Penalized Mask Loss must be enable to use this option." +msgstr "" + +#: plugins/train/_config.py:589 +msgid "" +"Image loss function is weighted by mask presence. For areas of the image " +"without the facial mask, reconstruction errors will be ignored while the " +"masked face area is prioritized. May increase overall quality by focusing " +"attention on the core face area." +msgstr "" + +#: plugins/train/_config.py:600 plugins/train/_config.py:643 +#: plugins/train/_config.py:656 plugins/train/_config.py:671 +#: plugins/train/_config.py:680 +msgid "mask" +msgstr "" + +#: plugins/train/_config.py:603 +msgid "" +"The mask to be used for training. If you have selected 'Learn Mask' or " +"'Penalized Mask Loss' you must select a value other than 'none'. The " +"required mask should have been selected as part of the Extract process. If " +"it does not exist in the alignments file then it will be generated prior to " +"training commencing.\n" +"\tnone: Don't use a mask.\n" +"\tbisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"\tbisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"\tcomponents: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"\tcustom_face: Custom user created, face centered mask.\n" +"\tcustom_head: Custom user created, head centered mask.\n" +"\textended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"\tvgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"\tvgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"\tunet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance." +msgstr "" + +#: plugins/train/_config.py:645 +msgid "" +"Dilate or erode the mask. Negative values erode the mask (make it smaller). " +"Positive values dilate the mask (make it larger). The value given is a " +"percentage of the total mask size." +msgstr "" + +#: plugins/train/_config.py:658 +msgid "" +"Apply gaussian blur to the mask input. This has the effect of smoothing the " +"edges of the mask, which can help with poorly calculated masks and give less " +"of a hard edge to the predicted mask. The size is in pixels (calculated from " +"a 128px mask). Set to 0 to not apply gaussian blur. This value should be " +"odd, if an even number is passed in then it will be rounded to the next odd " +"number." +msgstr "" + +#: plugins/train/_config.py:673 +msgid "" +"Sets pixels that are near white to white and near black to black. Set to 0 " +"for off." +msgstr "" + +#: plugins/train/_config.py:682 +msgid "" +"Dedicate a portion of the model to learning how to duplicate the input mask. " +"Increases VRAM usage in exchange for learning a quick ability to try to " +"replicate more complex mask models." +msgstr "" diff --git a/locales/ru/LC_MESSAGES/faceswap.mo b/locales/ru/LC_MESSAGES/faceswap.mo index f02a1b0ff3..db2449a789 100644 Binary files a/locales/ru/LC_MESSAGES/faceswap.mo and b/locales/ru/LC_MESSAGES/faceswap.mo differ diff --git a/locales/ru/LC_MESSAGES/faceswap.po b/locales/ru/LC_MESSAGES/faceswap.po index 2e0f297baf..7c37585784 100644 --- a/locales/ru/LC_MESSAGES/faceswap.po +++ b/locales/ru/LC_MESSAGES/faceswap.po @@ -6,28 +6,28 @@ msgid "" msgstr "" "Project-Id-Version: \n" "POT-Creation-Date: 2021-02-18 23:48-0000\n" -"PO-Revision-Date: 2021-02-19 21:30+0300\n" +"PO-Revision-Date: 2023-04-11 12:56+0700\n" +"Last-Translator: \n" "Language-Team: \n" +"Language: ru\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.2\n" -"Last-Translator: \n" "Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" -"Language: ru\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.2.2\n" #: faceswap.py:43 msgid "Extract the faces from pictures or a video" -msgstr "Извлечь лица из фотографий или видео" +msgstr "Извлечение лиц из картинок или видео" #: faceswap.py:44 msgid "Train a model for the two faces A and B" -msgstr "Обучить модель при помощи лиц A и B" +msgstr "Обучить модель для двух лиц A и B" #: faceswap.py:47 msgid "Convert source pictures or video to a new one with the face swapped" -msgstr "Преобразование исходных изображений или видео в новое с замененным лицом" +msgstr "Преобразование исходных изображений или видео в новое с заменой лиц" #: faceswap.py:48 msgid "Launch the Faceswap Graphical User Interface" diff --git a/locales/ru/LC_MESSAGES/gui.menu.mo b/locales/ru/LC_MESSAGES/gui.menu.mo new file mode 100644 index 0000000000..15df09a5cd Binary files /dev/null and b/locales/ru/LC_MESSAGES/gui.menu.mo differ diff --git a/locales/ru/LC_MESSAGES/gui.menu.po b/locales/ru/LC_MESSAGES/gui.menu.po new file mode 100644 index 0000000000..581dfde23f --- /dev/null +++ b/locales/ru/LC_MESSAGES/gui.menu.po @@ -0,0 +1,156 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-07 13:54+0100\n" +"PO-Revision-Date: 2023-06-07 20:29+0700\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/gui/menu.py:37 +msgid "faceswap.dev - Guides and Forum" +msgstr "faceswap.dev - Руководства и Форум" + +#: lib/gui/menu.py:38 +msgid "Patreon - Support this project" +msgstr "Patreon - Поддержите этот проект" + +#: lib/gui/menu.py:39 +msgid "Discord - The FaceSwap Discord server" +msgstr "Discord - Discord сервер Faceswap" + +#: lib/gui/menu.py:40 +msgid "Github - Our Source Code" +msgstr "Github - Наш исходный код" + +#: lib/gui/menu.py:60 +msgid "File" +msgstr "Файл" + +#: lib/gui/menu.py:61 +msgid "Settings" +msgstr "Настройки" + +#: lib/gui/menu.py:62 +msgid "Help" +msgstr "Помощь" + +#: lib/gui/menu.py:85 +msgid "Configure Settings..." +msgstr "Настройки..." + +#: lib/gui/menu.py:116 +msgid "New Project..." +msgstr "Новый проект..." + +#: lib/gui/menu.py:121 +msgid "Open Project..." +msgstr "Открыть проект..." + +#: lib/gui/menu.py:126 +msgid "Save Project" +msgstr "Сохранить проект" + +#: lib/gui/menu.py:131 +msgid "Save Project as..." +msgstr "Сохранить проект как..." + +#: lib/gui/menu.py:136 +msgid "Reload Project from Disk" +msgstr "Перезагрузить Проект из диска" + +#: lib/gui/menu.py:141 +msgid "Close Project" +msgstr "Закрыть проект" + +#: lib/gui/menu.py:147 +msgid "Open Task..." +msgstr "Открыть задачу..." + +#: lib/gui/menu.py:154 +msgid "Open recent" +msgstr "Открытые недавно" + +#: lib/gui/menu.py:156 +msgid "Quit" +msgstr "Выход" + +#: lib/gui/menu.py:211 +msgid "{} Task" +msgstr "{} Задача" + +#: lib/gui/menu.py:223 +msgid "Clear recent files" +msgstr "Очистить недавние файлы" + +#: lib/gui/menu.py:391 +msgid "Check for updates..." +msgstr "Проверить обновления..." + +#: lib/gui/menu.py:394 +msgid "Update Faceswap..." +msgstr "Обновить Faceswap..." + +#: lib/gui/menu.py:398 +msgid "Switch Branch" +msgstr "Сменить ветку" + +#: lib/gui/menu.py:401 +msgid "Resources" +msgstr "Ресурсы" + +#: lib/gui/menu.py:404 +msgid "Output System Information" +msgstr "Вывести информацию о системе" + +#: lib/gui/menu.py:589 +msgid "currently selected Task" +msgstr "текущую выбранную задачу" + +#: lib/gui/menu.py:589 +msgid "Project" +msgstr "Проект" + +#: lib/gui/menu.py:591 +msgid "Reload {} from disk" +msgstr "Перезагрузить {} из диска" + +#: lib/gui/menu.py:593 +msgid "Create a new {}..." +msgstr "Создать новый {}..." + +#: lib/gui/menu.py:595 +msgid "Reset {} to default" +msgstr "Сбросить {} по умолчанию" + +#: lib/gui/menu.py:597 +msgid "Save {}" +msgstr "Сохранить {}" + +#: lib/gui/menu.py:599 +msgid "Save {} as..." +msgstr "Сохранить {} как..." + +#: lib/gui/menu.py:603 +msgid " from a task or project file" +msgstr " из файла задачи или проекта" + +#: lib/gui/menu.py:604 +msgid "Load {}..." +msgstr "Загрузить {}..." + +#: lib/gui/menu.py:659 +msgid "Configure {} settings..." +msgstr "Настройка параметров {}..." diff --git a/locales/ru/LC_MESSAGES/gui.tooltips.mo b/locales/ru/LC_MESSAGES/gui.tooltips.mo new file mode 100644 index 0000000000..390771565d Binary files /dev/null and b/locales/ru/LC_MESSAGES/gui.tooltips.mo differ diff --git a/locales/ru/LC_MESSAGES/gui.tooltips.po b/locales/ru/LC_MESSAGES/gui.tooltips.po new file mode 100644 index 0000000000..c3b59f0cc3 --- /dev/null +++ b/locales/ru/LC_MESSAGES/gui.tooltips.po @@ -0,0 +1,210 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"POT-Creation-Date: 2021-03-22 18:37+0000\n" +"PO-Revision-Date: 2023-06-07 20:31+0700\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/gui/command.py:184 +msgid "Output command line options to the console" +msgstr "Вывод опций командной строки в консоль" + +#: lib/gui/command.py:195 +msgid "Run the {} script" +msgstr "Запуск сценария {}" + +#: lib/gui/control_helper.py:1234 +msgid "Select a folder..." +msgstr "Выбрать папку..." + +#: lib/gui/control_helper.py:1235 lib/gui/control_helper.py:1236 +msgid "Select a file..." +msgstr "Выбрать файл..." + +#: lib/gui/control_helper.py:1237 +msgid "Select a folder of images..." +msgstr "Выбрать папку с изображениями..." + +#: lib/gui/control_helper.py:1238 +msgid "Select a video..." +msgstr "Выбрать видео..." + +#: lib/gui/control_helper.py:1239 +msgid "Select a model folder..." +msgstr "Выбрать папку с моделью..." + +#: lib/gui/control_helper.py:1240 +msgid "Select one or more files..." +msgstr "Выбрать один или несколько файлов..." + +#: lib/gui/control_helper.py:1241 +msgid "Select a file or folder..." +msgstr "Выбрать файл или папку..." + +#: lib/gui/control_helper.py:1242 +msgid "Select a save location..." +msgstr "Выбрать место сохранения..." + +#: lib/gui/display.py:71 +msgid "Summary statistics for each training session" +msgstr "Сводная статистика для каждой тренировки" + +#: lib/gui/display.py:113 +msgid "Preview updates every 5 seconds" +msgstr "Предпросмотр обновляется каждые 5 секунд" + +#: lib/gui/display.py:122 +msgid "Graph showing Loss vs Iterations" +msgstr "График зависимости потерь от количества итераций" + +#: lib/gui/display.py:125 +msgid "Training preview. Updated on every save iteration" +msgstr "Предпросмотр тренировки. Обновляется каждую сохраняющую итерацию" + +#: lib/gui/display_analysis.py:342 +msgid "Load/Refresh stats for the currently training session" +msgstr "Загрузить/обновить статистику для текущей тренировки" + +#: lib/gui/display_analysis.py:344 +msgid "Clear currently displayed session stats" +msgstr "Очистить отображаемую статистику сессии" + +#: lib/gui/display_analysis.py:346 +msgid "Save session stats to csv" +msgstr "Сохранить статистику сессии в csv файл" + +#: lib/gui/display_analysis.py:348 +msgid "Load saved session stats" +msgstr "Загрузить сохраненную статистику" + +#: lib/gui/display_command.py:94 +msgid "Preview updates at every model save. Click to refresh now." +msgstr "" +"Предпросмотр обновляется при каждом сохранении модели. Нажмите, чтобы " +"обновить сейчас." + +#: lib/gui/display_command.py:261 +msgid "Graph updates at every model save. Click to refresh now." +msgstr "" +"График обновляется при каждом сохранении модели. Нажмите, чтобы обновить " +"сейчас." + +#: lib/gui/display_command.py:275 +msgid "Display the raw loss data" +msgstr "Показать необработанные данные о потерях" + +#: lib/gui/display_command.py:287 +msgid "Display the smoothed loss data" +msgstr "Показать сглаженные данные о потерях" + +#: lib/gui/display_command.py:294 +msgid "Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing." +msgstr "" +"Установите величину сглаживания. 0 - нет сглаживания, 0.99 - максимальное " +"сглаживание." + +#: lib/gui/display_command.py:324 +msgid "Set the number of iterations to display. 0 displays the full session." +msgstr "" +"Установите количество итераций для отображения. 0 отображает полный сеанс." + +#: lib/gui/display_page.py:238 +msgid "Save {}(s) to file" +msgstr "Сохранить {}(ы) в файл" + +#: lib/gui/display_page.py:250 +msgid "Enable or disable {} display" +msgstr "Включить или выключить отображение {}" + +#: lib/gui/popup_configure.py:209 +msgid "Close without saving" +msgstr "Закрыть без сохранения" + +#: lib/gui/popup_configure.py:210 +msgid "Save this page's config" +msgstr "Сохранить конфигурацию этой страницы" + +#: lib/gui/popup_configure.py:211 +msgid "Reset this page's config to default values" +msgstr "Сбросить конфигурацию этой страницы до заводских значений" + +#: lib/gui/popup_configure.py:213 +msgid "Save all settings for the currently selected config" +msgstr "Сохранить все настройки для текущей выбранной конфигурации" + +#: lib/gui/popup_configure.py:216 +msgid "Reset all settings for the currently selected config to default values" +msgstr "" +"Сбросить все настройки для текущей выбранной конфигурации до заводских " +"значений" + +#: lib/gui/popup_configure.py:538 +msgid "Select a plugin to configure:" +msgstr "Выбрать плагин для настройки:" + +#: lib/gui/popup_session.py:191 +msgid "Display {}" +msgstr "Показать {}" + +#: lib/gui/popup_session.py:342 +msgid "Refresh graph" +msgstr "Обновить график" + +#: lib/gui/popup_session.py:344 +msgid "Save display data to csv" +msgstr "Сохранить данные дисплея в csv файл" + +#: lib/gui/popup_session.py:346 +msgid "Number of data points to sample for rolling average" +msgstr "Количество точек данных для выборки среднего значения" + +#: lib/gui/popup_session.py:348 +msgid "Set the smoothing amount. 0 is no smoothing, 0.99 is maximum smoothing" +msgstr "" +"Установите величину сглаживания. 0 - нет сглаживания, 0.99 - максимальное " +"сглаживание" + +#: lib/gui/popup_session.py:350 +msgid "" +"Flatten data points that fall more than 1 standard deviation from the mean " +"to the mean value." +msgstr "" +"Сглаживание точек данных, которые отклоняются от среднего значения более чем " +"на 1 стандартное отклонение, до среднего значения." + +#: lib/gui/popup_session.py:353 +msgid "Display rolling average of the data" +msgstr "Показать среднее значение данных" + +#: lib/gui/popup_session.py:355 +msgid "Smooth the data" +msgstr "Сгладить данные" + +#: lib/gui/popup_session.py:357 +msgid "Display raw data" +msgstr "Показать необработанные данные" + +#: lib/gui/popup_session.py:359 +msgid "Display polynormal data trend" +msgstr "Отображение полинормальной тенденции данных" + +#: lib/gui/popup_session.py:361 +msgid "Set the data to display" +msgstr "Указать данные для отображения" + +#: lib/gui/popup_session.py:363 +msgid "Change y-axis scale" +msgstr "Изменить масштаб оси y" diff --git a/locales/ru/LC_MESSAGES/lib.cli.args.mo b/locales/ru/LC_MESSAGES/lib.cli.args.mo index 509595df4b..5b51831151 100644 Binary files a/locales/ru/LC_MESSAGES/lib.cli.args.mo and b/locales/ru/LC_MESSAGES/lib.cli.args.mo differ diff --git a/locales/ru/LC_MESSAGES/lib.cli.args.po b/locales/ru/LC_MESSAGES/lib.cli.args.po old mode 100644 new mode 100755 index 666d12ff23..2f6a55435f --- a/locales/ru/LC_MESSAGES/lib.cli.args.po +++ b/locales/ru/LC_MESSAGES/lib.cli.args.po @@ -1,970 +1,63 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # msgid "" msgstr "" "Project-Id-Version: \n" -"POT-Creation-Date: 2021-05-17 18:04+0100\n" -"PO-Revision-Date: 2021-05-17 18:11+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:06+0000\n" +"PO-Revision-Date: 2024-03-28 18:23+0000\n" "Last-Translator: \n" "Language-Team: \n" "Language: ru\n" "MIME-Version: 1.0\n" "Content-Type: text/plain; charset=UTF-8\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" -"X-Generator: Poedit 2.4.3\n" -"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && n" -"%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" -#: lib/cli/args.py:177 lib/cli/args.py:187 lib/cli/args.py:195 -#: lib/cli/args.py:205 +#: lib/cli/args.py:188 lib/cli/args.py:199 lib/cli/args.py:208 +#: lib/cli/args.py:219 msgid "Global Options" -msgstr "Общие настройки" +msgstr "Глобальные Настройки" -#: lib/cli/args.py:178 +#: lib/cli/args.py:190 msgid "" "R|Exclude GPUs from use by Faceswap. Select the number(s) which correspond " "to any GPU(s) that you do not wish to be made available to Faceswap. " "Selecting all GPUs here will force Faceswap into CPU mode.\n" "L|{}" msgstr "" -"R|Не использовать GPU для Faceswap. Выберите номер(а), которые соответствуют " -"тем GPU, которые вы не хотите использовать в Faceswap. При отключении всех " -"GPU Faceswap будет работать в режиме CPU.\n" +"R|Исключить GPU из использования Faceswap. Выберите номер (номера), " +"соответствующие любому GPU, который вы не хотите предоставлять Faceswap. " +"Если выбрать здесь все GPU, Faceswap перейдет в режим CPU.\n" "L|{}" -#: lib/cli/args.py:188 +#: lib/cli/args.py:201 msgid "" "Optionally overide the saved config with the path to a custom config file." msgstr "" -"Переназначить путь к файлу конфигурации пользовательским. (Необязательно)" +"Опционально переопределите сохраненную конфигурацию, указав путь к " +"пользовательскому файлу конфигурации." -#: lib/cli/args.py:196 +#: lib/cli/args.py:210 msgid "" "Log level. Stick with INFO or VERBOSE unless you need to file an error " "report. Be careful with TRACE as it will generate a lot of data" msgstr "" -"Уровень записи журнала. Придерживайтесь уровней INFO или VERBOSE, кроме " -"случаев когда вам нужно отправить отчёт об ошибке. Будьте осторожнее при " -"указании уровня TRACE, так как будет сгенерировано очень много данных" +"Уровень логирования. Придерживайтесь INFO или VERBOSE, если только вам не " +"нужно отправить отчет об ошибке. Будьте осторожны с TRACE, поскольку он " +"генерирует много данных" -#: lib/cli/args.py:206 +#: lib/cli/args.py:220 msgid "Path to store the logfile. Leave blank to store in the faceswap folder" msgstr "" -"Путь для сохранения файла журнала. Оставьте пустым, чтобы сохранить в папке " -"с faceswap" +"Путь для хранения файла журнала. Оставьте пустым, чтобы хранить в папке " +"faceswap" -#: lib/cli/args.py:299 lib/cli/args.py:308 lib/cli/args.py:316 -#: lib/cli/args.py:630 lib/cli/args.py:639 -msgid "Data" -msgstr "Данные" - -#: lib/cli/args.py:300 -msgid "" -"Input directory or video. Either a directory containing the image files you " -"wish to process or path to a video file. NB: This should be the source video/" -"frames NOT the source faces." -msgstr "" -"Входная папка либо видео файл. Папка с набором фотографий для обработки либо " -"видео файл. Примечание: должно указывать на исходное видео либо набор " -"извлеченных кадров, а НЕ уже извлеченных лица." - -#: lib/cli/args.py:309 -msgid "Output directory. This is where the converted files will be saved." -msgstr "Папка для сохранения преобразованных файлов." - -#: lib/cli/args.py:317 -msgid "" -"Optional path to an alignments file. Leave blank if the alignments file is " -"at the default location." -msgstr "Путь к файлу выравнивания. Оставьте пустым, для пути по умолчанию." - -#: lib/cli/args.py:340 -msgid "" -"Extract faces from image or video sources.\n" -"Extraction plugins can be configured in the 'Settings' Menu" -msgstr "" -"Извлечь лица из изображений или видео источников.\n" -"Плагины извлечения можно настроить в меню 'Настройки'" - -#: lib/cli/args.py:365 lib/cli/args.py:381 lib/cli/args.py:393 -#: lib/cli/args.py:428 lib/cli/args.py:446 lib/cli/args.py:458 -#: lib/cli/args.py:649 lib/cli/args.py:676 lib/cli/args.py:712 -msgid "Plugins" -msgstr "Плагины" - -#: lib/cli/args.py:366 -msgid "" -"R|Detector to use. Some of these have configurable settings in '/config/" -"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" -"L|cv2-dnn: A CPU only extractor which is the least reliable and least " -"resource intensive. Use this if not using a GPU and time is important.\n" -"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " -"than other GPU detectors but can often return more false positives.\n" -"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " -"fewer false positives than other GPU detectors, but is a lot more resource " -"intensive." -msgstr "" -"R|Тип детектора. Некоторые могут быть настроенны через '/config/extract.ini' " -"либо 'Settings > Configure Extract 'Plugins':\n" -"L|cv2-dnn: Работает только на CPU, наименее надежный и наименее требователен " -"к ресурсам. Используйте если для вас очень важна скорость, а также не " -"использовать GPU .\n" -"L|mtcnn: Хороший детектор. Быстрый на CPU, ещё быстрее на GPU. Использует " -"меньше ресурсов, нежели другие GPU детекторы, но может производить больше " -"ложных положительных детектирований.\n" -"L|s3fd: Лучший детектор. Медленный на CPU, быстре на GPU. Может " -"детектировать лицо в большем кол-ве ситуация и меньшим кол-вом ошибок, чем " -"другие GPU, но значительно более требователен к ресурсам." - -#: lib/cli/args.py:382 -msgid "" -"R|Aligner to use.\n" -"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " -"but less accurate. Only use this if not using a GPU and time is important.\n" -"L|fan: Best aligner. Fast on GPU, slow on CPU." -msgstr "" -"R|Выравнивание лица.\n" -"L|cv2-dnn: Детектор меток лица, только для CPU. Быстрый, не требователен к " -"ресурсам, но менее точный. Используйте только если вам необходимо не " -"использовать GPU.\n" -"L|fan: Лучший выравниватель. Быстрый на GPU, медленный на CPU." - -#: lib/cli/args.py:394 -msgid "" -"R|Additional Masker(s) to use. The masks generated here will all take up GPU " -"RAM. You can select none, one or multiple masks, but the extraction may take " -"longer the more you select. NB: The Extended and Components (landmark based) " -"masks are automatically generated on extraction.\n" -"L|bisenet-fp: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked including full head masking " -"(configurable in mask settings).\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " -"faces clear of obstructions. Profile faces and obstructions may result in " -"sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " -"frontal faces. The mask model has been specifically trained to recognize " -"some facial obstructions (hands and eyeglasses). Profile faces may result in " -"sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " -"faces. The mask model has been trained by community members and will need " -"testing for further description. Profile faces may result in sub-par " -"performance.\n" -"The auto generated masks are as follows:\n" -"L|components: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks and the mask is extended upwards onto the " -"forehead.\n" -"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" -msgstr "" -"R|Создание доп. масок. Генерация масок требует дополнительной памяти GPU. Вы " -"можете выбрать none, одну, или несколько масок, но процес извлечение может " -"занять больше времени в зависимости от выбора. Прим.: Маски Extended и " -"Components (на основе меток лица) всегда создаются автоматически при " -"извлечении лиц.\n" -"L|bisenet-fp: Относительно легкая маска на основе NN, которая обеспечивает " -"более точный контроль над маскируемой областью, включая полное маскирование " -"головы (настраивается в настройках маски).\n" -"L|vgg-clear: Маска предназначена для умной сегментации преимущественно " -"фронтальных лиц без препятствий. Фотографии в профиль могут быть обработаны " -"посредственно.\n" -"L|vgg-obstructed: Маска предназначена для умной сегментации преимущественно " -"фронтальных лиц. Эта маска была обучена распознавать некоторые препятствия, " -"такие как руки и очки. Фотографии в профиль могут быть обработаны " -"посредственно.\n" -"L|unet-dfl: Маска предназначена для умной сегментации преимущественно " -"фронтальных лиц. Маска была обучена силами участников сообщества и нуждается " -"в тестировании. Фотографии в профиль могут быть обработаны посредственно.\n" -"Следующие маски создаются автоматически:\n" -"L|components: Маска предназначена для сегментации лица на основе ориентиров " -"лица. Маска создается путем построения выпуклого полигона вокруг внешних " -"ориентиров лица.\n" -"L|extended: Маска предназначена для сегментации лица на основе ориентиров " -"лица. Маска создается путем построения выпуклого полигона вокруг внешних " -"ориентиров лица и расширяется вверх на лоб.\n" -"(пример: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" - -#: lib/cli/args.py:429 -msgid "" -"R|Performing normalization can help the aligner better align faces with " -"difficult lighting conditions at an extraction speed cost. Different methods " -"will yield different results on different sets. NB: This does not impact the " -"output face, just the input to the aligner.\n" -"L|none: Don't perform normalization on the face.\n" -"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " -"face.\n" -"L|hist: Equalize the histograms on the RGB channels.\n" -"L|mean: Normalize the face colors to the mean." -msgstr "" -"R|Нормализация может помочь выравниванию лиц при сложных условиях освещения, " -"ценой снижения скорости. Различные методы дают разные результаты в " -"зависимости от набора лиц. Прим.: Не влияет на вывод лица, только на " -"выравнивание.\n" -"L|none: Не производить нормализацию картинки лица.\n" -"L|clahe: Производить нормализацию методом CLAHE.\n" -"L|hist: Выравнивание гистограммы каналов RGB каналов.\n" -"L|mean: Усреднение цветов лица." - -#: lib/cli/args.py:447 -msgid "" -"The number of times to re-feed the detected face into the aligner. Each time " -"the face is re-fed into the aligner the bounding box is adjusted by a small " -"amount. The final landmarks are then averaged from each iteration. Helps to " -"remove 'micro-jitter' but at the cost of slower extraction speed. The more " -"times the face is re-fed into the aligner, the less micro-jitter should " -"occur but the longer extraction will take." -msgstr "" -"Кол-во проходов выравнивания после обнаружения лица. Каждый раз при " -"повторном выравнивании рамка лица немного корректируется. Окончательные " -"ориентиры затем усредняются. Помогает устранить «микроджиттер», но за счет " -"замедления скорости извлечения. Чем больше проходов выравнивания, тем меньше " -"микродрожание, но тем дольше идет извлечение." - -#: lib/cli/args.py:459 -msgid "" -"If a face isn't found, rotate the images to try to find a face. Can find " -"more faces at the cost of extraction speed. Pass in a single number to use " -"increments of that size up to 360, or pass in a list of numbers to enumerate " -"exactly what angles to check." -msgstr "" -"Если лицо не найдено, поворачивает картинку, чтобы попытаться найти лицо. " -"Может найти больше лиц ценой скорости извлечения. Укажите число, чтобы " -"использовать приращения этого размера до 360, либо передайте список чисел, " -"чтобы точно указать, какие углы проверять." - -#: lib/cli/args.py:471 lib/cli/args.py:481 lib/cli/args.py:494 -#: lib/cli/args.py:508 lib/cli/args.py:749 lib/cli/args.py:763 -#: lib/cli/args.py:776 lib/cli/args.py:790 -msgid "Face Processing" -msgstr "Обработка лиц" - -#: lib/cli/args.py:472 -msgid "" -"Filters out faces detected below this size. Length, in pixels across the " -"diagonal of the bounding box. Set to 0 for off" -msgstr "" -"Отбрасывает лица ниже указанного размера. Длина указывается в пикселях по " -"диагонали. Установите в 0 для отключения" - -#: lib/cli/args.py:482 lib/cli/args.py:764 -msgid "" -"Optionally filter out people who you do not wish to process by passing in an " -"image of that person. Should be a front portrait with a single person in the " -"image. Multiple images can be added space separated. NB: Using face filter " -"will significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Дополнительно вы можете отфильтровать лица людей, которых вы не хотите " -"обрабатывать указав изображение этого человека. На изображении должен быть " -"фронтальный портрет одного человека . Можно указать несколько файлов через " -"пробел. Прим.: Фильтрация лиц существенно снижает скорость извлечения, при " -"этом точность не гарантируется." - -#: lib/cli/args.py:495 lib/cli/args.py:777 -msgid "" -"Optionally select people you wish to process by passing in an image of that " -"person. Should be a front portrait with a single person in the image. " -"Multiple images can be added space separated. NB: Using face filter will " -"significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Дополнительно вы можете выбрать людей, которых вы хотели бы включить в " -"обработку путем указания изображения этого человека. Должен быть фронтальный " -"портрет с лишь одним человеком на картинке. Можно выбрать несколько " -"изображений через пробел. Прим.: Использование фильтра существенно замедлит " -"скорость извлечения. Также точность не гарантируется." - -#: lib/cli/args.py:509 lib/cli/args.py:791 -msgid "" -"For use with the optional nfilter/filter files. Threshold for positive face " -"recognition. Lower values are stricter. NB: Using face filter will " -"significantly decrease extraction speed and its accuracy cannot be " -"guaranteed." -msgstr "" -"Только при использовании файлов nfilter/filter. Порог для распознавания " -"лица. Чем ниже значения, тем строже. Прим.: Использование фильтра лиц " -"существенно замедлит скорость извлечения. Также точность не гарантируется." - -#: lib/cli/args.py:520 lib/cli/args.py:532 lib/cli/args.py:544 -#: lib/cli/args.py:556 -msgid "output" -msgstr "вывод" - -#: lib/cli/args.py:521 -msgid "" -"The output size of extracted faces. Make sure that the model you intend to " -"train supports your required size. This will only need to be changed for hi-" -"res models." -msgstr "" -"Размер извлекаемых лиц в пикселях. Убедитесь, что выбранная Вами модель " -"поддерживает такой входной размер. Стоит изменять только для моделей " -"высокого разрешения." - -#: lib/cli/args.py:533 -msgid "" -"Extract every 'nth' frame. This option will skip frames when extracting " -"faces. For example a value of 1 will extract faces from every frame, a value " -"of 10 will extract faces from every 10th frame." -msgstr "" -"Обрабатывать каждые N кадров. Эта опция будет пропускать лица при " -"извлечении. Например, значение 1 будет искать лица в каждом кадре, а " -"значение 10 в каждом 10том кадре." - -#: lib/cli/args.py:545 -msgid "" -"Automatically save the alignments file after a set amount of frames. By " -"default the alignments file is only saved at the end of the extraction " -"process. NB: If extracting in 2 passes then the alignments file will only " -"start to be saved out during the second pass. WARNING: Don't interrupt the " -"script when writing the file because it might get corrupted. Set to 0 to " -"turn off" -msgstr "" -"Автоматически сохранять файл выравнивания после указанного кол-ва кадров. По " -"умолчанию файл выравнивания сохраняется только в конце процедуры извлечения. " -"Прим.: При извлечении в 2 прохода, файл выравниваний начнёт сохранение " -"только во время второго прохода. ВНИМАНИЕ: Не прерывайте выполнение во время " -"записи, так как это может повлечь порчу файла. Установите в 0 для выключения" - -#: lib/cli/args.py:557 -msgid "Draw landmarks on the ouput faces for debugging purposes." -msgstr "Рисовать ландмарки на выходных лицах для нужд отладки." - -#: lib/cli/args.py:563 lib/cli/args.py:572 lib/cli/args.py:580 -#: lib/cli/args.py:587 lib/cli/args.py:803 lib/cli/args.py:814 -#: lib/cli/args.py:822 lib/cli/args.py:841 lib/cli/args.py:847 -msgid "settings" -msgstr "настройки" - -#: lib/cli/args.py:564 -msgid "" -"Don't run extraction in parallel. Will run each part of the extraction " -"process separately (one after the other) rather than all at the smae time. " -"Useful if VRAM is at a premium." -msgstr "" -"Не проводить параллельное извлечение. Вместо одновременного запуска, каждая " -"стадия извлечения будет запущена отдельно (одна, за другой). Полезно при " -"нехватке VRAM." - -#: lib/cli/args.py:573 -msgid "" -"Skips frames that have already been extracted and exist in the alignments " -"file" -msgstr "" -"Пропускать кадры, которые уже были извлечены и существуют в файле " -"выравнивания" - -#: lib/cli/args.py:581 -msgid "Skip frames that already have detected faces in the alignments file" -msgstr "Пропускать кадры, для которых в файле выравнивания есть найденные лица" - -#: lib/cli/args.py:588 -msgid "Skip saving the detected faces to disk. Just create an alignments file" -msgstr "" -"Не сохранять найденные лица на носитель. Просто создать файл выравнивания" - -#: lib/cli/args.py:610 -msgid "" -"Swap the original faces in a source video/images to your final faces.\n" -"Conversion plugins can be configured in the 'Settings' Menu" -msgstr "" -"Заменить оригиналы лица в исходном видео/фотографиях новыми.\n" -"Плагины конвертации могут быть настроены в меню 'Настройки'" - -#: lib/cli/args.py:631 -msgid "" -"Only required if converting from images to video. Provide The original video " -"that the source frames were extracted from (for extracting the fps and " -"audio)." -msgstr "" -"Нужно указывать лишь при конвертации из набора картинок в видео. " -"Предоставьте исходное видео, из которого были извлечены кадры (для настройки " -"частоты кадров, а также аудио)." - -#: lib/cli/args.py:640 -msgid "" -"Model directory. The directory containing the trained model you wish to use " -"for conversion." -msgstr "" -"Папка с моделью. Папка, содержащая обученную модель, которую вы хотите " -"использовать для преобразования." - -#: lib/cli/args.py:650 -msgid "" -"R|Performs color adjustment to the swapped face. Some of these options have " -"configurable settings in '/config/convert.ini' or 'Settings > Configure " -"Convert Plugins':\n" -"L|avg-color: Adjust the mean of each color channel in the swapped " -"reconstruction to equal the mean of the masked area in the original image.\n" -"L|color-transfer: Transfers the color distribution from the source to the " -"target image using the mean and standard deviations of the L*a*b* color " -"space.\n" -"L|manual-balance: Manually adjust the balance of the image in a variety of " -"color spaces. Best used with the Preview tool to set correct values.\n" -"L|match-hist: Adjust the histogram of each color channel in the swapped " -"reconstruction to equal the histogram of the masked area in the original " -"image.\n" -"L|seamless-clone: Use cv2's seamless clone function to remove extreme " -"gradients at the mask seam by smoothing colors. Generally does not give very " -"satisfactory results.\n" -"L|none: Don't perform color adjustment." -msgstr "" -"R|Производит подгонку цветов в измененном лице. Некоторые из этих опций " -"имеют настройки в файле '/config/convert.ini' либо 'Настройки > Настроить " -"Плагины Конверсии':\n" -"L|avg-color: Подогнать среднее значение каждого цветового канала в " -"замененном лице так, чтобы оно равнялось среднему значению области маски " -"исходного изображения.\n" -"L|color-transfer: Переносит распределение цвета от источника к целевому " -"изображению с использованием среднего и стандартного отклонения цветового " -"пространства L * a * b *.\n" -"L|manual-balance: Ручная настройка баланса изображения в различных цветовых " -"пространствах. Лучше всего использовать с инструментом предварительного " -"просмотра для установки правильных значений.\n" -"L|match-hist: Подгонять гистограмму каждого цветового канала нового лица, " -"гистограммой области маски исходного изображения\n" -"L|seamless-clone: Исп. фунцю cv2's незаметного переноса чтобы убрать " -"экстремальные градиенты на краях маски путём сглаживания цветов. Обычно не " -"дает удовлетворительных результатов.\n" -"L|none: Не производить подгонку цвета." - -#: lib/cli/args.py:677 -msgid "" -"R|Masker to use. NB: The mask you require must exist within the alignments " -"file. You can add additional masks with the Mask Tool.\n" -"L|none: Don't use a mask.\n" -"L|bisenet-fp-face: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked (configurable in mask settings). " -"Use this version of bisenet-fp if your model is trained with 'face' or " -"'legacy' centering.\n" -"L|bisenet-fp-head: Relatively lightweight NN based mask that provides more " -"refined control over the area to be masked (configurable in mask settings). " -"Use this version of bisenet-fp if your model is trained with 'head' " -"centering.\n" -"L|components: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the " -"positioning of landmark locations. A convex hull is constructed around the " -"exterior of the landmarks and the mask is extended upwards onto the " -"forehead.\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " -"faces clear of obstructions. Profile faces and obstructions may result in " -"sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " -"frontal faces. The mask model has been specifically trained to recognize " -"some facial obstructions (hands and eyeglasses). Profile faces may result in " -"sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " -"faces. The mask model has been trained by community members and will need " -"testing for further description. Profile faces may result in sub-par " -"performance.\n" -"L|predicted: If the 'Learn Mask' option was enabled during training, this " -"will use the mask that was created by the trained model." -msgstr "" -"R|Использовать маску. Прим.: Требуемая маска должна наличествовать в файле " -"выравнивания. Доп. маски можно добавить через Инструмент Создания Масок.\n" -"L|none: Не использовать маску.\n" -"L|bisenet-fp-face: Относительно легкая маска на основе NN, которая " -"обеспечивает более точный контроль над маскируемой областью (настраивается в " -"настройках маски). Используйте эту версию bisenet-fp, если ваша модель " -"обучена с центрированием «face» или «legacy» центрирование.\n" -"L|bisenet-fp-head: Относительно легкая маска на основе NN, которая " -"обеспечивает более точный контроль над маскируемой областью (настраивается в " -"настройках маски). Используйте эту версию bisenet-fp, если ваша модель " -"обучена с центрированием «head».\n" -"L| components: маска, предназначенная для сегментации лица на основе " -"найденных ориентиров. Маска создается построением выпуклого многоугольника " -"вокруг внешних ориентиров лица.\n" -"L| extended: маска, предназначенная для сегментации лица на основе " -"расположения ориентиров. Маска создается построением выпуклого " -"многоугольника вокруг внешних ориентиров лица и продолжается вверх на лоб.\n" -"L| vgg-clear: маска, предназначенная для умной сегментации преимущественно " -"фронтальных лиц без препятствий. Лица в профиль и препятствия могут привести " -"к некачественным результатам.\n" -"L| vgg-obstructed: маска, предназначенная для умной сегментации " -"преимущественно фронтальных лиц. Модель маски специально обучена " -"распознавать некоторые лицевые препятствия (руки и очки). Лица в профиль " -"могут привести к некачественным результатам..\n" -"L| unet-dfl: маска, предназначенная для умной сегментации преимущественно " -"фронтальных лиц. Модель маски была обучена членами сообщества и потребует " -"тестирования для дальнейшего описания. Лица в профиль могут привести к " -"некачественным результатам..\n" -"L| predicted: Если во время обучения была включена опция «Learn Mask», будет " -"использоваться маска, созданная обученной моделью." - -#: lib/cli/args.py:713 -msgid "" -"R|The plugin to use to output the converted images. The writers are " -"configurable in '/config/convert.ini' or 'Settings > Configure Convert " -"Plugins:'\n" -"L|ffmpeg: [video] Writes out the convert straight to video. When the input " -"is a series of images then the '-ref' (--reference-video) parameter must be " -"set.\n" -"L|gif: [animated image] Create an animated gif.\n" -"L|opencv: [images] The fastest image writer, but less options and formats " -"than other plugins.\n" -"L|pillow: [images] Slower than opencv, but has more options and supports " -"more formats." -msgstr "" -"R|Тип плагина для вывода конвертированных изображений. Записывающие плагины " -"можно настроить в '/config/convert.ini' либо 'Настройки > Настроить Плагины " -"Конверсии:'\n" -"L|ffmpeg: [видео] Записывает результат конверсии сразу в видео файл. Если " -"входом является серий изображений, то нужно также указать параметр '-ref' (--" -"reference-video).\n" -"L|gif: [анимированное изображение] Создает анимированный gif.\n" -"L|opencv: [изображения] Наибыстрейший способ записи, но с меньшим кол-вом " -"опций и форматов вывода.\n" -"L|pillow: [изображения] Более медленный, чем opencv, но имеет больше опций и " -"поддерживает больше форматов." - -#: lib/cli/args.py:732 lib/cli/args.py:739 lib/cli/args.py:833 -msgid "Frame Processing" -msgstr "Обработка кадров" - -#: lib/cli/args.py:733 -msgid "" -"Scale the final output frames by this amount. 100%% will output the frames " -"at source dimensions. 50%% at half size 200%% at double size" -msgstr "" -"Масштабировать оконечные кадры до указанного процента. 100%% будет выводить " -"кадры в исходном размере. 50%% половина от размера, а 200%% в удвоенном " -"размере" - -#: lib/cli/args.py:740 -msgid "" -"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " -"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " -"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " -"converting from images, then the filenames must end with the frame-number!" -msgstr "" -"Диапазон кадров к которым применять перенос, например, для кадров от 10 до " -"50, и 90 до 100 укажите: --frame-ranges 10-50 90-100. Кадры попадающие вне " -"выбранного диапазона будут отброшены если не указано '-k' (--keep-" -"unchanged). Прим.: Если при конверсии используются изображения, то имена " -"файлов должны заканчиваться номером кадра!" - -#: lib/cli/args.py:750 -msgid "" -"If you have not cleansed your alignments file, then you can filter out faces " -"by defining a folder here that contains the faces extracted from your input " -"files/video. If this folder is defined, then only faces that exist within " -"your alignments file and also exist within the specified folder will be " -"converted. Leaving this blank will convert all faces that exist within the " -"alignments file." -msgstr "" -"Если вы не вычистили ваш файл выравниваний, то вы можете отфильтровать лица " -"указав здесь папку, которая содержит лица извлеченные из входных файлов/" -"видео. Если эта папка указана, то, только лица, которые существуют в файле " -"выравниваний и ТАКЖЕ существуют в указанной папке будут сконвертированы. " -"Если оставить это поле пустым, то все лица, которые существуют в файле " -"выравниваний будут сконвертированы." - -#: lib/cli/args.py:804 -msgid "" -"The maximum number of parallel processes for performing conversion. " -"Converting images is system RAM heavy so it is possible to run out of memory " -"if you have a lot of processes and not enough RAM to accommodate them all. " -"Setting this to 0 will use the maximum available. No matter what you set " -"this to, it will never attempt to use more processes than are available on " -"your system. If singleprocess is enabled this setting will be ignored." -msgstr "" -"Максимальное количество параллельных процессов для выполнения " -"преобразования. Преобразование изображений требует большого объема системной " -"памяти, поэтому возможна ее нехватка, если у вас много процессов и не " -"хватает памяти для их всех. Установка этого значения на 0 будет использовать " -"максимально доступное значение. Независимо от ваших установок, никогда не " -"будет использоваться больше процессов, чем доступно в вашей системе. Если " -"включен одиночный процесс, этот параметр будет проигнорирован." - -#: lib/cli/args.py:815 -msgid "" -"[LEGACY] This only needs to be selected if a legacy model is being loaded or " -"if there are multiple models in the model folder" -msgstr "" -"[СОВМЕСТИМОСТЬ] Это нужно выбирать только в том случае, если загружается " -"устаревшая модель или если в папке сохранения есть несколько моделей" - -#: lib/cli/args.py:823 -msgid "" -"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " -"alignments file for your destination video. However, if you wish you can " -"generate the alignments on-the-fly by enabling this option. This will use an " -"inferior extraction pipeline and will lead to substandard results. If an " -"alignments file is found, this option will be ignored." -msgstr "" -"Включить преобразование на лету. НЕ рекомендуется. Вам стоит создать чистый " -"файл выравнивания для вашего целевого видео. Однако, если вы хотите, вы " -"можете сгенерировать выравнивания на лету, включив эту опцию. Это приведет к " -"использованию улучшенного конвейера экстракции и некачественных результатов. " -"Если файл выравниваний найден, этот параметр будет проигнорирован." - -#: lib/cli/args.py:834 -msgid "" -"When used with --frame-ranges outputs the unchanged frames that are not " -"processed instead of discarding them." -msgstr "" -"При использовании с --frame-range кадры не попавшие в диапазон выводятся " -"неизменными, вместо их пропуска." - -#: lib/cli/args.py:842 -msgid "Swap the model. Instead converting from of A -> B, converts B -> A" -msgstr "" -"Поменять модели местами. Вместо преобразования из A -> B, преобразует B -> A" - -#: lib/cli/args.py:848 -msgid "Disable multiprocessing. Slower but less resource intensive." -msgstr "Отключить многопроцессорность. Медленнее, но менее ресурсоемко." - -#: lib/cli/args.py:864 -msgid "" -"Train a model on extracted original (A) and swap (B) faces.\n" -"Training models can take a long time. Anything from 24hrs to over a week\n" -"Model plugins can be configured in the 'Settings' Menu" -msgstr "" -"Начать обучение модели используя наборы лиц: (A) - исходное лицо и (B) - " -"новое лицо.\n" -"Обучение моделей может занять долгое время: от 24 часов до недели\n" -"Каждую модель можно отдельно настроить в меню «Настройки»" - -#: lib/cli/args.py:883 lib/cli/args.py:892 -msgid "faces" -msgstr "лица" - -#: lib/cli/args.py:884 -msgid "" -"Input directory. A directory containing training images for face A. This is " -"the original face, i.e. the face that you want to remove and replace with " -"face B." -msgstr "" -"Входная папка. Папка содержащая изображения для тренировки лица A. Это " -"исходное лицо т.е. лицо, которое вы хотите убрать, заменив лицом B." - -#: lib/cli/args.py:893 -msgid "" -"Input directory. A directory containing training images for face B. This is " -"the swap face, i.e. the face that you want to place onto the head of person " -"A." -msgstr "" -"Входная папка. Папка содержащая изображения для тренировки лица B. Это новое " -"лицо т.е. лицо, которое вы хотите поместить на голову человека A." - -#: lib/cli/args.py:901 lib/cli/args.py:913 lib/cli/args.py:929 -#: lib/cli/args.py:954 lib/cli/args.py:964 -msgid "model" -msgstr "модель" - -#: lib/cli/args.py:902 -msgid "" -"Model directory. This is where the training data will be stored. You should " -"always specify a new folder for new models. If starting a new model, select " -"either an empty folder, or a folder which does not exist (which will be " -"created). If continuing to train an existing model, specify the location of " -"the existing model." -msgstr "" -"Папка сохранений модели. Здесь сохраняется прогресс тренировки. Следует " -"всегда создавать новую папку для новых моделей. При начале тренировки новой " -"модели, выберите пустую либо несуществующую папку (во втором случае она " -"будет создана). Если вы хотите продолжить тренировку, выберите папку с уже " -"существующими сохранениями." - -#: lib/cli/args.py:914 -msgid "" -"R|Load the weights from a pre-existing model into a newly created model. For " -"most models this will load weights from the Encoder of the given model into " -"the encoder of the newly created model. Some plugins may have specific " -"configuration options allowing you to load weights from other layers. " -"Weights will only be loaded when creating a new model. This option will be " -"ignored if you are resuming an existing model. Generally you will also want " -"to 'freeze-weights' whilst the rest of your model catches up with your " -"Encoder.\n" -"NB: Weights can only be loaded from models of the same plugin as you intend " -"to train." -msgstr "" -"R|Загрузите веса из уже существующей модели во вновь созданную модель. Для " -"большинства моделей это загрузит веса из кодировщика данной модели в " -"кодировщик вновь созданной модели. Некоторые плагины могут иметь " -"определенные параметры конфигурации, позволяющие загружать веса из других " -"слоев. Вес будет загружен только при создании новой модели. Этот параметр " -"будет проигнорирован, если вы возобновите работу с существующей моделью. Как " -"правило, вы также захотите «заморозить вес», пока остальная часть вашей " -"модели догонит ваш кодировщик.\n" -"NB: Вес можно загружать только из моделей того же плагина, который вы " -"собираетесь тренировать." - -#: lib/cli/args.py:930 -msgid "" -"R|Select which trainer to use. Trainers can be configured from the Settings " -"menu or the config folder.\n" -"L|original: The original model created by /u/deepfakes.\n" -"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " -"for full dfaker method.\n" -"L|dfl-h128: 128px in/out model from deepfacelab\n" -"L|dfl-sae: Adaptable model from deepfacelab\n" -"L|dlight: A lightweight, high resolution DFaker variant.\n" -"L|iae: A model that uses intermediate layers to try to get better details\n" -"L|lightweight: A lightweight model for low-end cards. Don't expect great " -"results. Can train as low as 1.6GB with batch size 8.\n" -"L|realface: A high detail, dual density model based on DFaker, with " -"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " -"won't work so well. By andenixa et al. Very configurable.\n" -"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " -"unbalanced so B>A swaps won't work so well. Very configurable.\n" -"L|villain: 128px in/out model from villainguy. Very resource hungry (You " -"will require a GPU with a fair amount of VRAM). Good for details, but more " -"susceptible to color differences." -msgstr "" -"R|Выберите тренера для использования. Тренеры могут быть настроенны через " -"меню Настройки либо в папке config.\n" -"L|original: Оригинальная модель созданная /u/deepfakes.\n" -"L|dfaker: модель с 64px вход/128px выходом от dfaker. Включите 'warp-to-" -"landmarks' для полного соответствия методу dfaker.\n" -"L|dfl-h128: 128px вход/выход модель от deepfacelab\n" -"L|dfl-sae: Адаптивная модель от deepfacelab\n" -"L|dlight: Легковесная модель высокого разрешения. Один из вариантов DFaker.\n" -"L|iae: Модель использующая промежуточные слои, для достижения лучшей " -"детализции\n" -"L|lightweight: Легковесная модель для младшей линейки видеокарт. Не ожидайте " -"хороших результатов. Может тренировать на картах с 1.6Гб памяти при размере " -"серии 8.\n" -"L|realface: Модель повышенной детализации, с двумя сложносоставными слоями, " -"базированная на DFaker, с настраиваемым разрешением входа/выхода. " -"Автоэнкодеры не сбалансированы, поэтому свапы B>A не дадут хорошего " -"качества. andenixa и другие. Очень настраиваемая.\n" -"L|unbalanced: Модель 128px вход/выход от andenixa. Автоэнкодеры не " -"сбалансированы, поэтому свапы B>A не будут очень хорошими. Очень " -"настраеваемая.\n" -"L|villain: Модель 128px вход/выход от villainguy. Очень требовательна к " -"ресурсам (Вам потребуется GPU с хорошим количеством видеопамяти). Хороша для " -"деталей, но подвержена к неправильной передаче цвета." - -#: lib/cli/args.py:955 -msgid "" -"Output a summary of the model and exit. If a model folder is provided then a " -"summary of the saved model is displayed. Otherwise a summary of the model " -"that would be created by the chosen plugin and configuration settings is " -"displayed." -msgstr "" -"Выведите сводку модели и выйдите. Если предоставлена папка модели, " -"отображается сводка сохраненной модели. В противном случае отображается " -"сводная информация о модели, которая будет создана выбранным плагином, и " -"параметрами конфигурации." - -#: lib/cli/args.py:965 -msgid "" -"Freeze the weights of the model. Freezing weights means that some of the " -"parameters in the model will no longer continue to learn, but those that are " -"not frozen will continue to learn. For most models, this will freeze the " -"encoder, but some models may have configuration options for freezing other " -"layers." -msgstr "" -"Зафиксируйте веса модели. Замораживание весов означает, что некоторые " -"параметры в модели больше не будут изучаться, но те, которые не заморожены, " -"продолжат обучение. Для большинства моделей это заморозит кодировщик, но " -"некоторые модели могут иметь параметры конфигурации для замораживания других " -"слоев." - -#: lib/cli/args.py:978 lib/cli/args.py:990 lib/cli/args.py:1001 -#: lib/cli/args.py:1087 -msgid "training" -msgstr "тренировка" - -#: lib/cli/args.py:979 -msgid "" -"Batch size. This is the number of images processed through the model for " -"each side per iteration. NB: As the model is fed 2 sides at a time, the " -"actual number of images within the model at any one time is double the " -"number that you set here. Larger batches require more GPU RAM." -msgstr "" -"Размер партии. Это количество изображений для каждой стороны, которые " -"обрабатываются моделью за одну итерацию. Примечание: Поскольку в модель " -"передается сразу две стороны за раз, реальное количество загружаемых " -"изображений в два раза больше этого числа. Увеличение размера партии требует " -"больше памяти GPU." - -#: lib/cli/args.py:991 -msgid "" -"Length of training in iterations. This is only really used for automation. " -"There is no 'correct' number of iterations a model should be trained for. " -"You should stop training when you are happy with the previews. However, if " -"you want the model to stop automatically at a set number of iterations, you " -"can set that value here." -msgstr "" -"Кол-во итераций для тренировки. Используется только для автоматизирования. " -"Не существует \"правильного\" кол-ва итераций для любой выбранной модели. " -"Тренировку стоит завершать только когда вы довольны кадрами на превью. " -"Однако, если вы хотите, чтобы тренировка прервалась после указанного кол-ва " -"итерация, вы можете ввести это здесь." - -#: lib/cli/args.py:1002 -msgid "" -"Use the Tensorflow Mirrored Distrubution Strategy to train on multiple GPUs." -msgstr "" -"Использовать стратегию зеркального распределения Tensorflow для совместной " -"тренировки сразу на нескольких GPU." - -#: lib/cli/args.py:1012 lib/cli/args.py:1022 -msgid "Saving" -msgstr "Сохранение" - -#: lib/cli/args.py:1013 -msgid "Sets the number of iterations between each model save." -msgstr "Установка количества итераций между сохранениями модели." - -#: lib/cli/args.py:1023 -msgid "" -"Sets the number of iterations before saving a backup snapshot of the model " -"in it's current state. Set to 0 for off." -msgstr "" -"Устанавливает кол-во итераций перед созданием резервной копии модели. " -"Установите в 0 для отключения." - -#: lib/cli/args.py:1030 lib/cli/args.py:1041 lib/cli/args.py:1052 -msgid "timelapse" -msgstr "таймлапс" - -#: lib/cli/args.py:1031 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. " -"This should be the input folder of 'A' faces that you would like to use for " -"creating the timelapse. You must also supply a --timelapse-output and a --" -"timelapse-input-B parameter." -msgstr "" -"Только при создании таймлапсов. Сохраняет предварительный просмотр выбранных " -"лиц в папку timelapse-output при каждом сохранении. Следует указать входную " -"папку лиц набора 'A' для использования при создании таймлапса. Вам также " -"нужно указать параметры--timelapse-output и --timelapse-input-B." - -#: lib/cli/args.py:1042 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. " -"This should be the input folder of 'B' faces that you would like to use for " -"creating the timelapse. You must also supply a --timelapse-output and a --" -"timelapse-input-A parameter." -msgstr "" -"Только при создании таймлапса. Таймлапс будет сохранять изображения " -"выбранных лиц в папке таймлапсов при каждой итерации сохранения. Это должна " -"быть папка для ввода лиц из набора 'B', для использования в создании " -"таймлапса. Вы также должны указать параметр --timelapse-output и --timelapse-" -"input-A." - -#: lib/cli/args.py:1053 -msgid "" -"Optional for creating a timelapse. Timelapse will save an image of your " -"selected faces into the timelapse-output folder at every save iteration. If " -"the input folders are supplied but no output folder, it will default to your " -"model folder /timelapse/" -msgstr "" -"Опционально, при создании таймлапса. Создаст картинку текущего таймлапса " -"выбранных лиц в папке timelapse-output при каждом сохранении модели. Если " -"указаны только входные папки, то по умолчанию вывод будет сохранен вместе с " -"моделью в подкаталог /timelapse/" - -#: lib/cli/args.py:1065 lib/cli/args.py:1072 lib/cli/args.py:1079 -msgid "preview" -msgstr "предварительный просмотр" - -#: lib/cli/args.py:1066 -msgid "" -"Percentage amount to scale the preview by. 100%% is the model output size." -msgstr "" -"Величина в процентах, на которую требуется масштабировать предварительный " -"просмотр. 100 %% - размер вывода модели." - -#: lib/cli/args.py:1073 -msgid "Show training preview output. in a separate window." -msgstr "Показывать предварительный просмотр в отдельном окне." - -#: lib/cli/args.py:1080 -msgid "" -"Writes the training result to a file. The image will be stored in the root " -"of your FaceSwap folder." -msgstr "" -"Записывает результат тренировки в файл. Файл будет сохранен в коренной папке " -"FaceSwap." - -#: lib/cli/args.py:1088 -msgid "" -"Disables TensorBoard logging. NB: Disabling logs means that you will not be " -"able to use the graph or analysis for this session in the GUI." -msgstr "" -"Отключает журнал TensorBoard. Примечание: Отключение журналов означает, что " -"вы не сможете использовать графики или анализ сессии внутри GUI." - -#: lib/cli/args.py:1095 lib/cli/args.py:1104 lib/cli/args.py:1113 -#: lib/cli/args.py:1122 -msgid "augmentation" -msgstr "аугментация" - -#: lib/cli/args.py:1096 -msgid "" -"Warps training faces to closely matched Landmarks from the opposite face-set " -"rather than randomly warping the face. This is the 'dfaker' way of doing " -"warping." -msgstr "" -"Вместо случайного искажения лица, деформирует лица в соответствии с " -"Ориентирами/Landmarks противоположного набора лиц. Этот способ используется " -"пакетом \"dfaker\"." - -#: lib/cli/args.py:1105 -msgid "" -"To effectively learn, a random set of images are flipped horizontally. " -"Sometimes it is desirable for this not to occur. Generally this should be " -"left off except for during 'fit training'." -msgstr "" -"Для повышения эффективности обучения, некоторые изображения случайным " -"образом переворачивается по горизонтали. Иногда желательно, чтобы этого не " -"происходило. Как правило, эту настройку не стоит трогать, за исключением " -"периода «финальной шлифовки»." - -#: lib/cli/args.py:1114 -msgid "" -"Color augmentation helps make the model less susceptible to color " -"differences between the A and B sets, at an increased training time cost. " -"Enable this option to disable color augmentation." -msgstr "" -"Цветовая аугментация помогает модели быть менее чувствительной к разнице " -"цвета между наборами A and B ценой некоторого замедления скорости " -"тренировки. Включите эту опцию для отключения цветовой аугментации." - -#: lib/cli/args.py:1123 -msgid "" -"Warping is integral to training the Neural Network. This option should only " -"be enabled towards the very end of training to try to bring out more detail. " -"Think of it as 'fine-tuning'. Enabling this option from the beginning is " -"likely to kill a model and lead to terrible results." -msgstr "" -"Внесение случайных искажение является неотъемлемой частью обучения нейронной " -"сети. Эту опцию следует включать только в самом конце обучения, чтобы " -"попытаться выявить больше деталей. Думайте об этом как о «стадии шлифовки». " -"Включение этой опции с самого начала может убить модель и привести к ужасным " -"результатам." - -#: lib/cli/args.py:1148 +#: lib/cli/args.py:319 msgid "Output to Shell console instead of GUI console" -msgstr "Вывод в системную консоль вместо GUI" - -#~ msgid "" -#~ "DEPRECATED - This option will be removed in a future update. Path to " -#~ "alignments file for training set A. Defaults to /alignments.json " -#~ "if not provided." -#~ msgstr "" -#~ "УСТАРЕЛО - Эта настройка будет удалена в будущих обновлениях. Путь к " -#~ "файлу выравнивания для обучающего набора A. По умолчанию используется " -#~ " /alignments.json, если он не указан." - -#~ msgid "" -#~ "DEPRECATED - This option will be removed in a future update. Path to " -#~ "alignments file for training set B. Defaults to /alignments.json " -#~ "if not provided." -#~ msgstr "" -#~ "УСТАРЕЛО - Эта настройка будет удалена в будущих обновлениях. Путь к " -#~ "файлу выравнивания для обучающего набора B. По умолчанию используется " -#~ " /alignments.json, если он не указан." +msgstr "Вывод в консоль Shell вместо консоли GUI" diff --git a/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.mo b/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.mo new file mode 100644 index 0000000000..51d7f9676f Binary files /dev/null and b/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.mo differ diff --git a/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.po b/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.po new file mode 100755 index 0000000000..e95bf84dd7 --- /dev/null +++ b/locales/ru/LC_MESSAGES/lib.cli.args_extract_convert.po @@ -0,0 +1,710 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-12 11:56+0100\n" +"PO-Revision-Date: 2024-04-12 11:59+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_extract_convert.py:46 lib/cli/args_extract_convert.py:56 +#: lib/cli/args_extract_convert.py:64 lib/cli/args_extract_convert.py:122 +#: lib/cli/args_extract_convert.py:483 lib/cli/args_extract_convert.py:492 +msgid "Data" +msgstr "Данные" + +#: lib/cli/args_extract_convert.py:48 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file. NB: This should be the source video/" +"frames NOT the source faces." +msgstr "" +"Входная папка или видео. Либо каталог, содержащий файлы изображений, которые " +"вы хотите обработать, либо путь к видеофайлу. ПРИМЕЧАНИЕ: Это должно быть " +"исходное видео/кадры, а не исходные лица." + +#: lib/cli/args_extract_convert.py:57 +msgid "Output directory. This is where the converted files will be saved." +msgstr "Выходная папка. Здесь будут сохранены преобразованные файлы." + +#: lib/cli/args_extract_convert.py:66 +msgid "" +"Optional path to an alignments file. Leave blank if the alignments file is " +"at the default location." +msgstr "" +"Необязательный путь к файлу выравниваний. Оставьте пустым, если файл " +"выравнивания находится в месте по умолчанию." + +#: lib/cli/args_extract_convert.py:97 +msgid "" +"Extract faces from image or video sources.\n" +"Extraction plugins can be configured in the 'Settings' Menu" +msgstr "" +"Извлечение лиц из источников изображений или видео.\n" +"Плагины извлечения можно настроить в меню \"Настройки\"" + +#: lib/cli/args_extract_convert.py:124 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple videos and/or folders of images you wish to extract from. The faces " +"will be output to separate sub-folders in the output_dir." +msgstr "" +"R|Если выбрано, то input_dir должен быть родительской папкой, содержащей " +"несколько видео и/или папок с изображениями, из которых вы хотите извлечь " +"изображение. Лица будут выведены в отдельные вложенные папки в output_dir." + +#: lib/cli/args_extract_convert.py:133 lib/cli/args_extract_convert.py:152 +#: lib/cli/args_extract_convert.py:167 lib/cli/args_extract_convert.py:206 +#: lib/cli/args_extract_convert.py:224 lib/cli/args_extract_convert.py:237 +#: lib/cli/args_extract_convert.py:247 lib/cli/args_extract_convert.py:257 +#: lib/cli/args_extract_convert.py:503 lib/cli/args_extract_convert.py:529 +#: lib/cli/args_extract_convert.py:568 +msgid "Plugins" +msgstr "Плагины" + +#: lib/cli/args_extract_convert.py:135 +msgid "" +"R|Detector to use. Some of these have configurable settings in '/config/" +"extract.ini' or 'Settings > Configure Extract 'Plugins':\n" +"L|cv2-dnn: A CPU only extractor which is the least reliable and least " +"resource intensive. Use this if not using a GPU and time is important.\n" +"L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " +"than other GPU detectors but can often return more false positives.\n" +"L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces and " +"fewer false positives than other GPU detectors, but is a lot more resource " +"intensive.\n" +"L|external: Import a face detection bounding box from a json file. " +"(configurable in Detect settings)" +msgstr "" +"R|Детектор для использования. Некоторые из них имеют настраиваемые параметры " +"в '/config/extract.ini' или 'Settings > Configure Extract 'Plugins':\n" +"L|cv2-dnn: Экстрактор только для процессора, который является наименее " +"надежным и наименее ресурсоемким. Используйте его, если не используется GPU " +"и важно время.\n" +"L|mtcnn: Хороший детектор. Быстрый на CPU, еще быстрее на GPU. Использует " +"меньше ресурсов, чем другие детекторы на GPU, но часто может давать больше " +"ложных срабатываний.\n" +"L|s3fd: Лучший детектор. Медленный на CPU, более быстрый на GPU. Может " +"обнаружить больше лиц и меньше ложных срабатываний, чем другие детекторы на " +"GPU, но требует гораздо больше ресурсов.\n" +"L|external: импортируйте ограничивающую коробку обнаружения лица из файла " +"JSON. (настраивается в настройках обнаружения)" + +#: lib/cli/args_extract_convert.py:154 +msgid "" +"R|Aligner to use.\n" +"L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " +"but less accurate. Only use this if not using a GPU and time is important.\n" +"L|fan: Best aligner. Fast on GPU, slow on CPU.\n" +"L|external: Import 68 point 2D landmarks or an aligned bounding box from a " +"json file. (configurable in Align settings)" +msgstr "" +"R|Выравниватель для использования.\n" +"L|cv2-dnn: Детектор ориентиров только для процессора. Быстрее, менее " +"ресурсоемкий, но менее точный. Используйте его, только если не используется " +"GPU и важно время.\n" +"L|fan: Лучший выравниватель. Быстрый на GPU, медленный на CPU.\n" +"L|external: импорт 68 баллов 2D достопримечательности или выровненная " +"ограничивающая коробка из файла JSON. (настраивается в настройках " +"выравнивания)" + +#: lib/cli/args_extract_convert.py:169 +msgid "" +"R|Additional Masker(s) to use. The masks generated here will all take up GPU " +"RAM. You can select none, one or multiple masks, but the extraction may take " +"longer the more you select. NB: The Extended and Components (landmark based) " +"masks are automatically generated on extraction.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU so will not use any additional VRAM.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"The auto generated masks are as follows:\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" +msgstr "" +"R|Дополнительный маскер(ы) для использования. Все маски, созданные здесь, " +"будут занимать видеопамять GPU. Вы можете выбрать ни одной, одну или " +"несколько масок, но извлечение может занять больше времени, чем больше масок " +"вы выберете. Примечание: Расширенные маски и маски компонентов (на основе " +"ориентиров) генерируются автоматически при извлечении.\n" +"L|bisenet-fp: Относительно легкая маска на основе NN, которая обеспечивает " +"более точный контроль над маскируемой областью, включая полное маскирование " +"головы (настраивается в настройках маски).\n" +"L|custom: Фиктивная маска, которая заполняет область маски всеми 1 или 0 " +"(настраивается в настройках). Она необходима только в том случае, если вы " +"собираетесь вручную редактировать пользовательские маски в ручном " +"инструменте. Эта маска не задействует GPU, поэтому не будет использовать " +"дополнительную память VRAM.\n" +"L|vgg-clear: Маска предназначена для интеллектуальной сегментации " +"преимущественно фронтальных лиц без препятствий. Профильные лица и " +"препятствия могут привести к снижению производительности.\n" +"L|vgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была специально обучена " +"распознавать некоторые препятствия на лице (руки и очки). Лица в профиль " +"могут иметь низкую производительность.\n" +"L|unet-dfl: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была обучена членами " +"сообщества и для дальнейшего описания нуждается в тестировании. Профильные " +"лица могут привести к низкой производительности.\n" +"Автоматически сгенерированные маски выглядят следующим образом:\n" +"L|components: Маска, разработанная для сегментации лица на основе " +"расположения ориентиров. Для создания маски вокруг внешних ориентиров " +"строится выпуклая оболочка.\n" +"L|extended: Маска, предназначенная для сегментации лица на основе " +"расположения ориентиров. Выпуклый корпус строится вокруг внешних ориентиров, " +"и маска расширяется вверх на лоб.\n" +"(например: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" + +#: lib/cli/args_extract_convert.py:208 +msgid "" +"R|Performing normalization can help the aligner better align faces with " +"difficult lighting conditions at an extraction speed cost. Different methods " +"will yield different results on different sets. NB: This does not impact the " +"output face, just the input to the aligner.\n" +"L|none: Don't perform normalization on the face.\n" +"L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"L|hist: Equalize the histograms on the RGB channels.\n" +"L|mean: Normalize the face colors to the mean." +msgstr "" +"R|Проведение нормализации может помочь выравнивателю лучше выравнивать лица " +"со сложными условиями освещения при затратах на скорость извлечения. " +"Различные методы дают разные результаты на разных наборах. NB: Это не влияет " +"на выходное лицо, только на вход выравнивателя.\n" +"L|none: Не выполнять нормализацию лица.\n" +"L|clahe: Выполнить для лица адаптивную гистограммную эквализацию с " +"ограничением контраста.\n" +"L|hist: Уравнять гистограммы в каналах RGB.\n" +"L|mean: Нормализовать цвета лица к среднему значению." + +#: lib/cli/args_extract_convert.py:226 +msgid "" +"The number of times to re-feed the detected face into the aligner. Each time " +"the face is re-fed into the aligner the bounding box is adjusted by a small " +"amount. The final landmarks are then averaged from each iteration. Helps to " +"remove 'micro-jitter' but at the cost of slower extraction speed. The more " +"times the face is re-fed into the aligner, the less micro-jitter should " +"occur but the longer extraction will take." +msgstr "" +"Количество повторных подач обнаруженной области лица в выравниватель. При " +"каждой повторной подаче лица в выравниватель ограничивающая рамка " +"корректируется на небольшую величину. Затем конечные ориентиры усредняются " +"по результатам каждой итерации. Это помогает устранить \"микро-дрожание\", " +"но ценой снижения скорости извлечения. Чем больше раз лицо повторно подается " +"в выравниватель, тем меньше микро-дрожание, но тем больше времени займет " +"извлечение." + +#: lib/cli/args_extract_convert.py:239 +msgid "" +"Re-feed the initially found aligned face through the aligner. Can help " +"produce better alignments for faces that are rotated beyond 45 degrees in " +"the frame or are at extreme angles. Slows down extraction." +msgstr "" +"Повторная подача первоначально найденной выровненной области лица через " +"выравниватель. Может помочь получить лучшее выравнивание для лиц, повернутых " +"в кадре более чем на 45 градусов или расположенных под экстремальными " +"углами. Замедляет извлечение." + +#: lib/cli/args_extract_convert.py:249 +msgid "" +"If a face isn't found, rotate the images to try to find a face. Can find " +"more faces at the cost of extraction speed. Pass in a single number to use " +"increments of that size up to 360, or pass in a list of numbers to enumerate " +"exactly what angles to check." +msgstr "" +"Если лицо не найдено, поворачивает изображения, чтобы попытаться найти лицо. " +"Может найти больше лиц ценой снижения скорости извлечения. Передайте одно " +"число, чтобы использовать приращения этого размера до 360, или передайте " +"список чисел, чтобы перечислить, какие именно углы нужно проверить." + +#: lib/cli/args_extract_convert.py:259 +msgid "" +"Obtain and store face identity encodings from VGGFace2. Slows down extract a " +"little, but will save time if using 'sort by face'" +msgstr "" +"Получение и хранение кодировок идентификации лица из VGGFace2. Немного " +"замедляет извлечение, но экономит время при использовании \"сортировки по " +"лицам\"." + +#: lib/cli/args_extract_convert.py:269 lib/cli/args_extract_convert.py:280 +#: lib/cli/args_extract_convert.py:293 lib/cli/args_extract_convert.py:307 +#: lib/cli/args_extract_convert.py:614 lib/cli/args_extract_convert.py:623 +#: lib/cli/args_extract_convert.py:638 lib/cli/args_extract_convert.py:651 +#: lib/cli/args_extract_convert.py:665 +msgid "Face Processing" +msgstr "Обработка лиц" + +#: lib/cli/args_extract_convert.py:271 +msgid "" +"Filters out faces detected below this size. Length, in pixels across the " +"diagonal of the bounding box. Set to 0 for off" +msgstr "" +"Отфильтровывает лица, обнаруженные ниже этого размера. Длина в пикселях по " +"диагонали ограничивающего поля. Установите значение 0, чтобы выключить" + +#: lib/cli/args_extract_convert.py:282 +msgid "" +"Optionally filter out people who you do not wish to extract by passing in " +"images of those people. Should be a small variety of images at different " +"angles and in different conditions. A folder containing the required images " +"or multiple image files, space separated, can be selected." +msgstr "" +"По желанию отфильтруйте людей, которых вы не хотите извлекать, передав " +"изображения этих людей. Должно быть небольшое разнообразие изображений под " +"разными углами и в разных условиях. Можно выбрать папку, содержащую " +"необходимые изображения, или несколько файлов изображений, разделенных " +"пробелами." + +#: lib/cli/args_extract_convert.py:295 +msgid "" +"Optionally select people you wish to extract by passing in images of that " +"person. Should be a small variety of images at different angles and in " +"different conditions A folder containing the required images or multiple " +"image files, space separated, can be selected." +msgstr "" +"По желанию выберите людей, которых вы хотите извлечь, передав изображения " +"этого человека. Должно быть небольшое разнообразие изображений под разными " +"углами и в разных условиях. Можно выбрать папку, содержащую необходимые " +"изображения, или несколько файлов изображений, разделенных пробелами." + +#: lib/cli/args_extract_convert.py:309 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Higher values are stricter." +msgstr "" +"Для использования с дополнительными файлами nfilter/filter. Порог для " +"положительного распознавания лица. Более высокие значения являются более " +"строгими." + +#: lib/cli/args_extract_convert.py:318 lib/cli/args_extract_convert.py:331 +#: lib/cli/args_extract_convert.py:344 lib/cli/args_extract_convert.py:356 +msgid "output" +msgstr "вывод" + +#: lib/cli/args_extract_convert.py:320 +msgid "" +"The output size of extracted faces. Make sure that the model you intend to " +"train supports your required size. This will only need to be changed for hi-" +"res models." +msgstr "" +"Выходной размер извлеченных лиц. Убедитесь, что модель, которую вы " +"собираетесь тренировать, поддерживает требуемый размер. Это необходимо " +"изменить только для моделей высокого разрешения." + +#: lib/cli/args_extract_convert.py:333 +msgid "" +"Extract every 'nth' frame. This option will skip frames when extracting " +"faces. For example a value of 1 will extract faces from every frame, a value " +"of 10 will extract faces from every 10th frame." +msgstr "" +"Извлекать каждый 'n-й' кадр. Этот параметр пропускает кадры при извлечении " +"лиц. Например, значение 1 будет извлекать лица из каждого кадра, значение 10 " +"будет извлекать лица из каждого 10-го кадра." + +#: lib/cli/args_extract_convert.py:346 +msgid "" +"Automatically save the alignments file after a set amount of frames. By " +"default the alignments file is only saved at the end of the extraction " +"process. NB: If extracting in 2 passes then the alignments file will only " +"start to be saved out during the second pass. WARNING: Don't interrupt the " +"script when writing the file because it might get corrupted. Set to 0 to " +"turn off" +msgstr "" +"Автоматическое сохранение файла выравнивания после заданного количества " +"кадров. По умолчанию файл выравнивания сохраняется только в конце процесса " +"извлечения. Примечание: Если извлечение выполняется в 2 прохода, то файл " +"выравнивания начнет сохраняться только во время второго прохода. " +"ПРЕДУПРЕЖДЕНИЕ: Не прерывайте работу скрипта при записи файла, так как он " +"может быть поврежден. Установите значение 0, чтобы отключить" + +#: lib/cli/args_extract_convert.py:357 +msgid "Draw landmarks on the ouput faces for debugging purposes." +msgstr "Нарисуйте ориентиры на выходящих гранях для отладки." + +#: lib/cli/args_extract_convert.py:363 lib/cli/args_extract_convert.py:373 +#: lib/cli/args_extract_convert.py:381 lib/cli/args_extract_convert.py:388 +#: lib/cli/args_extract_convert.py:678 lib/cli/args_extract_convert.py:691 +#: lib/cli/args_extract_convert.py:712 lib/cli/args_extract_convert.py:718 +msgid "settings" +msgstr "настройки" + +#: lib/cli/args_extract_convert.py:365 +msgid "" +"Don't run extraction in parallel. Will run each part of the extraction " +"process separately (one after the other) rather than all at the same time. " +"Useful if VRAM is at a premium." +msgstr "" +"Не запускать извлечение параллельно. Каждая часть процесса извлечения будет " +"выполняться отдельно (одна за другой), а не одновременно. Полезно, если " +"память VRAM ограничена." + +#: lib/cli/args_extract_convert.py:375 +msgid "" +"Skips frames that have already been extracted and exist in the alignments " +"file" +msgstr "" +"Пропускает кадры, которые уже были извлечены и существуют в файле " +"выравнивания" + +#: lib/cli/args_extract_convert.py:382 +msgid "Skip frames that already have detected faces in the alignments file" +msgstr "" +"Пропустить кадры, в которых уже есть обнаруженные лица в файле выравнивания" + +#: lib/cli/args_extract_convert.py:389 +msgid "Skip saving the detected faces to disk. Just create an alignments file" +msgstr "" +"Не сохранять обнаруженные лица на диск. Просто создать файл выравнивания" + +#: lib/cli/args_extract_convert.py:463 +msgid "" +"Swap the original faces in a source video/images to your final faces.\n" +"Conversion plugins can be configured in the 'Settings' Menu" +msgstr "" +"Поменять исходные лица в исходном видео/изображении на ваши конечные лица.\n" +"Плагины конвертирования можно настроить в меню \"Настройки\"" + +#: lib/cli/args_extract_convert.py:485 +msgid "" +"Only required if converting from images to video. Provide The original video " +"that the source frames were extracted from (for extracting the fps and " +"audio)." +msgstr "" +"Требуется только при преобразовании из изображений в видео. Предоставьте " +"исходное видео, из которого были извлечены исходные кадры (для извлечения " +"кадров в секунду и звука)." + +#: lib/cli/args_extract_convert.py:494 +msgid "" +"Model directory. The directory containing the trained model you wish to use " +"for conversion." +msgstr "" +"Папка модели. Папка, содержащая обученную модель, которую вы хотите " +"использовать для преобразования." + +#: lib/cli/args_extract_convert.py:505 +msgid "" +"R|Performs color adjustment to the swapped face. Some of these options have " +"configurable settings in '/config/convert.ini' or 'Settings > Configure " +"Convert Plugins':\n" +"L|avg-color: Adjust the mean of each color channel in the swapped " +"reconstruction to equal the mean of the masked area in the original image.\n" +"L|color-transfer: Transfers the color distribution from the source to the " +"target image using the mean and standard deviations of the L*a*b* color " +"space.\n" +"L|manual-balance: Manually adjust the balance of the image in a variety of " +"color spaces. Best used with the Preview tool to set correct values.\n" +"L|match-hist: Adjust the histogram of each color channel in the swapped " +"reconstruction to equal the histogram of the masked area in the original " +"image.\n" +"L|seamless-clone: Use cv2's seamless clone function to remove extreme " +"gradients at the mask seam by smoothing colors. Generally does not give very " +"satisfactory results.\n" +"L|none: Don't perform color adjustment." +msgstr "" +"R|Производит корректировку цвета поменявшегося лица. Некоторые из этих " +"параметров настраиваются в '/config/convert.ini' или 'Настройки > Настроить " +"плагины конвертации':\n" +"L|avg-color: корректирует среднее значение каждого цветового канала в " +"реконструкции, чтобы оно было равно среднему значению маскированной области " +"в исходном изображении.\n" +"L|color-transfer: Переносит распределение цветов с исходного изображения на " +"целевое, используя среднее и стандартные отклонения цветового пространства " +"L*a*b*.\n" +"L|manual-balance: Ручная настройка баланса изображения в различных цветовых " +"пространствах. Лучше всего использовать с инструментом предварительного " +"просмотра для установки правильных значений.\n" +"L|match-hist: Настроить гистограмму каждого цветового канала в измененном " +"восстановлении так, чтобы она соответствовала гистограмме маскированной " +"области исходного изображения.\n" +"L|seamless-clone: Используйте функцию бесшовного клонирования cv2 для " +"удаления экстремальных градиентов на шве маски путем сглаживания цветов. " +"Обычно дает не очень удовлетворительные результаты.\n" +"L|none: Не выполнять коррекцию цвета." + +#: lib/cli/args_extract_convert.py:531 +msgid "" +"R|Masker to use. NB: The mask you require must exist within the alignments " +"file. You can add additional masks with the Mask Tool.\n" +"L|none: Don't use a mask.\n" +"L|bisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"L|bisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"L|custom_face: Custom user created, face centered mask.\n" +"L|custom_head: Custom user created, head centered mask.\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance.\n" +"L|predicted: If the 'Learn Mask' option was enabled during training, this " +"will use the mask that was created by the trained model." +msgstr "" +"R|Маскер для использования. Примечание: Нужная маска должна существовать в " +"файле выравнивания. Вы можете добавить дополнительные маски с помощью " +"инструмента Mask Tool.\n" +"L|none: Не использовать маску.\n" +"L|bisenet-fp_face: Относительно легкая маска на основе NN, которая " +"обеспечивает более точный контроль над маскируемой областью (настраивается в " +"настройках маски). Используйте эту версию bisenet-fp, если ваша модель " +"обучена с центрированием 'face' или 'legacy'.\n" +"L|bisenet-fp_head: Относительно легкая маска на основе NN, которая " +"обеспечивает более точный контроль над маскируемой областью (настраивается в " +"настройках маски). Используйте эту версию bisenet-fp, если ваша модель " +"обучена с центрированием по \"голове\".\n" +"L|custom_face: Пользовательская маска, созданная пользователем и " +"центрированная по лицу.\n" +"L|custom_head: Созданная пользователем маска, центрированная по голове.\n" +"L|components: Маска, разработанная для сегментации лица на основе " +"расположения ориентиров. Для создания маски вокруг внешних ориентиров " +"строится выпуклая оболочка.\n" +"L|extended: Маска, предназначенная для сегментации лица на основе " +"расположения ориентиров. Выпуклый корпус строится вокруг внешних ориентиров, " +"и маска расширяется вверх на лоб.\n" +"L|vgg-clear: Маска предназначена для интеллектуальной сегментации " +"преимущественно фронтальных лиц без препятствий. Профильные лица и " +"препятствия могут привести к снижению производительности.\n" +"L|vgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была специально обучена " +"распознавать некоторые препятствия на лице (руки и очки). Лица в профиль " +"могут иметь низкую производительность.\n" +"L|unet-dfl: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была обучена членами " +"сообщества и для дальнейшего описания нуждается в тестировании. Профильные " +"лица могут привести к низкой производительности.\n" +"L|predicted: Если во время обучения была включена опция 'Изучить Маску', то " +"будет использоваться маска, созданная обученной моделью." + +#: lib/cli/args_extract_convert.py:570 +msgid "" +"R|The plugin to use to output the converted images. The writers are " +"configurable in '/config/convert.ini' or 'Settings > Configure Convert " +"Plugins:'\n" +"L|ffmpeg: [video] Writes out the convert straight to video. When the input " +"is a series of images then the '-ref' (--reference-video) parameter must be " +"set.\n" +"L|gif: [animated image] Create an animated gif.\n" +"L|opencv: [images] The fastest image writer, but less options and formats " +"than other plugins.\n" +"L|patch: [images] Outputs the raw swapped face patch, along with the " +"transformation matrix required to re-insert the face back into the original " +"frame. Use this option if you wish to post-process and composite the final " +"face within external tools.\n" +"L|pillow: [images] Slower than opencv, but has more options and supports " +"more formats." +msgstr "" +"R|Плагин, который нужно использовать для вывода преобразованных изображений. " +"Записи настраиваются в '/config/convert.ini' или 'Настройки > Настроить " +"плагины конвертации:'\n" +"L|ffmpeg: [видео] Записывает конвертацию прямо в видео. Если на вход " +"подается серия изображений, необходимо установить параметр '-ref' (--" +"reference-video).\n" +"L|gif: [анимированное изображение] Создает анимированный gif.\n" +"L|opencv: [изображения] Самый быстрый редактор изображений, но имеет меньше " +"опций и форматов, чем другие плагины.\n" +"L|patch: [изображения] Выводит необработанный фрагмент измененного лица " +"вместе с матрицей преобразования, необходимой для повторной вставки лица " +"обратно в исходный кадр.\n" +"L|pillow: [изображения] Медленнее, чем opencv, но имеет больше опций и " +"поддерживает больше форматов." + +#: lib/cli/args_extract_convert.py:591 lib/cli/args_extract_convert.py:600 +#: lib/cli/args_extract_convert.py:703 +msgid "Frame Processing" +msgstr "Обработка лиц" + +#: lib/cli/args_extract_convert.py:593 +#, python-format +msgid "" +"Scale the final output frames by this amount. 100%% will output the frames " +"at source dimensions. 50%% at half size 200%% at double size" +msgstr "" +"Масштабирование конечных выходных кадров на эту величину. 100%% выводит " +"кадры в исходном размере. 50%% при половинном размере 200%% при двойном " +"размере" + +#: lib/cli/args_extract_convert.py:602 +msgid "" +"Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 use " +"--frame-ranges 10-50 90-100. Frames falling outside of the selected range " +"will be discarded unless '-k' (--keep-unchanged) is selected. NB: If you are " +"converting from images, then the filenames must end with the frame-number!" +msgstr "" +"Диапазоны кадров для применения переноса, например, для кадров с 10 по 50 и " +"с 90 по 100 используйте --frame-ranges 10-50 90-100. Кадры, выходящие за " +"пределы выбранного диапазона, будут отброшены, если не выбрана опция '-k' (--" +"keep-unchanged). Примечание: Если вы конвертируете из изображений, то имена " +"файлов должны заканчиваться номером кадра!" + +#: lib/cli/args_extract_convert.py:616 +msgid "" +"Scale the swapped face by this percentage. Positive values will enlarge the " +"face, Negative values will shrink the face." +msgstr "" +"Увеличить масштаб нового лица на этот процент. Положительные значения " +"увеличат лицо, в то время как отрицательные значения уменьшат его." + +#: lib/cli/args_extract_convert.py:625 +msgid "" +"If you have not cleansed your alignments file, then you can filter out faces " +"by defining a folder here that contains the faces extracted from your input " +"files/video. If this folder is defined, then only faces that exist within " +"your alignments file and also exist within the specified folder will be " +"converted. Leaving this blank will convert all faces that exist within the " +"alignments file." +msgstr "" +"Если вы не очистили свой файл выравнивания, то вы можете отфильтровать лица, " +"определив здесь папку, содержащую лица, извлеченные из ваших входных файлов/" +"видео. Если эта папка определена, то будут преобразованы только те лица, " +"которые существуют в вашем файле выравнивания, а также в указанной папке. " +"Если оставить этот параметр пустым, будут преобразованы все лица, " +"существующие в файле выравнивания." + +#: lib/cli/args_extract_convert.py:640 +msgid "" +"Optionally filter out people who you do not wish to process by passing in an " +"image of that person. Should be a front portrait with a single person in the " +"image. Multiple images can be added space separated. NB: Using face filter " +"will significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"По желанию отфильтровать людей, которых вы не хотите обрабатывать, передав " +"изображение этого человека. Это должен быть фронтальный портрет с " +"изображением одного человека. Можно добавить несколько изображений, " +"разделенных пробелами. Примечание: Использование фильтра лиц значительно " +"снизит скорость извлечения, а его точность не гарантируется." + +#: lib/cli/args_extract_convert.py:653 +msgid "" +"Optionally select people you wish to process by passing in an image of that " +"person. Should be a front portrait with a single person in the image. " +"Multiple images can be added space separated. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"По желанию выберите людей, которых вы хотите обработать, передав изображение " +"этого человека. Это должен быть фронтальный портрет с изображением одного " +"человека. Можно добавить несколько изображений, разделенных пробелами. " +"Примечание: Использование фильтра лиц значительно снизит скорость " +"извлечения, а его точность не гарантируется." + +#: lib/cli/args_extract_convert.py:667 +msgid "" +"For use with the optional nfilter/filter files. Threshold for positive face " +"recognition. Lower values are stricter. NB: Using face filter will " +"significantly decrease extraction speed and its accuracy cannot be " +"guaranteed." +msgstr "" +"Для использования с дополнительными файлами nfilter/filter. Порог для " +"положительного распознавания лиц. Более низкие значения являются более " +"строгими. Примечание: Использование фильтра лиц значительно снизит скорость " +"извлечения, а его точность не гарантируется." + +#: lib/cli/args_extract_convert.py:680 +msgid "" +"The maximum number of parallel processes for performing conversion. " +"Converting images is system RAM heavy so it is possible to run out of memory " +"if you have a lot of processes and not enough RAM to accommodate them all. " +"Setting this to 0 will use the maximum available. No matter what you set " +"this to, it will never attempt to use more processes than are available on " +"your system. If singleprocess is enabled this setting will be ignored." +msgstr "" +"Максимальное количество параллельных процессов для выполнения конвертации. " +"Конвертирование изображений занимает много системной оперативной памяти, " +"поэтому может закончиться память, если у вас много процессов и недостаточно " +"оперативной памяти для их размещения. Если установить значение 0, будет " +"использован максимум доступной памяти. Независимо от того, какое значение вы " +"установите, программа никогда не будет пытаться использовать больше " +"процессов, чем доступно в вашей системе. Если включена однопоточная " +"обработка, этот параметр будет проигнорирован." + +#: lib/cli/args_extract_convert.py:693 +msgid "" +"Enable On-The-Fly Conversion. NOT recommended. You should generate a clean " +"alignments file for your destination video. However, if you wish you can " +"generate the alignments on-the-fly by enabling this option. This will use an " +"inferior extraction pipeline and will lead to substandard results. If an " +"alignments file is found, this option will be ignored." +msgstr "" +"Включить преобразование \"на лету\". НЕ рекомендуется. Вы должны " +"сгенерировать чистый файл выравнивания для конечного видео. Однако при " +"желании вы можете генерировать выравнивания \"на лету\", включив эту опцию. " +"При этом будет использоваться некачественный конвейер извлечения, что " +"приведет к некачественным результатам. Если файл выравнивания найден, этот " +"параметр будет проигнорирован." + +#: lib/cli/args_extract_convert.py:705 +msgid "" +"When used with --frame-ranges outputs the unchanged frames that are not " +"processed instead of discarding them." +msgstr "" +"При использовании с --frame-ranges выводит неизмененные кадры, которые не " +"были обработаны, вместо того, чтобы отбрасывать их." + +#: lib/cli/args_extract_convert.py:713 +msgid "Swap the model. Instead converting from of A -> B, converts B -> A" +msgstr "" +"Поменять модель местами. Вместо преобразования из A -> B, преобразуется B -> " +"A" + +#: lib/cli/args_extract_convert.py:719 +msgid "Disable multiprocessing. Slower but less resource intensive." +msgstr "Отключение многопоточной обработки. Медленнее, но менее ресурсоемко." + +#~ msgid "" +#~ "[LEGACY] This only needs to be selected if a legacy model is being loaded " +#~ "or if there are multiple models in the model folder" +#~ msgstr "" +#~ "[ОТБРОШЕН] Этот параметр необходимо выбрать только в том случае, если " +#~ "загружается устаревшая модель или если в папке моделей имеется несколько " +#~ "моделей" diff --git a/locales/ru/LC_MESSAGES/lib.cli.args_train.mo b/locales/ru/LC_MESSAGES/lib.cli.args_train.mo new file mode 100644 index 0000000000..f300361161 Binary files /dev/null and b/locales/ru/LC_MESSAGES/lib.cli.args_train.mo differ diff --git a/locales/ru/LC_MESSAGES/lib.cli.args_train.po b/locales/ru/LC_MESSAGES/lib.cli.args_train.po new file mode 100755 index 0000000000..4e41e6a3f2 --- /dev/null +++ b/locales/ru/LC_MESSAGES/lib.cli.args_train.po @@ -0,0 +1,1045 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 18:04+0000\n" +"PO-Revision-Date: 2024-03-28 18:18+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: lib/cli/args_train.py:30 +msgid "" +"Train a model on extracted original (A) and swap (B) faces.\n" +"Training models can take a long time. Anything from 24hrs to over a week\n" +"Model plugins can be configured in the 'Settings' Menu" +msgstr "" +"Обучение модели на извлеченных оригинальных (A) и подмененных (B) лицах.\n" +"Обучение моделей может занять много времени. От 24 часов до недели.\n" +"Плагины для моделей можно настроить в меню \"Настройки\"" + +#: lib/cli/args_train.py:49 lib/cli/args_train.py:58 +msgid "faces" +msgstr "лица" + +#: lib/cli/args_train.py:51 +msgid "" +"Input directory. A directory containing training images for face A. This is " +"the original face, i.e. the face that you want to remove and replace with " +"face B." +msgstr "" +"Входная папка. Папка, содержащая обучающие изображения для лица A. Это " +"исходное лицо, т.е. лицо, которое вы хотите удалить и заменить лицом B." + +#: lib/cli/args_train.py:60 +msgid "" +"Input directory. A directory containing training images for face B. This is " +"the swap face, i.e. the face that you want to place onto the head of person " +"A." +msgstr "" +"Входная папка. Папка, содержащая обучающие изображения для лица B. Это " +"подменное лицо, т.е. лицо, которое вы хотите поместить на голову человека A." + +#: lib/cli/args_train.py:67 lib/cli/args_train.py:80 lib/cli/args_train.py:97 +#: lib/cli/args_train.py:123 lib/cli/args_train.py:133 +msgid "model" +msgstr "модель" + +#: lib/cli/args_train.py:69 +msgid "" +"Model directory. This is where the training data will be stored. You should " +"always specify a new folder for new models. If starting a new model, select " +"either an empty folder, or a folder which does not exist (which will be " +"created). If continuing to train an existing model, specify the location of " +"the existing model." +msgstr "" +"Папка модели. Здесь будут храниться данные для обучения. Для новых моделей " +"всегда следует указывать новую папку. Если вы начинаете новую модель, " +"выберите либо пустую папку, либо несуществующую папку (которая будет " +"создана). Если вы продолжаете обучение существующей модели, укажите " +"местоположение существующей модели." + +#: lib/cli/args_train.py:82 +msgid "" +"R|Load the weights from a pre-existing model into a newly created model. For " +"most models this will load weights from the Encoder of the given model into " +"the encoder of the newly created model. Some plugins may have specific " +"configuration options allowing you to load weights from other layers. " +"Weights will only be loaded when creating a new model. This option will be " +"ignored if you are resuming an existing model. Generally you will also want " +"to 'freeze-weights' whilst the rest of your model catches up with your " +"Encoder.\n" +"NB: Weights can only be loaded from models of the same plugin as you intend " +"to train." +msgstr "" +"R|Загрузить веса из уже существующей модели во вновь созданную модель. Для " +"большинства моделей это означает загрузку весов из кодировщика данной модели " +"в кодировщик вновь создаваемой модели. Некоторые плагины могут иметь " +"специальные параметры конфигурации, позволяющие загружать веса из других " +"слоев. Веса будут загружаться только при создании новой модели. Эта опция " +"будет проигнорирована, если вы возобновляете существующую модель. Обычно " +"также требуется \"заморозить\" веса, пока остальная часть модели догоняет " +"кодировщик.\n" +"Примечание: Веса могут быть загружены только из моделей того же плагина, " +"который вы собираетесь обучать." + +#: lib/cli/args_train.py:99 +msgid "" +"R|Select which trainer to use. Trainers can be configured from the Settings " +"menu or the config folder.\n" +"L|original: The original model created by /u/deepfakes.\n" +"L|dfaker: 64px in/128px out model from dfaker. Enable 'warp-to-landmarks' " +"for full dfaker method.\n" +"L|dfl-h128: 128px in/out model from deepfacelab\n" +"L|dfl-sae: Adaptable model from deepfacelab\n" +"L|dlight: A lightweight, high resolution DFaker variant.\n" +"L|iae: A model that uses intermediate layers to try to get better details\n" +"L|lightweight: A lightweight model for low-end cards. Don't expect great " +"results. Can train as low as 1.6GB with batch size 8.\n" +"L|realface: A high detail, dual density model based on DFaker, with " +"customizable in/out resolution. The autoencoders are unbalanced so B>A swaps " +"won't work so well. By andenixa et al. Very configurable.\n" +"L|unbalanced: 128px in/out model from andenixa. The autoencoders are " +"unbalanced so B>A swaps won't work so well. Very configurable.\n" +"L|villain: 128px in/out model from villainguy. Very resource hungry (You " +"will require a GPU with a fair amount of VRAM). Good for details, but more " +"susceptible to color differences." +msgstr "" +"R|Выберите, какой тренажер использовать. Тренажеры можно настроить в меню " +"\"Настройки\" или в папке config.\n" +"L|original: Оригинальная модель, созданная /u/deepfakes.\n" +"L|dfaker: модель 64px вход/ 128px выход от dfaker. Включите 'warp-to-" +"landmarks' для полного метода dfaker.\n" +"L|dfl-h128: модель 128px вход/выход от deepfacelab\n" +"L|dfl-sae: Адаптируемая модель от deepfacelab\n" +"L|dlight: Легкий вариант DFaker с высоким разрешением.\n" +"L|iae: Модель, использующая промежуточные слои для получения лучших " +"деталей.\n" +"L|lightweight: Облегченная модель для карт низкого класса. Не ожидайте " +"высоких результатов. Может обучаться на 1,6 ГБ при размере пачки 8.\n" +"L|realface: Модель с высокой детализацией и двойной плотностью, основанная " +"на DFaker, с настраиваемым разрешением входа/выхода. Автоэнкодеры " +"несбалансированы, поэтому замены B>A не будут работать так хорошо. Автор " +"andenixa и др. Очень настраиваемая.\n" +"L|unbalanced: модель 128px вход/выход от andenixa. Автокодировщики " +"несбалансированы, поэтому замены B>A не будут работать так хорошо. Очень " +"настраиваемая.\n" +"L|villain: модель 128px вход/выход от villainguy. Очень требовательна к " +"ресурсам (вам потребуется GPU с достаточным количеством VRAM). Хороша для " +"детализации, но более восприимчива к цветовым различиям." + +#: lib/cli/args_train.py:125 +msgid "" +"Output a summary of the model and exit. If a model folder is provided then a " +"summary of the saved model is displayed. Otherwise a summary of the model " +"that would be created by the chosen plugin and configuration settings is " +"displayed." +msgstr "" +"Вывести сводку модели и выйти. Если указана папка модели, то выводится " +"сводка сохраненной модели. В противном случае отображается сводка модели, " +"которая будет создана выбранным плагином и настройками конфигурации." + +#: lib/cli/args_train.py:135 +msgid "" +"Freeze the weights of the model. Freezing weights means that some of the " +"parameters in the model will no longer continue to learn, but those that are " +"not frozen will continue to learn. For most models, this will freeze the " +"encoder, but some models may have configuration options for freezing other " +"layers." +msgstr "" +"Заморозить веса модели. Замораживание весов означает, что некоторые " +"параметры в модели больше не будут продолжать обучение, но те, которые не " +"заморожены, будут продолжать обучение. Для большинства моделей это означает " +"замораживание кодера, но некоторые модели могут иметь опции конфигурации для " +"замораживания других слоев." + +#: lib/cli/args_train.py:147 lib/cli/args_train.py:160 +#: lib/cli/args_train.py:175 lib/cli/args_train.py:191 +#: lib/cli/args_train.py:200 +msgid "training" +msgstr "тренировка" + +#: lib/cli/args_train.py:149 +msgid "" +"Batch size. This is the number of images processed through the model for " +"each side per iteration. NB: As the model is fed 2 sides at a time, the " +"actual number of images within the model at any one time is double the " +"number that you set here. Larger batches require more GPU RAM." +msgstr "" +"Размер пачки. Это количество изображений, обрабатываемых моделью для каждой " +"стороны за итерацию. Примечание: Поскольку модель обрабатывает 2 стороны " +"одновременно, фактическое количество изображений в модели в любой момент " +"времени будет вдвое больше, чем заданное здесь. Большие партии требуют " +"больше оперативной памяти GPU." + +#: lib/cli/args_train.py:162 +msgid "" +"Length of training in iterations. This is only really used for automation. " +"There is no 'correct' number of iterations a model should be trained for. " +"You should stop training when you are happy with the previews. However, if " +"you want the model to stop automatically at a set number of iterations, you " +"can set that value here." +msgstr "" +"Продолжительность обучения в итерациях. Этот параметр действительно " +"используется только для автоматизации. Не существует \"правильного\" " +"количества итераций, за которое следует обучить модель. Вы должны прекратить " +"обучение, когда будете удовлетворены предварительным просмотром. Однако если " +"вы хотите, чтобы модель автоматически останавливалась при определенном " +"количестве итераций, вы можете задать это значение здесь." + +#: lib/cli/args_train.py:177 +msgid "" +"R|Select the distribution stategy to use.\n" +"L|default: Use Tensorflow's default distribution strategy.\n" +"L|central-storage: Centralizes variables on the CPU whilst operations are " +"performed on 1 or more local GPUs. This can help save some VRAM at the cost " +"of some speed by not storing variables on the GPU. Note: Mixed-Precision is " +"not supported on multi-GPU setups.\n" +"L|mirrored: Supports synchronous distributed training across multiple local " +"GPUs. A copy of the model and all variables are loaded onto each GPU with " +"batches distributed to each GPU at each iteration." +msgstr "" +"R|Выберите стратегию распределения для использования.\n" +"L|default: Использовать стратегию распространения Tensorflow по умолчанию.\n" +"L|central-storage: Централизует переменные на CPU, в то время как операции " +"выполняются на 1 или более локальных GPU. Это может помочь сэкономить " +"немного VRAM за счет некоторой скорости, поскольку переменные не хранятся на " +"GPU. Примечание: Mixed-Precision не поддерживается на многопроцессорных " +"установках.\n" +"L|mirrored: Поддерживает синхронное распределенное обучение на нескольких " +"локальных GPU. Копия модели и все переменные загружаются на каждый GPU с " +"распределением партий на каждый GPU на каждой итерации." + +#: lib/cli/args_train.py:193 +msgid "" +"Disables TensorBoard logging. NB: Disabling logs means that you will not be " +"able to use the graph or analysis for this session in the GUI." +msgstr "" +"Отключает ведение журналов TensorBoard. Примечание: Отключение ведения " +"журналов означает, что вы не сможете использовать график или анализ для этой " +"сессии в графическом интерфейсе." + +#: lib/cli/args_train.py:202 +msgid "" +"Use the Learning Rate Finder to discover the optimal learning rate for " +"training. For new models, this will calculate the optimal learning rate for " +"the model. For existing models this will use the optimal learning rate that " +"was discovered when initializing the model. Setting this option will ignore " +"the manually configured learning rate (configurable in train settings)." +msgstr "" +"Используйте инструмент поиска коэффициента обучения, чтобы найти оптимальную " +"скорость обучения вашей модели. Для новых моделей это позволит рассчитать " +"оптимальный коэффициент обучения для модели. Для существующих моделей будет " +"использован оптимальный коэффициент обучения, найденный при инициализации " +"модели. Установка этой опции приведет к игнорированию вручную настроенного " +"коэффициента обучения (настраиваемого в параметрах обучения)." + +#: lib/cli/args_train.py:215 lib/cli/args_train.py:225 +msgid "Saving" +msgstr "Сохранение" + +#: lib/cli/args_train.py:216 +msgid "Sets the number of iterations between each model save." +msgstr "Устанавливает количество итераций между каждым сохранением модели." + +#: lib/cli/args_train.py:227 +msgid "" +"Sets the number of iterations before saving a backup snapshot of the model " +"in it's current state. Set to 0 for off." +msgstr "" +"Устанавливает количество итераций между каждым сохранением модели. " +"Устанавливает количество итераций перед сохранением резервного снимка модели " +"в текущем состоянии. Установите значение 0 для выключения." + +#: lib/cli/args_train.py:234 lib/cli/args_train.py:246 +#: lib/cli/args_train.py:258 +msgid "timelapse" +msgstr "таймлапс" + +#: lib/cli/args_train.py:236 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'A' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-B parameter." +msgstr "" +"Опционально для создания таймлапса. Timelapse будет сохранять изображение " +"выбранных лиц в папку timelapse-output на каждой итерации сохранения. Это " +"должна быть входная папка с лицами 'A', которые вы хотите использовать для " +"создания timelapse. Вы также должны указать параметры --timelapse-output и --" +"timelapse-input-B." + +#: lib/cli/args_train.py:248 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. " +"This should be the input folder of 'B' faces that you would like to use for " +"creating the timelapse. You must also supply a --timelapse-output and a --" +"timelapse-input-A parameter." +msgstr "" +"Опционально для создания таймлапса. Timelapse будет сохранять изображение " +"выбранных лиц в папку timelapse-output на каждой итерации сохранения. Это " +"должна быть входная папка с лицами 'B', которые вы хотите использовать для " +"создания timelapse. Вы также должны указать параметры --timelapse-output и --" +"timelapse-input-A." + +#: lib/cli/args_train.py:260 +msgid "" +"Optional for creating a timelapse. Timelapse will save an image of your " +"selected faces into the timelapse-output folder at every save iteration. If " +"the input folders are supplied but no output folder, it will default to your " +"model folder/timelapse/" +msgstr "" +"Опционально для создания таймлапса. Timelapse будет сохранять изображение " +"выбранных лиц в папку timelapse-output на каждой итерации сохранения. Если " +"указаны входные папки, но нет выходной папки, то по умолчанию будет выбрана " +"папка модели/timelapse/" + +#: lib/cli/args_train.py:269 lib/cli/args_train.py:276 +msgid "preview" +msgstr "предпросмотр" + +#: lib/cli/args_train.py:270 +msgid "Show training preview output. in a separate window." +msgstr "Показать вывод предварительного просмотра тренировки в отдельном окне." + +#: lib/cli/args_train.py:278 +msgid "" +"Writes the training result to a file. The image will be stored in the root " +"of your FaceSwap folder." +msgstr "" +"Записывает результат обучения в файл. Изображение будет сохранено в корне " +"папки Faceswap." + +#: lib/cli/args_train.py:285 lib/cli/args_train.py:295 +#: lib/cli/args_train.py:305 lib/cli/args_train.py:315 +msgid "augmentation" +msgstr "аугментация" + +#: lib/cli/args_train.py:287 +msgid "" +"Warps training faces to closely matched Landmarks from the opposite face-set " +"rather than randomly warping the face. This is the 'dfaker' way of doing " +"warping." +msgstr "" +"Искажает обучаемые лица до близко подходящих ориентиров из противоположного " +"набора лиц вместо случайного искажения лица. Это способ выполнения искажения " +"от \"dfaker\" ." + +#: lib/cli/args_train.py:297 +msgid "" +"To effectively learn, a random set of images are flipped horizontally. " +"Sometimes it is desirable for this not to occur. Generally this should be " +"left off except for during 'fit training'." +msgstr "" +"Для эффективного обучения случайный набор изображений переворачивается по " +"горизонтали. Иногда желательно, чтобы этого не происходило. Как правило, это " +"не нужно делать, за исключением случаев \"тренировки подгонки\"." + +#: lib/cli/args_train.py:307 +msgid "" +"Color augmentation helps make the model less susceptible to color " +"differences between the A and B sets, at an increased training time cost. " +"Enable this option to disable color augmentation." +msgstr "" +"Аугментация цвета помогает сделать модель менее восприимчивой к цветовым " +"различиям между наборами A и B, что влечет за собой увеличение затрат " +"времени на обучение. Включите этот параметр для отключения цветовой " +"аугментации." + +#: lib/cli/args_train.py:317 +msgid "" +"Warping is integral to training the Neural Network. This option should only " +"be enabled towards the very end of training to try to bring out more detail. " +"Think of it as 'fine-tuning'. Enabling this option from the beginning is " +"likely to kill a model and lead to terrible results." +msgstr "" +"Искажение является неотъемлемой частью обучения нейронной сети. Эту опцию " +"следует включать только в самом конце обучения, чтобы попытаться получить " +"больше деталей. Считайте это \"тонкой настройкой\". Включение этой опции в " +"самом начале, скорее всего, погубит модель и приведет к ужасным результатам." + +#~ msgid "Global Options" +#~ msgstr "Глобальные Настройки" + +#~ msgid "" +#~ "R|Exclude GPUs from use by Faceswap. Select the number(s) which " +#~ "correspond to any GPU(s) that you do not wish to be made available to " +#~ "Faceswap. Selecting all GPUs here will force Faceswap into CPU mode.\n" +#~ "L|{}" +#~ msgstr "" +#~ "R|Исключить GPU из использования Faceswap. Выберите номер (номера), " +#~ "соответствующие любому GPU, который вы не хотите предоставлять Faceswap. " +#~ "Если выбрать здесь все GPU, Faceswap перейдет в режим CPU.\n" +#~ "L|{}" + +#~ msgid "" +#~ "Optionally overide the saved config with the path to a custom config file." +#~ msgstr "" +#~ "Опционально переопределите сохраненную конфигурацию, указав путь к " +#~ "пользовательскому файлу конфигурации." + +#~ msgid "" +#~ "Log level. Stick with INFO or VERBOSE unless you need to file an error " +#~ "report. Be careful with TRACE as it will generate a lot of data" +#~ msgstr "" +#~ "Уровень логирования. Придерживайтесь INFO или VERBOSE, если только вам не " +#~ "нужно отправить отчет об ошибке. Будьте осторожны с TRACE, поскольку он " +#~ "генерирует много данных" + +#~ msgid "" +#~ "Path to store the logfile. Leave blank to store in the faceswap folder" +#~ msgstr "" +#~ "Путь для хранения файла журнала. Оставьте пустым, чтобы хранить в папке " +#~ "faceswap" + +#~ msgid "Data" +#~ msgstr "Данные" + +#~ msgid "" +#~ "Input directory or video. Either a directory containing the image files " +#~ "you wish to process or path to a video file. NB: This should be the " +#~ "source video/frames NOT the source faces." +#~ msgstr "" +#~ "Входная папка или видео. Либо каталог, содержащий файлы изображений, " +#~ "которые вы хотите обработать, либо путь к видеофайлу. ПРИМЕЧАНИЕ: Это " +#~ "должно быть исходное видео/кадры, а не исходные лица." + +#~ msgid "Output directory. This is where the converted files will be saved." +#~ msgstr "Выходная папка. Здесь будут сохранены преобразованные файлы." + +#~ msgid "" +#~ "Optional path to an alignments file. Leave blank if the alignments file " +#~ "is at the default location." +#~ msgstr "" +#~ "Необязательный путь к файлу выравниваний. Оставьте пустым, если файл " +#~ "выравнивания находится в месте по умолчанию." + +#~ msgid "" +#~ "Extract faces from image or video sources.\n" +#~ "Extraction plugins can be configured in the 'Settings' Menu" +#~ msgstr "" +#~ "Извлечение лиц из источников изображений или видео.\n" +#~ "Плагины извлечения можно настроить в меню \"Настройки\"" + +#~ msgid "" +#~ "R|If selected then the input_dir should be a parent folder containing " +#~ "multiple videos and/or folders of images you wish to extract from. The " +#~ "faces will be output to separate sub-folders in the output_dir." +#~ msgstr "" +#~ "R|Если выбрано, то input_dir должен быть родительской папкой, содержащей " +#~ "несколько видео и/или папок с изображениями, из которых вы хотите извлечь " +#~ "изображение. Лица будут выведены в отдельные вложенные папки в output_dir." + +#~ msgid "Plugins" +#~ msgstr "Плагины" + +#~ msgid "" +#~ "R|Detector to use. Some of these have configurable settings in '/config/" +#~ "extract.ini' or 'Settings > Configure Extract 'Plugins':\n" +#~ "L|cv2-dnn: A CPU only extractor which is the least reliable and least " +#~ "resource intensive. Use this if not using a GPU and time is important.\n" +#~ "L|mtcnn: Good detector. Fast on CPU, faster on GPU. Uses fewer resources " +#~ "than other GPU detectors but can often return more false positives.\n" +#~ "L|s3fd: Best detector. Slow on CPU, faster on GPU. Can detect more faces " +#~ "and fewer false positives than other GPU detectors, but is a lot more " +#~ "resource intensive." +#~ msgstr "" +#~ "R|Детектор для использования. Некоторые из них имеют настраиваемые " +#~ "параметры в '/config/extract.ini' или 'Settings > Configure Extract " +#~ "'Plugins':\n" +#~ "L|cv2-dnn: Экстрактор только для процессора, который является наименее " +#~ "надежным и наименее ресурсоемким. Используйте его, если не используется " +#~ "GPU и важно время.\n" +#~ "L|mtcnn: Хороший детектор. Быстрый на CPU, еще быстрее на GPU. Использует " +#~ "меньше ресурсов, чем другие детекторы на GPU, но часто может давать " +#~ "больше ложных срабатываний.\n" +#~ "L|s3fd: Лучший детектор. Медленный на CPU, более быстрый на GPU. Может " +#~ "обнаружить больше лиц и меньше ложных срабатываний, чем другие детекторы " +#~ "на GPU, но требует гораздо больше ресурсов." + +#~ msgid "" +#~ "R|Aligner to use.\n" +#~ "L|cv2-dnn: A CPU only landmark detector. Faster, less resource intensive, " +#~ "but less accurate. Only use this if not using a GPU and time is " +#~ "important.\n" +#~ "L|fan: Best aligner. Fast on GPU, slow on CPU." +#~ msgstr "" +#~ "R|Выравниватель для использования.\n" +#~ "L|cv2-dnn: Детектор ориентиров только для процессора. Быстрее, менее " +#~ "ресурсоемкий, но менее точный. Используйте его, только если не " +#~ "используется GPU и важно время.\n" +#~ "L|fan: Лучший выравниватель. Быстрый на GPU, медленный на CPU." + +#~ msgid "" +#~ "R|Additional Masker(s) to use. The masks generated here will all take up " +#~ "GPU RAM. You can select none, one or multiple masks, but the extraction " +#~ "may take longer the more you select. NB: The Extended and Components " +#~ "(landmark based) masks are automatically generated on extraction.\n" +#~ "L|bisenet-fp: Relatively lightweight NN based mask that provides more " +#~ "refined control over the area to be masked including full head masking " +#~ "(configurable in mask settings).\n" +#~ "L|custom: A dummy mask that fills the mask area with all 1s or 0s " +#~ "(configurable in settings). This is only required if you intend to " +#~ "manually edit the custom masks yourself in the manual tool. This mask " +#~ "does not use the GPU so will not use any additional VRAM.\n" +#~ "L|vgg-clear: Mask designed to provide smart segmentation of mostly " +#~ "frontal faces clear of obstructions. Profile faces and obstructions may " +#~ "result in sub-par performance.\n" +#~ "L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +#~ "frontal faces. The mask model has been specifically trained to recognize " +#~ "some facial obstructions (hands and eyeglasses). Profile faces may result " +#~ "in sub-par performance.\n" +#~ "L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +#~ "faces. The mask model has been trained by community members and will need " +#~ "testing for further description. Profile faces may result in sub-par " +#~ "performance.\n" +#~ "The auto generated masks are as follows:\n" +#~ "L|components: Mask designed to provide facial segmentation based on the " +#~ "positioning of landmark locations. A convex hull is constructed around " +#~ "the exterior of the landmarks to create a mask.\n" +#~ "L|extended: Mask designed to provide facial segmentation based on the " +#~ "positioning of landmark locations. A convex hull is constructed around " +#~ "the exterior of the landmarks and the mask is extended upwards onto the " +#~ "forehead.\n" +#~ "(eg: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" +#~ msgstr "" +#~ "R|Дополнительный маскер(ы) для использования. Все маски, созданные здесь, " +#~ "будут занимать видеопамять GPU. Вы можете выбрать ни одной, одну или " +#~ "несколько масок, но извлечение может занять больше времени, чем больше " +#~ "масок вы выберете. Примечание: Расширенные маски и маски компонентов (на " +#~ "основе ориентиров) генерируются автоматически при извлечении.\n" +#~ "L|bisenet-fp: Относительно легкая маска на основе NN, которая " +#~ "обеспечивает более точный контроль над маскируемой областью, включая " +#~ "полное маскирование головы (настраивается в настройках маски).\n" +#~ "L|custom: Фиктивная маска, которая заполняет область маски всеми 1 или 0 " +#~ "(настраивается в настройках). Она необходима только в том случае, если вы " +#~ "собираетесь вручную редактировать пользовательские маски в ручном " +#~ "инструменте. Эта маска не задействует GPU, поэтому не будет использовать " +#~ "дополнительную память VRAM.\n" +#~ "L|vgg-clear: Маска предназначена для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц без препятствий. Профильные лица и " +#~ "препятствия могут привести к снижению производительности.\n" +#~ "L|vgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц. Модель маски была специально обучена " +#~ "распознавать некоторые препятствия на лице (руки и очки). Лица в профиль " +#~ "могут иметь низкую производительность.\n" +#~ "L|unet-dfl: Маска, разработанная для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц. Модель маски была обучена членами " +#~ "сообщества и для дальнейшего описания нуждается в тестировании. " +#~ "Профильные лица могут привести к низкой производительности.\n" +#~ "Автоматически сгенерированные маски выглядят следующим образом:\n" +#~ "L|components: Маска, разработанная для сегментации лица на основе " +#~ "расположения ориентиров. Для создания маски вокруг внешних ориентиров " +#~ "строится выпуклая оболочка.\n" +#~ "L|extended: Маска, предназначенная для сегментации лица на основе " +#~ "расположения ориентиров. Выпуклый корпус строится вокруг внешних " +#~ "ориентиров, и маска расширяется вверх на лоб.\n" +#~ "(например: `-M unet-dfl vgg-clear`, `--masker vgg-obstructed`)" + +#~ msgid "" +#~ "R|Performing normalization can help the aligner better align faces with " +#~ "difficult lighting conditions at an extraction speed cost. Different " +#~ "methods will yield different results on different sets. NB: This does not " +#~ "impact the output face, just the input to the aligner.\n" +#~ "L|none: Don't perform normalization on the face.\n" +#~ "L|clahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +#~ "face.\n" +#~ "L|hist: Equalize the histograms on the RGB channels.\n" +#~ "L|mean: Normalize the face colors to the mean." +#~ msgstr "" +#~ "R|Проведение нормализации может помочь выравнивателю лучше выравнивать " +#~ "лица со сложными условиями освещения при затратах на скорость извлечения. " +#~ "Различные методы дают разные результаты на разных наборах. NB: Это не " +#~ "влияет на выходное лицо, только на вход выравнивателя.\n" +#~ "L|none: Не выполнять нормализацию лица.\n" +#~ "L|clahe: Выполнить для лица адаптивную гистограммную эквализацию с " +#~ "ограничением контраста.\n" +#~ "L|hist: Уравнять гистограммы в каналах RGB.\n" +#~ "L|mean: Нормализовать цвета лица к среднему значению." + +#~ msgid "" +#~ "The number of times to re-feed the detected face into the aligner. Each " +#~ "time the face is re-fed into the aligner the bounding box is adjusted by " +#~ "a small amount. The final landmarks are then averaged from each " +#~ "iteration. Helps to remove 'micro-jitter' but at the cost of slower " +#~ "extraction speed. The more times the face is re-fed into the aligner, the " +#~ "less micro-jitter should occur but the longer extraction will take." +#~ msgstr "" +#~ "Количество повторных подач обнаруженной области лица в выравниватель. При " +#~ "каждой повторной подаче лица в выравниватель ограничивающая рамка " +#~ "корректируется на небольшую величину. Затем конечные ориентиры " +#~ "усредняются по результатам каждой итерации. Это помогает устранить " +#~ "\"микро-дрожание\", но ценой снижения скорости извлечения. Чем больше раз " +#~ "лицо повторно подается в выравниватель, тем меньше микро-дрожание, но тем " +#~ "больше времени займет извлечение." + +#~ msgid "" +#~ "Re-feed the initially found aligned face through the aligner. Can help " +#~ "produce better alignments for faces that are rotated beyond 45 degrees in " +#~ "the frame or are at extreme angles. Slows down extraction." +#~ msgstr "" +#~ "Повторная подача первоначально найденной выровненной области лица через " +#~ "выравниватель. Может помочь получить лучшее выравнивание для лиц, " +#~ "повернутых в кадре более чем на 45 градусов или расположенных под " +#~ "экстремальными углами. Замедляет извлечение." + +#~ msgid "" +#~ "If a face isn't found, rotate the images to try to find a face. Can find " +#~ "more faces at the cost of extraction speed. Pass in a single number to " +#~ "use increments of that size up to 360, or pass in a list of numbers to " +#~ "enumerate exactly what angles to check." +#~ msgstr "" +#~ "Если лицо не найдено, поворачивает изображения, чтобы попытаться найти " +#~ "лицо. Может найти больше лиц ценой снижения скорости извлечения. " +#~ "Передайте одно число, чтобы использовать приращения этого размера до 360, " +#~ "или передайте список чисел, чтобы перечислить, какие именно углы нужно " +#~ "проверить." + +#~ msgid "" +#~ "Obtain and store face identity encodings from VGGFace2. Slows down " +#~ "extract a little, but will save time if using 'sort by face'" +#~ msgstr "" +#~ "Получение и хранение кодировок идентификации лица из VGGFace2. Немного " +#~ "замедляет извлечение, но экономит время при использовании \"сортировки по " +#~ "лицам\"." + +#~ msgid "Face Processing" +#~ msgstr "Обработка лиц" + +#~ msgid "" +#~ "Filters out faces detected below this size. Length, in pixels across the " +#~ "diagonal of the bounding box. Set to 0 for off" +#~ msgstr "" +#~ "Отфильтровывает лица, обнаруженные ниже этого размера. Длина в пикселях " +#~ "по диагонали ограничивающего поля. Установите значение 0, чтобы выключить" + +#~ msgid "" +#~ "Optionally filter out people who you do not wish to extract by passing in " +#~ "images of those people. Should be a small variety of images at different " +#~ "angles and in different conditions. A folder containing the required " +#~ "images or multiple image files, space separated, can be selected." +#~ msgstr "" +#~ "По желанию отфильтруйте людей, которых вы не хотите извлекать, передав " +#~ "изображения этих людей. Должно быть небольшое разнообразие изображений " +#~ "под разными углами и в разных условиях. Можно выбрать папку, содержащую " +#~ "необходимые изображения, или несколько файлов изображений, разделенных " +#~ "пробелами." + +#~ msgid "" +#~ "Optionally select people you wish to extract by passing in images of that " +#~ "person. Should be a small variety of images at different angles and in " +#~ "different conditions A folder containing the required images or multiple " +#~ "image files, space separated, can be selected." +#~ msgstr "" +#~ "По желанию выберите людей, которых вы хотите извлечь, передав изображения " +#~ "этого человека. Должно быть небольшое разнообразие изображений под " +#~ "разными углами и в разных условиях. Можно выбрать папку, содержащую " +#~ "необходимые изображения, или несколько файлов изображений, разделенных " +#~ "пробелами." + +#~ msgid "" +#~ "For use with the optional nfilter/filter files. Threshold for positive " +#~ "face recognition. Higher values are stricter." +#~ msgstr "" +#~ "Для использования с дополнительными файлами nfilter/filter. Порог для " +#~ "положительного распознавания лица. Более высокие значения являются более " +#~ "строгими." + +#~ msgid "output" +#~ msgstr "вывод" + +#~ msgid "" +#~ "The output size of extracted faces. Make sure that the model you intend " +#~ "to train supports your required size. This will only need to be changed " +#~ "for hi-res models." +#~ msgstr "" +#~ "Выходной размер извлеченных лиц. Убедитесь, что модель, которую вы " +#~ "собираетесь тренировать, поддерживает требуемый размер. Это необходимо " +#~ "изменить только для моделей высокого разрешения." + +#~ msgid "" +#~ "Extract every 'nth' frame. This option will skip frames when extracting " +#~ "faces. For example a value of 1 will extract faces from every frame, a " +#~ "value of 10 will extract faces from every 10th frame." +#~ msgstr "" +#~ "Извлекать каждый 'n-й' кадр. Этот параметр пропускает кадры при " +#~ "извлечении лиц. Например, значение 1 будет извлекать лица из каждого " +#~ "кадра, значение 10 будет извлекать лица из каждого 10-го кадра." + +#~ msgid "" +#~ "Automatically save the alignments file after a set amount of frames. By " +#~ "default the alignments file is only saved at the end of the extraction " +#~ "process. NB: If extracting in 2 passes then the alignments file will only " +#~ "start to be saved out during the second pass. WARNING: Don't interrupt " +#~ "the script when writing the file because it might get corrupted. Set to 0 " +#~ "to turn off" +#~ msgstr "" +#~ "Автоматическое сохранение файла выравнивания после заданного количества " +#~ "кадров. По умолчанию файл выравнивания сохраняется только в конце " +#~ "процесса извлечения. Примечание: Если извлечение выполняется в 2 прохода, " +#~ "то файл выравнивания начнет сохраняться только во время второго прохода. " +#~ "ПРЕДУПРЕЖДЕНИЕ: Не прерывайте работу скрипта при записи файла, так как он " +#~ "может быть поврежден. Установите значение 0, чтобы отключить" + +#~ msgid "Draw landmarks on the ouput faces for debugging purposes." +#~ msgstr "Нарисуйте ориентиры на выходящих гранях для отладки." + +#~ msgid "settings" +#~ msgstr "настройки" + +#~ msgid "" +#~ "Don't run extraction in parallel. Will run each part of the extraction " +#~ "process separately (one after the other) rather than all at the same " +#~ "time. Useful if VRAM is at a premium." +#~ msgstr "" +#~ "Не запускать извлечение параллельно. Каждая часть процесса извлечения " +#~ "будет выполняться отдельно (одна за другой), а не одновременно. Полезно, " +#~ "если память VRAM ограничена." + +#~ msgid "" +#~ "Skips frames that have already been extracted and exist in the alignments " +#~ "file" +#~ msgstr "" +#~ "Пропускает кадры, которые уже были извлечены и существуют в файле " +#~ "выравнивания" + +#~ msgid "Skip frames that already have detected faces in the alignments file" +#~ msgstr "" +#~ "Пропустить кадры, в которых уже есть обнаруженные лица в файле " +#~ "выравнивания" + +#~ msgid "" +#~ "Skip saving the detected faces to disk. Just create an alignments file" +#~ msgstr "" +#~ "Не сохранять обнаруженные лица на диск. Просто создать файл выравнивания" + +#~ msgid "" +#~ "Swap the original faces in a source video/images to your final faces.\n" +#~ "Conversion plugins can be configured in the 'Settings' Menu" +#~ msgstr "" +#~ "Поменять исходные лица в исходном видео/изображении на ваши конечные " +#~ "лица.\n" +#~ "Плагины конвертирования можно настроить в меню \"Настройки\"" + +#~ msgid "" +#~ "Only required if converting from images to video. Provide The original " +#~ "video that the source frames were extracted from (for extracting the fps " +#~ "and audio)." +#~ msgstr "" +#~ "Требуется только при преобразовании из изображений в видео. Предоставьте " +#~ "исходное видео, из которого были извлечены исходные кадры (для извлечения " +#~ "кадров в секунду и звука)." + +#~ msgid "" +#~ "Model directory. The directory containing the trained model you wish to " +#~ "use for conversion." +#~ msgstr "" +#~ "Папка модели. Папка, содержащая обученную модель, которую вы хотите " +#~ "использовать для преобразования." + +#~ msgid "" +#~ "R|Performs color adjustment to the swapped face. Some of these options " +#~ "have configurable settings in '/config/convert.ini' or 'Settings > " +#~ "Configure Convert Plugins':\n" +#~ "L|avg-color: Adjust the mean of each color channel in the swapped " +#~ "reconstruction to equal the mean of the masked area in the original " +#~ "image.\n" +#~ "L|color-transfer: Transfers the color distribution from the source to the " +#~ "target image using the mean and standard deviations of the L*a*b* color " +#~ "space.\n" +#~ "L|manual-balance: Manually adjust the balance of the image in a variety " +#~ "of color spaces. Best used with the Preview tool to set correct values.\n" +#~ "L|match-hist: Adjust the histogram of each color channel in the swapped " +#~ "reconstruction to equal the histogram of the masked area in the original " +#~ "image.\n" +#~ "L|seamless-clone: Use cv2's seamless clone function to remove extreme " +#~ "gradients at the mask seam by smoothing colors. Generally does not give " +#~ "very satisfactory results.\n" +#~ "L|none: Don't perform color adjustment." +#~ msgstr "" +#~ "R|Производит корректировку цвета поменявшегося лица. Некоторые из этих " +#~ "параметров настраиваются в '/config/convert.ini' или 'Настройки > " +#~ "Настроить плагины конвертации':\n" +#~ "L|avg-color: корректирует среднее значение каждого цветового канала в " +#~ "реконструкции, чтобы оно было равно среднему значению маскированной " +#~ "области в исходном изображении.\n" +#~ "L|color-transfer: Переносит распределение цветов с исходного изображения " +#~ "на целевое, используя среднее и стандартные отклонения цветового " +#~ "пространства L*a*b*.\n" +#~ "L|manual-balance: Ручная настройка баланса изображения в различных " +#~ "цветовых пространствах. Лучше всего использовать с инструментом " +#~ "предварительного просмотра для установки правильных значений.\n" +#~ "L|match-hist: Настроить гистограмму каждого цветового канала в измененном " +#~ "восстановлении так, чтобы она соответствовала гистограмме маскированной " +#~ "области исходного изображения.\n" +#~ "L|seamless-clone: Используйте функцию бесшовного клонирования cv2 для " +#~ "удаления экстремальных градиентов на шве маски путем сглаживания цветов. " +#~ "Обычно дает не очень удовлетворительные результаты.\n" +#~ "L|none: Не выполнять коррекцию цвета." + +#~ msgid "" +#~ "R|Masker to use. NB: The mask you require must exist within the " +#~ "alignments file. You can add additional masks with the Mask Tool.\n" +#~ "L|none: Don't use a mask.\n" +#~ "L|bisenet-fp_face: Relatively lightweight NN based mask that provides " +#~ "more refined control over the area to be masked (configurable in mask " +#~ "settings). Use this version of bisenet-fp if your model is trained with " +#~ "'face' or 'legacy' centering.\n" +#~ "L|bisenet-fp_head: Relatively lightweight NN based mask that provides " +#~ "more refined control over the area to be masked (configurable in mask " +#~ "settings). Use this version of bisenet-fp if your model is trained with " +#~ "'head' centering.\n" +#~ "L|custom_face: Custom user created, face centered mask.\n" +#~ "L|custom_head: Custom user created, head centered mask.\n" +#~ "L|components: Mask designed to provide facial segmentation based on the " +#~ "positioning of landmark locations. A convex hull is constructed around " +#~ "the exterior of the landmarks to create a mask.\n" +#~ "L|extended: Mask designed to provide facial segmentation based on the " +#~ "positioning of landmark locations. A convex hull is constructed around " +#~ "the exterior of the landmarks and the mask is extended upwards onto the " +#~ "forehead.\n" +#~ "L|vgg-clear: Mask designed to provide smart segmentation of mostly " +#~ "frontal faces clear of obstructions. Profile faces and obstructions may " +#~ "result in sub-par performance.\n" +#~ "L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +#~ "frontal faces. The mask model has been specifically trained to recognize " +#~ "some facial obstructions (hands and eyeglasses). Profile faces may result " +#~ "in sub-par performance.\n" +#~ "L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +#~ "faces. The mask model has been trained by community members and will need " +#~ "testing for further description. Profile faces may result in sub-par " +#~ "performance.\n" +#~ "L|predicted: If the 'Learn Mask' option was enabled during training, this " +#~ "will use the mask that was created by the trained model." +#~ msgstr "" +#~ "R|Маскер для использования. Примечание: Нужная маска должна существовать " +#~ "в файле выравнивания. Вы можете добавить дополнительные маски с помощью " +#~ "инструмента Mask Tool.\n" +#~ "L|none: Не использовать маску.\n" +#~ "L|bisenet-fp_face: Относительно легкая маска на основе NN, которая " +#~ "обеспечивает более точный контроль над маскируемой областью " +#~ "(настраивается в настройках маски). Используйте эту версию bisenet-fp, " +#~ "если ваша модель обучена с центрированием 'face' или 'legacy'.\n" +#~ "L|bisenet-fp_head: Относительно легкая маска на основе NN, которая " +#~ "обеспечивает более точный контроль над маскируемой областью " +#~ "(настраивается в настройках маски). Используйте эту версию bisenet-fp, " +#~ "если ваша модель обучена с центрированием по \"голове\".\n" +#~ "L|custom_face: Пользовательская маска, созданная пользователем и " +#~ "центрированная по лицу.\n" +#~ "L|custom_head: Созданная пользователем маска, центрированная по голове.\n" +#~ "L|components: Маска, разработанная для сегментации лица на основе " +#~ "расположения ориентиров. Для создания маски вокруг внешних ориентиров " +#~ "строится выпуклая оболочка.\n" +#~ "L|extended: Маска, предназначенная для сегментации лица на основе " +#~ "расположения ориентиров. Выпуклый корпус строится вокруг внешних " +#~ "ориентиров, и маска расширяется вверх на лоб.\n" +#~ "L|vgg-clear: Маска предназначена для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц без препятствий. Профильные лица и " +#~ "препятствия могут привести к снижению производительности.\n" +#~ "L|vgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц. Модель маски была специально обучена " +#~ "распознавать некоторые препятствия на лице (руки и очки). Лица в профиль " +#~ "могут иметь низкую производительность.\n" +#~ "L|unet-dfl: Маска, разработанная для интеллектуальной сегментации " +#~ "преимущественно фронтальных лиц. Модель маски была обучена членами " +#~ "сообщества и для дальнейшего описания нуждается в тестировании. " +#~ "Профильные лица могут привести к низкой производительности.\n" +#~ "L|predicted: Если во время обучения была включена опция 'Изучить Маску', " +#~ "то будет использоваться маска, созданная обученной моделью." + +#~ msgid "" +#~ "R|The plugin to use to output the converted images. The writers are " +#~ "configurable in '/config/convert.ini' or 'Settings > Configure Convert " +#~ "Plugins:'\n" +#~ "L|ffmpeg: [video] Writes out the convert straight to video. When the " +#~ "input is a series of images then the '-ref' (--reference-video) parameter " +#~ "must be set.\n" +#~ "L|gif: [animated image] Create an animated gif.\n" +#~ "L|opencv: [images] The fastest image writer, but less options and formats " +#~ "than other plugins.\n" +#~ "L|patch: [images] Outputs the raw swapped face patch, along with the " +#~ "transformation matrix required to re-insert the face back into the " +#~ "original frame. Use this option if you wish to post-process and composite " +#~ "the final face within external tools.\n" +#~ "L|pillow: [images] Slower than opencv, but has more options and supports " +#~ "more formats." +#~ msgstr "" +#~ "R|Плагин, который нужно использовать для вывода преобразованных " +#~ "изображений. Записи настраиваются в '/config/convert.ini' или 'Настройки " +#~ "> Настроить плагины конвертации:'\n" +#~ "L|ffmpeg: [видео] Записывает конвертацию прямо в видео. Если на вход " +#~ "подается серия изображений, необходимо установить параметр '-ref' (--" +#~ "reference-video).\n" +#~ "L|gif: [анимированное изображение] Создает анимированный gif.\n" +#~ "L|opencv: [изображения] Самый быстрый редактор изображений, но имеет " +#~ "меньше опций и форматов, чем другие плагины.\n" +#~ "L|patch: [изображения] Выводит необработанный фрагмент измененного лица " +#~ "вместе с матрицей преобразования, необходимой для повторной вставки лица " +#~ "обратно в исходный кадр.\n" +#~ "L|pillow: [изображения] Медленнее, чем opencv, но имеет больше опций и " +#~ "поддерживает больше форматов." + +#~ msgid "Frame Processing" +#~ msgstr "Обработка лиц" + +#, python-format +#~ msgid "" +#~ "Scale the final output frames by this amount. 100%% will output the " +#~ "frames at source dimensions. 50%% at half size 200%% at double size" +#~ msgstr "" +#~ "Масштабирование конечных выходных кадров на эту величину. 100%% выводит " +#~ "кадры в исходном размере. 50%% при половинном размере 200%% при двойном " +#~ "размере" + +#~ msgid "" +#~ "Frame ranges to apply transfer to e.g. For frames 10 to 50 and 90 to 100 " +#~ "use --frame-ranges 10-50 90-100. Frames falling outside of the selected " +#~ "range will be discarded unless '-k' (--keep-unchanged) is selected. NB: " +#~ "If you are converting from images, then the filenames must end with the " +#~ "frame-number!" +#~ msgstr "" +#~ "Диапазоны кадров для применения переноса, например, для кадров с 10 по 50 " +#~ "и с 90 по 100 используйте --frame-ranges 10-50 90-100. Кадры, выходящие " +#~ "за пределы выбранного диапазона, будут отброшены, если не выбрана опция '-" +#~ "k' (--keep-unchanged). Примечание: Если вы конвертируете из изображений, " +#~ "то имена файлов должны заканчиваться номером кадра!" + +#~ msgid "" +#~ "Scale the swapped face by this percentage. Positive values will enlarge " +#~ "the face, Negative values will shrink the face." +#~ msgstr "" +#~ "Увеличить масштаб нового лица на этот процент. Положительные значения " +#~ "увеличат лицо, в то время как отрицательные значения уменьшат его." + +#~ msgid "" +#~ "If you have not cleansed your alignments file, then you can filter out " +#~ "faces by defining a folder here that contains the faces extracted from " +#~ "your input files/video. If this folder is defined, then only faces that " +#~ "exist within your alignments file and also exist within the specified " +#~ "folder will be converted. Leaving this blank will convert all faces that " +#~ "exist within the alignments file." +#~ msgstr "" +#~ "Если вы не очистили свой файл выравнивания, то вы можете отфильтровать " +#~ "лица, определив здесь папку, содержащую лица, извлеченные из ваших " +#~ "входных файлов/видео. Если эта папка определена, то будут преобразованы " +#~ "только те лица, которые существуют в вашем файле выравнивания, а также в " +#~ "указанной папке. Если оставить этот параметр пустым, будут преобразованы " +#~ "все лица, существующие в файле выравнивания." + +#~ msgid "" +#~ "Optionally filter out people who you do not wish to process by passing in " +#~ "an image of that person. Should be a front portrait with a single person " +#~ "in the image. Multiple images can be added space separated. NB: Using " +#~ "face filter will significantly decrease extraction speed and its accuracy " +#~ "cannot be guaranteed." +#~ msgstr "" +#~ "По желанию отфильтровать людей, которых вы не хотите обрабатывать, " +#~ "передав изображение этого человека. Это должен быть фронтальный портрет с " +#~ "изображением одного человека. Можно добавить несколько изображений, " +#~ "разделенных пробелами. Примечание: Использование фильтра лиц значительно " +#~ "снизит скорость извлечения, а его точность не гарантируется." + +#~ msgid "" +#~ "Optionally select people you wish to process by passing in an image of " +#~ "that person. Should be a front portrait with a single person in the " +#~ "image. Multiple images can be added space separated. NB: Using face " +#~ "filter will significantly decrease extraction speed and its accuracy " +#~ "cannot be guaranteed." +#~ msgstr "" +#~ "По желанию выберите людей, которых вы хотите обработать, передав " +#~ "изображение этого человека. Это должен быть фронтальный портрет с " +#~ "изображением одного человека. Можно добавить несколько изображений, " +#~ "разделенных пробелами. Примечание: Использование фильтра лиц значительно " +#~ "снизит скорость извлечения, а его точность не гарантируется." + +#~ msgid "" +#~ "For use with the optional nfilter/filter files. Threshold for positive " +#~ "face recognition. Lower values are stricter. NB: Using face filter will " +#~ "significantly decrease extraction speed and its accuracy cannot be " +#~ "guaranteed." +#~ msgstr "" +#~ "Для использования с дополнительными файлами nfilter/filter. Порог для " +#~ "положительного распознавания лиц. Более низкие значения являются более " +#~ "строгими. Примечание: Использование фильтра лиц значительно снизит " +#~ "скорость извлечения, а его точность не гарантируется." + +#~ msgid "" +#~ "The maximum number of parallel processes for performing conversion. " +#~ "Converting images is system RAM heavy so it is possible to run out of " +#~ "memory if you have a lot of processes and not enough RAM to accommodate " +#~ "them all. Setting this to 0 will use the maximum available. No matter " +#~ "what you set this to, it will never attempt to use more processes than " +#~ "are available on your system. If singleprocess is enabled this setting " +#~ "will be ignored." +#~ msgstr "" +#~ "Максимальное количество параллельных процессов для выполнения " +#~ "конвертации. Конвертирование изображений занимает много системной " +#~ "оперативной памяти, поэтому может закончиться память, если у вас много " +#~ "процессов и недостаточно оперативной памяти для их размещения. Если " +#~ "установить значение 0, будет использован максимум доступной памяти. " +#~ "Независимо от того, какое значение вы установите, программа никогда не " +#~ "будет пытаться использовать больше процессов, чем доступно в вашей " +#~ "системе. Если включена однопоточная обработка, этот параметр будет " +#~ "проигнорирован." + +#~ msgid "" +#~ "[LEGACY] This only needs to be selected if a legacy model is being loaded " +#~ "or if there are multiple models in the model folder" +#~ msgstr "" +#~ "[ОТБРОШЕН] Этот параметр необходимо выбрать только в том случае, если " +#~ "загружается устаревшая модель или если в папке моделей имеется несколько " +#~ "моделей" + +#~ msgid "" +#~ "Enable On-The-Fly Conversion. NOT recommended. You should generate a " +#~ "clean alignments file for your destination video. However, if you wish " +#~ "you can generate the alignments on-the-fly by enabling this option. This " +#~ "will use an inferior extraction pipeline and will lead to substandard " +#~ "results. If an alignments file is found, this option will be ignored." +#~ msgstr "" +#~ "Включить преобразование \"на лету\". НЕ рекомендуется. Вы должны " +#~ "сгенерировать чистый файл выравнивания для конечного видео. Однако при " +#~ "желании вы можете генерировать выравнивания \"на лету\", включив эту " +#~ "опцию. При этом будет использоваться некачественный конвейер извлечения, " +#~ "что приведет к некачественным результатам. Если файл выравнивания найден, " +#~ "этот параметр будет проигнорирован." + +#~ msgid "" +#~ "When used with --frame-ranges outputs the unchanged frames that are not " +#~ "processed instead of discarding them." +#~ msgstr "" +#~ "При использовании с --frame-ranges выводит неизмененные кадры, которые не " +#~ "были обработаны, вместо того, чтобы отбрасывать их." + +#~ msgid "Swap the model. Instead converting from of A -> B, converts B -> A" +#~ msgstr "" +#~ "Поменять модель местами. Вместо преобразования из A -> B, преобразуется B " +#~ "-> A" + +#~ msgid "Disable multiprocessing. Slower but less resource intensive." +#~ msgstr "" +#~ "Отключение многопоточной обработки. Медленнее, но менее ресурсоемко." + +#~ msgid "Output to Shell console instead of GUI console" +#~ msgstr "Вывод в консоль Shell вместо консоли GUI" + +#~ msgid "" +#~ "[Deprecated - Use '-D, --distribution-strategy' instead] Use the " +#~ "Tensorflow Mirrored Distrubution Strategy to train on multiple GPUs." +#~ msgstr "" +#~ "[Устарело - Используйте '-D, --distribution-strategy' вместо этого] " +#~ "Используйте стратегию Tensorflow Mirrored Distrubution Strategy(Стратегия " +#~ "Зеркального Распределения Tensorflow) для обучения на нескольких GPU." diff --git a/locales/ru/LC_MESSAGES/lib.config.mo b/locales/ru/LC_MESSAGES/lib.config.mo new file mode 100644 index 0000000000..49787224ff Binary files /dev/null and b/locales/ru/LC_MESSAGES/lib.config.mo differ diff --git a/locales/ru/LC_MESSAGES/lib.config.po b/locales/ru/LC_MESSAGES/lib.config.po new file mode 100644 index 0000000000..f9c3d0d252 --- /dev/null +++ b/locales/ru/LC_MESSAGES/lib.config.po @@ -0,0 +1,76 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-11 23:28+0100\n" +"PO-Revision-Date: 2023-06-12 21:25+0700\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru_RU\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"X-Generator: Poedit 3.3.1\n" + +#: lib/config.py:393 +msgid "" +"\n" +"This option can be updated for existing models.\n" +msgstr "" +"\n" +"Эта настройка будет обновлена для существующих моделей.\n" + +#: lib/config.py:395 +msgid "" +"\n" +"If selecting multiple options then each option should be separated by a " +"space or a comma (e.g. item1, item2, item3)\n" +msgstr "" +"\n" +"Если выбираете несколько опций, тогда каждая опция должна быть " +"разделена пробелом или запятой (например: опция1, опция2, опция3)\n" + +#: lib/config.py:398 +msgid "" +"\n" +"Choose from: {}" +msgstr "" +"\n" +"Выберите из: {}" + +#: lib/config.py:400 +msgid "" +"\n" +"Choose from: True, False" +msgstr "" +"\n" +"Выберите из: True, False" + +#: lib/config.py:404 +msgid "" +"\n" +"Select an integer between {} and {}" +msgstr "" +"\n" +"Выберите число между {} и {}" + +#: lib/config.py:408 +msgid "" +"\n" +"Select a decimal number between {} and {}" +msgstr "" +"\n" +"Выберите десятичное число между {} и {}" + +#: lib/config.py:409 +msgid "" +"\n" +"[Default: {}]" +msgstr "" +"\n" +"[По умолчанию: {}]" diff --git a/locales/ru/LC_MESSAGES/plugins.extract._config.mo b/locales/ru/LC_MESSAGES/plugins.extract._config.mo new file mode 100644 index 0000000000..4ec5cb72f3 Binary files /dev/null and b/locales/ru/LC_MESSAGES/plugins.extract._config.mo differ diff --git a/locales/ru/LC_MESSAGES/plugins.extract._config.po b/locales/ru/LC_MESSAGES/plugins.extract._config.po new file mode 100644 index 0000000000..0500d72cda --- /dev/null +++ b/locales/ru/LC_MESSAGES/plugins.extract._config.po @@ -0,0 +1,166 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-06-08 16:43+0100\n" +"PO-Revision-Date: 2023-06-12 19:42+0700\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru_RU\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"X-Generator: Poedit 3.3.1\n" + +#: plugins/extract/_config.py:32 +msgid "Options that apply to all extraction plugins" +msgstr "Параметры, применимые ко всем плагинам извлечения" + +#: plugins/extract/_config.py:38 +msgid "settings" +msgstr "настройки" + +#: plugins/extract/_config.py:39 +msgid "" +"Enable the Tensorflow GPU `allow_growth` configuration " +"option. This option prevents Tensorflow from allocating all of the GPU VRAM " +"at launch but can lead to higher VRAM fragmentation and slower performance. " +"Should only be enabled if you are having problems running extraction." +msgstr "" +"Включите опцию конфигурации Tensorflow GPU " +"`allow_growth`. Эта опция не позволяет Tensorflow выделять всю видеопамять " +"видеокарты при запуске, но может привести к повышенной фрагментации " +"видеопамяти и снижению производительности. Следует включать только в том " +"случае, если у вас есть проблемы с запуском извлечения." + +#: plugins/extract/_config.py:50 plugins/extract/_config.py:64 +#: plugins/extract/_config.py:78 plugins/extract/_config.py:89 +#: plugins/extract/_config.py:99 plugins/extract/_config.py:108 +#: plugins/extract/_config.py:119 +msgid "filters" +msgstr "фильтры" + +#: plugins/extract/_config.py:51 +msgid "" +"Filters out faces below this size. This is a multiplier of the minimum " +"dimension of the frame (i.e. 1280x720 = 720). If the original face extract " +"box is smaller than the minimum dimension times this multiplier, it is " +"considered a false positive and discarded. Faces which are found to be " +"unusually smaller than the frame tend to be misaligned images, except in " +"extreme long-shots. These can be usually be safely discarded." +msgstr "" +"Отфильтровывает лица меньше этого размера. Это множитель минимального " +"размера кадра (т.е. 1280x720 = 720). Если исходное поле извлечения лица " +"меньше минимального размера, умноженного на этот множитель, оно считается " +"ложным срабатыванием и отбрасывается. Лица, которые оказываются необычно " +"меньшего размера, чем кадр, как правило, являются неправильно выровненными " +"изображениями, за исключением экстремально длинных снимков. Обычно их можно " +"смело отбрасывать." + +#: plugins/extract/_config.py:65 +msgid "" +"Filters out faces above this size. This is a multiplier of the minimum " +"dimension of the frame (i.e. 1280x720 = 720). If the original face extract " +"box is larger than the minimum dimension times this multiplier, it is " +"considered a false positive and discarded. Faces which are found to be " +"unusually larger than the frame tend to be misaligned images except in " +"extreme close-ups. These can be usually be safely discarded." +msgstr "" +"Отфильтровывает лица, превышающие этот размер. Это множитель минимального " +"размера кадра (т.е. 1280x720 = 720). Если исходный блок извлечения лица " +"больше, чем минимальный размер кадра, умноженный на этот множитель, он " +"считается ложным срабатыванием и отбрасывается. Лица, размер которых " +"необычно превышает размер кадра, как правило, являются несогласованными " +"изображениями, за исключением экстремальных крупных планов. Обычно их можно " +"смело отбрасывать." + +#: plugins/extract/_config.py:79 +msgid "" +"Filters out faces who's landmarks are above this distance from an 'average' " +"face. Values above 15 tend to be fairly safe. Values above 10 will remove " +"more false positives, but may also filter out some faces at extreme angles." +msgstr "" +"Отфильтровывает лица, ориентиры которых находятся на расстоянии, превышающем " +"это расстояние от 'среднего' лица. Значения выше 15, как правило, достаточно " +"безопасны. Значения выше 10 устраняют больше ложных срабатываний, но также " +"могут отфильтровать некоторые лица под экстремальными углами." + +#: plugins/extract/_config.py:90 +msgid "" +"Filters out faces who's calculated roll is greater than zero +/- this value " +"in degrees. Aligned faces should have a roll value close to zero. Values " +"that are a significant distance from 0 degrees tend to be misaligned images. " +"These can usually be safely disgarded." +msgstr "" +"Отфильтровывает лица, у которых расчетный угол наклона больше нуля +/- это " +"значение в градусах. Выровненные лица должны иметь значение угла наклона, " +"близкое к нулю. Значения, которые значительно удалены от 0 градусов, как " +"правило, представляют собой неправильно выровненные изображения. Обычно их " +"можно смело отбрасывать." + +#: plugins/extract/_config.py:100 +msgid "" +"Filters out faces where the lowest point of the aligned face's eye or " +"eyebrow is lower than the highest point of the aligned face's mouth. Any " +"faces where this occurs are misaligned and can be safely disgarded." +msgstr "" +"Отфильтровывает лица, у которых нижняя точка глаза или брови выровненного " +"лица находится ниже, чем верхняя точка рта выровненного лица. Все лица, на " +"которых это происходит, являются неправильно выровненными и могут быть смело " +"отброшены." + +#: plugins/extract/_config.py:109 +msgid "" +"If enabled, and 're-feed' has been selected for extraction, then interim " +"alignments will be filtered prior to averaging the final landmarks. This can " +"help improve the final alignments by removing any obvious misaligns from the " +"interim results, and may also help pick up difficult alignments. If " +"disabled, then all re-feed results will be averaged." +msgstr "" +"Если эта функция включена, и для извлечения выбрана 'повторная подача'('re-" +"feed'), то промежуточные выравнивания будут отфильтрованы перед усреднением " +"окончательных ориентиров. Это может помочь улучшить окончательное " +"выравнивание, удалив любые очевидные несоответствия из промежуточных " +"результатов, а также может помочь выявить сложные выравнивания. Если эта " +"функция отключена, то все результаты повторной подачи будут усреднены." + +#: plugins/extract/_config.py:120 +msgid "" +"If enabled, saves any filtered out images into a sub-folder during the " +"extraction process. If disabled, filtered faces are deleted. Note: The faces " +"will always be filtered out of the alignments file, regardless of whether " +"you keep the faces or not." +msgstr "" +"Если включена, то в процессе извлечения отфильтрованные изображения " +"сохраняются в подпапке. Если отключено, отфильтрованные лица удаляются. " +"Примечание: Лица всегда будут отфильтрованы из файла выравнивания, " +"независимо от того, сохраняете вы эти лица или нет." + +#: plugins/extract/_config.py:129 plugins/extract/_config.py:138 +msgid "re-align" +msgstr "повторное выравнивание" + +#: plugins/extract/_config.py:130 +msgid "" +"If enabled, and 're-align' has been selected for extraction, then all re-" +"feed iterations are re-aligned. If disabled, then only the final averaged " +"output from re-feed will be re-aligned." +msgstr "" +"Если включено, и для извлечения выбрано 'повторное выравнивание'('re-" +"align'), то все итерации повторной подачи выравниваются повторно. Если " +"отключено, то выравнивается только конечный усредненный результат повторной " +"подачи." + +#: plugins/extract/_config.py:139 +msgid "" +"If enabled, and 're-align' has been selected for extraction, then any " +"alignments which would be filtered out will not be re-aligned." +msgstr "" +"Если эта функция включена, и для извлечения выбрано 'повторное " +"выравнивание'('re-align'), то все выравнивания, которые будут отфильтрованы, " +"не будут повторно выравниваться." diff --git a/locales/ru/LC_MESSAGES/plugins.train._config.mo b/locales/ru/LC_MESSAGES/plugins.train._config.mo new file mode 100644 index 0000000000..a1032a91d1 Binary files /dev/null and b/locales/ru/LC_MESSAGES/plugins.train._config.mo differ diff --git a/locales/ru/LC_MESSAGES/plugins.train._config.po b/locales/ru/LC_MESSAGES/plugins.train._config.po new file mode 100644 index 0000000000..04191e5e70 --- /dev/null +++ b/locales/ru/LC_MESSAGES/plugins.train._config.po @@ -0,0 +1,973 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-26 17:37+0000\n" +"PO-Revision-Date: 2024-03-26 17:40+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru_RU\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"X-Generator: Poedit 3.4.2\n" + +#: plugins/train/_config.py:17 +msgid "" +"\n" +"NB: Unless specifically stated, values changed here will only take effect " +"when creating a new model." +msgstr "" +"\n" +"Примечание: До тех пор, пока об этом не сказано, значения, измененные здесь, " +"будут применены при создании новой модели." + +#: plugins/train/_config.py:22 +msgid "" +"Focal Frequency Loss. Analyzes the frequency spectrum of the images rather " +"than the images themselves. This loss function can be used on its own, but " +"the original paper found increased benefits when using it as a complementary " +"loss to another spacial loss function (e.g. MSE). Ref: Focal Frequency Loss " +"for Image Reconstruction and Synthesis https://arxiv.org/pdf/2012.12821.pdf " +"NB: This loss does not currently work on AMD cards." +msgstr "" +"Потеря фокальной частоты. Анализирует частотный спектр изображений, а не " +"сами изображения. Эта функция потерь может использоваться сама по себе, но в " +"оригинальной статье было обнаружено, что она дает больше преимуществ при " +"использовании в качестве дополнительной потери к другой пространственной " +"функции потерь (например, MSE). Ссылка: Focal Frequency Loss for Image " +"Reconstruction and Synthesis [ТОЛЬКО на английском] https://arxiv.org/" +"pdf/2012.12821.pdf NB: Эта потеря в настоящее время не работает на картах " +"AMD." + +#: plugins/train/_config.py:29 +msgid "" +"Nvidia FLIP. A perceptual loss measure that approximates the difference " +"perceived by humans as they alternate quickly (or flip) between two images. " +"Used on its own and this loss function creates a distinct grid on the " +"output. However it can be helpful when used as a complimentary loss " +"function. Ref: FLIP: A Difference Evaluator for Alternating Images: https://" +"research.nvidia.com/sites/default/files/node/3260/FLIP_Paper.pdf" +msgstr "" +"Nvidia FLIP. Мера потерь восприятия, которая приближает разницу, " +"воспринимаемую человеком при быстром чередовании (или перелистывании) двух " +"изображений. Используемая сама по себе, эта функция потерь создает на выходе " +"отчетливую сетку. Однако она может быть полезна при использовании в качестве " +"дополнительной функции потерь. Ссылка: FLIP: A Difference Evaluator for " +"Alternating Images [ТОЛЬКО на английском]: https://research.nvidia.com/sites/" +"default/files/node/3260/FLIP_Paper.pdf" + +#: plugins/train/_config.py:36 +msgid "" +"Gradient Magnitude Similarity Deviation seeks to match the global standard " +"deviation of the pixel to pixel differences between two images. Similar in " +"approach to SSIM. Ref: Gradient Magnitude Similarity Deviation: An Highly " +"Efficient Perceptual Image Quality Index https://arxiv.org/ftp/arxiv/" +"papers/1308/1308.3052.pdf" +msgstr "" +"Отклонение Схожести Магнитуды Градиентов(Gradient Magnitude Similarity " +"Deviation) пытается совместить глобальную стандартную девиацию различий " +"пикселя к пикселю между двумя изображениями. Подход похож на SSIM. Ссылка: " +"Gradient Magnitude Similarity Deviation: An Highly Efficient Perceptual " +"Image Quality Index [ТОЛЬКО на английском] https://arxiv.org/ftp/arxiv/" +"papers/1308/1308.3052.pdf" + +#: plugins/train/_config.py:41 +msgid "" +"The L_inf norm will reduce the largest individual pixel error in an image. " +"As each largest error is minimized sequentially, the overall error is " +"improved. This loss will be extremely focused on outliers." +msgstr "" +"Норма L_inf уменьшает наибольшую ошибку отдельного пикселя в изображении. По " +"мере последовательной минимизации каждой наибольшей ошибки улучшается общая " +"ошибка. Эта потеря будет чрезвычайно сосредоточена на выбросах." + +#: plugins/train/_config.py:45 +msgid "" +"Laplacian Pyramid Loss. Attempts to improve results by focussing on edges " +"using Laplacian Pyramids. As this loss function gives priority to edges over " +"other low-frequency information, like color, it should not be used on its " +"own. The original implementation uses this loss as a complimentary function " +"to MSE. Ref: Optimizing the Latent Space of Generative Networks https://" +"arxiv.org/abs/1707.05776" +msgstr "" +"Потеря пирамиды Лапласиана. Пытается улучшить результаты, концентрируясь на " +"краях с помощью пирамид Лапласиана. Поскольку эта функция потерь отдает " +"приоритет краям, а не другой низкочастотной информации, например, цвету, ее " +"не следует использовать самостоятельно. В оригинальной реализации эта потеря " +"используется как дополнительная функция к MSE. Ссылка: Optimizing the Latent " +"Space of Generative Networks [ТОЛЬКО на английском] https://arxiv.org/" +"abs/1707.05776" + +#: plugins/train/_config.py:52 +msgid "" +"LPIPS is a perceptual loss that uses the feature outputs of other pretrained " +"models as a loss metric. Be aware that this loss function will use more " +"VRAM. Used on its own and this loss will create a distinct moire pattern on " +"the output, however it can be helpful as a complimentary loss function. The " +"output of this function is strong, so depending on your chosen primary loss " +"function, you are unlikely going to want to set the weight above about 25%. " +"Ref: The Unreasonable Effectiveness of Deep Features as a Perceptual Metric " +"http://arxiv.org/abs/1801.03924\n" +"This variant uses the AlexNet backbone. A fairly light and old model which " +"performed best in the paper's original implementation.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" +"LPIPS - это перцептивная потеря, которая использует в качестве метрики " +"потерь выходные характеристики других предварительно обученных моделей. " +"Имейте в виду, что эта функция потерь использует больше VRAM. При " +"самостоятельном использовании эта потеря создает на выходе отчетливый " +"муаровый рисунок, однако она может быть полезна как дополнительная функция " +"потерь. Вывод этой функции является сильным, поэтому, в зависимости от " +"выбранной вами основной функции потерь, вы вряд ли захотите устанавливать " +"вес выше 25%. Ссылка: The Unreasonable Effectiveness of Deep Features as a " +"Perceptual Metric [ТОЛЬКО на английском] http://arxiv.org/abs/1801.03924.\n" +"Этот вариант использует основу AlexNet. Это довольно легкая и старая модель, " +"которая лучше всего показала себя в оригинальной реализации.\n" +"NB: Для пользователей AMD последний линейный слой не реализован." + +#: plugins/train/_config.py:62 +msgid "" +"Same as lpips_alex, but using the SqueezeNet backbone. A more lightweight " +"version of AlexNet.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" +"То же, что и lpips_alex, но использует основу SqueezeNet. Более облегченная " +"версия AlexNet.\n" +"NB: Для пользователей AMD последний линейный слой не реализован." + +#: plugins/train/_config.py:65 +msgid "" +"Same as lpips_alex, but using the VGG16 backbone. A more heavyweight model.\n" +"NB: For AMD Users the final linear layer is not implemented." +msgstr "" +"То же, что и lpips_alex, но использует основу VGG16. Более тяжелая модель.\n" +"NB: Для пользователей AMD последний линейный слой не реализован." + +#: plugins/train/_config.py:68 +msgid "" +"log(cosh(x)) acts similar to MSE for small errors and to MAE for large " +"errors. Like MSE, it is very stable and prevents overshoots when errors are " +"near zero. Like MAE, it is robust to outliers." +msgstr "" +"log(cosh(x)) действует аналогично MSE для малых ошибок и MAE для больших " +"ошибок. Как и MSE, он очень стабилен и предотвращает переборы, когда ошибки " +"близки к нулю. Как и MAE, он устойчив к выбросам." + +#: plugins/train/_config.py:72 +msgid "" +"Mean absolute error will guide reconstructions of each pixel towards its " +"median value in the training dataset. Robust to outliers but as a median, it " +"can potentially ignore some infrequent image types in the dataset." +msgstr "" +"Средняя абсолютная погрешность направляет реконструкцию каждого пикселя к " +"его медианному значению в обучающем наборе данных. Устойчив к выбросам, но в " +"качестве медианы может игнорировать некоторые редкие типы изображений в " +"наборе данных." + +#: plugins/train/_config.py:76 +msgid "" +"Mean squared error will guide reconstructions of each pixel towards its " +"average value in the training dataset. As an avg, it will be susceptible to " +"outliers and typically produces slightly blurrier results. Ref: Multi-Scale " +"Structural Similarity for Image Quality Assessment https://www.cns.nyu.edu/" +"pub/eero/wang03b.pdf" +msgstr "" +"Средняя квадратичная погрешность направляет реконструкцию каждого пикселя к " +"его среднему значению в наборе данных для обучения. Как среднее значение, " +"оно будет чувствительно к выбросам и обычно дает немного более размытые " +"результаты. Ссылка: Multi-Scale Structural Similarity for Image Quality " +"Assessment [ТОЛЬКО на английском]https://www.cns.nyu.edu/pub/eero/wang03b.pdf" + +#: plugins/train/_config.py:81 +msgid "" +"Multiscale Structural Similarity Index Metric is similar to SSIM except that " +"it performs the calculations along multiple scales of the input image." +msgstr "" +"Метрика Индекса Многомасштабного Структурного Сходства (Multiscale " +"Structural Similarity Index Metric) похожа на SSIM, за исключением того, что " +"она выполняет вычисления по нескольким масштабам входного изображения." + +#: plugins/train/_config.py:84 +msgid "" +"Smooth_L1 is a modification of the MAE loss to correct two of its " +"disadvantages. This loss has improved stability and guidance for small " +"errors. Ref: A General and Adaptive Robust Loss Function https://arxiv.org/" +"pdf/1701.03077.pdf" +msgstr "" +"Smooth_L1 - это модификация потери MAE для исправления двух ее недостатков. " +"Эта потеря улучшает стабильность и ориентирование при небольших " +"погрешностях. Ссылка: A General and Adaptive Robust Loss Function [ТОЛЬКО на " +"английском] https://arxiv.org/pdf/1701.03077.pdf" + +#: plugins/train/_config.py:88 +msgid "" +"Structural Similarity Index Metric is a perception-based loss that considers " +"changes in texture, luminance, contrast, and local spatial statistics of an " +"image. Potentially delivers more realistic looking images. Ref: Image " +"Quality Assessment: From Error Visibility to Structural Similarity http://" +"www.cns.nyu.edu/pub/eero/wang03-reprint.pdf" +msgstr "" +"Метрика индекса структурного сходства ('Structural Similarity Index Metric') " +"- это основанная на восприятии потеря, которая учитывает изменения в " +"текстуре, яркости, контрасте и локальной пространственной статистике " +"изображения. Потенциально обеспечивает более реалистичный вид изображений. " +"Ссылка: Image Quality Assessment: From Error Visibility to Structural " +"Similarity [ТОЛЬКО на английском] http://www.cns.nyu.edu/pub/eero/wang03-" +"reprint.pdf" + +#: plugins/train/_config.py:93 +msgid "" +"Instead of minimizing the difference between the absolute value of each " +"pixel in two reference images, compute the pixel to pixel spatial difference " +"in each image and then minimize that difference between two images. Allows " +"for large color shifts, but maintains the structure of the image." +msgstr "" +"Вместо того чтобы минимизировать разницу между абсолютным значением каждого " +"пикселя в двух образцовых изображениях, вычислить пространственную разницу " +"между пикселями в каждом изображении и затем минимизировать эту разницу " +"между двумя изображениями. Это позволяет получить большие цветовые сдвиги, " +"но сохраняет структуру изображения." + +#: plugins/train/_config.py:97 +msgid "Do not use an additional loss function." +msgstr "Не использовать функцию дополнительных потерь." + +#: plugins/train/_config.py:117 +msgid "Options that apply to all models" +msgstr "Настройки, применимые ко всем моделям" + +#: plugins/train/_config.py:126 plugins/train/_config.py:150 +msgid "face" +msgstr "лицо" + +#: plugins/train/_config.py:128 +msgid "" +"How to center the training image. The extracted images are centered on the " +"middle of the skull based on the face's estimated pose. A subsection of " +"these images are used for training. The centering used dictates how this " +"subsection will be cropped from the aligned images.\n" +"\tface: Centers the training image on the center of the face, adjusting for " +"pitch and yaw.\n" +"\thead: Centers the training image on the center of the head, adjusting for " +"pitch and yaw. NB: You should only select head centering if you intend to " +"include the full head (including hair) in the final swap. This may give " +"mixed results. Additionally, it is only worth choosing head centering if you " +"are training with a mask that includes the hair (e.g. BiSeNet-FP-Head).\n" +"\tlegacy: The 'original' extraction technique. Centers the training image " +"near the tip of the nose with no adjustment. Can result in the edges of the " +"face appearing outside of the training area." +msgstr "" +"Как централизовывать тренировочное изображение. Центр в извлеченных " +"изображениях находится в середине черепа, основанный на примерной позе лица. " +"Подсекция этих изображений используется для тренировки. Используемый центр " +"диктует то, как эта подсекция будет обрезана из выравненных изображений.\n" +"\tface: Центрирует учебное изображение по центру лица, регулируя угол " +"наклона и поворота.\n" +"\thead: Централизует тренировочное изображение в центре головы, регулируя " +"угол наклона и поворота. Примечание: Следует выбирать централизацию головы, " +"если вы планируете включать голову полностью (включая волосы) в финальную " +"замену. Может дать смешанные результаты. В дополнении, оно стоит того только " +"если вы тренируете с маской, что включает в себя волосы (к примеру: BiSeNet-" +"FP-Head).\n" +"\tlegacy: 'оригинальная' техника извлечения. Централизует тренировочное " +"изображение ближе к кончику носа без правок. Может привести к тому, что края " +"лица будут вне тренировочной зоны." + +#: plugins/train/_config.py:152 +msgid "" +"How much of the extracted image to train on. A lower coverage will limit the " +"model's scope to a zoomed-in central area while higher amounts can include " +"the entire face. A trade-off exists between lower amounts given more detail " +"versus higher amounts avoiding noticeable swap transitions. For 'Face' " +"centering you will want to leave this above 75%. For Head centering you will " +"most likely want to set this to 100%. Sensible values for 'Legacy' centering " +"are:\n" +"\t62.5% spans from eyebrow to eyebrow.\n" +"\t75.0% spans from temple to temple.\n" +"\t87.5% spans from ear to ear.\n" +"\t100.0% is a mugshot." +msgstr "" +"Сколько извлеченного изображения тренировать. Низкая покрытость ограничит " +"прицел модели к приближенной центральной зоне, в то время как большие " +"значения могут включать в себя целое лицо. Существует компромисс между " +"меньшими объемами, дающими больше деталей, и большими объемами, позволяющими " +"избежать заметных переходов замены. Для централизации 'Face', вам нужно " +"будет оставить значение выше 75%. Для централизации 'Head', вам скорее всего " +"нужно будет поставить значение 100%. Адекватные значения для 'Legacy':\n" +"\t62.5% охватывает от бровей до бровей.\n" +"\t75% охватывает от виска до виска.\n" +"\t87.5% охватывает от уха до уха.\n" +"\t100% - полный снимок." + +#: plugins/train/_config.py:168 plugins/train/_config.py:179 +msgid "initialization" +msgstr "инициализация" + +#: plugins/train/_config.py:170 +msgid "" +"Use ICNR to tile the default initializer in a repeating pattern. This " +"strategy is designed for pairing with sub-pixel / pixel shuffler to reduce " +"the 'checkerboard effect' in image reconstruction. \n" +"\t https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf" +msgstr "" +"Использовать ICNR для чередования инициализатора по умолчанию в " +"повторяющемся шаблоне. Эта стратегия предназначена для использования в паре " +"с субпиксельным/пиксельным перетасовщиком для уменьшения \"эффекта шахматной " +"доски\" при реконструкции изображения. \n" +"\t [ТОЛЬКО на английском] https://arxiv.org/ftp/arxiv/papers/1707/1707.02937." +"pdf" + +#: plugins/train/_config.py:181 +msgid "" +"Use Convolution Aware Initialization for convolutional layers. This can help " +"eradicate the vanishing and exploding gradient problem as well as lead to " +"higher accuracy, lower loss and faster convergence.\n" +"NB:\n" +"\t This can use more VRAM when creating a new model so you may want to lower " +"the batch size for the first run. The batch size can be raised again when " +"reloading the model. \n" +"\t Multi-GPU is not supported for this option, so you should start the model " +"on a single GPU. Once training has started, you can stop training, enable " +"multi-GPU and resume.\n" +"\t Building the model will likely take several minutes as the calculations " +"for this initialization technique are expensive. This will only impact " +"starting a new model." +msgstr "" +"Использовать инициализацию с учетом свертки для сверточных слоев. Это " +"поможет устранить проблему исчезающего и взрывающегося градиента, а также " +"повысить точность, снизить потери и ускорить сходимость.\n" +"Примечание:\n" +"\tПри создании новой модели может потребоваться больше видеопамяти, поэтому " +"для первого запуска лучше уменьшить размер пачки. Размер пачки может быть " +"увеличен при перезагрузке модели. \n" +"\tИспользование нескольких видеокарт не поддерживается, поэтому модель " +"следует запускать на одной видеокарте. После начала обучения вы можете " +"остановить обучение, включить несколько видеокарт и возобновить его.\n" +"\t Построение модели, скорее всего, займет несколько минут, поскольку " +"вычисления для этой техники инициализации являются дорогостоящими. Это " +"повлияет только на запуск новой модели." + +#: plugins/train/_config.py:198 plugins/train/_config.py:223 +#: plugins/train/_config.py:238 plugins/train/_config.py:256 +#: plugins/train/_config.py:337 +msgid "optimizer" +msgstr "оптимизатор" + +#: plugins/train/_config.py:202 +msgid "" +"The optimizer to use.\n" +"\t adabelief - Adapting Stepsizes by the Belief in Observed Gradients. An " +"optimizer with the aim to converge faster, generalize better and remain more " +"stable. (https://arxiv.org/abs/2010.07468). NB: Epsilon for AdaBelief needs " +"to be set to a smaller value than other Optimizers. Generally setting the " +"'Epsilon Exponent' to around '-16' should work.\n" +"\t adam - Adaptive Moment Optimization. A stochastic gradient descent method " +"that is based on adaptive estimation of first-order and second-order " +"moments.\n" +"\t nadam - Adaptive Moment Optimization with Nesterov Momentum. Much like " +"Adam but uses a different formula for calculating momentum.\n" +"\t rms-prop - Root Mean Square Propagation. Maintains a moving (discounted) " +"average of the square of the gradients. Divides the gradient by the root of " +"this average." +msgstr "" +"Используемый оптимизатор.\n" +"\t adabelief - Адаптация размеров шагов по убеждению в наблюдаемых " +"градиентах('Adapting Stepsizes by the Belief in Observed Gradients'). " +"Оптимизатор, цель которого - быстрее сходиться, лучше обобщаться и " +"оставаться более стабильным. ([ТОЛЬКО на английском] https://arxiv.org/" +"abs/2010.07468). Примечание: значение Epsilon для AdaBelief должно быть " +"меньше, чем для других оптимизаторов. Как правило, значение 'Epsilon " +"Exponent' должно быть около '-16'.\n" +"\t adam - Адаптивная оптимизация моментов('Adaptive Moment Optimization'). " +"Стохастический метод градиентного спуска, основанный на адаптивной оценке " +"моментов первого и второго порядка.\n" +"\t nadam - Адаптивная оптимизация моментов с моментумом Нестерова ('Adaptive " +"Moment Optimization with Nesterov Momentum'). Похож на Adam, но использует " +"другую формулу для вычисления момента.\n" +"rms-prop - Распространение корневого среднего квадрата ('Root Mean Square " +"Propagation'). Поддерживает скользящее (дисконтированное) среднее квадрата " +"градиентов. Делит градиент на корень из этого среднего." + +#: plugins/train/_config.py:225 +msgid "" +"Learning rate - how fast your network will learn (how large are the " +"modifications to the model weights after one batch of training). Values that " +"are too large might result in model crashes and the inability of the model " +"to find the best solution. Values that are too small might be unable to " +"escape from dead-ends and find the best global minimum." +msgstr "" +"Скорость обучения - насколько быстро ваша модель будет обучаться (насколько " +"огромны изменения весов модели после одной пачки тренировки). Слишком " +"большие значения могут привести к крахам модели и невозможности модели найти " +"лучшее решение. Слишком маленькие значения могут привести к невозможности " +"выбраться из тупиков и найти лучший глобальный минимум." + +#: plugins/train/_config.py:240 +msgid "" +"The epsilon adds a small constant to weight updates to attempt to avoid " +"'divide by zero' errors. Unless you are using the AdaBelief Optimizer, then " +"Generally this option should be left at default value, For AdaBelief, " +"setting this to around '-16' should work.\n" +"In all instances if you are getting 'NaN' loss values, and have been unable " +"to resolve the issue any other way (for example, increasing batch size, or " +"lowering learning rate), then raising the epsilon can lead to a more stable " +"model. It may, however, come at the cost of slower training and a less " +"accurate final result.\n" +"NB: The value given here is the 'exponent' to the epsilon. For example, " +"choosing '-7' will set the epsilon to 1e-7. Choosing '-3' will set the " +"epsilon to 0.001 (1e-3)." +msgstr "" +"Эпсилон добавляет небольшую константу к обновлениям веса, чтобы попытаться " +"избежать ошибок \"деления на ноль\". Если вы не используете оптимизатор " +"AdaBelief, то, как правило, этот параметр следует оставить по умолчанию. Для " +"AdaBelief подойдет значение около '-16'.\n" +"Во всех случаях, если вы получаете значения потерь 'NaN' и не смогли решить " +"проблему другим способом (например, увеличив размер пачки или уменьшив " +"скорость обучения), то увеличение эпсилона может привести к более стабильной " +"модели. Однако это может стоить более медленного обучения и менее точного " +"конечного результата.\n" +"Примечание: Значение, указанное здесь, является \"экспонентой\" к эпсилону. " +"Например, при выборе значения '-7' эпсилон будет равен 1e-7. При выборе " +"значения \"-3\" эпсилон будет равен 0,001 (1e-3)." + +#: plugins/train/_config.py:262 +msgid "" +"When to save the Optimizer Weights. Saving the optimizer weights is not " +"necessary and will increase the model file size 3x (and by extension the " +"amount of time it takes to save the model). However, it can be useful to " +"save these weights if you want to guarantee that a resumed model carries off " +"exactly from where it left off, rather than spending a few hundred " +"iterations catching up.\n" +"\t never - Don't save optimizer weights.\n" +"\t always - Save the optimizer weights at every save iteration. Model saving " +"will take longer, due to the increased file size, but you will always have " +"the last saved optimizer state in your model file.\n" +"\t exit - Only save the optimizer weights when explicitly terminating a " +"model. This can be when the model is actively stopped or when the target " +"iterations are met. Note: If the training session ends because of another " +"reason (e.g. power outage, Out of Memory Error, NaN detected) then the " +"optimizer weights will NOT be saved." +msgstr "" +"Когда сохранять веса оптимизатора. Сохранение весов оптимизатора не является " +"необходимым и увеличит размер файла модели в 3 раза (и соответственно время, " +"необходимое для сохранения модели). Однако может быть полезно сохранить эти " +"веса, если вы хотите гарантировать, что возобновленная модель продолжит " +"работу именно с того места, где она остановилась, а не тратит несколько " +"сотен итераций на догонялки.\n" +"\t never - не сохранять веса оптимизатора.\n" +"\t always - сохранять веса оптимизатора при каждой итерации сохранения. " +"Сохранение модели займет больше времени из-за увеличенного размера файла, но " +"в файле модели всегда будет последнее сохраненное состояние оптимизатора.\n" +"\t exit - сохранять веса оптимизатора только при явном завершении модели. " +"Это может быть, когда модель активно останавливается или когда выполняются " +"целевые итерации. Примечание. Если сеанс обучения завершается по другой " +"причине (например, отключение питания, ошибка нехватки памяти, обнаружение " +"NaN), веса оптимизатора НЕ будут сохранены." + +#: plugins/train/_config.py:285 plugins/train/_config.py:297 +#: plugins/train/_config.py:314 +msgid "Learning Rate Finder" +msgstr "Инструмент поиска оптимального коэффициента обучения" + +#: plugins/train/_config.py:287 +msgid "" +"The number of iterations to process to find the optimal learning rate. " +"Higher values will take longer, but will be more accurate." +msgstr "" +"Количество итераций для поиска оптимального коэффициента обучения. Большие " +"значения займут больше времени, но будут более точными." + +#: plugins/train/_config.py:299 +msgid "" +"The operation mode for the learning rate finder. Only applicable to new " +"models. For existing models this will always default to 'set'.\n" +"\tset - Train with the discovered optimal learning rate.\n" +"\tgraph_and_set - Output a graph in the training folder showing the " +"discovered learning rates and train with the optimal learning rate.\n" +"\tgraph_and_exit - Output a graph in the training folder with the discovered " +"learning rates and exit." +msgstr "" +"Режим работы для поиска коэффициента обучения. Применимо только для новых " +"моделей. Для уже существующих моделей режим будет автоматически выставлен в " +"'set'.\n" +"\tset - Обучение с найденным оптимальным коэффициентом обучения.\n" +"\tgraph_and_set - Вывод графика в папку обучения, показывающего найденные " +"коэффициенты обучения, и обучение с оптимальным коэффициентом.\n" +"\tgraph_and_exit - Вывод графика в папку обучения с найденными " +"коэффициентами обучения с последующим выходом из программы." + +#: plugins/train/_config.py:316 +msgid "" +"How aggressively to set the Learning Rate. More aggressive can learn faster, " +"but is more likely to lead to exploding gradients.\n" +"\tdefault - The default optimal learning rate. A safe choice for nearly all " +"use cases.\n" +"\taggressive - Set's a higher learning rate than the default. May learn " +"faster but with a higher chance of exploding gradients.\n" +"\textreme - The highest optimal learning rate. A much higher risk of " +"exploding gradients." +msgstr "" +"Насколько агрессивно устанавливать коэффициент обучения. Более агрессивный " +"подход может обучать быстрее, но с большей вероятностью может привести к " +"взрыву градиентов.\n" +"\tdefault - Оптимальный коэффициент обучения по умолчанию. Безопасный выбор " +"для почти всех случаев использования.\n" +"\taggressive - Устанавливает коэффициент обучения выше, чем по умолчанию. " +"Может обучать быстрее, но с большей вероятностью взрыва градиента.\n" +"\textreme - Наивысший оптимальный коэффициент обучения. Гораздо выше риск " +"взрыва градиента." + +#: plugins/train/_config.py:330 +msgid "" +"Apply AutoClipping to the gradients. AutoClip analyzes the gradient weights " +"and adjusts the normalization value dynamically to fit the data. Can help " +"prevent NaNs and improve model optimization at the expense of VRAM. Ref: " +"AutoClip: Adaptive Gradient Clipping for Source Separation Networks https://" +"arxiv.org/abs/2007.14469" +msgstr "" +"Применить AutoClipping к градиентам. AutoClip анализирует веса градиентов и " +"динамически корректирует значение нормализации, чтобы оно подходило к " +"данным. Может помочь избежать NaN('не число') и улучшить оптимизацию модели " +"ценой видеопамяти. Ссылка: AutoClip: Adaptive Gradient Clipping for Source " +"Separation Networks [ТОЛЬКО на английском] https://arxiv.org/abs/2007.14469" + +#: plugins/train/_config.py:343 plugins/train/_config.py:355 +#: plugins/train/_config.py:369 plugins/train/_config.py:386 +msgid "network" +msgstr "сеть" + +#: plugins/train/_config.py:345 +msgid "" +"Use reflection padding rather than zero padding with convolutions. Each " +"convolution must pad the image boundaries to maintain the proper sizing. " +"More complex padding schemes can reduce artifacts at the border of the " +"image.\n" +"\t http://www-cs.engr.ccny.cuny.edu/~wolberg/cs470/hw/hw2_pad.txt" +msgstr "" +"Используйте для сверток не нулевую, а отражающую подкладку. Каждая свертка " +"должна заполнять границы изображения для поддержания правильного размера. " +"Более сложные схемы вставки могут уменьшить артефакты на границе " +"изображения.\n" +"\t http://www-cs.engr.ccny.cuny.edu/~wolberg/cs470/hw/hw2_pad.txt" + +#: plugins/train/_config.py:358 +msgid "" +"Enable the Tensorflow GPU 'allow_growth' configuration option. This option " +"prevents Tensorflow from allocating all of the GPU VRAM at launch but can " +"lead to higher VRAM fragmentation and slower performance. Should only be " +"enabled if you are receiving errors regarding 'cuDNN fails to initialize' " +"when commencing training." +msgstr "" +"[Только для Nvidia]. Включите опцию конфигурации Tensorflow GPU " +"`allow_growth`. Эта опция не позволяет Tensorflow выделять всю видеопамять " +"видеокарты при запуске, но может привести к повышенной фрагментации " +"видеопамяти и снижению производительности. Следует включать только в том " +"случае, если у вас появляются ошибки, рода 'cuDNN fails to initialize'(cuDNN " +"не может инициализироваться) при начале тренировки." + +#: plugins/train/_config.py:371 +msgid "" +"NVIDIA GPUs can run operations in float16 faster than in float32. Mixed " +"precision allows you to use a mix of float16 with float32, to get the " +"performance benefits from float16 and the numeric stability benefits from " +"float32.\n" +"\n" +"This is untested on DirectML backend, but will run on most Nvidia models. it " +"will only speed up training on more recent GPUs. Those with compute " +"capability 7.0 or higher will see the greatest performance benefit from " +"mixed precision because they have Tensor Cores. Older GPUs offer no math " +"performance benefit for using mixed precision, however memory and bandwidth " +"savings can enable some speedups. Generally RTX GPUs and later will offer " +"the most benefit." +msgstr "" +"Видеокарты от NVIDIA могут оперировать в 'float16' быстрее, чем в 'float32'. " +"Смешанная точность позволяет вам использовать микс float16 с float32, чтобы " +"получить улучшение производительности от float16 и числовую стабильность от " +"float32.\n" +"\n" +"Это не было проверено на DirectML, но будет работать на большенстве моделей " +"Nvidia. Оно только ускорит тренировку на более недавних видеокартах. Те, что " +"имеют возможность вычислений('Compute Capability') 7.0 и выше, получат самое " +"большое ускорение от смешанной точности, потому что у них имеются тензор " +"ядра. Старые видеокарты предлагают никакого ускорения от смешанной точности, " +"однако экономия памяти и (хз, честно, словаря нет) могут дать небольшое " +"ускорение. В основном RTX видеокарты и позже предлагают самое большое " +"ускорение." + +#: plugins/train/_config.py:388 +msgid "" +"If a 'NaN' is generated in the model, this means that the model has " +"corrupted and the model is likely to start deteriorating from this point on. " +"Enabling NaN protection will stop training immediately in the event of a " +"NaN. The last save will not contain the NaN, so you may still be able to " +"rescue your model." +msgstr "" +"Если 'Не число'(далее, NaN) сгенерировано в модели - это значит, что модель " +"повреждена и с этого момента, скорее всего, начнет деградировать. Включение " +"защиты от NaN немедленно остановит тренировку, в случае, если был обнаружен " +"NaN. Последнее сохранение не будет содержать в себе NaN, так что у вас будет " +"возможность спасти вашу модель." + +#: plugins/train/_config.py:401 +msgid "convert" +msgstr "конвертирование" + +#: plugins/train/_config.py:403 +msgid "" +"[GPU Only]. The number of faces to feed through the model at once when " +"running the Convert process.\n" +"\n" +"NB: Increasing this figure is unlikely to improve convert speed, however, if " +"you are getting Out of Memory errors, then you may want to reduce the batch " +"size." +msgstr "" +"[Только для видеокарт] Количество лиц, проходящих через модель в одно время " +"во время конвертирования\n" +"\n" +"Примечание: Увеличение этого значения вряд ли повлечет за собой ускорение " +"конвертирования, однако, если у вас появляются ошибки 'Out of Memory', тогда " +"стоит снизить размер пачки." + +#: plugins/train/_config.py:422 +msgid "" +"Loss configuration options\n" +"Loss is the mechanism by which a Neural Network judges how well it thinks " +"that it is recreating a face." +msgstr "" +"Настройки потерь\n" +"Потеря - механизм, по которому Нейронная Сеть судит, насколько хорошо она " +"воспроизводит лицо." + +#: plugins/train/_config.py:429 plugins/train/_config.py:441 +#: plugins/train/_config.py:454 plugins/train/_config.py:474 +#: plugins/train/_config.py:486 plugins/train/_config.py:506 +#: plugins/train/_config.py:518 plugins/train/_config.py:538 +#: plugins/train/_config.py:554 plugins/train/_config.py:570 +#: plugins/train/_config.py:587 +msgid "loss" +msgstr "потери" + +#: plugins/train/_config.py:433 +msgid "The loss function to use." +msgstr "Какую функцию потерь стоит использовать." + +#: plugins/train/_config.py:445 +msgid "" +"The second loss function to use. If using a structural based loss (such as " +"SSIM, MS-SSIM or GMSD) it is common to add an L1 regularization(MAE) or L2 " +"regularization (MSE) function. You can adjust the weighting of this loss " +"function with the loss_weight_2 option." +msgstr "" +"Вторая используемая функция потерь. При использовании потерь, основанных на " +"структуре (таких как SSIM, MS-SSIM или GMSD), обычно добавляется функция " +"регуляризации L1 (MAE) или регуляризации L2 (MSE). Вы можете настроить вес " +"этой функции потерь с помощью параметра loss_weight_2." + +#: plugins/train/_config.py:460 +msgid "" +"The amount of weight to apply to the second loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the second loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the second loss function will be reduced by " +"a quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the second loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the second loss function altogether." +msgstr "" +"Величина веса, применяемая ко второй функции потерь.\n" +"\n" +"\n" +"\n" +"Значение задается в процентах и показывает, какой вклад выбранная функция " +"должна внести в общую стоимость потерь модели. Например:\n" +"\t 100 - Потери, рассчитанные для четвертой функции потерь, будут применены " +"в полном объеме к общей стоимости потерь. \n" +"\t25 - Потери, рассчитанные для четвертой функции потерь, будут уменьшены на " +"четверть перед добавлением к общей стоимости потерь. \n" +"\t400 - Потери, рассчитанные для четвертой функции потерь, будут умножены в " +"4 раза перед добавлением к общей оценке потерь. \n" +"\t 0 - Полностью отключает четвертую функцию потерь." + +#: plugins/train/_config.py:478 +msgid "" +"The third loss function to use. You can adjust the weighting of this loss " +"function with the loss_weight_3 option." +msgstr "" +"Третья используемая функция потерь. Вы можете настроить вес этой функции " +"потерь с помощью параметра loss_weight_3." + +#: plugins/train/_config.py:492 +msgid "" +"The amount of weight to apply to the third loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the third loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the third loss function will be reduced by a " +"quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the third loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the third loss function altogether." +msgstr "" +"Величина веса, применяемая к третьей функции потерь.\n" +"\n" +"\n" +"\n" +"Значение задается в процентах и показывает, какой вклад выбранная функция " +"должна внести в общую стоимость потерь модели. Например:\n" +"\t 100 - Потери, рассчитанные для четвертой функции потерь, будут применены " +"в полном объеме к общей стоимости потерь. \n" +"\t25 - Потери, рассчитанные для четвертой функции потерь, будут уменьшены на " +"четверть перед добавлением к общей стоимости потерь. \n" +"\t400 - Потери, рассчитанные для четвертой функции потерь, будут умножены в " +"4 раза перед добавлением к общей оценке потерь. \n" +"\t 0 - Полностью отключает четвертую функцию потерь." + +#: plugins/train/_config.py:510 +msgid "" +"The fourth loss function to use. You can adjust the weighting of this loss " +"function with the loss_weight_3 option." +msgstr "" +"Четвертая используемая функция потерь. Вы можете настроить вес этой функции " +"потерь с помощью параметра 'loss_weight_4'." + +#: plugins/train/_config.py:524 +msgid "" +"The amount of weight to apply to the fourth loss function.\n" +"\n" +"\n" +"\n" +"The value given here is as a percentage denoting how much the selected " +"function should contribute to the overall loss cost of the model. For " +"example:\n" +"\t 100 - The loss calculated for the fourth loss function will be applied at " +"its full amount towards the overall loss score. \n" +"\t 25 - The loss calculated for the fourth loss function will be reduced by " +"a quarter prior to adding to the overall loss score. \n" +"\t 400 - The loss calculated for the fourth loss function will be mulitplied " +"4 times prior to adding to the overall loss score. \n" +"\t 0 - Disables the fourth loss function altogether." +msgstr "" +"Величина веса, применяемая к четвертой функции потерь.\n" +"\n" +"\n" +"\n" +"Значение задается в процентах и показывает, какой вклад выбранная функция " +"должна внести в общую стоимость потерь модели. Например:\n" +"\t 100 - Потери, рассчитанные для четвертой функции потерь, будут применены " +"в полном объеме к общей стоимости потерь. \n" +"\t25 - Потери, рассчитанные для четвертой функции потерь, будут уменьшены на " +"четверть перед добавлением к общей стоимости потерь. \n" +"\t400 - Потери, рассчитанные для четвертой функции потерь, будут умножены в " +"4 раза перед добавлением к общей оценке потерь. \n" +"\t 0 - Полностью отключает четвертую функцию потерь." + +#: plugins/train/_config.py:543 +msgid "" +"The loss function to use when learning a mask.\n" +"\t MAE - Mean absolute error will guide reconstructions of each pixel " +"towards its median value in the training dataset. Robust to outliers but as " +"a median, it can potentially ignore some infrequent image types in the " +"dataset.\n" +"\t MSE - Mean squared error will guide reconstructions of each pixel towards " +"its average value in the training dataset. As an average, it will be " +"susceptible to outliers and typically produces slightly blurrier results." +msgstr "" +"Функция потерь, используемая при обучении маски.\n" +"\tMAE - средняя абсолютная погрешность('Mean absolute error') направляет " +"реконструкцию каждого пикселя к его срединному значению в обучающем наборе " +"данных. Устойчива к выбросам, но как медиана может игнорировать некоторые " +"редкие типы изображений в наборе данных.\n" +"\tMSE - средняя квадратичная погрешность('Mean squared error') направляет " +"реконструкцию каждого пикселя к его срединному значению в обучающем наборе " +"данных. Как среднее значение, оно чувствительно к выбросам и обычно дает " +"немного более размытые результаты." + +#: plugins/train/_config.py:560 +msgid "" +"The amount of priority to give to the eyes.\n" +"\n" +"The value given here is as a multiplier of the main loss score. For " +"example:\n" +"\t 1 - The eyes will receive the same priority as the rest of the face. \n" +"\t 10 - The eyes will be given a score 10 times higher than the rest of the " +"face.\n" +"\n" +"NB: Penalized Mask Loss must be enable to use this option." +msgstr "" +"Величина приоритета, которую следует придать глазам.\n" +"\n" +"Значение дается как множитель основного показателя потерь. Например:\n" +"\t 1 - Глаза получат тот же приоритет, что и остальное лицо. \n" +"\t 10 - глаза получат оценку в 10 раз выше, чем остальные части лица.\n" +"\n" +"NB: Penalized Mask Loss должен быть включен, чтобы использовать эту опцию." + +#: plugins/train/_config.py:576 +msgid "" +"The amount of priority to give to the mouth.\n" +"\n" +"The value given here is as a multiplier of the main loss score. For " +"Example:\n" +"\t 1 - The mouth will receive the same priority as the rest of the face. \n" +"\t 10 - The mouth will be given a score 10 times higher than the rest of the " +"face.\n" +"\n" +"NB: Penalized Mask Loss must be enable to use this option." +msgstr "" +"Величина приоритета, которую следует придать рту.\n" +"\n" +"Значение дается как множитель основного показателя потерь. Например:\n" +"\t 1 - Рот получит тот же приоритет, что и остальное лицо. \n" +"\t 10 - Рот получит оценку в 10 раз выше, чем остальные части лица.\n" +"\n" +"NB: Penalized Mask Loss должен быть включен, чтобы использовать эту опцию." + +#: plugins/train/_config.py:589 +msgid "" +"Image loss function is weighted by mask presence. For areas of the image " +"without the facial mask, reconstruction errors will be ignored while the " +"masked face area is prioritized. May increase overall quality by focusing " +"attention on the core face area." +msgstr "" +"Функция потерь изображения взвешивается по наличию маски. Для областей " +"изображения без маски лица погрешности реконструкции игнорируются, в то " +"время как область лица с маской является приоритетной. Может повысить общее " +"качество за счет концентрации внимания на основной области лица." + +#: plugins/train/_config.py:600 plugins/train/_config.py:643 +#: plugins/train/_config.py:656 plugins/train/_config.py:671 +#: plugins/train/_config.py:680 +msgid "mask" +msgstr "маска" + +#: plugins/train/_config.py:603 +msgid "" +"The mask to be used for training. If you have selected 'Learn Mask' or " +"'Penalized Mask Loss' you must select a value other than 'none'. The " +"required mask should have been selected as part of the Extract process. If " +"it does not exist in the alignments file then it will be generated prior to " +"training commencing.\n" +"\tnone: Don't use a mask.\n" +"\tbisenet-fp_face: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'face' or " +"'legacy' centering.\n" +"\tbisenet-fp_head: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked (configurable in mask settings). " +"Use this version of bisenet-fp if your model is trained with 'head' " +"centering.\n" +"\tcomponents: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"\tcustom_face: Custom user created, face centered mask.\n" +"\tcustom_head: Custom user created, head centered mask.\n" +"\textended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"\tvgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"\tvgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"\tunet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members and will need " +"testing for further description. Profile faces may result in sub-par " +"performance." +msgstr "" +"Маска, которая будет использоваться для обучения. Если вы выбрали 'Learn " +"Mask' или 'Penalized Mask Loss', вы должны выбрать значение, отличное от " +"'none'. Необходимая маска должна быть выбрана в процессе извлечения. Если " +"она не существует в файле выравниваний, то она будет создана до начала " +"обучения.\n" +"\tnone: Не использовать маску.\n" +"\tbisenet-fp_face: Относительно легкая маска на основе NN, которая " +"обеспечивает более точный контроль над маскируемой областью (настраивается в " +"настройках маски). Используйте эту версию bisenet-fp, если ваша модель " +"обучена с центрированием 'face' или 'legacy'.\n" +"\tbisenet-fp_head: Относительно легкая маска на основе NN, которая " +"обеспечивает более точный контроль над маскируемой областью (настраивается в " +"параметрах маски). Используйте эту версию bisenet-fp, если ваша модель " +"обучена с центрированием 'head'.\n" +"\tcomponents: Маска, разработанная для сегментации лица на основе " +"расположения ориентиров. Для создания маски вокруг внешних ориентиров " +"строится выпуклая оболочка.\n" +"\tcustom_face: Пользовательская маска, созданная пользователем и " +"центрированная по лицу.\n" +"\tcustom_head: Созданная пользователем маска, центрированная по голове.\n" +"\textended: Маска, разработанная для сегментации лица на основе расположения " +"ориентиров. Выпуклый корпус строится вокруг внешних ориентиров, и маска " +"расширяется вверх на лоб.\n" +"\tvgg-clear: Маска предназначена для интеллектуальной сегментации " +"преимущественно фронтальных лиц без препятствий. Профильные лица и " +"препятствия могут привести к снижению производительности.\n" +"\tvgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была специально обучена " +"распознавать некоторые препятствия на лице (руки и очки). Профильные лица " +"могут иметь низкую производительность.\n" +"\tunet-dfl: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была обучена членами " +"сообщества и для дальнейшего описания нуждается в тестировании. Профильные " +"лица могут иметь низкую производительность." + +#: plugins/train/_config.py:645 +msgid "" +"Dilate or erode the mask. Negative values erode the mask (make it smaller). " +"Positive values dilate the mask (make it larger). The value given is a " +"percentage of the total mask size." +msgstr "" +"Расширяет или сужает маску. Отрицательные значения сужают маску (делают её " +"меньше). Положительные значения расширяют маску (делают её больше). " + +#: plugins/train/_config.py:658 +msgid "" +"Apply gaussian blur to the mask input. This has the effect of smoothing the " +"edges of the mask, which can help with poorly calculated masks and give less " +"of a hard edge to the predicted mask. The size is in pixels (calculated from " +"a 128px mask). Set to 0 to not apply gaussian blur. This value should be " +"odd, if an even number is passed in then it will be rounded to the next odd " +"number." +msgstr "" +"Применить размытие по Гауссу на входную маску. Дает эффект сглаживания краев " +"маски, что может помочь с плохо вычисленными масками и дает менее резкий " +"край предугаданной маске. Размер в пикселях (вычисленно из маски на 128 " +"пикселей). Установите 0, чтобы не применять размытие по Гауссу. Это значение " +"должно быть нечетным, если передано четное число, то оно будет округлено до " +"следующего нечетного числа." + +#: plugins/train/_config.py:673 +msgid "" +"Sets pixels that are near white to white and near black to black. Set to 0 " +"for off." +msgstr "" +"Устанавливает пиксели, которые почти белые - в белые и которые почти черные " +"- в черные. Установите 0, чтобы выключить." + +#: plugins/train/_config.py:682 +msgid "" +"Dedicate a portion of the model to learning how to duplicate the input mask. " +"Increases VRAM usage in exchange for learning a quick ability to try to " +"replicate more complex mask models." +msgstr "" +"Выделить частичку модели обучению тому, как дублировать входную маску. " +"Увеличивает использование видеопамяти в обмен на обучение быстрой " +"способности попытки переделывать более сложные маски." diff --git a/locales/ru/LC_MESSAGES/tools.alignments.cli.mo b/locales/ru/LC_MESSAGES/tools.alignments.cli.mo new file mode 100644 index 0000000000..5277c793d2 Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.alignments.cli.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.alignments.cli.po b/locales/ru/LC_MESSAGES/tools.alignments.cli.po new file mode 100644 index 0000000000..3f68a44acf --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.alignments.cli.po @@ -0,0 +1,276 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-04-19 11:28+0100\n" +"PO-Revision-Date: 2024-04-19 11:31+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/alignments/cli.py:16 +msgid "" +"This command lets you perform various tasks pertaining to an alignments file." +msgstr "" +"Эта команда позволяет выполнять различные задачи, относящиеся к файлу " +"выравнивания." + +#: tools/alignments/cli.py:31 +msgid "" +"Alignments tool\n" +"This tool allows you to perform numerous actions on or using an alignments " +"file against its corresponding faceset/frame source." +msgstr "" +"Инструмент выравнивания\n" +"Этот инструмент позволяет выполнять многочисленные действия с файлом " +"выравнивания или с его использованием против соответствующего набора лиц/" +"кадров." + +#: tools/alignments/cli.py:43 +msgid " Must Pass in a frames folder/source video file (-r)." +msgstr " Должен проходить в папке с кадрами/исходным видеофайлом (-r)." + +#: tools/alignments/cli.py:44 +msgid " Must Pass in a faces folder (-c)." +msgstr " Должен проходить в папке с лицами (-c)." + +#: tools/alignments/cli.py:45 +msgid "" +" Must Pass in either a frames folder/source video file OR a faces folder (-r " +"or -c)." +msgstr "" +" Должно передаваться либо в папку с кадрами/исходным видеофайлом, либо в " +"папку с лицами (-r или -c)." + +#: tools/alignments/cli.py:47 +msgid "" +" Must Pass in a frames folder/source video file AND a faces folder (-r and -" +"c)." +msgstr "" +" Должно передаваться либо в папку с кадрами/исходным видеофайлом И в папку с " +"лицами (-r и -c)." + +#: tools/alignments/cli.py:49 +msgid " Use the output option (-o) to process results." +msgstr " Используйте опцию вывода (-o) для обработки результатов." + +#: tools/alignments/cli.py:58 tools/alignments/cli.py:104 +msgid "processing" +msgstr "обработка" + +#: tools/alignments/cli.py:61 +#, python-brace-format +msgid "" +"R|Choose which action you want to perform. NB: All actions require an " +"alignments file (-a) to be passed in.\n" +"L|'draw': Draw landmarks on frames in the selected folder/video. A subfolder " +"will be created within the frames folder to hold the output.{0}\n" +"L|'export': Export the contents of an alignments file to a json file. Can be " +"used for editing alignment information in external tools and then re-" +"importing by using Faceswap's Extract 'Import' plugins. Note: masks and " +"identity vectors will not be included in the exported file, so will be re-" +"generated when the json file is imported back into Faceswap. All data is " +"exported with the origin (0, 0) at the top left of the canvas.\n" +"L|'extract': Re-extract faces from the source frames/video based on " +"alignment data. This is a lot quicker than re-detecting faces. Can pass in " +"the '-een' (--extract-every-n) parameter to only extract every nth frame." +"{1}\n" +"L|'from-faces': Generate alignment file(s) from a folder of extracted faces. " +"if the folder of faces comes from multiple sources, then multiple alignments " +"files will be created. NB: for faces which have been extracted from folders " +"of source images, rather than a video, a single alignments file will be " +"created as there is no way for the process to know how many folders of " +"images were originally used. You do not need to provide an alignments file " +"path to run this job. {3}\n" +"L|'missing-alignments': Identify frames that do not exist in the alignments " +"file.{2}{0}\n" +"L|'missing-frames': Identify frames in the alignments file that do not " +"appear within the frames folder/video.{2}{0}\n" +"L|'multi-faces': Identify where multiple faces exist within the alignments " +"file.{2}{4}\n" +"L|'no-faces': Identify frames that exist within the alignment file but no " +"faces were detected.{2}{0}\n" +"L|'remove-faces': Remove deleted faces from an alignments file. The original " +"alignments file will be backed up.{3}\n" +"L|'rename' - Rename faces to correspond with their parent frame and position " +"index in the alignments file (i.e. how they are named after running extract)." +"{3}\n" +"L|'sort': Re-index the alignments from left to right. For alignments with " +"multiple faces this will ensure that the left-most face is at index 0.\n" +"L|'spatial': Perform spatial and temporal filtering to smooth alignments " +"(EXPERIMENTAL!)" +msgstr "" +"R|Выберите действие, которое вы хотите выполнить. Примечание: Все действия " +"требуют передачи файла выравнивания (-a).\n" +"L|'draw': Нарисовать ориентиры на кадрах в выбранной папке/видео. В папке " +"frames будет создана подпапка для хранения результатов.\n" +"L|'export': экспортировать содержимое файла выравнивания в файл JSON. Может " +"использоваться для редактирования информации о выравнивании во внешних " +"инструментах, а затем повторно импортируется с помощью плагинов Faceswap " +"Extract 'Import'. ПРИМЕЧАНИЕ. Маски и векторы идентификации не будут " +"включены в экспортированный файл, поэтому будут повторно сгенерированы, " +"когда файл JSON будет импортирован обратно в Faceswap. Все данные " +"экспортируются с началом координат (0, 0) в верхнем левом углу холста.\n" +"L|'extract': Повторное извлечение лиц из исходных кадров/видео на основе " +"данных о выравнивании. Это намного быстрее, чем повторное обнаружение лиц. " +"Можно передать параметр '-een' (--extract-every-n), чтобы извлекать только " +"каждый n-й кадр.{1}\n" +"L|'from-faces': Создать файл(ы) выравнивания из папки с извлеченными лицами. " +"Если папка с лицами получена из нескольких источников, то будет создано " +"несколько файлов выравнивания. Примечание: для лиц, которые были извлечены " +"из папок с исходными изображениями, а не из видео, будет создан один файл " +"выравнивания, поскольку процесс не может знать, сколько папок с " +"изображениями было использовано изначально. Для выполнения этого задания не " +"нужно указывать путь к файлу выравнивания. {3}\n" +"L|'missing-alignments': Определить кадры, которых нет в файле выравнивания." +"{2}{0}\n" +"L|'missing-frames': Определить кадры в файле выравнивания, которые не " +"появляются в папке frames/video.{2}{0}\n" +"L|'multi-faces': Определить, где в файле выравнивания существует несколько " +"лиц.{2}{4}\n" +"L|'no-faces': Идентифицировать кадры, которые существуют в файле " +"выравнивания, но лица не были обнаружены.{2}{0}\n" +"L|'remove-faces': Удалить удаленные лица из файла выравнивания. Оригинальный " +"файл выравнивания будет сохранен.{3}\n" +"L|'rename' - Переименовать лица в соответствии с их родительским кадром и " +"индексом позиции в файле выравниваний (т.е. как они будут названы после " +"запуска extract).{3}\n" +"L|'sort': Переиндексирует выравнивания слева направо. Для выравниваний с " +"несколькими гранями это гарантирует, что самое левое лицо будет иметь индекс " +"0.\n" +"L|'spatial': Выполнить пространственную и временную фильтрацию для " +"сглаживания выравниваний (ЭКСПЕРИМЕНТАЛЬНО!)." + +#: tools/alignments/cli.py:107 +msgid "" +"R|How to output discovered items ('faces' and 'frames' only):\n" +"L|'console': Print the list of frames to the screen. (DEFAULT)\n" +"L|'file': Output the list of frames to a text file (stored within the source " +"directory).\n" +"L|'move': Move the discovered items to a sub-folder within the source " +"directory." +msgstr "" +"R|Как вывести обнаруженные элементы (только \"лица\" и \"кадры\"):\n" +"L|'console': Вывести список рамок на экран. (DEFAULT)\n" +"L|'file': Вывести список кадров в текстовый файл (хранящийся в исходном " +"каталоге).\n" +"L|'move': Переместить обнаруженные элементы в подпапку в исходном каталоге." + +#: tools/alignments/cli.py:118 tools/alignments/cli.py:141 +#: tools/alignments/cli.py:148 +msgid "data" +msgstr "данные" + +#: tools/alignments/cli.py:125 +msgid "" +"Full path to the alignments file to be processed. If you have input a " +"'frames_dir' and don't provide this option, the process will try to find the " +"alignments file at the default location. All jobs require an alignments file " +"with the exception of 'from-faces' when the alignments file will be " +"generated in the specified faces folder." +msgstr "" +"Полный путь к обрабатываемому файлу выравниваний. Если вы ввели 'frames_dir' " +"и не указали этот параметр, процесс попытается найти файл выравнивания в " +"месте по умолчанию. Все задания требуют файл выравнивания, за исключением " +"задания 'from-faces', когда файл выравнивания будет создан в указанной папке " +"с лицами." + +#: tools/alignments/cli.py:142 +msgid "Directory containing source frames that faces were extracted from." +msgstr "Папка, содержащая исходные кадры, из которых были извлечены лица." + +#: tools/alignments/cli.py:150 +msgid "" +"R|Run the aligmnents tool on multiple sources. The following jobs support " +"batch mode:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"If batch mode is selected then the other options should be set as follows:\n" +"L|alignments_file: For 'sort' and 'spatial' this should point to the parent " +"folder containing the alignments files to be processed. For all other jobs " +"this option is ignored, and the alignments files must exist at their default " +"location relative to the original frames folder/video.\n" +"L|faces_dir: For 'from-faces' this should be a parent folder, containing sub-" +"folders of extracted faces from which to generate alignments files. For " +"'extract' this should be a parent folder where sub-folders will be created " +"for each extraction to be run. For all other jobs this option is ignored.\n" +"L|frames_dir: For 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"and 'no-faces' this should be a parent folder containing video files or sub-" +"folders of images to perform the alignments job on. The alignments file " +"should exist at the default location. For all other jobs this option is " +"ignored." +msgstr "" +"R|Запуск инструмента выравнивания на нескольких источниках. Следующие " +"задания поддерживают пакетный режим:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"Если выбран пакетный режим, то остальные опции должны быть установлены " +"следующим образом:\n" +"L|alignments_file: Для заданий 'sort' и 'spatial' этот параметр должен " +"указывать на родительскую папку, содержащую файлы выравниваний, которые " +"будут обрабатываться. Для всех остальных заданий этот параметр игнорируется, " +"и файлы выравнивания должны существовать в их расположении по умолчанию " +"относительно исходной папки кадров/видео.\n" +"L|faces_dir: Для 'from-faces' это должна быть родительская папка, содержащая " +"вложенные папки с извлеченными лицами, из которых будут сгенерированы файлы " +"выравнивания. Для 'extract' это должна быть родительская папка, в которой " +"будут создаваться вложенные папки для каждой выполняемой экстракции. Для " +"всех остальных заданий этот параметр игнорируется.\n" +"L|frames_dir: Для 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"и 'no-faces' это должна быть родительская папка, содержащая видеофайлы или " +"вложенные папки изображений для выполнения задания выравнивания. Файл " +"выравнивания должен существовать в месте по умолчанию. Для всех остальных " +"заданий этот параметр игнорируется." + +#: tools/alignments/cli.py:176 tools/alignments/cli.py:188 +#: tools/alignments/cli.py:198 +msgid "extract" +msgstr "извлечение" + +#: tools/alignments/cli.py:178 +msgid "" +"[Extract only] Extract every 'nth' frame. This option will skip frames when " +"extracting faces. For example a value of 1 will extract faces from every " +"frame, a value of 10 will extract faces from every 10th frame." +msgstr "" +"[Только извлечение] Извлекать каждый \"n-й\" кадр. Этот параметр пропускает " +"кадры при извлечении лиц. Например, значение 1 будет извлекать лица из " +"каждого кадра, значение 10 будет извлекать лица из каждого 10-го кадра." + +#: tools/alignments/cli.py:189 +msgid "[Extract only] The output size of extracted faces." +msgstr "[Только извлечение] Выходной размер извлеченных лиц." + +#: tools/alignments/cli.py:200 +msgid "" +"[Extract only] Only extract faces that have been resized by this percent or " +"more to meet the specified extract size (`-sz`, `--size`). Useful for " +"excluding low-res images from a training set. Set to 0 to extract all faces. " +"Eg: For an extract size of 512px, A setting of 50 will only include faces " +"that have been resized from 256px or above. Setting to 100 will only extract " +"faces that have been resized from 512px or above. A setting of 200 will only " +"extract faces that have been downscaled from 1024px or above." +msgstr "" +"[Только извлечение] Извлекать только те лица, размер которых был изменен на " +"данный процент или более, чтобы соответствовать заданному размеру извлечения " +"(`-sz`, `--size`). Полезно для исключения изображений с низким разрешением " +"из обучающего набора. Установите значение 0, чтобы извлечь все лица. " +"Например: Для размера экстракта 512px, при установке значения 50 будут " +"извлечены только лица, размер которых был изменен с 256px или выше. При " +"значении 100 будут извлечены только лица, размер которых был изменен с 512px " +"или выше. При значении 200 будут извлечены только лица, уменьшенные с 1024px " +"или выше." + +#~ msgid "Directory containing extracted faces." +#~ msgstr "Папка, содержащая извлеченные лица." diff --git a/locales/ru/LC_MESSAGES/tools.effmpeg.cli.mo b/locales/ru/LC_MESSAGES/tools.effmpeg.cli.mo new file mode 100644 index 0000000000..b47b17d08d Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.effmpeg.cli.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.effmpeg.cli.po b/locales/ru/LC_MESSAGES/tools.effmpeg.cli.po new file mode 100644 index 0000000000..8322e90571 --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.effmpeg.cli.po @@ -0,0 +1,191 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:50+0000\n" +"PO-Revision-Date: 2024-03-29 00:08+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/effmpeg/cli.py:15 +msgid "This command allows you to easily execute common ffmpeg tasks." +msgstr "Эта команда позволяет легко выполнять общие задачи ffmpeg." + +#: tools/effmpeg/cli.py:52 +msgid "A wrapper for ffmpeg for performing image <> video converting." +msgstr "Обертка для ffmpeg для выполнения конвертации изображений <> видео." + +#: tools/effmpeg/cli.py:64 +msgid "" +"R|Choose which action you want ffmpeg ffmpeg to do.\n" +"L|'extract': turns videos into images \n" +"L|'gen-vid': turns images into videos \n" +"L|'get-fps' returns the chosen video's fps.\n" +"L|'get-info' returns information about a video.\n" +"L|'mux-audio' add audio from one video to another.\n" +"L|'rescale' resize video.\n" +"L|'rotate' rotate video.\n" +"L|'slice' cuts a portion of the video into a separate video file." +msgstr "" +"R|Выберите, какое действие вы хотите, чтобы выполнял ffmpeg.\n" +"L|'extract': превращает видео в изображения \n" +"L|'gen-vid': превращает изображения в видео. \n" +"L|'get-fps' возвращает частоту кадров в секунду выбранного видео.\n" +"L|'get-info': возвращает информацию о видео.\n" +"L|'mux-audio' добавляет звук из одного видео в другое.\n" +"L|'rescale' изменить размер видео.\n" +"L|'rotate' вращение видео.\n" +"L|'slice' вырезает часть видео в отдельный видеофайл." + +#: tools/effmpeg/cli.py:78 +msgid "Input file." +msgstr "Входной файл." + +#: tools/effmpeg/cli.py:79 tools/effmpeg/cli.py:86 tools/effmpeg/cli.py:100 +msgid "data" +msgstr "данные" + +#: tools/effmpeg/cli.py:89 +msgid "" +"Output file. If no output is specified then: if the output is meant to be a " +"video then a video called 'out.mkv' will be created in the input directory; " +"if the output is meant to be a directory then a directory called 'out' will " +"be created inside the input directory. Note: the chosen output file " +"extension will determine the file encoding." +msgstr "" +"Выходной файл. Если выходной файл не указан, то: если выходным файлом " +"является видео, то в каталоге ввода будет создан видеофайл с именем 'out." +"mkv'; если выходным файлом является каталог, то внутри каталога ввода будет " +"создан каталог с именем 'out'. Примечание: выбранное расширение выходного " +"файла определяет кодировку файла." + +#: tools/effmpeg/cli.py:102 +msgid "Path to reference video if 'input' was not a video." +msgstr "Путь к опорному видео, если 'input' не является видео." + +#: tools/effmpeg/cli.py:108 tools/effmpeg/cli.py:118 tools/effmpeg/cli.py:156 +#: tools/effmpeg/cli.py:185 +msgid "output" +msgstr "выход" + +#: tools/effmpeg/cli.py:110 +msgid "" +"Provide video fps. Can be an integer, float or fraction. Negative values " +"will will make the program try to get the fps from the input or reference " +"videos." +msgstr "" +"Предоставляет количество кадров в секунду. Может быть целым числом, " +"плавающей цифрой или дробью. Отрицательные значения заставят программу " +"попытаться получить fps из входного или опорного видео." + +#: tools/effmpeg/cli.py:120 +msgid "" +"Image format that extracted images should be saved as. '.bmp' will offer the " +"fastest extraction speed, but will take the most storage space. '.png' will " +"be slower but will take less storage." +msgstr "" +"Формат изображения, в котором должны быть сохранены извлеченные изображения. " +"'.bmp' обеспечивает самую высокую скорость извлечения, но занимает больше " +"всего места в памяти. '.png' будет медленнее, но займет меньше места." + +#: tools/effmpeg/cli.py:127 tools/effmpeg/cli.py:136 tools/effmpeg/cli.py:145 +msgid "clip" +msgstr "клип" + +#: tools/effmpeg/cli.py:129 +msgid "" +"Enter the start time from which an action is to be applied. Default: " +"00:00:00, in HH:MM:SS format. You can also enter the time with or without " +"the colons, e.g. 00:0000 or 026010." +msgstr "" +"Введите время начала, с которого будет применяться действие. По умолчанию: " +"00:00:00, в формате ЧЧ:ММ:СС. Вы также можете ввести время с двоеточием или " +"без него, например, 00:0000 или 026010." + +#: tools/effmpeg/cli.py:138 +msgid "" +"Enter the end time to which an action is to be applied. If both an end time " +"and duration are set, then the end time will be used and the duration will " +"be ignored. Default: 00:00:00, in HH:MM:SS." +msgstr "" +"Введите время окончания, до которого будет применяться действие. Если заданы " +"и время окончания, и продолжительность, то будет использоваться время " +"окончания, а продолжительность будет игнорироваться. По умолчанию: 00:00:00, " +"в формате ЧЧ:ММ:СС." + +#: tools/effmpeg/cli.py:147 +msgid "" +"Enter the duration of the chosen action, for example if you enter 00:00:10 " +"for slice, then the first 10 seconds after and including the start time will " +"be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. You can " +"also enter the time with or without the colons, e.g. 00:0000 or 026010." +msgstr "" +"Введите продолжительность выбранного действия, например, если вы введете " +"00:00:10 для нарезки, то первые 10 секунд после начала и включая время " +"начала будут вырезаны в новое видео. По умолчанию: 00:00:00, в формате ЧЧ:ММ:" +"СС. Вы также можете ввести время с двоеточием или без него, например, " +"00:0000 или 026010." + +#: tools/effmpeg/cli.py:158 +msgid "" +"Mux the audio from the reference video into the input video. This option is " +"only used for the 'gen-vid' action. 'mux-audio' action has this turned on " +"implicitly." +msgstr "" +"Mux аудио из опорного видео во входное видео. Эта опция используется только " +"для действия 'gen-vid'. Действие 'mux-audio' включает эту опцию неявно." + +#: tools/effmpeg/cli.py:169 tools/effmpeg/cli.py:179 +msgid "rotate" +msgstr "поворот" + +#: tools/effmpeg/cli.py:171 +msgid "" +"Transpose the video. If transpose is set, then degrees will be ignored. For " +"cli you can enter either the number or the long command name, e.g. to use " +"(1, 90Clockwise) -tr 1 or -tr 90Clockwise" +msgstr "" +"Транспонировать видео. Если задано транспонирование, то градусы будут " +"игнорироваться. Для командой строки вы можете ввести либо число, либо " +"длинное имя команды, например, для использования (1, 90 по часовой стрелке) -" +"tr 1 или -tr 90 по часовой стрелке" + +#: tools/effmpeg/cli.py:180 +msgid "Rotate the video clockwise by the given number of degrees." +msgstr "Поверните видео по часовой стрелке на заданное количество градусов." + +#: tools/effmpeg/cli.py:187 +msgid "Set the new resolution scale if the chosen action is 'rescale'." +msgstr "Установите новый масштаб разрешения, если выбрано действие 'rescale'." + +#: tools/effmpeg/cli.py:192 tools/effmpeg/cli.py:200 +msgid "settings" +msgstr "настройки" + +#: tools/effmpeg/cli.py:194 +msgid "" +"Reduces output verbosity so that only serious errors are printed. If both " +"quiet and verbose are set, verbose will override quiet." +msgstr "" +"Уменьшает многословность вывода, чтобы выводились только серьезные ошибки. " +"Если заданы и quiet, и verbose, то verbose будет преобладать над quiet." + +#: tools/effmpeg/cli.py:202 +msgid "" +"Increases output verbosity. If both quiet and verbose are set, verbose will " +"override quiet." +msgstr "" +"Повышает точность вывода. Если заданы и quiet, и verbose, то verbose будет " +"преобладать над quiet." diff --git a/locales/ru/LC_MESSAGES/tools.manual.mo b/locales/ru/LC_MESSAGES/tools.manual.mo new file mode 100644 index 0000000000..6e724e6f9d Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.manual.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.manual.po b/locales/ru/LC_MESSAGES/tools.manual.po new file mode 100644 index 0000000000..2c74501b46 --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.manual.po @@ -0,0 +1,294 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR ORGANIZATION +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:55+0000\n" +"PO-Revision-Date: 2024-03-29 00:07+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"Generated-By: pygettext.py 1.5\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/manual/cli.py:13 +msgid "" +"This command lets you perform various actions on frames, faces and " +"alignments files using visual tools." +msgstr "" +"Эта команда позволяет выполнять различные действия с кадрами, гранями и " +"файлами выравнивания с помощью визуальных инструментов." + +#: tools/manual/cli.py:23 +msgid "" +"A tool to perform various actions on frames, faces and alignments files " +"using visual tools" +msgstr "" +"Инструмент для выполнения различных действий с кадрами, лицами и файлами " +"выравнивания с помощью визуальных инструментов" + +#: tools/manual/cli.py:35 tools/manual/cli.py:44 +msgid "data" +msgstr "данные" + +#: tools/manual/cli.py:38 +msgid "" +"Path to the alignments file for the input, if not at the default location" +msgstr "" +"Путь к файлу выравниваний для входных данных, если он не находится в месте " +"по умолчанию" + +#: tools/manual/cli.py:46 +msgid "" +"Video file or directory containing source frames that faces were extracted " +"from." +msgstr "" +"Видеофайл или папка, содержащая исходные кадры, из которых были извлечены " +"лица." + +#: tools/manual/cli.py:53 tools/manual/cli.py:62 +msgid "options" +msgstr "опции" + +#: tools/manual/cli.py:55 +msgid "" +"Force regeneration of the low resolution jpg thumbnails in the alignments " +"file." +msgstr "" +"Принудительное восстановление миниатюр jpg низкого разрешения в файле " +"выравнивания." + +#: tools/manual/cli.py:64 +msgid "" +"The process attempts to speed up generation of thumbnails by extracting from " +"the video in parallel threads. For some videos, this causes the caching " +"process to hang. If this happens, then set this option to generate the " +"thumbnails in a slower, but more stable single thread." +msgstr "" +"Процесс пытается ускорить генерацию эскизов путем извлечения из видео в " +"параллельных потоках. Для некоторых видео это приводит к зависанию процесса " +"кэширования. Если это происходит, установите этот параметр, чтобы " +"генерировать эскизы в более медленном, но более стабильном одном потоке." + +#: tools/manual\faceviewer\frame.py:163 +msgid "Display the landmarks mesh" +msgstr "Отображение сетки ориентиров" + +#: tools/manual\faceviewer\frame.py:164 +msgid "Display the mask" +msgstr "Отображение маски" + +#: tools/manual\frameviewer\editor\_base.py:628 +#: tools/manual\frameviewer\editor\landmarks.py:44 +#: tools/manual\frameviewer\editor\mask.py:75 +msgid "Magnify/Demagnify the View" +msgstr "Увеличение/уменьшение изображения" + +#: tools/manual\frameviewer\editor\bounding_box.py:33 +#: tools/manual\frameviewer\editor\extract_box.py:32 +msgid "Delete Face" +msgstr "Удалить лицо" + +#: tools/manual\frameviewer\editor\bounding_box.py:36 +msgid "" +"Bounding Box Editor\n" +"Edit the bounding box being fed into the aligner to recalculate the " +"landmarks.\n" +"\n" +" - Grab the corner anchors to resize the bounding box.\n" +" - Click and drag the bounding box to relocate.\n" +" - Click in empty space to create a new bounding box.\n" +" - Right click a bounding box to delete a face." +msgstr "" +"Редактор ограничительных рамок\n" +"Отредактируйте ограничивающую рамку, подаваемую в выравниватель, чтобы " +"пересчитать ориентиры.\n" +"\n" +"- Захватите угловые опоры, чтобы изменить размер ограничивающей рамки.\n" +" - Щелкните и перетащите ограничивающую рамку для перемещения.\n" +" - Щелкните в пустом пространстве, чтобы создать новую ограничивающую " +"рамку.\n" +"- Щелкните правой кнопкой мыши ограничительную рамку, чтобы удалить лицо." + +#: tools/manual\frameviewer\editor\bounding_box.py:70 +msgid "" +"Aligner to use. FAN will obtain better alignments, but cv2-dnn can be useful " +"if FAN cannot get decent alignments and you want to set a base to edit from." +msgstr "" +"Выравниватель для использования. FAN получит лучшие выравнивания, но cv2-dnn " +"может быть полезен, если FAN не может получить достойные выравнивания, и вы " +"хотите установить базу для редактирования." + +#: tools/manual\frameviewer\editor\bounding_box.py:83 +msgid "" +"Normalization method to use for feeding faces to the aligner. This can help " +"the aligner better align faces with difficult lighting conditions. Different " +"methods will yield different results on different sets. NB: This does not " +"impact the output face, just the input to the aligner.\n" +"\tnone: Don't perform normalization on the face.\n" +"\tclahe: Perform Contrast Limited Adaptive Histogram Equalization on the " +"face.\n" +"\thist: Equalize the histograms on the RGB channels.\n" +"\tmean: Normalize the face colors to the mean." +msgstr "" +"Метод нормализации, используемый для подачи лиц в выравниватель. Это может " +"помочь выравнивателю лучше выравнивать лица при сложных условиях освещения. " +"Различные методы дают разные результаты на разных наборах. Примечание: Это " +"не влияет на выходное лицо, только на входное в выравниватель.\n" +"\tnone: Не выполнять нормализацию лица.\n" +"\tclahe: Выполнить для лица адаптивную гистограммную эквализацию с " +"ограничением контраста.\n" +"\thist: Выравнивание гистограмм по каналам RGB.\n" +"\tmean: Нормализовать цвета лица к среднему значению." + +#: tools/manual\frameviewer\editor\extract_box.py:35 +msgid "" +"Extract Box Editor\n" +"Move the extract box that has been generated by the aligner. Click and " +"drag:\n" +"\n" +" - Inside the bounding box to relocate the landmarks.\n" +" - The corner anchors to resize the landmarks.\n" +" - Outside of the corners to rotate the landmarks." +msgstr "" +"Редактор поля извлечения\n" +"Переместите поле извлечения, созданное выравнивателем. Нажмите и " +"перетащите:\n" +"\n" +" - Внутри ограничивающей рамки для перемещения опорных точек.\n" +"- По угловым опорам для изменения размера опорных точек.\n" +"- За пределами углов, чтобы повернуть опорные точки." + +#: tools/manual\frameviewer\editor\landmarks.py:27 +msgid "" +"Landmark Point Editor\n" +"Edit the individual landmark points.\n" +"\n" +" - Click and drag individual points to relocate.\n" +" - Draw a box to select multiple points to relocate." +msgstr "" +"Редактор точек ориентира\n" +"Редактирование отдельных опорных точек.\n" +"\n" +" - Щелкните и перетащите отдельные точки для перемещения.\n" +" - Нарисуйте рамку, чтобы выбрать несколько точек для перемещения." + +#: tools/manual\frameviewer\editor\mask.py:33 +msgid "" +"Mask Editor\n" +"Edit the mask.\n" +" - NB: For Landmark based masks (e.g. components/extended) it is better to " +"make sure the landmarks are correct rather than editing the mask directly. " +"Any change to the landmarks after editing the mask will override your manual " +"edits." +msgstr "" +"Редактор маски\n" +"Отредактировать маску.\n" +" - Примечание: Для масок, основанных на ориентирах (например, компоненты/" +"расширенные), лучше убедиться в правильности ориентиров, а не редактировать " +"маску напрямую. Любое изменение ориентиров после редактирования маски " +"отменит ваши ручные правки." + +#: tools/manual\frameviewer\editor\mask.py:77 +msgid "Draw Tool" +msgstr "Инструмент рисования" + +#: tools/manual\frameviewer\editor\mask.py:78 +msgid "Erase Tool" +msgstr "Инструмент \"Ластик\"" + +#: tools/manual\frameviewer\editor\mask.py:97 +msgid "Select which mask to edit" +msgstr "Выбрать, какую маску редактировать" + +#: tools/manual\frameviewer\editor\mask.py:104 +msgid "Set the brush size. ([ - decrease, ] - increase)" +msgstr "Установить размер кисти. ([ - уменьшение, ] - увеличение)" + +#: tools/manual\frameviewer\editor\mask.py:111 +msgid "Select the brush cursor color." +msgstr "Установить цвет курсора кисти." + +#: tools/manual\frameviewer\frame.py:78 +msgid "Play/Pause (SPACE)" +msgstr "Воспроизвести/Приостановить (ПРОБЕЛ)" + +#: tools/manual\frameviewer\frame.py:79 +msgid "Go to First Frame (HOME)" +msgstr "Перейти к первому кадру (HOME)" + +#: tools/manual\frameviewer\frame.py:80 +msgid "Go to Previous Frame (Z)" +msgstr "Перейти к предыдущему кадру (Z/Я)" + +#: tools/manual\frameviewer\frame.py:81 +msgid "Go to Next Frame (X)" +msgstr "Перейти к следующему кадру (X/Ч)" + +#: tools/manual\frameviewer\frame.py:82 +msgid "Go to Last Frame (END)" +msgstr "Перейти к последнему кадру (END)" + +#: tools/manual\frameviewer\frame.py:83 +msgid "Extract the faces to a folder... (Ctrl+E)" +msgstr "Извлечь лица в папку... (Ctrl+E)" + +#: tools/manual\frameviewer\frame.py:84 +msgid "Save the Alignments file (Ctrl+S)" +msgstr "Сохранить файл выравнивания (Ctrl+S)" + +#: tools/manual\frameviewer\frame.py:85 +msgid "Filter Frames to only those Containing the Selected Item (F)" +msgstr "Отфильтровать кадры, содержащие только выбранный элемент (F/А)" + +#: tools/manual\frameviewer\frame.py:86 +msgid "" +"Set the distance from an 'average face' to be considered misaligned. Higher " +"distances are more restrictive" +msgstr "" +"Установить расстояние от \"среднего лица\", на котором оно будет считаться " +"смещенным. Большие расстояния являются более ограничительными" + +#: tools/manual\frameviewer\frame.py:391 +msgid "View alignments" +msgstr "Просмотреть выравнивания" + +#: tools/manual\frameviewer\frame.py:392 +msgid "Bounding box editor" +msgstr "Редактор ограничительных рамок" + +#: tools/manual\frameviewer\frame.py:393 +msgid "Location editor" +msgstr "Редактор расположения" + +#: tools/manual\frameviewer\frame.py:394 +msgid "Mask editor" +msgstr "Редактор маски" + +#: tools/manual\frameviewer\frame.py:395 +msgid "Landmark point editor" +msgstr "Редактор точек ориентира" + +#: tools/manual\frameviewer\frame.py:470 +msgid "Next" +msgstr "Следующий" + +#: tools/manual\frameviewer\frame.py:470 +msgid "Previous" +msgstr "Предыдущий" + +#: tools/manual\frameviewer\frame.py:481 +msgid "Revert to saved Alignments ({})" +msgstr "Откатить до сохраненных выравниваний ({})" + +#: tools/manual\frameviewer\frame.py:487 +msgid "Copy {} Alignments ({})" +msgstr "Копировать {} выравнивания ({})" diff --git a/locales/ru/LC_MESSAGES/tools.mask.cli.mo b/locales/ru/LC_MESSAGES/tools.mask.cli.mo new file mode 100644 index 0000000000..89631d0cd9 Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.mask.cli.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.mask.cli.po b/locales/ru/LC_MESSAGES/tools.mask.cli.po new file mode 100644 index 0000000000..6cabf81c53 --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.mask.cli.po @@ -0,0 +1,331 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-06-28 13:45+0100\n" +"PO-Revision-Date: 2024-06-28 13:48+0100\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.4\n" + +#: tools/mask/cli.py:15 +msgid "" +"This tool allows you to generate, import, export or preview masks for " +"existing alignments." +msgstr "" +"Этот инструмент позволяет создавать, импортировать, экспортировать или " +"просматривать маски для существующих трасс." + +#: tools/mask/cli.py:25 +msgid "" +"Mask tool\n" +"Generate, import, export or preview masks for existing alignments files." +msgstr "" +"Инструмент \"Маска\"\n" +"Создавайте, импортируйте, экспортируйте или просматривайте маски для " +"существующих файлов трасс." + +#: tools/mask/cli.py:35 tools/mask/cli.py:47 tools/mask/cli.py:58 +#: tools/mask/cli.py:69 +msgid "data" +msgstr "данные" + +#: tools/mask/cli.py:39 +msgid "" +"Full path to the alignments file that contains the masks if not at the " +"default location. NB: If the input-type is faces and you wish to update the " +"corresponding alignments file, then you must provide a value here as the " +"location cannot be automatically detected." +msgstr "" +"Полный путь к файлу выравниваний для добавления маски, если он не находится " +"в месте по умолчанию. Примечание: Если input-type - лица, и вы хотите " +"обновить соответствующий файл выравнивания, то вы должны указать значение " +"здесь, так как местоположение не может быть определено автоматически." + +#: tools/mask/cli.py:51 +msgid "Directory containing extracted faces, source frames, or a video file." +msgstr "Папка, содержащая извлеченные лица, исходные кадры или видеофайл." + +#: tools/mask/cli.py:61 +msgid "" +"R|Whether the `input` is a folder of faces or a folder frames/video\n" +"L|faces: The input is a folder containing extracted faces.\n" +"L|frames: The input is a folder containing frames or is a video" +msgstr "" +"R|Выбирается ли \"вход\" как папка лиц или как папка кадров/видео\n" +"L|faces: Входом является папка, содержащая извлеченные лица.\n" +"L|frames: Входом является папка с кадрами или видео" + +#: tools/mask/cli.py:71 +msgid "" +"R|Run the mask tool on multiple sources. If selected then the other options " +"should be set as follows:\n" +"L|input: A parent folder containing either all of the video files to be " +"processed, or containing sub-folders of frames/faces.\n" +"L|output-folder: If provided, then sub-folders will be created within the " +"given location to hold the previews for each input.\n" +"L|alignments: Alignments field will be ignored for batch processing. The " +"alignments files must exist at the default location (for frames). For batch " +"processing of masks with 'faces' as the input type, then only the PNG header " +"within the extracted faces will be updated." +msgstr "" +"R|Запустить инструмент маски на нескольких источниках. Если выбрано, то " +"остальные параметры должны быть установлены следующим образом:\n" +"L|input: Родительская папка, содержащая либо все видеофайлы для обработки, " +"либо содержащая вложенные папки кадров/лиц.\n" +"L|output-folder: Если указано, то в заданном месте будут созданы вложенные " +"папки для хранения превью для каждого входа.\n" +"L|alignments: Поле выравнивания будет игнорироваться при пакетной обработке. " +"Файлы выравнивания должны существовать в месте по умолчанию (для кадров). " +"При пакетной обработке масок с типом входа \"лица\" будут обновлены только " +"заголовки PNG в извлеченных лицах." + +#: tools/mask/cli.py:87 tools/mask/cli.py:119 +msgid "process" +msgstr "обработка" + +#: tools/mask/cli.py:89 +msgid "" +"R|Masker to use.\n" +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members. Profile faces " +"may result in sub-par performance." +msgstr "" +"R|Маскер для использования.\n" +"L|bisenet-fp: Относительно легкая маска на основе NN, которая обеспечивает " +"более точный контроль над маскируемой областью, включая полное маскирование " +"головы (настраивается в настройках маски).\n" +"L|components: Маска, разработанная для сегментации лица на основе " +"расположения ориентиров. Для создания маски вокруг внешних ориентиров " +"строится выпуклая оболочка.\n" +"L|custom (пользовательская): Фиктивная маска, которая заполняет область " +"маски всеми 1 или 0 (настраивается в настройках). Она необходима только в " +"том случае, если вы собираетесь вручную редактировать пользовательские маски " +"в ручном инструменте. Эта маска не использует GPU.\n" +"L|extended: Маска предназначена для сегментации лица на основе расположения " +"ориентиров. Выпуклая оболочка строится вокруг внешних ориентиров, и маска " +"расширяется вверх на лоб.\n" +"L|vgg-clear: Маска предназначена для интеллектуальной сегментации " +"преимущественно фронтальных лиц без препятствий. Профильные лица и " +"препятствия могут привести к снижению производительности.\n" +"L|vgg-obstructed: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была специально обучена " +"распознавать некоторые препятствия на лице (руки и очки). Лица в профиль " +"могут иметь низкую производительность.\n" +"L|unet-dfl: Маска, разработанная для интеллектуальной сегментации " +"преимущественно фронтальных лиц. Модель маски была обучена членами " +"сообщества и для дальнейшего описания нуждается в тестировании. Профильные " +"лица могут иметь низкую производительность." + +#: tools/mask/cli.py:121 +msgid "" +"R|The Mask tool process to perform.\n" +"L|all: Update the mask for all faces in the alignments file for the selected " +"'masker'.\n" +"L|missing: Create a mask for all faces in the alignments file where a mask " +"does not previously exist for the selected 'masker'.\n" +"L|output: Don't update the masks, just output the selected 'masker' for " +"review/editing in external tools to the given output folder.\n" +"L|import: Import masks that have been edited outside of faceswap into the " +"alignments file. Note: 'custom' must be the selected 'masker' and the masks " +"must be in the same format as the 'input-type' (frames or faces)" +msgstr "" +"R|El proceso de la herramienta Máscara a realizar.\n" +"L|all: actualiza la máscara de todas las caras en el archivo de alineaciones " +"para el 'masker' seleccionado.\n" +"L|missing: crea una máscara para todas las caras en el archivo de " +"alineaciones donde no existe previamente una máscara para el 'masker' " +"seleccionado.\n" +"L|output: no actualice las máscaras, simplemente envíe el 'masker' " +"seleccionado para su revisión/edición en herramientas externas a la carpeta " +"de salida proporcionada.\n" +"L|import: importa máscaras que se han editado fuera de faceswap al archivo " +"de alineaciones. Nota: 'custom' debe ser el 'masker' seleccionado y las " +"máscaras deben tener el mismo formato que el 'input-type' (frames o faces)" + +#: tools/mask/cli.py:135 tools/mask/cli.py:154 tools/mask/cli.py:176 +msgid "import" +msgstr "Импортировать" + +#: tools/mask/cli.py:137 +msgid "" +"R|Import only. The path to the folder that contains masks to be imported.\n" +"L|How the masks are provided is not important, but they will be stored, " +"internally, as 8-bit grayscale images.\n" +"L|If the input are images, then the masks must be named exactly the same as " +"input frames/faces (excluding the file extension).\n" +"L|If the input is a video file, then the filename of the masks is not " +"important but should contain the frame number at the end of the filename " +"(but before the file extension). The frame number can be separated from the " +"rest of the filename by any non-numeric character and can be padded by any " +"number of zeros. The frame number must correspond correctly to the frame " +"number in the original video (starting from frame 1)." +msgstr "" +"R|Только импорт. Путь к папке, содержащей маски для импорта.\n" +"L|Как предоставляются маски, не важно, но они будут храниться внутри как 8-" +"битные изображения в оттенках серого.\n" +"L|Если входными данными являются изображения, то имена масок должны быть " +"точно такими же, как у входных кадров/лиц (за исключением расширения " +"файла).\n" +"L|Если входной файл представляет собой видеофайл, то имя файла масок не " +"важно, но должно содержать номер кадра в конце имени файла (но перед " +"расширением файла). Номер кадра может быть отделен от остальной части имени " +"файла любым нечисловым символом и дополнен любым количеством нулей. Номер " +"кадра должен правильно соответствовать номеру кадра в исходном видео " +"(начиная с кадра 1)." + +#: tools/mask/cli.py:156 +msgid "" +"R|Import/Output only. When importing masks, this is the centering to use. " +"For output this is only used for outputting custom imported masks, and " +"should correspond to the centering used when importing the mask. Note: For " +"any job other than 'import' and 'output' this option is ignored as mask " +"centering is handled internally.\n" +"L|face: Centers the mask on the center of the face, adjusting for pitch and " +"yaw. Outside of requirements for full head masking/training, this is likely " +"to be the best choice.\n" +"L|head: Centers the mask on the center of the head, adjusting for pitch and " +"yaw. Note: You should only select head centering if you intend to include " +"the full head (including hair) within the mask and are looking to train a " +"full head model.\n" +"L|legacy: The 'original' extraction technique. Centers the mask near the of " +"the nose with and crops closely to the face. Can result in the edges of the " +"mask appearing outside of the training area." +msgstr "" +"R|Только импорт/вывод. При импорте масок это центрирование для " +"использования. Для вывода это используется только для вывода " +"пользовательских импортированных масок и должно соответствовать " +"центрированию, используемому при импорте маски. Примечание: для любого " +"задания, кроме «импорта» и «вывода», эта опция игнорируется, поскольку " +"центрирование маски обрабатывается внутренне.\n" +"L|face: центрирует маску по центру лица с регулировкой угла наклона и " +"отклонения от курса. Помимо требований к полной маскировке/тренировке " +"головы, это, вероятно, будет лучшим выбором.\n" +"L|head: центрирует маску по центру головы с регулировкой угла наклона и " +"отклонения от курса. Примечание. Выбирать центрирование головы следует " +"только в том случае, если вы собираетесь включить в маску всю голову " +"(включая волосы) и хотите обучить модель полной головы.\n" +"L|legacy: «Оригинальная» техника извлечения. Центрирует маску возле носа и " +"приближает ее к лицу. Это может привести к тому, что края маски окажутся за " +"пределами тренировочной зоны." + +#: tools/mask/cli.py:181 +msgid "" +"Import only. The size, in pixels to internally store the mask at.\n" +"The default is 128 which is fine for nearly all usecases. Larger sizes will " +"result in larger alignments files and longer processing." +msgstr "" +"Только импорт. Размер в пикселях для внутреннего хранения маски.\n" +"Значение по умолчанию — 128, что подходит практически для всех случаев " +"использования. Большие размеры приведут к увеличению размера файлов " +"выравниваний и более длительной обработке." + +#: tools/mask/cli.py:189 tools/mask/cli.py:197 tools/mask/cli.py:211 +#: tools/mask/cli.py:225 tools/mask/cli.py:235 +msgid "output" +msgstr "вывод" + +#: tools/mask/cli.py:191 +msgid "" +"Optional output location. If provided, a preview of the masks created will " +"be output in the given folder." +msgstr "" +"Необязательное местоположение вывода. Если указано, предварительный просмотр " +"созданных масок будет выведен в указанную папку." + +#: tools/mask/cli.py:202 +msgid "" +"Apply gaussian blur to the mask output. Has the effect of smoothing the " +"edges of the mask giving less of a hard edge. the size is in pixels. This " +"value should be odd, if an even number is passed in then it will be rounded " +"to the next odd number. NB: Only effects the output preview. Set to 0 for off" +msgstr "" +"Применяет гауссово размытие к выходу маски. Сглаживает края маски, делая их " +"менее жесткими. размер в пикселях. Это значение должно быть нечетным, если " +"передано четное число, то оно будет округлено до следующего нечетного числа. " +"Примечание: влияет только на предварительный просмотр. Установите значение 0 " +"для выключения" + +#: tools/mask/cli.py:216 +msgid "" +"Helps reduce 'blotchiness' on some masks by making light shades white and " +"dark shades black. Higher values will impact more of the mask. NB: Only " +"effects the output preview. Set to 0 for off" +msgstr "" +"Помогает уменьшить \"пятнистость\" на некоторых масках, делая светлые " +"оттенки белыми, а темные - черными. Более высокие значения влияют на большую " +"часть маски. Примечание: влияет только на предварительный просмотр. " +"Установите значение 0 для выключения" + +#: tools/mask/cli.py:227 +msgid "" +"R|How to format the output when processing is set to 'output'.\n" +"L|combined: The image contains the face/frame, face mask and masked face.\n" +"L|masked: Output the face/frame as rgba image with the face masked.\n" +"L|mask: Only output the mask as a single channel image." +msgstr "" +"R|Как форматировать вывод, когда обработка установлена на 'output'.\n" +"L|combined: Изображение содержит лицо/кадр, маску лица и маскированное " +"лицо.\n" +"L|masked: Вывести лицо/кадр как изображение rgba с маскированным лицом.\n" +"L|mask: Выводить только маску как одноканальное изображение." + +#: tools/mask/cli.py:237 +msgid "" +"R|Whether to output the whole frame or only the face box when using output " +"processing. Only has an effect when using frames as input." +msgstr "" +"R|Выводить ли весь кадр или только поле лица при использовании выходной " +"обработки. Имеет значение только при использовании кадров в качестве входных " +"данных." + +#~ msgid "" +#~ "R|Whether to update all masks in the alignments files, only those faces " +#~ "that do not already have a mask of the given `mask type` or just to " +#~ "output the masks to the `output` location.\n" +#~ "L|all: Update the mask for all faces in the alignments file.\n" +#~ "L|missing: Create a mask for all faces in the alignments file where a " +#~ "mask does not previously exist.\n" +#~ "L|output: Don't update the masks, just output them for review in the " +#~ "given output folder." +#~ msgstr "" +#~ "R|Обновлять ли все маски в файлах выравнивания, только те лица, которые " +#~ "еще не имеют маски заданного `mask type` или просто выводить маски в " +#~ "место `output`.\n" +#~ "L|all: Обновить маску для всех лиц в файле выравнивания.\n" +#~ "L|missing: Создать маску для всех лиц в файле выравнивания, для которых " +#~ "маска ранее не существовала.\n" +#~ "L|output: Не обновлять маски, а просто вывести их для просмотра в " +#~ "указанную выходную папку." diff --git a/locales/ru/LC_MESSAGES/tools.model.cli.mo b/locales/ru/LC_MESSAGES/tools.model.cli.mo new file mode 100644 index 0000000000..37b7545821 Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.model.cli.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.model.cli.po b/locales/ru/LC_MESSAGES/tools.model.cli.po new file mode 100644 index 0000000000..bef71ab233 --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.model.cli.po @@ -0,0 +1,92 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:51+0000\n" +"PO-Revision-Date: 2024-03-29 00:07+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/model/cli.py:13 +msgid "This tool lets you perform actions on saved Faceswap models." +msgstr "" +"Этот инструмент позволяет выполнять действия над сохраненными моделями " +"Faceswap." + +#: tools/model/cli.py:22 +msgid "A tool for performing actions on Faceswap trained model files" +msgstr "" +"Инструмент для выполнения действий над файлами обученных моделей Faceswap" + +#: tools/model/cli.py:34 +msgid "" +"Model directory. A directory containing the model you wish to perform an " +"action on." +msgstr "" +"Папка модели. Папка, содержащая модель, над которой вы хотите выполнить " +"действие." + +#: tools/model/cli.py:43 +msgid "" +"R|Choose which action you want to perform.\n" +"L|'inference' - Create an inference only copy of the model. Strips any " +"layers from the model which are only required for training. NB: This is for " +"exporting the model for use in external applications. Inference generated " +"models cannot be used within Faceswap. See the 'format' option for " +"specifying the model output format.\n" +"L|'nan-scan' - Scan the model file for NaNs or Infs (invalid data).\n" +"L|'restore' - Restore a model from backup." +msgstr "" +"R|Выберите действие, которое вы хотите выполнить.\n" +"L|'inference' - Создать копию модели только для проведения расчетов. Удаляет " +"из модели все слои, которые нужны только для обучения. Примечание: Эта " +"функция предназначена для экспорта модели для использования во внешних " +"приложениях. Модели, созданные в режиме вывода, не могут быть использованы в " +"Faceswap. См. опцию 'format' для указания формата вывода модели.\n" +"L|'nan-scan' - Проверить файл модели на наличие NaNs или Infs (недопустимых " +"данных).\n" +"L|'restore' - Восстановить модель из резервной копии." + +#: tools/model/cli.py:57 tools/model/cli.py:69 +msgid "inference" +msgstr "вывод" + +#: tools/model/cli.py:59 +msgid "" +"R|The format to save the model as. Note: Only used for 'inference' job.\n" +"L|'h5' - Standard Keras H5 format. Does not store any custom layer " +"information. Layers will need to be loaded from Faceswap to use.\n" +"L|'saved-model' - Tensorflow's Saved Model format. Contains all information " +"required to load the model outside of Faceswap." +msgstr "" +"R|Формат для сохранения модели. Примечание: Используется только для задания " +"'inference'.\n" +"L||'h5' - Стандартный формат Keras H5. Не хранит никакой информации о " +"пользовательских слоях. Для использования слои должны быть загружены из " +"Faceswap.\n" +"L|'saved-model' - формат сохраненной модели Tensorflow. Содержит всю " +"информацию, необходимую для загрузки модели вне Faceswap." + +#: tools/model/cli.py:71 +#, fuzzy +#| msgid "" +#| "Only used for 'inference' job. Generate the inference model for B -> A " +#| "instead of A -> B." +msgid "" +"Only used for 'inference' job. Generate the inference model for B -> A " +"instead of A -> B." +msgstr "" +"Используется только для задания 'inference'. Создайте модель вывода для B -> " +"A вместо A -> B." diff --git a/locales/ru/LC_MESSAGES/tools.preview.mo b/locales/ru/LC_MESSAGES/tools.preview.mo new file mode 100644 index 0000000000..780e7173eb Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.preview.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.preview.po b/locales/ru/LC_MESSAGES/tools.preview.po new file mode 100644 index 0000000000..ebcaea18d8 --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.preview.po @@ -0,0 +1,93 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:06+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/preview/cli.py:15 +msgid "This command allows you to preview swaps to tweak convert settings." +msgstr "" +"Эта команда позволяет просматривать замены для настройки параметров " +"конвертирования." + +#: tools/preview/cli.py:30 +msgid "" +"Preview tool\n" +"Allows you to configure your convert settings with a live preview" +msgstr "" +"Инструмент предпросмотра\n" +"Позволяет настраивать параметры конвертации с помощью предварительного " +"просмотра в реальном времени" + +#: tools/preview/cli.py:47 tools/preview/cli.py:57 tools/preview/cli.py:65 +msgid "data" +msgstr "данные" + +#: tools/preview/cli.py:50 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file." +msgstr "" +"Входная папка или видео. Либо папка, содержащая файлы изображений, которые " +"необходимо обработать, либо путь к видеофайлу." + +#: tools/preview/cli.py:60 +msgid "" +"Path to the alignments file for the input, if not at the default location" +msgstr "" +"Путь к файлу выравниваний для входных данных, если он не находится в месте " +"по умолчанию" + +#: tools/preview/cli.py:68 +msgid "" +"Model directory. A directory containing the trained model you wish to " +"process." +msgstr "" +"Папка модели. Папка, содержащая обученную модель, которую вы хотите " +"обработать." + +#: tools/preview/cli.py:74 +msgid "Swap the model. Instead of A -> B, swap B -> A" +msgstr "Поменять местами модели. Вместо A -> B заменить B -> A" + +#: tools/preview/control_panels.py:510 +msgid "Save full config" +msgstr "Сохранить полную конфигурацию" + +#: tools/preview/control_panels.py:513 +msgid "Reset full config to default values" +msgstr "Сбросить полную конфигурацию до заводских значений" + +#: tools/preview/control_panels.py:516 +msgid "Reset full config to saved values" +msgstr "Сбросить полную конфигурацию до сохраненных значений" + +#: tools/preview/control_panels.py:667 +#, python-brace-format +msgid "Save {title} config" +msgstr "Сохранить конфигурацию {title}" + +#: tools/preview/control_panels.py:670 +#, python-brace-format +msgid "Reset {title} config to default values" +msgstr "Сбросить полную конфигурацию {title} до заводских значений" + +#: tools/preview/control_panels.py:673 +#, python-brace-format +msgid "Reset {title} config to saved values" +msgstr "Сбросить полную конфигурацию {title} до сохраненных значений" diff --git a/locales/ru/LC_MESSAGES/tools.sort.cli.mo b/locales/ru/LC_MESSAGES/tools.sort.cli.mo new file mode 100644 index 0000000000..6b832be91e Binary files /dev/null and b/locales/ru/LC_MESSAGES/tools.sort.cli.mo differ diff --git a/locales/ru/LC_MESSAGES/tools.sort.cli.po b/locales/ru/LC_MESSAGES/tools.sort.cli.po new file mode 100644 index 0000000000..9b76d494ae --- /dev/null +++ b/locales/ru/LC_MESSAGES/tools.sort.cli.po @@ -0,0 +1,414 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: \n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" +"PO-Revision-Date: 2024-03-29 00:06+0000\n" +"Last-Translator: \n" +"Language-Team: \n" +"Language: ru\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=(n%10==1 && n%100!=11 ? 0 : n%10>=2 && " +"n%10<=4 && (n%100<12 || n%100>14) ? 1 : 2);\n" +"X-Generator: Poedit 3.4.2\n" + +#: tools/sort/cli.py:15 +msgid "This command lets you sort images using various methods." +msgstr "Эта команда позволяет сортировать изображения различными методами." + +#: tools/sort/cli.py:21 +msgid "" +" Adjust the '-t' ('--threshold') parameter to control the strength of " +"grouping." +msgstr "" +" Настройте параметр '-t' ('--threshold') для контроля силы группировки." + +#: tools/sort/cli.py:22 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the percentage of color pixels " +"that appear in the image." +msgstr "" +" Настройте параметр '-b' ('--bins') для управления количеством корзинок для " +"группировки. Каждое изображение распределяется по корзинкам в зависимости от " +"процента цветных пикселей, присутствующих в изображении." + +#: tools/sort/cli.py:25 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the number of degrees the face " +"is orientated from center." +msgstr "" +" Настройте параметр '-b' ('--bins') для управления количеством корзинок для " +"группировки. Каждое изображение распределяется по корзинам по количеству " +"градусов, на которые лицо ориентировано от центра." + +#: tools/sort/cli.py:28 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. The minimum and maximum values are taken for the chosen sort " +"metric. The bins are then populated with the results from the group sorting." +msgstr "" +" Настройте параметр '-b' ('--bins') для управления количеством корзинок для " +"группировки. Для выбранной метрики сортировки берутся минимальное и " +"максимальное значения. Затем корзины заполняются результатами групповой " +"сортировки." + +#: tools/sort/cli.py:32 +msgid "faces by blurriness." +msgstr "лица по размытости." + +#: tools/sort/cli.py:33 +msgid "faces by fft filtered blurriness." +msgstr "лица по размытости с фильтрацией fft." + +#: tools/sort/cli.py:34 +msgid "" +"faces by the estimated distance of the alignments from an 'average' face. " +"This can be useful for eliminating misaligned faces. Sorts from most like an " +"average face to least like an average face." +msgstr "" +"лица по оценочному расстоянию выравнивания от \"среднего\" лица. Это может " +"быть полезно для устранения неправильно расположенных лиц. Сортирует от " +"наиболее похожего на среднее лицо к наименее похожему на среднее лицо." + +#: tools/sort/cli.py:37 +msgid "" +"faces using VGG Face2 by face similarity. This uses a pairwise clustering " +"algorithm to check the distances between 512 features on every face in your " +"set and order them appropriately." +msgstr "" +"лиц с помощью VGG Face2 по сходству лиц. При этом используется алгоритм " +"парной кластеризации для проверки расстояний между 512 признаками на каждом " +"лице в вашем наборе и их упорядочивания соответствующим образом." + +#: tools/sort/cli.py:40 +msgid "faces by their landmarks." +msgstr "лица по их ориентирам." + +#: tools/sort/cli.py:41 +msgid "Like 'face-cnn' but sorts by dissimilarity." +msgstr "Как 'face-cnn', но сортирует по непохожести." + +#: tools/sort/cli.py:42 +msgid "faces by Yaw (rotation left to right)." +msgstr "лица по Yaw (вращение слева направо)." + +#: tools/sort/cli.py:43 +msgid "faces by Pitch (rotation up and down)." +msgstr "лица по Pitch (вращение вверх и вниз)." + +#: tools/sort/cli.py:44 +msgid "" +"faces by Roll (rotation). Aligned faces should have a roll value close to " +"zero. The further the Roll value from zero the higher liklihood the face is " +"misaligned." +msgstr "" +"грани по Roll (повороту). Выровненные грани должны иметь значение Roll, " +"близкое к нулю. Чем дальше значение Roll от нуля, тем выше вероятность того, " +"что лицо неправильно выровнено." + +#: tools/sort/cli.py:46 +msgid "faces by their color histogram." +msgstr "лица по их цветовой гистограмме." + +#: tools/sort/cli.py:47 +msgid "Like 'hist' but sorts by dissimilarity." +msgstr "Как 'hist', но сортирует по непохожести." + +#: tools/sort/cli.py:48 +msgid "" +"images by the average intensity of the converted grayscale color channel." +msgstr "" +"изображения по средней интенсивности преобразованного полутонового цветового " +"канала." + +#: tools/sort/cli.py:49 +msgid "" +"images by their number of black pixels. Useful when faces are near borders " +"and a large part of the image is black." +msgstr "" +"изображения по количеству черных пикселей. Полезно, когда лица находятся " +"вблизи границ и большая часть изображения черная." + +#: tools/sort/cli.py:51 +msgid "" +"images by the average intensity of the converted Y color channel. Bright " +"lighting and oversaturated images will be ranked first." +msgstr "" +"изображений по средней интенсивности преобразованного цветового канала Y. " +"Яркое освещение и перенасыщенные изображения будут ранжироваться в первую " +"очередь." + +#: tools/sort/cli.py:53 +msgid "" +"images by the average intensity of the converted Cg color channel. Green " +"images will be ranked first and red images will be last." +msgstr "" +"изображений по средней интенсивности преобразованного цветового канала Cg. " +"Зеленые изображения занимают первое место, а красные - последнее." + +#: tools/sort/cli.py:55 +msgid "" +"images by the average intensity of the converted Co color channel. Orange " +"images will be ranked first and blue images will be last." +msgstr "" +"изображений по средней интенсивности преобразованного цветового канала Co. " +"Оранжевые изображения занимают первое место, а синие - последнее." + +#: tools/sort/cli.py:57 +msgid "" +"images by their size in the original frame. Faces further from the camera " +"and from lower resolution sources will be sorted first, whilst faces closer " +"to the camera and from higher resolution sources will be sorted last." +msgstr "" +"изображения по их размеру в исходном кадре. Лица, расположенные дальше от " +"камеры и полученные из источников с низким разрешением, будут отсортированы " +"первыми, а лица, расположенные ближе к камере и полученные из источников с " +"высоким разрешением, будут отсортированы последними." + +#: tools/sort/cli.py:81 +msgid "Sort faces using a number of different techniques" +msgstr "Сортировка лиц с использованием различных методов" + +#: tools/sort/cli.py:91 tools/sort/cli.py:98 tools/sort/cli.py:110 +#: tools/sort/cli.py:150 +msgid "data" +msgstr "данные" + +#: tools/sort/cli.py:92 +msgid "Input directory of aligned faces." +msgstr "Входная папка соотнесенных лиц." + +#: tools/sort/cli.py:100 +msgid "" +"Output directory for sorted aligned faces. If not provided and 'keep' is " +"selected then a new folder called 'sorted' will be created within the input " +"folder to house the output. If not provided and 'keep' is not selected then " +"the images will be sorted in-place, overwriting the original contents of the " +"'input_dir'" +msgstr "" +"Выходная папка для отсортированных выровненных лиц. Если не указано и " +"выбрано 'keep', то в папке input будет создана новая папка под названием " +"'sorted' для размещения выходных данных. Если не указано и не выбрано " +"'keep', то изображения будут отсортированы на месте, перезаписывая исходное " +"содержимое 'input_dir'." + +#: tools/sort/cli.py:112 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple folders of faces you wish to sort. The faces will be output to " +"separate sub-folders in the output_dir" +msgstr "" +"R|Если выбрано, то input_dir должен быть родительской папкой, содержащей " +"несколько папок с лицами, которые вы хотите отсортировать. Лица будут " +"выведены в отдельные вложенные папки в output_dir" + +#: tools/sort/cli.py:121 +msgid "sort settings" +msgstr "настройки сортировки" + +#: tools/sort/cli.py:124 +msgid "" +"R|Choose how images are sorted. Selecting a sort method gives the images a " +"new filename based on the order the image appears within the given method.\n" +"L|'none': Don't sort the images. When a 'group-by' method is selected, " +"selecting 'none' means that the files will be moved/copied into their " +"respective bins, but the files will keep their original filenames. Selecting " +"'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|Выбор способа сортировки изображений. При выборе метода сортировки " +"изображениям присваивается новое имя файла, основанное на порядке появления " +"изображения в данном методе.\n" +"L|'none': Не сортировать изображения. Если выбран метод 'group-by', выбор " +"'none' означает, что файлы будут перемещены/скопированы в соответствующие " +"корзины, но файлы сохранят свои оригинальные имена. Выбор значения 'none' " +"как для 'sort-by', так и для 'group-by' ничего не даст" + +#: tools/sort/cli.py:136 tools/sort/cli.py:164 tools/sort/cli.py:184 +msgid "group settings" +msgstr "настройки группировки" + +#: tools/sort/cli.py:139 +#, fuzzy +#| msgid "" +#| "R|Selecting a group by method will move/copy files into numbered bins " +#| "based on the selected method.\n" +#| "L|'none': Don't bin the images. Folders will be sorted by the selected " +#| "'sort-by' but will not be binned, instead they will be sorted into a " +#| "single folder. Selecting 'none' for both 'sort-by' and 'group-by' will " +#| "do nothing" +msgid "" +"R|Selecting a group by method will move/copy files into numbered bins based " +"on the selected method.\n" +"L|'none': Don't bin the images. Folders will be sorted by the selected 'sort-" +"by' but will not be binned, instead they will be sorted into a single " +"folder. Selecting 'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" +"R|Выбор группы по методу приведет к перемещению/копированию файлов в " +"пронумерованные корзины в соответствии с выбранным методом.\n" +"L|'none': Не сортировать изображения. Папки будут отсортированы по " +"выбранному \"sort-by\", но не будут разбиты на папки, вместо этого они будут " +"отсортированы в одну папку. Выбор значения 'none' как для 'sort-by', так и " +"для 'group-by' ничего не даст" + +#: tools/sort/cli.py:152 +msgid "" +"Whether to keep the original files in their original location. Choosing a " +"'sort-by' method means that the files have to be renamed. Selecting 'keep' " +"means that the original files will be kept, and the renamed files will be " +"created in the specified output folder. Unselecting keep means that the " +"original files will be moved and renamed based on the selected sort/group " +"criteria." +msgstr "" +"Сохранять ли исходные файлы в их первоначальном расположении. Выбор метода " +"\"сортировать по\" означает, что файлы должны быть переименованы. Выбор " +"'keep' означает, что исходные файлы будут сохранены, а переименованные файлы " +"будут созданы в указанной выходной папке. Отмена выбора \"keep\" означает, " +"что исходные файлы будут перемещены и переименованы в соответствии с " +"выбранными критериями сортировки/группировки." + +#: tools/sort/cli.py:167 +msgid "" +"R|Float value. Minimum threshold to use for grouping comparison with 'face-" +"cnn' 'hist' and 'face' methods.\n" +"The lower the value the more discriminating the grouping is. Leaving -1.0 " +"will allow Faceswap to choose the default value.\n" +"L|For 'face-cnn' 7.2 should be enough, with 4 being very discriminating. \n" +"L|For 'hist' 0.3 should be enough, with 0.2 being very discriminating. \n" +"L|For 'face' between 0.1 (more bins) to 0.5 (fewer bins) should be about " +"right.\n" +"Be careful setting a value that's too extrene in a directory with many " +"images, as this could result in a lot of folders being created. Defaults: " +"face-cnn 7.2, hist 0.3, face 0.25" +msgstr "" +"R|Плавающее значение. Минимальный порог, используемый для сравнения " +"группировок с методами 'face-cnn' 'hist' и 'face'.\n" +"Чем меньше значение, тем более дискриминационной является группировка. Если " +"оставить значение -1.0, Faceswap сможет выбрать значение по умолчанию.\n" +"L|Для 'face-cnn' 7,2 должно быть достаточно, при этом 4 будет очень " +"дискриминационным. \n" +"L|Для 'hist' 0.3 должно быть достаточно, при этом 0.2 очень хорошо " +"различает. \n" +"L|For 'face' от 0,1 (больше бинов) до 0,5 (меньше бинов) должно быть " +"достаточно.\n" +"Будьте осторожны, устанавливая слишком большое значение в каталоге с большим " +"количеством изображений, так как это может привести к созданию большого " +"количества папок. По умолчанию: face-cnn 7.2, hist 0.3, face 0.25" + +#: tools/sort/cli.py:187 +#, fuzzy, python-format +#| msgid "" +#| "R|Integer value. Used to control the number of bins created for grouping " +#| "by: any 'blur' methods, 'color' methods or 'face metric' methods " +#| "('distance', 'size') and 'orientation; methods ('yaw', 'pitch'). For any " +#| "other grouping methods see the '-t' ('--threshold') option.\n" +#| "L|For 'face metric' methods the bins are filled, according the the " +#| "distribution of faces between the minimum and maximum chosen metric.\n" +#| "L|For 'color' methods the number of bins represents the divider of the " +#| "percentage of colored pixels. Eg. For a bin number of '5': The first " +#| "folder will have the faces with 0%% to 20%% colored pixels, second 21%% " +#| "to 40%%, etc. Any empty bins will be deleted, so you may end up with " +#| "fewer bins than selected.\n" +#| "L|For 'blur' methods folder 0 will be the least blurry, while the last " +#| "folder will be the blurriest.\n" +#| "L|For 'orientation' methods the number of bins is dictated by how much " +#| "180 degrees is divided. Eg. If 18 is selected, then each folder will be a " +#| "10 degree increment. Folder 0 will contain faces looking the most to the " +#| "left/down whereas the last folder will contain the faces looking the most " +#| "to the right/up. NB: Some bins may be empty if faces do not fit the " +#| "criteria.\n" +#| "Default value: 5" +msgid "" +"R|Integer value. Used to control the number of bins created for grouping by: " +"any 'blur' methods, 'color' methods or 'face metric' methods ('distance', " +"'size') and 'orientation; methods ('yaw', 'pitch'). For any other grouping " +"methods see the '-t' ('--threshold') option.\n" +"L|For 'face metric' methods the bins are filled, according the the " +"distribution of faces between the minimum and maximum chosen metric.\n" +"L|For 'color' methods the number of bins represents the divider of the " +"percentage of colored pixels. Eg. For a bin number of '5': The first folder " +"will have the faces with 0%% to 20%% colored pixels, second 21%% to 40%%, " +"etc. Any empty bins will be deleted, so you may end up with fewer bins than " +"selected.\n" +"L|For 'blur' methods folder 0 will be the least blurry, while the last " +"folder will be the blurriest.\n" +"L|For 'orientation' methods the number of bins is dictated by how much 180 " +"degrees is divided. Eg. If 18 is selected, then each folder will be a 10 " +"degree increment. Folder 0 will contain faces looking the most to the left/" +"down whereas the last folder will contain the faces looking the most to the " +"right/up. NB: Some bins may be empty if faces do not fit the criteria. \n" +"Default value: 5" +msgstr "" +"R| Целочисленное значение. Используется для управления количеством бинов, " +"создаваемых для группировки: любыми методами 'размытия', 'цвета' или " +"методами 'метрики лица' ('расстояние', 'размер') и 'ориентации; методы " +"('yaw', 'pitch'). Для любых других методов группировки смотрите опцию '-" +"t' ('--threshold').\n" +"L|Для методов 'face metric' бины заполняются в соответствии с распределением " +"лиц между минимальной и максимальной выбранной метрикой.\n" +"L|Для методов 'color' количество бинов представляет собой делитель процента " +"цветных пикселей. Например, для числа бинов \"5\": В первой папке будут лица " +"с 0%% - 20%% цветных пикселей, во второй 21%% - 40%% и т.д. Все пустые папки " +"будут удалены, поэтому в итоге у вас может оказаться меньше папок, чем было " +"выбрано.\n" +"L|Для методов 'blur' папка 0 будет наименее размытой, а последняя папка " +"будет самой размытой.\n" +"L|Для методов \"orientation\" количество бинов диктуется тем, на сколько " +"делится 180 градусов. Например, если выбрано 18, то каждая папка будет иметь " +"шаг в 10 градусов. Папка 0 будет содержать лица, направленные больше всего " +"влево/вниз, а последняя папка будет содержать лица, направленные больше " +"всего вправо/вверх. Примечание: Некоторые папки могут быть пустыми, если " +"лица не соответствуют критериям.\n" +"Значение по умолчанию: 5" + +#: tools/sort/cli.py:207 tools/sort/cli.py:217 +msgid "settings" +msgstr "настройки" + +#: tools/sort/cli.py:210 +msgid "" +"Logs file renaming changes if grouping by renaming, or it logs the file " +"copying/movement if grouping by folders. If no log file is specified with " +"'--log-file', then a 'sort_log.json' file will be created in the input " +"directory." +msgstr "" +"Ведет журнал изменений переименования файлов при группировке по " +"переименованию, или журнал копирования/перемещения файлов при группировке по " +"папкам. Если файл журнала не указан с помощью '--log-file', то в каталоге " +"ввода будет создан файл 'sort_log.json'." + +#: tools/sort/cli.py:221 +msgid "" +"Specify a log file to use for saving the renaming or grouping information. " +"If specified extension isn't 'json' or 'yaml', then json will be used as the " +"serializer, with the supplied filename. Default: sort_log.json" +msgstr "" +"Укажите файл журнала, который будет использоваться для сохранения информации " +"о переименовании или группировке. Если указанное расширение не 'json' или " +"'yaml', то в качестве сериализатора будет использоваться json, с указанным " +"именем файла. По умолчанию: sort_log.json" + +#~ msgid " option is deprecated. Use 'yaw'" +#~ msgstr " является устаревшей. Используйте 'yaw'" + +#~ msgid " option is deprecated. Use 'color-black'" +#~ msgstr " является устаревшей. Используйте 'color-black'" + +#~ msgid "output" +#~ msgstr "вывод" + +#~ msgid "" +#~ "Deprecated and no longer used. The final processing will be dictated by " +#~ "the sort/group by methods and whether 'keep_original' is selected." +#~ msgstr "" +#~ "Устарело и больше не используется. Окончательная обработка будет " +#~ "диктоваться методами sort/group by и тем, выбрана ли опция " +#~ "'keep_original'." diff --git a/locales/tools.alignments.cli.pot b/locales/tools.alignments.cli.pot index 985132cecd..4f1e02ae15 100644 --- a/locales/tools.alignments.cli.pot +++ b/locales/tools.alignments.cli.pot @@ -8,7 +8,7 @@ msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" "Report-Msgid-Bugs-To: \n" -"POT-Creation-Date: 2022-05-24 12:38+0100\n" +"POT-Creation-Date: 2024-04-19 11:28+0100\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" @@ -17,64 +17,70 @@ msgstr "" "Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -#: tools/alignments/cli.py:15 +#: tools/alignments/cli.py:16 msgid "" "This command lets you perform various tasks pertaining to an alignments file." msgstr "" -#: tools/alignments/cli.py:30 +#: tools/alignments/cli.py:31 msgid "" "Alignments tool\n" "This tool allows you to perform numerous actions on or using an alignments " "file against its corresponding faceset/frame source." msgstr "" -#: tools/alignments/cli.py:41 -msgid " Must Pass in a frames folder/source video file (-fr)." +#: tools/alignments/cli.py:43 +msgid " Must Pass in a frames folder/source video file (-r)." msgstr "" -#: tools/alignments/cli.py:42 -msgid " Must Pass in a faces folder (-fc)." +#: tools/alignments/cli.py:44 +msgid " Must Pass in a faces folder (-c)." msgstr "" -#: tools/alignments/cli.py:43 +#: tools/alignments/cli.py:45 msgid "" -" Must Pass in either a frames folder/source video file OR afaces folder (-fr " -"or -fc)." +" Must Pass in either a frames folder/source video file OR a faces folder (-r " +"or -c)." msgstr "" -#: tools/alignments/cli.py:45 +#: tools/alignments/cli.py:47 msgid "" -" Must Pass in a frames folder/source video file AND a faces folder (-fr and -" -"fc)." +" Must Pass in a frames folder/source video file AND a faces folder (-r and -" +"c)." msgstr "" -#: tools/alignments/cli.py:47 +#: tools/alignments/cli.py:49 msgid " Use the output option (-o) to process results." msgstr "" -#: tools/alignments/cli.py:55 tools/alignments/cli.py:94 +#: tools/alignments/cli.py:58 tools/alignments/cli.py:104 msgid "processing" msgstr "" -#: tools/alignments/cli.py:57 +#: tools/alignments/cli.py:61 #, python-brace-format msgid "" "R|Choose which action you want to perform. NB: All actions require an " "alignments file (-a) to be passed in.\n" "L|'draw': Draw landmarks on frames in the selected folder/video. A subfolder " "will be created within the frames folder to hold the output.{0}\n" +"L|'export': Export the contents of an alignments file to a json file. Can be " +"used for editing alignment information in external tools and then re-" +"importing by using Faceswap's Extract 'Import' plugins. Note: masks and " +"identity vectors will not be included in the exported file, so will be re-" +"generated when the json file is imported back into Faceswap. All data is " +"exported with the origin (0, 0) at the top left of the canvas.\n" "L|'extract': Re-extract faces from the source frames/video based on " "alignment data. This is a lot quicker than re-detecting faces. Can pass in " "the '-een' (--extract-every-n) parameter to only extract every nth frame." "{1}\n" "L|'from-faces': Generate alignment file(s) from a folder of extracted faces. " "if the folder of faces comes from multiple sources, then multiple alignments " -"files will be created. NB: for faces which have been extracted folders of " -"source images, rather than a video, a single alignments file will be created " -"as there is no way for the process to know how many folders of images were " -"originally used. You do not need to provide an alignments file path to run " -"this job. {3}\n" +"files will be created. NB: for faces which have been extracted from folders " +"of source images, rather than a video, a single alignments file will be " +"created as there is no way for the process to know how many folders of " +"images were originally used. You do not need to provide an alignments file " +"path to run this job. {3}\n" "L|'missing-alignments': Identify frames that do not exist in the alignments " "file.{2}{0}\n" "L|'missing-frames': Identify frames in the alignments file that do not " @@ -94,7 +100,7 @@ msgid "" "(EXPERIMENTAL!)" msgstr "" -#: tools/alignments/cli.py:96 +#: tools/alignments/cli.py:107 msgid "" "R|How to output discovered items ('faces' and 'frames' only):\n" "L|'console': Print the list of frames to the screen. (DEFAULT)\n" @@ -104,43 +110,63 @@ msgid "" "directory." msgstr "" -#: tools/alignments/cli.py:107 tools/alignments/cli.py:118 -#: tools/alignments/cli.py:125 +#: tools/alignments/cli.py:118 tools/alignments/cli.py:141 +#: tools/alignments/cli.py:148 msgid "data" msgstr "" -#: tools/alignments/cli.py:111 +#: tools/alignments/cli.py:125 msgid "" -"Full path to the alignments file to be processed. This is required for all " -"jobs except for 'from-faces' when the alignments file will be generated in " -"the specified faces folder." +"Full path to the alignments file to be processed. If you have input a " +"'frames_dir' and don't provide this option, the process will try to find the " +"alignments file at the default location. All jobs require an alignments file " +"with the exception of 'from-faces' when the alignments file will be " +"generated in the specified faces folder." msgstr "" -#: tools/alignments/cli.py:119 -msgid "Directory containing extracted faces." -msgstr "" - -#: tools/alignments/cli.py:126 +#: tools/alignments/cli.py:142 msgid "Directory containing source frames that faces were extracted from." msgstr "" -#: tools/alignments/cli.py:135 tools/alignments/cli.py:146 -#: tools/alignments/cli.py:156 +#: tools/alignments/cli.py:150 +msgid "" +"R|Run the aligmnents tool on multiple sources. The following jobs support " +"batch mode:\n" +"L|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " +"sort, spatial.\n" +"If batch mode is selected then the other options should be set as follows:\n" +"L|alignments_file: For 'sort' and 'spatial' this should point to the parent " +"folder containing the alignments files to be processed. For all other jobs " +"this option is ignored, and the alignments files must exist at their default " +"location relative to the original frames folder/video.\n" +"L|faces_dir: For 'from-faces' this should be a parent folder, containing sub-" +"folders of extracted faces from which to generate alignments files. For " +"'extract' this should be a parent folder where sub-folders will be created " +"for each extraction to be run. For all other jobs this option is ignored.\n" +"L|frames_dir: For 'draw', 'extract', 'missing-alignments', 'missing-frames' " +"and 'no-faces' this should be a parent folder containing video files or sub-" +"folders of images to perform the alignments job on. The alignments file " +"should exist at the default location. For all other jobs this option is " +"ignored." +msgstr "" + +#: tools/alignments/cli.py:176 tools/alignments/cli.py:188 +#: tools/alignments/cli.py:198 msgid "extract" msgstr "" -#: tools/alignments/cli.py:136 +#: tools/alignments/cli.py:178 msgid "" "[Extract only] Extract every 'nth' frame. This option will skip frames when " "extracting faces. For example a value of 1 will extract faces from every " "frame, a value of 10 will extract faces from every 10th frame." msgstr "" -#: tools/alignments/cli.py:147 +#: tools/alignments/cli.py:189 msgid "[Extract only] The output size of extracted faces." msgstr "" -#: tools/alignments/cli.py:157 +#: tools/alignments/cli.py:200 msgid "" "[Extract only] Only extract faces that have been resized by this percent or " "more to meet the specified extract size (`-sz`, `--size`). Useful for " diff --git a/locales/tools.effmpeg.cli.pot b/locales/tools.effmpeg.cli.pot index 83c4ac8f05..72ab831efa 100644 --- a/locales/tools.effmpeg.cli.pot +++ b/locales/tools.effmpeg.cli.pot @@ -1,29 +1,31 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-02-18 23:34-0000\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:50+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - #: tools/effmpeg/cli.py:15 msgid "This command allows you to easily execute common ffmpeg tasks." msgstr "" -#: tools/effmpeg/cli.py:24 +#: tools/effmpeg/cli.py:52 msgid "A wrapper for ffmpeg for performing image <> video converting." msgstr "" -#: tools/effmpeg/cli.py:51 +#: tools/effmpeg/cli.py:64 msgid "" "R|Choose which action you want ffmpeg ffmpeg to do.\n" "L|'extract': turns videos into images \n" @@ -36,80 +38,110 @@ msgid "" "L|'slice' cuts a portion of the video into a separate video file." msgstr "" -#: tools/effmpeg/cli.py:65 +#: tools/effmpeg/cli.py:78 msgid "Input file." msgstr "" -#: tools/effmpeg/cli.py:66 tools/effmpeg/cli.py:73 tools/effmpeg/cli.py:87 +#: tools/effmpeg/cli.py:79 tools/effmpeg/cli.py:86 tools/effmpeg/cli.py:100 msgid "data" msgstr "" -#: tools/effmpeg/cli.py:76 -msgid "Output file. If no output is specified then: if the output is meant to be a video then a video called 'out.mkv' will be created in the input directory; if the output is meant to be a directory then a directory called 'out' will be created inside the input directory. Note: the chosen output file extension will determine the file encoding." +#: tools/effmpeg/cli.py:89 +msgid "" +"Output file. If no output is specified then: if the output is meant to be a " +"video then a video called 'out.mkv' will be created in the input directory; " +"if the output is meant to be a directory then a directory called 'out' will " +"be created inside the input directory. Note: the chosen output file " +"extension will determine the file encoding." msgstr "" -#: tools/effmpeg/cli.py:89 +#: tools/effmpeg/cli.py:102 msgid "Path to reference video if 'input' was not a video." msgstr "" -#: tools/effmpeg/cli.py:95 tools/effmpeg/cli.py:105 tools/effmpeg/cli.py:142 -#: tools/effmpeg/cli.py:171 +#: tools/effmpeg/cli.py:108 tools/effmpeg/cli.py:118 tools/effmpeg/cli.py:156 +#: tools/effmpeg/cli.py:185 msgid "output" msgstr "" -#: tools/effmpeg/cli.py:97 -msgid "Provide video fps. Can be an integer, float or fraction. Negative values will will make the program try to get the fps from the input or reference videos." +#: tools/effmpeg/cli.py:110 +msgid "" +"Provide video fps. Can be an integer, float or fraction. Negative values " +"will will make the program try to get the fps from the input or reference " +"videos." msgstr "" -#: tools/effmpeg/cli.py:107 -msgid "Image format that extracted images should be saved as. '.bmp' will offer the fastest extraction speed, but will take the most storage space. '.png' will be slower but will take less storage." +#: tools/effmpeg/cli.py:120 +msgid "" +"Image format that extracted images should be saved as. '.bmp' will offer the " +"fastest extraction speed, but will take the most storage space. '.png' will " +"be slower but will take less storage." msgstr "" -#: tools/effmpeg/cli.py:114 tools/effmpeg/cli.py:123 tools/effmpeg/cli.py:132 +#: tools/effmpeg/cli.py:127 tools/effmpeg/cli.py:136 tools/effmpeg/cli.py:145 msgid "clip" msgstr "" -#: tools/effmpeg/cli.py:116 -msgid "Enter the start time from which an action is to be applied. Default: 00:00:00, in HH:MM:SS format. You can also enter the time with or without the colons, e.g. 00:0000 or 026010." +#: tools/effmpeg/cli.py:129 +msgid "" +"Enter the start time from which an action is to be applied. Default: " +"00:00:00, in HH:MM:SS format. You can also enter the time with or without " +"the colons, e.g. 00:0000 or 026010." msgstr "" -#: tools/effmpeg/cli.py:125 -msgid "Enter the end time to which an action is to be applied. If both an end time and duration are set, then the end time will be used and the duration will be ignored. Default: 00:00:00, in HH:MM:SS." +#: tools/effmpeg/cli.py:138 +msgid "" +"Enter the end time to which an action is to be applied. If both an end time " +"and duration are set, then the end time will be used and the duration will " +"be ignored. Default: 00:00:00, in HH:MM:SS." msgstr "" -#: tools/effmpeg/cli.py:134 -msgid "Enter the duration of the chosen action, for example if you enter 00:00:10 for slice, then the first 10 seconds after and including the start time will be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. You can also enter the time with or without the colons, e.g. 00:0000 or 026010." +#: tools/effmpeg/cli.py:147 +msgid "" +"Enter the duration of the chosen action, for example if you enter 00:00:10 " +"for slice, then the first 10 seconds after and including the start time will " +"be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. You can " +"also enter the time with or without the colons, e.g. 00:0000 or 026010." msgstr "" -#: tools/effmpeg/cli.py:144 -msgid "Mux the audio from the reference video into the input video. This option is only used for the 'gen-vid' action. 'mux-audio' action has this turned on implicitly." +#: tools/effmpeg/cli.py:158 +msgid "" +"Mux the audio from the reference video into the input video. This option is " +"only used for the 'gen-vid' action. 'mux-audio' action has this turned on " +"implicitly." msgstr "" -#: tools/effmpeg/cli.py:155 tools/effmpeg/cli.py:165 +#: tools/effmpeg/cli.py:169 tools/effmpeg/cli.py:179 msgid "rotate" msgstr "" -#: tools/effmpeg/cli.py:157 -msgid "Transpose the video. If transpose is set, then degrees will be ignored. For cli you can enter either the number or the long command name, e.g. to use (1, 90Clockwise) -tr 1 or -tr 90Clockwise" +#: tools/effmpeg/cli.py:171 +msgid "" +"Transpose the video. If transpose is set, then degrees will be ignored. For " +"cli you can enter either the number or the long command name, e.g. to use " +"(1, 90Clockwise) -tr 1 or -tr 90Clockwise" msgstr "" -#: tools/effmpeg/cli.py:166 +#: tools/effmpeg/cli.py:180 msgid "Rotate the video clockwise by the given number of degrees." msgstr "" -#: tools/effmpeg/cli.py:173 +#: tools/effmpeg/cli.py:187 msgid "Set the new resolution scale if the chosen action is 'rescale'." msgstr "" -#: tools/effmpeg/cli.py:178 tools/effmpeg/cli.py:186 +#: tools/effmpeg/cli.py:192 tools/effmpeg/cli.py:200 msgid "settings" msgstr "" -#: tools/effmpeg/cli.py:180 -msgid "Reduces output verbosity so that only serious errors are printed. If both quiet and verbose are set, verbose will override quiet." +#: tools/effmpeg/cli.py:194 +msgid "" +"Reduces output verbosity so that only serious errors are printed. If both " +"quiet and verbose are set, verbose will override quiet." msgstr "" -#: tools/effmpeg/cli.py:188 -msgid "Increases output verbosity. If both quiet and verbose are set, verbose will override quiet." +#: tools/effmpeg/cli.py:202 +msgid "" +"Increases output verbosity. If both quiet and verbose are set, verbose will " +"override quiet." msgstr "" - diff --git a/locales/tools.manual.pot b/locales/tools.manual.pot index a6cff24f13..4e3fe2e9ab 100644 --- a/locales/tools.manual.pot +++ b/locales/tools.manual.pot @@ -1,50 +1,65 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-06-08 19:24+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:55+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - -#: tools/manual\cli.py:13 -msgid "This command lets you perform various actions on frames, faces and alignments files using visual tools." +#: tools/manual/cli.py:13 +msgid "" +"This command lets you perform various actions on frames, faces and " +"alignments files using visual tools." msgstr "" -#: tools/manual\cli.py:23 -msgid "A tool to perform various actions on frames, faces and alignments files using visual tools" +#: tools/manual/cli.py:23 +msgid "" +"A tool to perform various actions on frames, faces and alignments files " +"using visual tools" msgstr "" -#: tools/manual\cli.py:35 tools/manual\cli.py:43 +#: tools/manual/cli.py:35 tools/manual/cli.py:44 msgid "data" msgstr "" -#: tools/manual\cli.py:37 -msgid "Path to the alignments file for the input, if not at the default location" +#: tools/manual/cli.py:38 +msgid "" +"Path to the alignments file for the input, if not at the default location" msgstr "" -#: tools/manual\cli.py:44 -msgid "Video file or directory containing source frames that faces were extracted from." +#: tools/manual/cli.py:46 +msgid "" +"Video file or directory containing source frames that faces were extracted " +"from." msgstr "" -#: tools/manual\cli.py:51 tools/manual\cli.py:59 +#: tools/manual/cli.py:53 tools/manual/cli.py:62 msgid "options" msgstr "" -#: tools/manual\cli.py:52 -msgid "Force regeneration of the low resolution jpg thumbnails in the alignments file." +#: tools/manual/cli.py:55 +msgid "" +"Force regeneration of the low resolution jpg thumbnails in the alignments " +"file." msgstr "" -#: tools/manual\cli.py:60 -msgid "The process attempts to speed up generation of thumbnails by extracting from the video in parallel threads. For some videos, this causes the caching process to hang. If this happens, then set this option to generate the thumbnails in a slower, but more stable single thread." +#: tools/manual/cli.py:64 +msgid "" +"The process attempts to speed up generation of thumbnails by extracting from " +"the video in parallel threads. For some videos, this causes the caching " +"process to hang. If this happens, then set this option to generate the " +"thumbnails in a slower, but more stable single thread." msgstr "" #: tools/manual\faceviewer\frame.py:163 @@ -207,4 +222,3 @@ msgstr "" #: tools/manual\frameviewer\frame.py:487 msgid "Copy {} Alignments ({})" msgstr "" - diff --git a/locales/tools.mask.cli.pot b/locales/tools.mask.cli.pot index 0e3d2c1c09..f8024c88ee 100644 --- a/locales/tools.mask.cli.pot +++ b/locales/tools.mask.cli.pot @@ -1,90 +1,191 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-05-17 18:17+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-06-28 13:45+0100\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - #: tools/mask/cli.py:15 -msgid "This command lets you generate masks for existing alignments." +msgid "" +"This tool allows you to generate, import, export or preview masks for " +"existing alignments." msgstr "" -#: tools/mask/cli.py:24 +#: tools/mask/cli.py:25 msgid "" "Mask tool\n" -"Generate masks for existing alignments files." +"Generate, import, export or preview masks for existing alignments files." msgstr "" -#: tools/mask/cli.py:32 tools/mask/cli.py:41 tools/mask/cli.py:51 +#: tools/mask/cli.py:35 tools/mask/cli.py:47 tools/mask/cli.py:58 +#: tools/mask/cli.py:69 msgid "data" msgstr "" -#: tools/mask/cli.py:35 -msgid "Full path to the alignments file to add the mask to. NB: if the mask already exists in the alignments file it will be overwritten." +#: tools/mask/cli.py:39 +msgid "" +"Full path to the alignments file that contains the masks if not at the " +"default location. NB: If the input-type is faces and you wish to update the " +"corresponding alignments file, then you must provide a value here as the " +"location cannot be automatically detected." msgstr "" -#: tools/mask/cli.py:44 +#: tools/mask/cli.py:51 msgid "Directory containing extracted faces, source frames, or a video file." msgstr "" -#: tools/mask/cli.py:53 +#: tools/mask/cli.py:61 msgid "" "R|Whether the `input` is a folder of faces or a folder frames/video\n" "L|faces: The input is a folder containing extracted faces.\n" "L|frames: The input is a folder containing frames or is a video" msgstr "" -#: tools/mask/cli.py:62 tools/mask/cli.py:90 +#: tools/mask/cli.py:71 +msgid "" +"R|Run the mask tool on multiple sources. If selected then the other options " +"should be set as follows:\n" +"L|input: A parent folder containing either all of the video files to be " +"processed, or containing sub-folders of frames/faces.\n" +"L|output-folder: If provided, then sub-folders will be created within the " +"given location to hold the previews for each input.\n" +"L|alignments: Alignments field will be ignored for batch processing. The " +"alignments files must exist at the default location (for frames). For batch " +"processing of masks with 'faces' as the input type, then only the PNG header " +"within the extracted faces will be updated." +msgstr "" + +#: tools/mask/cli.py:87 tools/mask/cli.py:119 msgid "process" msgstr "" -#: tools/mask/cli.py:63 +#: tools/mask/cli.py:89 msgid "" "R|Masker to use.\n" -"L|bisenet-fp: Relatively lightweight NN based mask that provides more refined control over the area to be masked including full head masking (configurable in mask settings).\n" -"L|components: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks to create a mask.\n" -"L|extended: Mask designed to provide facial segmentation based on the positioning of landmark locations. A convex hull is constructed around the exterior of the landmarks and the mask is extended upwards onto the forehead.\n" -"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal faces clear of obstructions. Profile faces and obstructions may result in sub-par performance.\n" -"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been specifically trained to recognize some facial obstructions (hands and eyeglasses). Profile faces may result in sub-par performance.\n" -"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal faces. The mask model has been trained by community members and will need testing for further description. Profile faces may result in sub-par performance." +"L|bisenet-fp: Relatively lightweight NN based mask that provides more " +"refined control over the area to be masked including full head masking " +"(configurable in mask settings).\n" +"L|components: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks to create a mask.\n" +"L|custom: A dummy mask that fills the mask area with all 1s or 0s " +"(configurable in settings). This is only required if you intend to manually " +"edit the custom masks yourself in the manual tool. This mask does not use " +"the GPU.\n" +"L|extended: Mask designed to provide facial segmentation based on the " +"positioning of landmark locations. A convex hull is constructed around the " +"exterior of the landmarks and the mask is extended upwards onto the " +"forehead.\n" +"L|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " +"faces clear of obstructions. Profile faces and obstructions may result in " +"sub-par performance.\n" +"L|vgg-obstructed: Mask designed to provide smart segmentation of mostly " +"frontal faces. The mask model has been specifically trained to recognize " +"some facial obstructions (hands and eyeglasses). Profile faces may result in " +"sub-par performance.\n" +"L|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " +"faces. The mask model has been trained by community members. Profile faces " +"may result in sub-par performance." +msgstr "" + +#: tools/mask/cli.py:121 +msgid "" +"R|The Mask tool process to perform.\n" +"L|all: Update the mask for all faces in the alignments file for the selected " +"'masker'.\n" +"L|missing: Create a mask for all faces in the alignments file where a mask " +"does not previously exist for the selected 'masker'.\n" +"L|output: Don't update the masks, just output the selected 'masker' for " +"review/editing in external tools to the given output folder.\n" +"L|import: Import masks that have been edited outside of faceswap into the " +"alignments file. Note: 'custom' must be the selected 'masker' and the masks " +"must be in the same format as the 'input-type' (frames or faces)" +msgstr "" + +#: tools/mask/cli.py:135 tools/mask/cli.py:154 tools/mask/cli.py:176 +msgid "import" msgstr "" -#: tools/mask/cli.py:91 +#: tools/mask/cli.py:137 msgid "" -"R|Whether to update all masks in the alignments files, only those faces that do not already have a mask of the given `mask type` or just to output the masks to the `output` location.\n" -"L|all: Update the mask for all faces in the alignments file.\n" -"L|missing: Create a mask for all faces in the alignments file where a mask does not previously exist.\n" -"L|output: Don't update the masks, just output them for review in the given output folder." +"R|Import only. The path to the folder that contains masks to be imported.\n" +"L|How the masks are provided is not important, but they will be stored, " +"internally, as 8-bit grayscale images.\n" +"L|If the input are images, then the masks must be named exactly the same as " +"input frames/faces (excluding the file extension).\n" +"L|If the input is a video file, then the filename of the masks is not " +"important but should contain the frame number at the end of the filename " +"(but before the file extension). The frame number can be separated from the " +"rest of the filename by any non-numeric character and can be padded by any " +"number of zeros. The frame number must correspond correctly to the frame " +"number in the original video (starting from frame 1)." +msgstr "" + +#: tools/mask/cli.py:156 +msgid "" +"R|Import/Output only. When importing masks, this is the centering to use. " +"For output this is only used for outputting custom imported masks, and " +"should correspond to the centering used when importing the mask. Note: For " +"any job other than 'import' and 'output' this option is ignored as mask " +"centering is handled internally.\n" +"L|face: Centers the mask on the center of the face, adjusting for pitch and " +"yaw. Outside of requirements for full head masking/training, this is likely " +"to be the best choice.\n" +"L|head: Centers the mask on the center of the head, adjusting for pitch and " +"yaw. Note: You should only select head centering if you intend to include " +"the full head (including hair) within the mask and are looking to train a " +"full head model.\n" +"L|legacy: The 'original' extraction technique. Centers the mask near the of " +"the nose with and crops closely to the face. Can result in the edges of the " +"mask appearing outside of the training area." +msgstr "" + +#: tools/mask/cli.py:181 +msgid "" +"Import only. The size, in pixels to internally store the mask at.\n" +"The default is 128 which is fine for nearly all usecases. Larger sizes will " +"result in larger alignments files and longer processing." msgstr "" -#: tools/mask/cli.py:104 tools/mask/cli.py:111 tools/mask/cli.py:124 -#: tools/mask/cli.py:137 tools/mask/cli.py:146 +#: tools/mask/cli.py:189 tools/mask/cli.py:197 tools/mask/cli.py:211 +#: tools/mask/cli.py:225 tools/mask/cli.py:235 msgid "output" msgstr "" -#: tools/mask/cli.py:105 -msgid "Optional output location. If provided, a preview of the masks created will be output in the given folder." +#: tools/mask/cli.py:191 +msgid "" +"Optional output location. If provided, a preview of the masks created will " +"be output in the given folder." msgstr "" -#: tools/mask/cli.py:115 -msgid "Apply gaussian blur to the mask output. Has the effect of smoothing the edges of the mask giving less of a hard edge. the size is in pixels. This value should be odd, if an even number is passed in then it will be rounded to the next odd number. NB: Only effects the output preview. Set to 0 for off" +#: tools/mask/cli.py:202 +msgid "" +"Apply gaussian blur to the mask output. Has the effect of smoothing the " +"edges of the mask giving less of a hard edge. the size is in pixels. This " +"value should be odd, if an even number is passed in then it will be rounded " +"to the next odd number. NB: Only effects the output preview. Set to 0 for off" msgstr "" -#: tools/mask/cli.py:128 -msgid "Helps reduce 'blotchiness' on some masks by making light shades white and dark shades black. Higher values will impact more of the mask. NB: Only effects the output preview. Set to 0 for off" +#: tools/mask/cli.py:216 +msgid "" +"Helps reduce 'blotchiness' on some masks by making light shades white and " +"dark shades black. Higher values will impact more of the mask. NB: Only " +"effects the output preview. Set to 0 for off" msgstr "" -#: tools/mask/cli.py:138 +#: tools/mask/cli.py:227 msgid "" "R|How to format the output when processing is set to 'output'.\n" "L|combined: The image contains the face/frame, face mask and masked face.\n" @@ -92,7 +193,8 @@ msgid "" "L|mask: Only output the mask as a single channel image." msgstr "" -#: tools/mask/cli.py:147 -msgid "R|Whether to output the whole frame or only the face box when using output processing. Only has an effect when using frames as input." +#: tools/mask/cli.py:237 +msgid "" +"R|Whether to output the whole frame or only the face box when using output " +"processing. Only has an effect when using frames as input." msgstr "" - diff --git a/locales/tools.model.cli.pot b/locales/tools.model.cli.pot new file mode 100644 index 0000000000..f5f2e9c690 --- /dev/null +++ b/locales/tools.model.cli.pot @@ -0,0 +1,63 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +#, fuzzy +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:51+0000\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: FULL NAME \n" +"Language-Team: LANGUAGE \n" +"Language: \n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=CHARSET\n" +"Content-Transfer-Encoding: 8bit\n" + +#: tools/model/cli.py:13 +msgid "This tool lets you perform actions on saved Faceswap models." +msgstr "" + +#: tools/model/cli.py:22 +msgid "A tool for performing actions on Faceswap trained model files" +msgstr "" + +#: tools/model/cli.py:34 +msgid "" +"Model directory. A directory containing the model you wish to perform an " +"action on." +msgstr "" + +#: tools/model/cli.py:43 +msgid "" +"R|Choose which action you want to perform.\n" +"L|'inference' - Create an inference only copy of the model. Strips any " +"layers from the model which are only required for training. NB: This is for " +"exporting the model for use in external applications. Inference generated " +"models cannot be used within Faceswap. See the 'format' option for " +"specifying the model output format.\n" +"L|'nan-scan' - Scan the model file for NaNs or Infs (invalid data).\n" +"L|'restore' - Restore a model from backup." +msgstr "" + +#: tools/model/cli.py:57 tools/model/cli.py:69 +msgid "inference" +msgstr "" + +#: tools/model/cli.py:59 +msgid "" +"R|The format to save the model as. Note: Only used for 'inference' job.\n" +"L|'h5' - Standard Keras H5 format. Does not store any custom layer " +"information. Layers will need to be loaded from Faceswap to use.\n" +"L|'saved-model' - Tensorflow's Saved Model format. Contains all information " +"required to load the model outside of Faceswap." +msgstr "" + +#: tools/model/cli.py:71 +msgid "" +"Only used for 'inference' job. Generate the inference model for B -> A " +"instead of A -> B." +msgstr "" diff --git a/locales/tools.pot b/locales/tools.pot deleted file mode 100644 index 78cf388e1e..0000000000 --- a/locales/tools.pot +++ /dev/null @@ -1,21 +0,0 @@ -# SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION -# FIRST AUTHOR , YEAR. -# -msgid "" -msgstr "" -"Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-02-18 23:49-0000\n" -"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" -"Last-Translator: FULL NAME \n" -"Language-Team: LANGUAGE \n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" -"Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - - -#: tools.py:46 -msgid "Please backup your data and/or test the tool you want to use with a smaller data set to make sure you understand how it works." -msgstr "" - diff --git a/locales/tools.preview.pot b/locales/tools.preview.pot index 3d98ad6363..1dac39da19 100644 --- a/locales/tools.preview.pot +++ b/locales/tools.preview.pot @@ -1,72 +1,80 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-03-10 16:51-0000\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - -#: ./tools/preview\cli.py:13 +#: tools/preview/cli.py:15 msgid "This command allows you to preview swaps to tweak convert settings." msgstr "" -#: ./tools/preview\cli.py:22 +#: tools/preview/cli.py:30 msgid "" "Preview tool\n" "Allows you to configure your convert settings with a live preview" msgstr "" -#: ./tools/preview\cli.py:32 ./tools/preview\cli.py:41 -#: ./tools/preview\cli.py:48 +#: tools/preview/cli.py:47 tools/preview/cli.py:57 tools/preview/cli.py:65 msgid "data" msgstr "" -#: ./tools/preview\cli.py:34 -msgid "Input directory or video. Either a directory containing the image files you wish to process or path to a video file." +#: tools/preview/cli.py:50 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file." msgstr "" -#: ./tools/preview\cli.py:43 -msgid "Path to the alignments file for the input, if not at the default location" +#: tools/preview/cli.py:60 +msgid "" +"Path to the alignments file for the input, if not at the default location" msgstr "" -#: ./tools/preview\cli.py:50 -msgid "Model directory. A directory containing the trained model you wish to process." +#: tools/preview/cli.py:68 +msgid "" +"Model directory. A directory containing the trained model you wish to " +"process." msgstr "" -#: ./tools/preview\cli.py:57 +#: tools/preview/cli.py:74 msgid "Swap the model. Instead of A -> B, swap B -> A" msgstr "" -#: ./tools/preview\preview.py:1303 +#: tools/preview/control_panels.py:510 msgid "Save full config" msgstr "" -#: ./tools/preview\preview.py:1306 +#: tools/preview/control_panels.py:513 msgid "Reset full config to default values" msgstr "" -#: ./tools/preview\preview.py:1309 +#: tools/preview/control_panels.py:516 msgid "Reset full config to saved values" msgstr "" -#: ./tools/preview\preview.py:1453 -msgid "Save {} config" +#: tools/preview/control_panels.py:667 +#, python-brace-format +msgid "Save {title} config" msgstr "" -#: ./tools/preview\preview.py:1456 -msgid "Reset {} config to default values" +#: tools/preview/control_panels.py:670 +#, python-brace-format +msgid "Reset {title} config to default values" msgstr "" -#: ./tools/preview\preview.py:1459 -msgid "Reset {} config to saved values" +#: tools/preview/control_panels.py:673 +#, python-brace-format +msgid "Reset {title} config to saved values" msgstr "" - diff --git a/locales/tools.restore.cli.pot b/locales/tools.restore.cli.pot deleted file mode 100644 index 95e331fb3a..0000000000 --- a/locales/tools.restore.cli.pot +++ /dev/null @@ -1,29 +0,0 @@ -# SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION -# FIRST AUTHOR , YEAR. -# -msgid "" -msgstr "" -"Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-02-18 23:06-0000\n" -"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" -"Last-Translator: FULL NAME \n" -"Language-Team: LANGUAGE \n" -"MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" -"Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - - -#: tools/restore/cli.py:13 -msgid "This command lets you restore models from backup." -msgstr "" - -#: tools/restore/cli.py:22 -msgid "A tool for restoring models from backup (.bk) files" -msgstr "" - -#: tools/restore/cli.py:33 -msgid "Model directory. A directory containing the model you wish to restore from backup." -msgstr "" - diff --git a/locales/tools.sort.cli.pot b/locales/tools.sort.cli.pot index ef03658792..8a963636d0 100644 --- a/locales/tools.sort.cli.pot +++ b/locales/tools.sort.cli.pot @@ -1,102 +1,262 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-08-07 12:34+0100\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2024-03-28 23:53+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - -#: tools/sort/cli.py:14 +#: tools/sort/cli.py:15 msgid "This command lets you sort images using various methods." msgstr "" -#: tools/sort/cli.py:23 -msgid "Sort faces using a number of different techniques" +#: tools/sort/cli.py:21 +msgid "" +" Adjust the '-t' ('--threshold') parameter to control the strength of " +"grouping." msgstr "" -#: tools/sort/cli.py:33 tools/sort/cli.py:40 -msgid "data" +#: tools/sort/cli.py:22 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the percentage of color pixels " +"that appear in the image." +msgstr "" + +#: tools/sort/cli.py:25 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. Each image is allocated to a bin by the number of degrees the face " +"is orientated from center." +msgstr "" + +#: tools/sort/cli.py:28 +msgid "" +" Adjust the '-b' ('--bins') parameter to control the number of bins for " +"grouping. The minimum and maximum values are taken for the chosen sort " +"metric. The bins are then populated with the results from the group sorting." +msgstr "" + +#: tools/sort/cli.py:32 +msgid "faces by blurriness." +msgstr "" + +#: tools/sort/cli.py:33 +msgid "faces by fft filtered blurriness." msgstr "" #: tools/sort/cli.py:34 -msgid "Input directory of aligned faces." +msgid "" +"faces by the estimated distance of the alignments from an 'average' face. " +"This can be useful for eliminating misaligned faces. Sorts from most like an " +"average face to least like an average face." +msgstr "" + +#: tools/sort/cli.py:37 +msgid "" +"faces using VGG Face2 by face similarity. This uses a pairwise clustering " +"algorithm to check the distances between 512 features on every face in your " +"set and order them appropriately." +msgstr "" + +#: tools/sort/cli.py:40 +msgid "faces by their landmarks." msgstr "" #: tools/sort/cli.py:41 -msgid "Output directory for sorted aligned faces." +msgid "Like 'face-cnn' but sorts by dissimilarity." msgstr "" -#: tools/sort/cli.py:50 tools/sort/cli.py:99 -msgid "sort settings" +#: tools/sort/cli.py:42 +msgid "faces by Yaw (rotation left to right)." msgstr "" -#: tools/sort/cli.py:52 +#: tools/sort/cli.py:43 +msgid "faces by Pitch (rotation up and down)." +msgstr "" + +#: tools/sort/cli.py:44 msgid "" -"R|Sort by method. Choose how images are sorted. \n" -"L|'blur': Sort faces by blurriness.\n" -"L|'blur-fft': Sort faces by fft filtered blurriness.\n" -"L|'distance' Sort faces by the estimated distance of the alignments from an 'average' face. This can be useful for eliminating misaligned faces.\n" -"L|'face': Use VGG Face to sort by face similarity. This uses a pairwise clustering algorithm to check the distances between 512 features on every face in your set and order them appropriately.\n" -"L|'face-cnn': Sort faces by their landmarks. You can adjust the threshold with the '-t' (--ref_threshold) option.\n" -"L|'face-cnn-dissim': Like 'face-cnn' but sorts by dissimilarity.\n" -"L|'face-yaw': Sort faces by Yaw (rotation left to right).\n" -"L|'hist': Sort faces by their color histogram. You can adjust the threshold with the '-t' (--ref_threshold) option.\n" -"L|'hist-dissim': Like 'hist' but sorts by dissimilarity.\n" -"L|'color-gray': Sort images by the average intensity of the converted grayscale color channel.\n" -"L|'color-luma': Sort images by the average intensity of the converted Y color channel. Bright lighting and oversaturated images will be ranked first.\n" -"L|'color-green': Sort images by the average intensity of the converted Cg color channel. Green images will be ranked first and red images will be last.\n" -"L|'color-orange': Sort images by the average intensity of the converted Co color channel. Orange images will be ranked first and blue images will be last.\n" -"L|'size': Sort images by their size in the original frame. Faces closer to the camera and from higher resolution sources will be sorted first, whilst faces further from the camera and from lower resolution sources will be sorted last.\n" -"L|'black-pixels': Sort images by their number of black pixels. Useful when faces are near borders and a large part of the image is black.\n" -"Default: face" +"faces by Roll (rotation). Aligned faces should have a roll value close to " +"zero. The further the Roll value from zero the higher liklihood the face is " +"misaligned." +msgstr "" + +#: tools/sort/cli.py:46 +msgid "faces by their color histogram." msgstr "" -#: tools/sort/cli.py:88 tools/sort/cli.py:115 tools/sort/cli.py:127 -#: tools/sort/cli.py:138 -msgid "output" +#: tools/sort/cli.py:47 +msgid "Like 'hist' but sorts by dissimilarity." msgstr "" -#: tools/sort/cli.py:89 -msgid "Keeps the original files in the input directory. Be careful when using this with rename grouping and no specified output directory as this would keep the original and renamed files in the same directory." +#: tools/sort/cli.py:48 +msgid "" +"images by the average intensity of the converted grayscale color channel." msgstr "" -#: tools/sort/cli.py:101 -msgid "Float value. Minimum threshold to use for grouping comparison with 'face-cnn' and 'hist' methods. The lower the value the more discriminating the grouping is. Leaving -1.0 will allow the program set the default value automatically. For face-cnn 7.2 should be enough, with 4 being very discriminating. For hist 0.3 should be enough, with 0.2 being very discriminating. Be careful setting a value that's too low in a directory with many images, as this could result in a lot of directories being created. Defaults: face-cnn 7.2, hist 0.3" +#: tools/sort/cli.py:49 +msgid "" +"images by their number of black pixels. Useful when faces are near borders " +"and a large part of the image is black." msgstr "" -#: tools/sort/cli.py:116 +#: tools/sort/cli.py:51 msgid "" -"R|Default: rename.\n" -"L|'folders': files are sorted using the -s/--sort-by method, then they are organized into folders using the -g/--group-by grouping method.\n" -"L|'rename': files are sorted using the -s/--sort-by then they are renamed." +"images by the average intensity of the converted Y color channel. Bright " +"lighting and oversaturated images will be ranked first." msgstr "" -#: tools/sort/cli.py:129 -msgid "Group by method. When -fp/--final-processing by folders choose the how the images are grouped after sorting. Default: hist" +#: tools/sort/cli.py:53 +msgid "" +"images by the average intensity of the converted Cg color channel. Green " +"images will be ranked first and red images will be last." msgstr "" -#: tools/sort/cli.py:140 -msgid "Integer value. Number of folders that will be used to group by blur, face-yaw and black-pixels. For blur folder 0 will be the least blurry, while the last folder will be the blurriest. For face-yaw the number of bins is by how much 180 degrees is divided. So if you use 18, then each folder will be a 10 degree increment. Folder 0 will contain faces looking the most to the left whereas the last folder will contain the faces looking the most to the right. If the number of images doesn't divide evenly into the number of bins, the remaining images get put in the last bin. For black-pixels it represents the divider of the percentage of black pixels. For 10, first folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, etc. Default value: 5" +#: tools/sort/cli.py:55 +msgid "" +"images by the average intensity of the converted Co color channel. Orange " +"images will be ranked first and blue images will be last." msgstr "" -#: tools/sort/cli.py:154 tools/sort/cli.py:164 -msgid "settings" +#: tools/sort/cli.py:57 +msgid "" +"images by their size in the original frame. Faces further from the camera " +"and from lower resolution sources will be sorted first, whilst faces closer " +"to the camera and from higher resolution sources will be sorted last." +msgstr "" + +#: tools/sort/cli.py:81 +msgid "Sort faces using a number of different techniques" msgstr "" -#: tools/sort/cli.py:156 -msgid "Logs file renaming changes if grouping by renaming, or it logs the file copying/movement if grouping by folders. If no log file is specified with '--log-file', then a 'sort_log.json' file will be created in the input directory." +#: tools/sort/cli.py:91 tools/sort/cli.py:98 tools/sort/cli.py:110 +#: tools/sort/cli.py:150 +msgid "data" +msgstr "" + +#: tools/sort/cli.py:92 +msgid "Input directory of aligned faces." +msgstr "" + +#: tools/sort/cli.py:100 +msgid "" +"Output directory for sorted aligned faces. If not provided and 'keep' is " +"selected then a new folder called 'sorted' will be created within the input " +"folder to house the output. If not provided and 'keep' is not selected then " +"the images will be sorted in-place, overwriting the original contents of the " +"'input_dir'" +msgstr "" + +#: tools/sort/cli.py:112 +msgid "" +"R|If selected then the input_dir should be a parent folder containing " +"multiple folders of faces you wish to sort. The faces will be output to " +"separate sub-folders in the output_dir" +msgstr "" + +#: tools/sort/cli.py:121 +msgid "sort settings" +msgstr "" + +#: tools/sort/cli.py:124 +msgid "" +"R|Choose how images are sorted. Selecting a sort method gives the images a " +"new filename based on the order the image appears within the given method.\n" +"L|'none': Don't sort the images. When a 'group-by' method is selected, " +"selecting 'none' means that the files will be moved/copied into their " +"respective bins, but the files will keep their original filenames. Selecting " +"'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" + +#: tools/sort/cli.py:136 tools/sort/cli.py:164 tools/sort/cli.py:184 +msgid "group settings" +msgstr "" + +#: tools/sort/cli.py:139 +msgid "" +"R|Selecting a group by method will move/copy files into numbered bins based " +"on the selected method.\n" +"L|'none': Don't bin the images. Folders will be sorted by the selected 'sort-" +"by' but will not be binned, instead they will be sorted into a single " +"folder. Selecting 'none' for both 'sort-by' and 'group-by' will do nothing" +msgstr "" + +#: tools/sort/cli.py:152 +msgid "" +"Whether to keep the original files in their original location. Choosing a " +"'sort-by' method means that the files have to be renamed. Selecting 'keep' " +"means that the original files will be kept, and the renamed files will be " +"created in the specified output folder. Unselecting keep means that the " +"original files will be moved and renamed based on the selected sort/group " +"criteria." msgstr "" #: tools/sort/cli.py:167 -msgid "Specify a log file to use for saving the renaming or grouping information. If specified extension isn't 'json' or 'yaml', then json will be used as the serializer, with the supplied filename. Default: sort_log.json" +msgid "" +"R|Float value. Minimum threshold to use for grouping comparison with 'face-" +"cnn' 'hist' and 'face' methods.\n" +"The lower the value the more discriminating the grouping is. Leaving -1.0 " +"will allow Faceswap to choose the default value.\n" +"L|For 'face-cnn' 7.2 should be enough, with 4 being very discriminating. \n" +"L|For 'hist' 0.3 should be enough, with 0.2 being very discriminating. \n" +"L|For 'face' between 0.1 (more bins) to 0.5 (fewer bins) should be about " +"right.\n" +"Be careful setting a value that's too extrene in a directory with many " +"images, as this could result in a lot of folders being created. Defaults: " +"face-cnn 7.2, hist 0.3, face 0.25" +msgstr "" + +#: tools/sort/cli.py:187 +#, python-format +msgid "" +"R|Integer value. Used to control the number of bins created for grouping by: " +"any 'blur' methods, 'color' methods or 'face metric' methods ('distance', " +"'size') and 'orientation; methods ('yaw', 'pitch'). For any other grouping " +"methods see the '-t' ('--threshold') option.\n" +"L|For 'face metric' methods the bins are filled, according the the " +"distribution of faces between the minimum and maximum chosen metric.\n" +"L|For 'color' methods the number of bins represents the divider of the " +"percentage of colored pixels. Eg. For a bin number of '5': The first folder " +"will have the faces with 0%% to 20%% colored pixels, second 21%% to 40%%, " +"etc. Any empty bins will be deleted, so you may end up with fewer bins than " +"selected.\n" +"L|For 'blur' methods folder 0 will be the least blurry, while the last " +"folder will be the blurriest.\n" +"L|For 'orientation' methods the number of bins is dictated by how much 180 " +"degrees is divided. Eg. If 18 is selected, then each folder will be a 10 " +"degree increment. Folder 0 will contain faces looking the most to the left/" +"down whereas the last folder will contain the faces looking the most to the " +"right/up. NB: Some bins may be empty if faces do not fit the criteria. \n" +"Default value: 5" msgstr "" +#: tools/sort/cli.py:207 tools/sort/cli.py:217 +msgid "settings" +msgstr "" + +#: tools/sort/cli.py:210 +msgid "" +"Logs file renaming changes if grouping by renaming, or it logs the file " +"copying/movement if grouping by folders. If no log file is specified with " +"'--log-file', then a 'sort_log.json' file will be created in the input " +"directory." +msgstr "" + +#: tools/sort/cli.py:221 +msgid "" +"Specify a log file to use for saving the renaming or grouping information. " +"If specified extension isn't 'json' or 'yaml', then json will be used as the " +"serializer, with the supplied filename. Default: sort_log.json" +msgstr "" diff --git a/plugins/convert/_config.py b/plugins/convert/_config.py index 5f5ad26c0f..3deb1857f4 100644 --- a/plugins/convert/_config.py +++ b/plugins/convert/_config.py @@ -6,7 +6,7 @@ from lib.config import FaceswapConfig -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Config(FaceswapConfig): diff --git a/plugins/convert/color/_base.py b/plugins/convert/color/_base.py index 1a5c4ebd72..7f58d45ead 100644 --- a/plugins/convert/color/_base.py +++ b/plugins/convert/color/_base.py @@ -6,7 +6,7 @@ from plugins.convert._config import Config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) def get_config(plugin_name, configfile=None): diff --git a/plugins/convert/color/avg_color.py b/plugins/convert/color/avg_color.py index 4483a3104b..f62024289e 100644 --- a/plugins/convert/color/avg_color.py +++ b/plugins/convert/color/avg_color.py @@ -8,11 +8,32 @@ class Color(Adjustment): """ Adjust the mean of the color channels to be the same for the swap and old frame """ - @staticmethod - def process(old_face, new_face, raw_mask): + def process(self, + old_face: np.ndarray, + new_face: np.ndarray, + raw_mask: np.ndarray) -> np.ndarray: + """ Adjust the mean of the original face and the new face to be the same + + Parameters + ---------- + old_face: :class:`numpy.ndarray` + The original face + new_face: :class:`numpy.ndarray` + The Faceswap generated face + raw_mask: :class:`numpy.ndarray` + A raw mask for including the face area only + + Returns + ------- + :class:`numpy.ndarray` + The adjusted face patch + """ for _ in [0, 1]: diff = old_face - new_face - avg_diff = np.sum(diff * raw_mask, axis=(0, 1)) - adjustment = avg_diff / np.sum(raw_mask, axis=(0, 1)) + if np.any(raw_mask): + avg_diff = np.sum(diff * raw_mask, axis=(0, 1)) + adjustment = avg_diff / np.sum(raw_mask, axis=(0, 1)) + else: + adjustment = diff new_face += adjustment return new_face diff --git a/plugins/convert/color/color_transfer.py b/plugins/convert/color/color_transfer.py index 17ae9d29ff..4b8080a433 100644 --- a/plugins/convert/color/color_transfer.py +++ b/plugins/convert/color/color_transfer.py @@ -40,8 +40,8 @@ class Color(Adjustment): def process(self, old_face, new_face, raw_mask): """ - Parameters: - ------- + Parameters + ---------- source: NumPy array OpenCV image in BGR color space (the source image) target: NumPy array @@ -59,7 +59,7 @@ def process(self, old_face, new_face, raw_mask): the scaling factor proposed in the paper. This method seems to produce more consistently aesthetically pleasing results - Returns: + Returns ------- transfer: NumPy array OpenCV image (w, h, 3) NumPy array (uint8) @@ -70,12 +70,12 @@ def process(self, old_face, new_face, raw_mask): # convert the images from the RGB to L*ab* color space, being # sure to utilizing the floating point data type (note: OpenCV # expects floats to be 32-bit, so use that instead of 64-bit) - source = cv2.cvtColor( # pylint: disable=no-member + source = cv2.cvtColor( # pylint:disable=no-member np.rint(old_face * raw_mask * 255.0).astype("uint8"), - cv2.COLOR_BGR2LAB).astype("float32") # pylint: disable=no-member - target = cv2.cvtColor( # pylint: disable=no-member + cv2.COLOR_BGR2LAB).astype("float32") # pylint:disable=no-member + target = cv2.cvtColor( # pylint:disable=no-member np.rint(new_face * raw_mask * 255.0).astype("uint8"), - cv2.COLOR_BGR2LAB).astype("float32") # pylint: disable=no-member + cv2.COLOR_BGR2LAB).astype("float32") # pylint:disable=no-member # compute color statistics for the source and target images (l_mean_src, l_std_src, a_mean_src, a_std_src, @@ -85,7 +85,7 @@ def process(self, old_face, new_face, raw_mask): b_mean_tar, b_std_tar) = self.image_stats(target) # subtract the means from the target image - (light, col_a, col_b) = cv2.split(target) # pylint: disable=no-member + (light, col_a, col_b) = cv2.split(target) # pylint:disable=no-member light -= l_mean_tar col_a -= a_mean_tar col_b -= b_mean_tar @@ -115,10 +115,10 @@ def process(self, old_face, new_face, raw_mask): # merge the channels together and convert back to the RGB color # space, being sure to utilize the 8-bit unsigned integer data # type - transfer = cv2.merge([light, col_a, col_b]) # pylint: disable=no-member - transfer = cv2.cvtColor( # pylint: disable=no-member + transfer = cv2.merge([light, col_a, col_b]) # pylint:disable=no-member + transfer = cv2.cvtColor( # pylint:disable=no-member transfer.astype("uint8"), - cv2.COLOR_LAB2BGR).astype("float32") / 255.0 # pylint: disable=no-member + cv2.COLOR_LAB2BGR).astype("float32") / 255.0 # pylint:disable=no-member background = new_face * (1 - raw_mask) merged = transfer + background # return the color transferred image @@ -127,18 +127,19 @@ def process(self, old_face, new_face, raw_mask): @staticmethod def image_stats(image): """ - Parameters: - ------- + Parameters + ---------- + image: NumPy array OpenCV image in L*a*b* color space - Returns: + Returns ------- Tuple of mean and standard deviations for the L*, a*, and b* channels, respectively """ # compute the mean and standard deviation of each channel - (light, col_a, col_b) = cv2.split(image) # pylint: disable=no-member + (light, col_a, col_b) = cv2.split(image) # pylint:disable=no-member (l_mean, l_std) = (light.mean(), light.std()) (a_mean, a_std) = (col_a.mean(), col_a.std()) (b_mean, b_std) = (col_b.mean(), col_b.std()) @@ -151,13 +152,13 @@ def _min_max_scale(arr, new_range=(0, 255)): """ Perform min-max scaling to a NumPy array - Parameters: - ------- + Parameters + ---------- arr: NumPy array to be scaled to [new_min, new_max] range new_range: tuple of form (min, max) specifying range of transformed array - Returns: + Returns ------- NumPy array that has been scaled to be in [new_range[0], new_range[1]] range @@ -182,14 +183,14 @@ def _scale_array(self, arr, clip=True): Trim NumPy array values to be in [0, 255] range with option of clipping or scaling. - Parameters: - ------- + Parameters + ---------- arr: array to be trimmed to [0, 255] range clip: should array be scaled by np.clip? if False then input array will be min-max scaled to range [max([arr.min(), 0]), min([arr.max(), 255])] - Returns: + Returns ------- NumPy array that has been scaled to be in [0, 255] range """ diff --git a/plugins/convert/color/manual_balance.py b/plugins/convert/color/manual_balance.py index 7acb30c95c..7dc6950bb6 100644 --- a/plugins/convert/color/manual_balance.py +++ b/plugins/convert/color/manual_balance.py @@ -43,7 +43,7 @@ def convert_colorspace(self, new_face, to_bgr=False): """ Convert colorspace based on mode or back to bgr """ mode = self.config["colorspace"].lower() colorspace = "YCrCb" if mode == "ycrcb" else mode.upper() - conversion = "{}2BGR".format(colorspace) if to_bgr else "BGR2{}".format(colorspace) - image = cv2.cvtColor(new_face.astype("uint8"), # pylint: disable=no-member - getattr(cv2, "COLOR_{}".format(conversion))).astype("float32") / 255.0 + conversion = f"{colorspace}2BGR" if to_bgr else f"BGR2{colorspace}" + image = cv2.cvtColor(new_face.astype("uint8"), # pylint:disable=no-member + getattr(cv2, f"COLOR_{conversion}")).astype("float32") / 255.0 return image diff --git a/plugins/convert/color/seamless_clone.py b/plugins/convert/color/seamless_clone.py index ccbc3bd19c..dc2f1fe21d 100644 --- a/plugins/convert/color/seamless_clone.py +++ b/plugins/convert/color/seamless_clone.py @@ -16,8 +16,7 @@ class Color(Adjustment): and does not have a natural home, so here for now. """ - @staticmethod - def process(old_face, new_face, raw_mask): + def process(self, old_face, new_face, raw_mask): height, width, _ = old_face.shape height = height // 2 width = width // 2 @@ -35,11 +34,11 @@ def process(old_face, new_face, raw_mask): ((height, height), (width, width), (0, 0)), 'constant')).astype("uint8") - blended = cv2.seamlessClone(insertion, # pylint: disable=no-member + blended = cv2.seamlessClone(insertion, # pylint:disable=no-member prior, insertion_mask, (x_center, y_center), - cv2.NORMAL_CLONE) # pylint: disable=no-member + cv2.NORMAL_CLONE) # pylint:disable=no-member blended = blended[height:-height, width:-width] return blended.astype("float32") / 255.0 diff --git a/plugins/convert/mask/mask_blend.py b/plugins/convert/mask/mask_blend.py index 83477884e8..6683bdf10b 100644 --- a/plugins/convert/mask/mask_blend.py +++ b/plugins/convert/mask/mask_blend.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 """ Plugin to blend the edges of the face between the swap and the original face. """ import logging -import sys -from typing import List, Optional, Tuple +import typing as T import cv2 import numpy as np @@ -11,16 +10,10 @@ from lib.config import FaceswapConfig from plugins.convert._config import Config -if sys.version_info < (3, 8): - from typing_extensions import Literal -else: - from typing import Literal - - logger = logging.getLogger(__name__) -class Mask(): # pylint:disable=too-few-public-methods +class Mask(): """ Manipulations to perform to the mask that is to be applied to the output of the Faceswap model. @@ -44,8 +37,8 @@ def __init__(self, mask_type: str, output_size: int, coverage_ratio: float, - configfile: Optional[str] = None, - config: Optional[FaceswapConfig] = None) -> None: + configfile: str | None = None, + config: FaceswapConfig | None = None) -> None: logger.debug("Initializing %s: (mask_type: '%s', output_size: %s, coverage_ratio: %s, " "configfile: %s, config: %s)", self.__class__.__name__, mask_type, coverage_ratio, output_size, configfile, config) @@ -61,8 +54,8 @@ def __init__(self, self._do_erode = any(amount != 0 for amount in self._erodes) def _set_config(self, - configfile: Optional[str], - config: Optional[FaceswapConfig]) -> dict: + configfile: str | None, + config: FaceswapConfig | None) -> dict: """ Set the correct configuration for the plugin based on whether a config file or pre-loaded config has been passed in. @@ -121,17 +114,20 @@ def _get_box(self, output_size: int) -> np.ndarray: def run(self, detected_face: DetectedFace, - sub_crop_offset: Optional[np.ndarray], - centering: Literal["legacy", "face", "head"], - predicted_mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]: + source_offset: np.ndarray, + target_offset: np.ndarray, + centering: T.Literal["legacy", "face", "head"], + predicted_mask: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]: """ Obtain the requested mask type and perform any defined mask manipulations. Parameters ---------- detected_face: :class:`lib.align.DetectedFace` The DetectedFace object as returned from :class:`scripts.convert.Predictor`. - sub_crop_offset: :class:`numpy.ndarray`, optional - The (x, y) offset to crop the mask from the center point. + source_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at its stored centering + target_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at the requested target centering centering: [`"legacy"`, `"face"`, `"head"`] The centering to obtain the mask for predicted_mask: :class:`numpy.ndarray`, optional @@ -146,15 +142,21 @@ def run(self, The mask with no erosion/dilation applied """ logger.trace("Performing mask adjustment: (detected_face: %s, " # type: ignore - "sub_crop_offset: %s, centering: '%s', predicted_mask: %s", - detected_face, sub_crop_offset, centering, predicted_mask is not None) - mask = self._get_mask(detected_face, predicted_mask, centering, sub_crop_offset) + "source_offset: %s, target_offset: %s, centering: '%s', predicted_mask: %s", + detected_face, source_offset, target_offset, centering, + predicted_mask is not None) + mask = self._get_mask(detected_face, + predicted_mask, + centering, + source_offset, + target_offset) raw_mask = mask.copy() if self._mask_type != "none": - out = self._erode(mask) if self._do_erode else mask out = np.minimum(out, self._box) + else: + out = mask logger.trace( # type: ignore "mask shape: %s, raw_mask shape: %s", mask.shape, raw_mask.shape) @@ -162,9 +164,10 @@ def run(self, def _get_mask(self, detected_face: DetectedFace, - predicted_mask: Optional[np.ndarray], - centering: Literal["legacy", "face", "head"], - sub_crop_offset: Optional[np.ndarray]) -> np.ndarray: + predicted_mask: np.ndarray | None, + centering: T.Literal["legacy", "face", "head"], + source_offset: np.ndarray, + target_offset: np.ndarray) -> np.ndarray: """ Return the requested mask with any requested blurring applied. Parameters @@ -176,9 +179,10 @@ def _get_mask(self, with a mask, otherwise ``None`` centering: [`"legacy"`, `"face"`, `"head"`] The centering to obtain the mask for - sub_crop_offset: :class:`numpy.ndarray` - The (x, y) offset to crop the mask from the center point. Set to `None` if the mask - does not need to be offset for alternative centering + source_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at its stored centering + target_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at the requested target centering Returns ------- @@ -190,7 +194,7 @@ def _get_mask(self, elif self._mask_type == "predicted" and predicted_mask is not None: mask = self._process_predicted_mask(predicted_mask) else: - mask = self._get_stored_mask(detected_face, centering, sub_crop_offset) + mask = self._get_stored_mask(detected_face, centering, source_offset, target_offset) logger.trace(mask.shape) # type: ignore return mask @@ -208,7 +212,7 @@ def _process_predicted_mask(self, mask: np.ndarray) -> np.ndarray: :class:`numpy.ndarray` The processed predicted mask """ - blur_type = self._config["type"] + blur_type = self._config["type"].lower() if blur_type is not None: mask = BlurMask(blur_type, mask, @@ -218,8 +222,9 @@ def _process_predicted_mask(self, mask: np.ndarray) -> np.ndarray: def _get_stored_mask(self, detected_face: DetectedFace, - centering: Literal["legacy", "face", "head"], - sub_crop_offset: Optional[np.ndarray]) -> np.ndarray: + centering: T.Literal["legacy", "face", "head"], + source_offset: np.ndarray, + target_offset: np.ndarray) -> np.ndarray: """ get the requested stored mask from the detected face object. Parameters @@ -228,9 +233,10 @@ def _get_stored_mask(self, The DetectedFace object as returned from :class:`scripts.convert.Predictor`. centering: [`"legacy"`, `"face"`, `"head"`] The centering to obtain the mask for - sub_crop_offset: :class:`numpy.ndarray` - The (x, y) offset to crop the mask from the center point. Set to `None` if the mask - does not need to be offset for alternative centering + source_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at its stored centering + target_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at the requested target centering Returns ------- @@ -242,42 +248,18 @@ def _get_stored_mask(self, blur_type=self._config["type"], blur_passes=self._config["passes"], threshold=self._config["threshold"]) - if sub_crop_offset is not None and np.any(sub_crop_offset): - mask.set_sub_crop(sub_crop_offset, centering) - mask = self._crop_to_coverage(mask.mask) - mask_size = mask.shape[0] + mask.set_sub_crop(source_offset, target_offset, centering, self._coverage_ratio) + face_mask = mask.mask + mask_size = face_mask.shape[0] face_size = self._box.shape[0] if mask_size != face_size: interp = cv2.INTER_CUBIC if mask_size < face_size else cv2.INTER_AREA - mask = cv2.resize(mask, - self._box.shape[:2], - interpolation=interp)[..., None].astype("float32") / 255. + face_mask = cv2.resize(face_mask, + self._box.shape[:2], + interpolation=interp)[..., None].astype("float32") / 255. else: - mask = np.float32(mask) / 255. - return mask - - def _crop_to_coverage(self, mask: np.ndarray) -> np.ndarray: - """ Crop the mask to the correct dimensions based on coverage ratio. - - Parameters - ---------- - mask: :class:`numpy.ndarray` - The original mask to be cropped - - Returns - ------- - :class:`numpy.ndarray` - The cropped mask - """ - if self._coverage_ratio == 1.0: - return mask - mask_size = mask.shape[0] - padding = round((mask_size * (1 - self._coverage_ratio)) / 2) - mask_slice = slice(padding, mask_size - padding) - mask = mask[mask_slice, mask_slice, :] - logger.trace("mask_size: %s, coverage: %s, padding: %s, final shape: %s", # type: ignore - mask_size, self._coverage_ratio, padding, mask.shape) - return mask + face_mask = face_mask.astype("float32") / 255. + return face_mask # MASK MANIPULATIONS def _erode(self, mask: np.ndarray) -> np.ndarray: @@ -314,7 +296,7 @@ def _erode(self, mask: np.ndarray) -> np.ndarray: return eroded[..., None] - def _get_erosion_kernels(self, mask: np.ndarray) -> List[np.ndarray]: + def _get_erosion_kernels(self, mask: np.ndarray) -> list[np.ndarray]: """ Get the erosion kernels for each of the center, left, top right and bottom erosions. An approximation is made based on the number of positive pixels within the mask to create diff --git a/plugins/convert/scaling/_base.py b/plugins/convert/scaling/_base.py index bee1321be9..036ddc1557 100644 --- a/plugins/convert/scaling/_base.py +++ b/plugins/convert/scaling/_base.py @@ -6,7 +6,7 @@ from plugins.convert._config import Config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) def get_config(plugin_name, configfile=None): diff --git a/plugins/convert/scaling/sharpen.py b/plugins/convert/scaling/sharpen.py index 1ecca3faad..0179de9fd6 100644 --- a/plugins/convert/scaling/sharpen.py +++ b/plugins/convert/scaling/sharpen.py @@ -34,15 +34,15 @@ def box(new_face, kernel_center, amount): kernel[center, center] = 1.0 box_filter = np.ones(kernel_size, dtype="float32") / kernel_size[0]**2 kernel = kernel + (kernel - box_filter) * amount - new_face = cv2.filter2D(new_face, -1, kernel) # pylint: disable=no-member + new_face = cv2.filter2D(new_face, -1, kernel) # pylint:disable=no-member return new_face @staticmethod def gaussian(new_face, kernel_center, amount): """ Sharpen using gaussian filter """ kernel_size = kernel_center[0] - blur = cv2.GaussianBlur(new_face, kernel_size, 0) # pylint: disable=no-member - new_face = cv2.addWeighted(new_face, # pylint: disable=no-member + blur = cv2.GaussianBlur(new_face, kernel_size, 0) # pylint:disable=no-member + new_face = cv2.addWeighted(new_face, # pylint:disable=no-member 1.0 + (0.5 * amount), blur, -(0.5 * amount), @@ -53,7 +53,7 @@ def unsharp_mask(self, new_face, kernel_center, amount): """ Sharpen using unsharp mask """ kernel_size = kernel_center[0] threshold = self.config["threshold"] / 255.0 - blur = cv2.GaussianBlur(new_face, kernel_size, 0) # pylint: disable=no-member + blur = cv2.GaussianBlur(new_face, kernel_size, 0) # pylint:disable=no-member low_contrast_mask = (abs(new_face - blur) < threshold).astype("float32") sharpened = (new_face * (1.0 + amount)) + (blur * -amount) new_face = (new_face * (1.0 - low_contrast_mask)) + (sharpened * low_contrast_mask) diff --git a/plugins/convert/writer/_base.py b/plugins/convert/writer/_base.py index 1be1e9c04f..a67ebee28e 100644 --- a/plugins/convert/writer/_base.py +++ b/plugins/convert/writer/_base.py @@ -4,15 +4,16 @@ import logging import os import re +import typing as T -from typing import Optional +import numpy as np from plugins.convert._config import Config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -def get_config(plugin_name: str, configfile: Optional[str] = None) -> dict: +def get_config(plugin_name: str, configfile: str | None = None) -> dict: """ Obtain the configuration settings for the writer plugin. Parameters @@ -42,7 +43,7 @@ class Output(): The full path to a custom configuration ini file. If ``None`` is passed then the file is loaded from the default location. Default: ``None``. """ - def __init__(self, output_folder: str, configfile: Optional[str] = None) -> None: + def __init__(self, output_folder: str, configfile: str | None = None) -> None: logger.debug("Initializing %s: (output_folder: '%s')", self.__class__.__name__, output_folder) self.config: dict = get_config(".".join(self.__module__.split(".")[-2:]), @@ -50,8 +51,11 @@ def __init__(self, output_folder: str, configfile: Optional[str] = None) -> None logger.debug("config: %s", self.config) self.output_folder: str = output_folder + # For creating subfolders when separate mask is selected + self._subfolders_created: bool = False + # Methods for making sure frames are written out in frame order - self.re_search: re.Pattern = re.compile(r"(\d+)(?=\.\w+$)") # Identify frame numbers + self.re_search = re.compile(r"(\d+)(?=\.\w+$)") # Identify frame numbers self.cache: dict = {} # Cache for when frames must be written in correct order logger.debug("Initialized %s", self.__class__.__name__) @@ -61,10 +65,40 @@ def is_stream(self) -> bool: Writers that write to a stream have a frame_order paramater to dictate the order in which frames should be written out (eg. gif/ffmpeg) """ - retval = hasattr(self, "frame_order") + retval = hasattr(self, "_frame_order") return retval - def output_filename(self, filename: str) -> str: + @classmethod + def _set_frame_order(cls, + total_count: int, + frame_ranges: list[tuple[int, int]] | None) -> list[int]: + """ Obtain the full list of frames to be converted in order. + + Used for FFMPEG and Gif writers to ensure correct frame order + + Parameters + ---------- + total_count: int + The total number of frames to be converted + frame_ranges: list or ``None`` + List of tuples for starting and end values of each frame range to be converted or + ``None`` if all frames are to be converted + + Returns + ------- + list + Full list of all frame indices to be converted + """ + if frame_ranges is None: + retval = list(range(1, total_count + 1)) + else: + retval = [] + for rng in frame_ranges: + retval.extend(list(range(rng[0], rng[1] + 1))) + logger.debug("frame_order: %s", retval) + return retval + + def output_filename(self, filename: str, separate_mask: bool = False) -> list[str]: """ Obtain the full path for the output file, including the correct extension, for the given input filename. @@ -75,19 +109,31 @@ def output_filename(self, filename: str) -> str: ---------- filename: str The input frame filename to generate the output file name for + separate_mask: bool, optional + ``True`` if the mask should be saved out to a sub-folder otherwise ``False`` Returns ------- - str - The full path for the output converted frame to be saved to. + list + The full path for the output converted frame to be saved to in position 1. The full + path for the mask to be output to in position 2 (if requested) """ filename = os.path.splitext(os.path.basename(filename))[0] out_filename = f"{filename}.{self.config['format']}" - out_filename = os.path.join(self.output_folder, out_filename) - logger.trace("in filename: '%s', out filename: '%s'", filename, out_filename) - return out_filename + retval = [os.path.join(self.output_folder, out_filename)] + if separate_mask: + retval.append(os.path.join(self.output_folder, "masks", out_filename)) + + if separate_mask and not self._subfolders_created: + locations = [os.path.dirname(loc) for loc in retval] + logger.debug("Creating sub-folders: %s", locations) + for location in locations: + os.makedirs(location, exist_ok=True) - def cache_frame(self, filename, image) -> None: + logger.trace("in filename: '%s', out filename: '%s'", filename, retval) # type:ignore + return retval + + def cache_frame(self, filename: str, image: np.ndarray) -> None: """ Add the incoming converted frame to the cache ready for writing out. Used for ffmpeg and gif writers to ensure that the frames are written out in the correct @@ -100,24 +146,27 @@ def cache_frame(self, filename, image) -> None: image: class:`numpy.ndarray` The converted frame corresponding to the given filename """ - frame_no = int(re.search(self.re_search, filename).group()) + re_frame = re.search(self.re_search, filename) + assert re_frame is not None + frame_no = int(re_frame.group()) self.cache[frame_no] = image - logger.trace("Added to cache. Frame no: %s", frame_no) - logger.trace("Current cache: %s", sorted(self.cache.keys())) + logger.trace("Added to cache. Frame no: %s", frame_no) # type: ignore + logger.trace("Current cache: %s", sorted(self.cache.keys())) # type:ignore - def write(self, filename: str, image) -> None: + def write(self, filename: str, image: T.Any) -> None: """ Override for specific frame writing method. Parameters ---------- filename: str The incoming frame filename. - image: :class:`numpy.ndarray` - The converted image to be written + image: Any + The converted image to be written. Could be a numpy array, a bytes encoded image or + any other plugin specific format """ raise NotImplementedError - def pre_encode(self, image) -> None: # pylint: disable=unused-argument,no-self-use + def pre_encode(self, image: np.ndarray, **kwargs) -> T.Any: # pylint:disable=unused-argument """ Some writer plugins support the pre-encoding of images prior to saving out. As patching is done in multiple threads, but writing is done in a single thread, it can speed up the process to do any pre-encoding as part of the converter process. @@ -132,9 +181,9 @@ def pre_encode(self, image) -> None: # pylint: disable=unused-argument,no-self- Returns ------- - python function or ``None`` - If ``None`` then the writer does not support pre-encoding, otherwise return the python - function that will pre-encode the image + Any or ``None`` + If ``None`` then the writer does not support pre-encoding, otherwise return output of + the plugin specific pre-enccode function """ return None diff --git a/plugins/convert/writer/ffmpeg.py b/plugins/convert/writer/ffmpeg.py index ca1f33bf02..92143a13bd 100644 --- a/plugins/convert/writer/ffmpeg.py +++ b/plugins/convert/writer/ffmpeg.py @@ -1,9 +1,11 @@ #!/usr/bin/env python3 """ Video output writer for faceswap.py converter """ +from __future__ import annotations import os +import typing as T + from math import ceil from subprocess import CalledProcessError, check_output, STDOUT -from typing import Optional, List, Tuple, Generator import imageio import imageio_ffmpeg as im_ffm @@ -11,6 +13,9 @@ from ._base import Output, logger +if T.TYPE_CHECKING: + from collections.abc import Generator + class Writer(Output): """ Video output writer using imageio-ffmpeg. @@ -32,7 +37,7 @@ class Writer(Output): def __init__(self, output_folder: str, total_count: int, - frame_ranges: Optional[List[Tuple[int]]], + frame_ranges: list[tuple[int, int]] | None, source_video: str, **kwargs) -> None: super().__init__(output_folder, **kwargs) @@ -40,11 +45,11 @@ def __init__(self, total_count, frame_ranges, source_video) self._source_video: str = source_video self._output_filename: str = self._get_output_filename() - self._frame_ranges: Optional[List[Tuple[int]]] = frame_ranges - self.frame_order: List[int] = self._set_frame_order(total_count) - self._output_dimensions: Optional[str] = None # Fix dims on 1st received frame + self._frame_ranges: list[tuple[int, int]] | None = frame_ranges + self._frame_order: list[int] = self._set_frame_order(total_count, frame_ranges) + self._output_dimensions: str | None = None # Fix dims on 1st received frame # Need to know dimensions of first frame, so set writer then - self._writer: Optional[Generator[None, np.ndarray, None]] = None + self._writer: Generator[None, np.ndarray, None] | None = None @property def _valid_tunes(self) -> dict: @@ -56,14 +61,14 @@ def _valid_tunes(self) -> dict: @property def _video_fps(self) -> float: """ float: The fps of the source video. """ - reader = imageio.get_reader(self._source_video, "ffmpeg") + reader = imageio.get_reader(self._source_video, "ffmpeg") # type:ignore[arg-type] retval = reader.get_meta_data()["fps"] reader.close() logger.debug(retval) return retval @property - def _output_params(self) -> List[str]: + def _output_params(self) -> list[str]: """ list: The FFMPEG Output parameters """ codec = self.config["codec"] tune = self.config["tune"] @@ -86,11 +91,11 @@ def _output_params(self) -> List[str]: return output_args @property - def _audio_codec(self) -> Optional[str]: - """ str or ``None``: The audio codec to use. This will either be ``"copy"`` (the default) or - ``None`` if skip muxing has been selected in configuration options, or if frame ranges have - been passed in the command line arguments. """ - retval = "copy" + def _audio_codec(self) -> str | None: + """ str or ``None``: The audio codec to use. This will either be ``"copy"`` (the default) + or ``None`` if skip muxing has been selected in configuration options, or if frame ranges + have been passed in the command line arguments. """ + retval: str | None = "copy" if self.config["skip_mux"]: logger.info("Skipping audio muxing due to configuration settings.") retval = None @@ -128,9 +133,9 @@ def _test_for_audio_stream(self) -> bool: try: out = check_output(cmd, stderr=STDOUT) except CalledProcessError as err: - out = err.output.decode(errors="ignore") - raise ValueError("Error checking audio stream. Status: " - f"{err.returncode}\n{out}") from err + err_out = err.output.decode(errors="ignore") + msg = f"Error checking audio stream. Status: {err.returncode}\n{err_out}" + raise ValueError(msg) from err retval = False for line in out.splitlines(): @@ -169,29 +174,7 @@ def _get_output_filename(self) -> str: logger.info("Outputting to: '%s'", retval) return retval - def _set_frame_order(self, total_count: int) -> List[int]: - """ Obtain the full list of frames to be converted in order. - - Parameters - ---------- - total_count: int - The total number of frames to be converted - - Returns - ------- - list - Full list of all frame indices to be converted - """ - if self._frame_ranges is None: - retval = list(range(1, total_count + 1)) - else: - retval = [] - for rng in self._frame_ranges: - retval.extend(list(range(rng[0], rng[1] + 1))) - logger.debug("frame_order: %s", retval) - return retval - - def _get_writer(self, frame_dims: Tuple[int]) -> Generator[None, np.ndarray, None]: + def _get_writer(self, frame_dims: tuple[int, int]) -> Generator[None, np.ndarray, None]: """ Add the requested encoding options and return the writer. Parameters @@ -235,15 +218,16 @@ def write(self, filename: str, image: np.ndarray) -> None: image: :class:`numpy.ndarray` The converted image to be written """ - logger.trace("Received frame: (filename: '%s', shape: %s", filename, image.shape) + logger.trace("Received frame: (filename: '%s', shape: %s", # type:ignore[attr-defined] + filename, image.shape) if not self._output_dimensions: - input_dims = image.shape[:2] + input_dims = T.cast(tuple[int, int], image.shape[:2]) self._set_dimensions(input_dims) self._writer = self._get_writer(input_dims) self.cache_frame(filename, image) self._save_from_cache() - def _set_dimensions(self, frame_dims: Tuple[int]) -> None: + def _set_dimensions(self, frame_dims: tuple[int, int]) -> None: """ Set the attribute :attr:`_output_dimensions` based on the first frame received. This protects against different sized images coming in and ensures all images are written to ffmpeg at the same size. Dimensions are mapped to a macro block size 8. @@ -261,16 +245,19 @@ def _set_dimensions(self, frame_dims: Tuple[int]) -> None: def _save_from_cache(self) -> None: """ Writes any consecutive frames to the video container that are ready to be output from the cache. """ - while self.frame_order: - if self.frame_order[0] not in self.cache: - logger.trace("Next frame not ready. Continuing") + assert self._writer is not None + while self._frame_order: + if self._frame_order[0] not in self.cache: + logger.trace("Next frame not ready. Continuing") # type:ignore[attr-defined] break - save_no = self.frame_order.pop(0) + save_no = self._frame_order.pop(0) save_image = self.cache.pop(save_no) - logger.trace("Rendering from cache. Frame no: %s", save_no) + logger.trace("Rendering from cache. Frame no: %s", # type:ignore[attr-defined] + save_no) self._writer.send(np.ascontiguousarray(save_image[:, :, ::-1])) - logger.trace("Current cache size: %s", len(self.cache)) + logger.trace("Current cache size: %s", len(self.cache)) # type:ignore[attr-defined] def close(self) -> None: """ Close the ffmpeg writer and mux the audio """ - self._writer.close() + if self._writer is not None: + self._writer.close() diff --git a/plugins/convert/writer/gif.py b/plugins/convert/writer/gif.py index fa9bfe0c52..7a75ec93f3 100644 --- a/plugins/convert/writer/gif.py +++ b/plugins/convert/writer/gif.py @@ -1,13 +1,17 @@ #!/usr/bin/env python3 """ Animated GIF writer for faceswap.py converter """ +from __future__ import annotations import os -from typing import Optional, List, Tuple +import typing as T import cv2 import imageio from ._base import Output, logger +if T.TYPE_CHECKING: + from imageio.core import format as im_format # noqa:F401 + class Writer(Output): """ GIF output writer using imageio. @@ -28,15 +32,16 @@ class Writer(Output): def __init__(self, output_folder: str, total_count: int, - frame_ranges: Optional[List[Tuple[int]]], + frame_ranges: list[tuple[int, int]] | None, **kwargs) -> None: logger.debug("total_count: %s, frame_ranges: %s", total_count, frame_ranges) super().__init__(output_folder, **kwargs) - self.frame_order: List[int] = self._set_frame_order(total_count, frame_ranges) - self._output_dimensions: Optional[str] = None # Fix dims on 1st received frame + self._frame_order: list[int] = self._set_frame_order(total_count, frame_ranges) + # Fix dims on 1st received frame + self._output_dimensions: tuple[int, int] | None = None # Need to know dimensions of first frame, so set writer then - self._writer: Optional[imageio.plugins.pillowmulti.GIFFormat.Writer] = None - self._gif_file: Optional[str] = None # Set filename based on first file seen + self._writer: imageio.plugins.pillowmulti.GIFFormat.Writer | None = None + self._gif_file: str | None = None # Set filename based on first file seen @property def _gif_params(self) -> dict: @@ -45,33 +50,7 @@ def _gif_params(self) -> dict: logger.debug(kwargs) return kwargs - @staticmethod - def _set_frame_order(total_count: int, frame_ranges: Optional[List[Tuple[int]]]) -> List[int]: - """ Obtain the full list of frames to be converted in order. - - Parameters - ---------- - total_count: int - The total number of frames to be converted - frame_ranges: list or ``None`` - List of tuples for starting and end values of each frame range to be converted or - ``None`` if all frames are to be converted - - Returns - ------- - list - Full list of all frame indices to be converted - """ - if frame_ranges is None: - retval = list(range(1, total_count + 1)) - else: - retval = [] - for rng in frame_ranges: - retval.extend(list(range(rng[0], rng[1] + 1))) - logger.debug("frame_order: %s", retval) - return retval - - def _get_writer(self) -> imageio.plugins.pillowmulti.GIFFormat.Writer: + def _get_writer(self) -> im_format.Format.Writer: """ Obtain the GIF writer with the requested GIF encoding options. Returns @@ -80,7 +59,7 @@ def _get_writer(self) -> imageio.plugins.pillowmulti.GIFFormat.Writer: The imageio GIF writer """ logger.debug("writer config: %s", self.config) - + assert self._gif_file is not None return imageio.get_writer(self._gif_file, mode="i", **self._gif_params) @@ -96,13 +75,14 @@ def write(self, filename: str, image) -> None: image: :class:`numpy.ndarray` The converted image to be written """ - logger.trace("Received frame: (filename: '%s', shape: %s", filename, image.shape) + logger.trace("Received frame: (filename: '%s', shape: %s", # type: ignore + filename, image.shape) if not self._gif_file: self._set_gif_filename(filename) self._set_dimensions(image.shape[:2]) self._writer = self._get_writer() if (image.shape[1], image.shape[0]) != self._output_dimensions: - image = cv2.resize(image, self._output_dimensions) # pylint: disable=no-member + image = cv2.resize(image, self._output_dimensions) # pylint:disable=no-member self.cache_frame(filename, image) self._save_from_cache() @@ -140,7 +120,7 @@ def _set_gif_filename(self, filename: str) -> None: self._gif_file = retval logger.info("Outputting to: '%s'", self._gif_file) - def _set_dimensions(self, frame_dims: str) -> None: + def _set_dimensions(self, frame_dims: tuple[int, int]) -> None: """ Set the attribute :attr:`_output_dimensions` based on the first frame received. This protects against different sized images coming in and ensure all images get written to the Gif at the sema dimensions. """ @@ -151,16 +131,18 @@ def _set_dimensions(self, frame_dims: str) -> None: def _save_from_cache(self) -> None: """ Writes any consecutive frames to the GIF container that are ready to be output from the cache. """ - while self.frame_order: - if self.frame_order[0] not in self.cache: - logger.trace("Next frame not ready. Continuing") + assert self._writer is not None + while self._frame_order: + if self._frame_order[0] not in self.cache: + logger.trace("Next frame not ready. Continuing") # type: ignore break - save_no = self.frame_order.pop(0) + save_no = self._frame_order.pop(0) save_image = self.cache.pop(save_no) - logger.trace("Rendering from cache. Frame no: %s", save_no) + logger.trace("Rendering from cache. Frame no: %s", save_no) # type: ignore self._writer.append_data(save_image[:, :, ::-1]) - logger.trace("Current cache size: %s", len(self.cache)) + logger.trace("Current cache size: %s", len(self.cache)) # type: ignore def close(self) -> None: """ Close the GIF writer on completion. """ - self._writer.close() + if self._writer is not None: + self._writer.close() diff --git a/plugins/convert/writer/opencv.py b/plugins/convert/writer/opencv.py index 17cafdc591..17b025bfd9 100644 --- a/plugins/convert/writer/opencv.py +++ b/plugins/convert/writer/opencv.py @@ -2,20 +2,31 @@ """ Image output writer for faceswap.py converter Uses cv2 for writing as in testing this was a lot faster than both Pillow and ImageIO """ - import cv2 +import numpy as np + from ._base import Output, logger class Writer(Output): - """ Images output writer using cv2 """ - def __init__(self, output_folder, **kwargs): + """ Images output writer using cv2 + + Parameters + ---------- + output_folder: str + The full path to the output folder where the converted media should be saved + configfile: str, optional + The full path to a custom configuration ini file. If ``None`` is passed + then the file is loaded from the default location. Default: ``None``. + """ + def __init__(self, output_folder: str, **kwargs) -> None: super().__init__(output_folder, **kwargs) - self.extension = ".{}".format(self.config["format"]) - self.check_transparency_format() - self.args = self.get_save_args() + self._extension = f".{self.config['format']}" + self._check_transparency_format() + self._separate_mask = self.config["draw_transparent"] and self.config["separate_mask"] + self._args = self._get_save_args() - def check_transparency_format(self): + def _check_transparency_format(self) -> None: """ Make sure that the output format is correct if draw_transparent is selected """ transparent = self.config["draw_transparent"] if not transparent or (transparent and self.config["format"] == "png"): @@ -24,34 +35,77 @@ def check_transparency_format(self): "transparency. Changing output format to 'png'") self.config["format"] = "png" - def get_save_args(self): - """ Return the save parameters for the file format """ + def _get_save_args(self) -> tuple[int, ...]: + """ Obtain the save parameters for the file format. + + Returns + ------- + tuple + The OpenCV specific arguments for the selected file format + """ filetype = self.config["format"] - args = list() + args: tuple[int, ...] = tuple() if filetype == "jpg" and self.config["jpg_quality"] > 0: - args = (cv2.IMWRITE_JPEG_QUALITY, # pylint: disable=no-member + args = (cv2.IMWRITE_JPEG_QUALITY, self.config["jpg_quality"]) if filetype == "png" and self.config["png_compress_level"] > -1: - args = (cv2.IMWRITE_PNG_COMPRESSION, # pylint: disable=no-member + args = (cv2.IMWRITE_PNG_COMPRESSION, self.config["png_compress_level"]) logger.debug(args) return args - def write(self, filename, image): - logger.trace("Outputting: (filename: '%s'", filename) - filename = self.output_filename(filename) - try: - with open(filename, "wb") as outfile: - outfile.write(image) - except Exception as err: # pylint: disable=broad-except - logger.error("Failed to save image '%s'. Original Error: %s", filename, err) - - def pre_encode(self, image): - """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker """ - logger.trace("Pre-encoding image") - image = cv2.imencode(self.extension, image, self.args)[1] # pylint: disable=no-member - return image - - def close(self): - """ Image writer does not need a close method """ + def write(self, filename: str, image: list[bytes]) -> None: + """ Write out the pre-encoded image to disk. If separate mask has been selected, write out + the encoded mask to a sub-folder in the output directory. + + Parameters + ---------- + filename: str + The full path to write out the image to. + image: list + List of :class:`bytes` objects of length 1 (containing just the image to write out) + or length 2 (containing the image and mask to write out) + """ + logger.trace("Outputting: (filename: '%s'", filename) # type:ignore + filenames = self.output_filename(filename, self._separate_mask) + for fname, img in zip(filenames, image): + try: + with open(fname, "wb") as outfile: + outfile.write(img) + except Exception as err: # pylint:disable=broad-except + logger.error("Failed to save image '%s'. Original Error: %s", filename, err) + + def pre_encode(self, image: np.ndarray, **kwargs) -> list[bytes]: + """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker. + + Parameters + ---------- + image: :class:`numpy.ndarray` + A 3 or 4 channel BGR swapped frame + + Returns + ------- + list + List of :class:`bytes` objects ready for writing. The list will be of length 1 with + image bytes object as the only member unless separate mask has been requested, in which + case it will be length 2 with the image in position 0 and mask in position 1 + """ + logger.trace("Pre-encoding image") # type:ignore + retval = [] + + if self._separate_mask: + mask = image[..., -1] + image = image[..., :3] + + retval.append(cv2.imencode(self._extension, + mask, + self._args)[1]) + + retval.insert(0, cv2.imencode(self._extension, + image, + self._args)[1]) + return retval + + def close(self) -> None: + """ Does nothing as OpenCV writer does not need a close method """ return diff --git a/plugins/convert/writer/opencv_defaults.py b/plugins/convert/writer/opencv_defaults.py index c1bd96368b..67022e1ae0 100755 --- a/plugins/convert/writer/opencv_defaults.py +++ b/plugins/convert/writer/opencv_defaults.py @@ -78,6 +78,21 @@ gui_radio=False, fixed=True, ), + separate_mask=dict( + default=False, + info="Seperate the mask into its own single channel image. This only applies when " + "'draw-transparent' is selected. If enabled, the RGB image will be saved into the " + "selected output folder whilst the masks will be saved into a sub-folder named " + "`masks`. If not enabled then the mask will be included in the alpha-channel of the " + "RGBA output.", + datatype=bool, + rounding=None, + min_max=None, + choices=[], + group="format", + gui_radio=False, + fixed=True, + ), jpg_quality=dict( default=75, info="[jpg only] Set the jpg quality. 1 is worst 95 is best. Higher quality leads to " diff --git a/plugins/convert/writer/patch.py b/plugins/convert/writer/patch.py new file mode 100644 index 0000000000..fe4a7f4bb1 --- /dev/null +++ b/plugins/convert/writer/patch.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +""" Face patch output writer for faceswap.py converter + Extracts the swapped Face Patch from faceswap rather than the final composited frame along with + the transformation matrix for re-inserting the face into the origial frame +""" +import json +import logging +import re + +import os +import cv2 +import numpy as np + +from lib.image import encode_image, png_read_meta, tiff_read_meta +from ._base import Output + +logger = logging.getLogger(__name__) + + +class Writer(Output): + """ Face patch writer for outputting swapped face patches and transformation matrices + + Parameters + ---------- + output_folder: str + The full path to the output folder where the face patches should besaved + patch_size: int + The size of the face patch output from the model + configfile: str, optional + The full path to a custom configuration ini file. If ``None`` is passed + then the file is loaded from the default location. Default: ``None``. + """ + def __init__(self, output_folder: str, patch_size: int, **kwargs) -> None: + logger.debug("patch_size: %s", patch_size) + super().__init__(output_folder, **kwargs) + self._extension = {"png": ".png", "tiff": ".tif"}[self.config["format"]] + self._separate_mask = self.config["separate_mask"] + self._fname_split = re.compile("[^0-9a-zA-Z]") + + if self._extension == ".png" and self.config["bit_depth"] not in ("8", "16"): + logger.warning("Patch Writer: Bit Depth '%s' is unsupported for format '%s'. " + "Updating to '16'", self.config["bit_depth"], self.config["format"]) + self.config["bit_depth"] = "16" + + self._dtype = {"8": np.uint8, "16": np.uint16, "32": np.float32}[self.config["bit_depth"]] + self._multiplier = {"8": 255., "16": 65535., "32": 1.}[self.config["bit_depth"]] + + self._dummy_patch = np.zeros((1, patch_size, patch_size, 4), dtype=np.float32) + + tl_box = np.array([[0, 0], [patch_size, 0], [patch_size, patch_size], [0, patch_size]], + dtype=np.float32) + self._patch_corner = {"top-left": tl_box[0], + "top-right": tl_box[1], + "bottom-right": tl_box[2], + "bottom-left": tl_box[3]}[self.config["origin"]].copy() + self._box = tl_box + if self.config["origin"] in ("top-right", "bottom-left"): + self._box[[1, 3], :] = self._box[[3, 1], :] # keep clockwise from 0,0 + + self._args = self._get_save_args() + self._matrices: dict[str, dict[str, list[list[float]]]] = {} + + def _get_save_args(self) -> tuple[int, ...]: + """ Obtain the save parameters for the file format. + + Returns + ------- + tuple + The OpenCV specific arguments for the selected file format + """ + args: tuple[int, ...] = tuple() + if self._extension == ".png" and self.config["png_compress_level"] > -1: + args = (cv2.IMWRITE_PNG_COMPRESSION, self.config["png_compress_level"]) + if self._extension == ".tif" and self.config["bit_depth"] != "32": + tiff_methods = {"none": 1, "lzw": 5, "deflate": 8} + method = self.config["tiff_compression_method"] + method = "none" if method is None else method + args = (cv2.IMWRITE_TIFF_COMPRESSION, tiff_methods[method]) + logger.debug(args) + return args + + def _get_new_filename(self, filename: str, face_index: int) -> str: + """ Obtain the filename for the output file based on the frame's filename and the user + selected naming options + + Parameters + ---------- + filename: str + The original frame's filename + face_index: int + The index of the face within the frame + + Returns + ------- + str + The new filename for naming the output face patch + """ + face_idx = str(face_index).rjust(2, "0") + fname, ext = os.path.splitext(filename) + fname = os.path.basename(fname) + + split_fname = self._fname_split.split(fname) + if split_fname and split_fname[-1].isdigit(): + i_frame_no = (int(split_fname[-1]) + + (int(self.config["start_index"]) - 1) + + self.config["index_offset"]) + frame_no = f".{str(i_frame_no).rjust(self.config['number_padding'], '0')}" + base_fname = fname[:-len(split_fname[-1]) - 1] + else: + frame_no = "" + base_fname = fname + + retval = "" + if self.config["include_filename"]: + retval += base_fname + if self.config["face_index_location"] == "before": + retval = f"{retval}_{face_idx}" + retval += frame_no + if self.config["face_index_location"] == "after": + retval = f"{retval}.{face_idx}" + retval += ext + logger.trace("source filename: '%s', output filename: '%s'", # type:ignore[attr-defined] + filename, retval) + return retval + + def write(self, filename: str, image: list[list[bytes]]) -> None: + """ Write out the pre-encoded image to disk. If separate mask has been selected, write out + the encoded mask to a sub-folder in the output directory. + + Parameters + ---------- + filename: str + The full path to write out the image to. + image: list[list[bytes]] + List of list of :class:`bytes` objects of containing all swapped faces from a frame to + write out. The inner list will be of length 1 (mask included in the alpha channel) or + length 2 (mask to write out separately) + """ + logger.trace("Outputting: (filename: '%s')", filename) # type:ignore[attr-defined] + + read_func = png_read_meta if self._extension == ".png" else tiff_read_meta + for idx, face in enumerate(image): + new_filename = self._get_new_filename(filename, idx) + filenames = self.output_filename(new_filename, self._separate_mask) + for fname, img in zip(filenames, face): + try: + with open(fname, "wb") as outfile: + outfile.write(img) + except Exception as err: # pylint:disable=broad-except + logger.error("Failed to save image '%s'. Original Error: %s", filename, err) + if not self.config["json_output"]: + continue + mat = read_func(img) + self._matrices[os.path.splitext(os.path.basename(fname))[0]] = mat + + @classmethod + def _get_inverse_matrices(cls, matrices: np.ndarray) -> np.ndarray: + """ Obtain the inverse matrices for the given matrices. If ``None`` is supplied return a + dummy transformation matrix that performs no action + + Parameters + ---------- + matrices : :class:`numpy.ndarray` + The original transform matrices that the inverse needs to be calculated for + + Returns + ------- + :class:`numpy.ndarray` + The inverse transformation matrices + """ + if not np.any(matrices): + return np.array([[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]], dtype=np.float32) + + identity = np.array([[[0., 0., 1.]]], dtype=np.float32) + mat = np.concatenate([matrices, np.repeat(identity, matrices.shape[0], axis=0)], axis=1) + retval = np.linalg.inv(mat) + logger.trace("matrix: %s, inverse: %s", mat, retval) # type:ignore[attr-defined] + return retval + + def _adjust_to_origin(self, matrices: np.ndarray, canvas_size: tuple[int, int]) -> None: + """ Adjust the transformation matrix to use the correct target coordinates system. The + matrix adjustment is done in place, so this does not return a value + + Parameters + ---------- + matrices: :class:`numpy.ndarray` + The transformation matrices to be adjusted + canvas_size: tuple[int, int] + The size of the canvas width, height) that the transformation matrix applies to. + """ + if self.config["origin"] == "top-left": + return + + for mat in matrices: + og_cnr = cv2.transform(self._patch_corner[None, None], mat[:2, ...]).squeeze() + x_shift, y_shift = og_cnr + if self.config["origin"].split("-")[-1] == "right": + x_shift = canvas_size[0] - x_shift + if self.config["origin"].split("-")[0] == "bottom": + y_shift = canvas_size[1] - y_shift + mat[:2, 2] = [x_shift, y_shift] + + if self.config["origin"] in ("top-right", "bottom-left"): + matrices[..., :2, :2] *= [[[1, -1], [-1, 1]]] # switch shear + + def _get_roi(self, matrices: np.ndarray) -> np.ndarray: + """ Obtain the (x, y) ROI points of the patch in the original frame. Points are returned + in clockwise order from the origin location + + Parameters + ---------- + matrices: :class:`numpy.ndarray` + The transformation matrices for the current frame + + Returns + ------- + np.ndarray + The ROI of the patches in original frame co-ordinates in clockwise order from the + origin point + """ + retval = [cv2.transform(np.expand_dims(self._box, axis=1), mat[:2, ...]).squeeze() + for mat in matrices] + return np.array(retval, dtype=np.float32) + + def pre_encode(self, image: np.ndarray, **kwargs) -> list[list[bytes]]: + """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker. + + Parameters + ---------- + image: :class:`numpy.ndarray` + A 3 or 4 channel BGR swapped face batch as float32 + canvas_size: tuple[int, int] + The size of the canvas (x, y) that the transformation matrix applies to. + matrices: :class:`numpy.ndarray`, optional + The transformation matrices for extracting the face patches from the original frame. + Must be provided if an image is provided, otherwise ``None`` to insert a dummy matrix + + Returns + ------- + list + List of :class:`bytes` objects ready for writing. The list will be of length 1 with + image bytes object as the only member unless separate mask has been requested, in which + case it will be length 2 with the image in position 0 and mask in position 1 + """ + logger.trace("Pre-encoding image") # type:ignore[attr-defined] + retval = [] + canvas_size: tuple[int, int] = kwargs.get("canvas_size", (1, 1)) + matrices: np.ndarray = kwargs.get("matrices", np.array([])) + + if not np.any(image) and self.config["empty_frames"] == "blank": + image = self._dummy_patch + + matrices = self._get_inverse_matrices(matrices) + self._adjust_to_origin(matrices, canvas_size) + rois = self._get_roi(matrices) + patches = (image * self._multiplier).astype(self._dtype) + + for patch, matrix, roi in zip(patches, matrices, rois): + this_face = [] + mat = json.dumps({"transform_matrix": matrix.tolist(), "roi": roi.tolist()}, + ensure_ascii=True).encode("ascii") + if self._separate_mask: + mask = patch[..., -1] + face = patch[..., :3] + + this_face.append(encode_image(mask, + self._extension, + encoding_args=self._args, + metadata=mat)) + else: + face = patch + + this_face.insert(0, encode_image(face, + self._extension, + encoding_args=self._args, + metadata=mat)) + retval.append(this_face) + return retval + + def close(self) -> None: + """ Outputs json file if requested """ + if not self.config["json_output"]: + return + fname = os.path.join(self.output_folder, "matrices.json") + with open(fname, "w", encoding="utf-8") as ofile: + json.dump(self._matrices, ofile, indent=2, sort_keys=True) + logger.info("Patch matrices written to: '%s'", fname) diff --git a/plugins/convert/writer/patch_defaults.py b/plugins/convert/writer/patch_defaults.py new file mode 100755 index 0000000000..76f3f4c6e2 --- /dev/null +++ b/plugins/convert/writer/patch_defaults.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" + The default options for the faceswap patch Writer plugin. + + Defaults files should be named _defaults.py + Any items placed into this file will automatically get added to the relevant config .ini files + within the faceswap/config folder. + + The following variables should be defined: + _HELPTEXT: A string describing what this plugin does + _DEFAULTS: A dictionary containing the options, defaults and meta information. The + dictionary should be defined as: + {: {}} + + should always be lower text. + dictionary requirements are listed below. + + The following keys are expected for the _DEFAULTS dict: + datatype: [required] A python type class. This limits the type of data that can be + provided in the .ini file and ensures that the value is returned in the + correct type to faceswap. Valid datatypes are: , , + , . + default: [required] The default value for this option. + info: [required] A string describing what this option does. + choices: [optional] If this option's datatype is of then valid + selections can be defined here. This validates the option and also enables + a combobox / radio option in the GUI. + gui_radio: [optional] If are defined, this indicates that the GUI should use + radio buttons rather than a combobox to display this option. + min_max: [partial] For and datatypes this is required + otherwise it is ignored. Should be a tuple of min and max accepted values. + This is used for controlling the GUI slider range. Values are not enforced. + rounding: [partial] For and datatypes this is + required otherwise it is ignored. Used for the GUI slider. For floats, this + is the number of decimal places to display. For ints this is the step size. + fixed: [optional] [train only]. Training configurations are fixed when the model is + created, and then reloaded from the state file. Marking an item as fixed=False + indicates that this value can be changed for existing models, and will override + the value saved in the state file with the updated value in config. If not + provided this will default to True. +""" + +_HELPTEXT = ( + "Options for outputting the raw converted face patches from faceswap\n" + "The raw face patches are output along with the transformation matrix, per face, to " + "transform the face back into the original frame in external tools" +) + +_DEFAULTS = { + "start_index": { + "default": "0", + "info": "The starting frame number for the first output frame.", + "datatype": str, + "choices": ["0", "1"], + "group": "file_naming", + "gui_radio": True, + }, + "index_offset": { + "default": 0, + "info": "How much to offset the frame numbering by.", + "datatype": int, + "rounding": 1, + "min_max": (0, 1000), + "group": "file_naming", + }, + "number_padding": { + "default": 6, + "info": "Length to pad the frame numbers by.", + "datatype": int, + "rounding": 6, + "min_max": (0, 10), + "group": "file_naming", + }, + "include_filename": { + "default": True, + "info": "Prefix the filename of the original frame to each face patch's output filename.", + "datatype": bool, + "group": "file_naming", + }, + "face_index_location": { + "default": "before", + "info": "For frames that contain multiple faces, where the face index should appear in " + "the filename:" + "\n\t before: places the face index before the frame number." + "\n\t after: places the face index after the frame number.", + "datatype": str, + "choices": ["before", "after"], + "group": "file_naming", + "gui_radio": True, + }, + "origin": { + "default": "bottom-left", + "info": "The origin (0, 0) location of the software that patches will be imported into. " + "This impacts the transformation matrix that is supplied with the image patch. " + "Setting the correct origin here will make importing into the external tool " + "simpler." + "\n\t top-left: The origin (0, 0) of the external canvas is at the top left " + "corner." + "\n\t bottom-left: The origin (0, 0) of the external canvas is at the bottom " + "left corner." + "\n\t top-right: The origin (0, 0) of the external canvas is at the top right " + "corner." + "\n\t bottom-right: The origin (0, 0) of the external canvas is at the bottom " + "right corner.", + "datatype": str, + "choices": ["top-left", "bottom-left", "top-right", "bottom-right"], + "group": "output", + "gui_radio": True + }, + "empty_frames": { + "default": "blank", + "info": "How to handle the output of frames without faces:" + "\n\t skip: skips any frames that do not have a face within it. This will lead to " + "gaps within the final image sequence." + "\n\t blank: outputs a blank (empty) face patch for any frames without faces. " + "There will be no gaps within the final image sequence, as those gaps will be " + "padded with empty face patches", + "datatype": str, + "choices": ["skip", "blank"], + "group": "output", + "gui_radio": True, + }, + "json_output": { + "default": False, + "info": "The transformation matrix, and other associated metadata, is output within the " + "face images EXIF fields. Some external tools can read this data, others cannot." + "enable this option to output a json file which contains this same metadata " + "mapped to each output face patch's filename.", + "datatype": bool, + "group": "output" + }, + "separate_mask": { + "default": False, + "info": "Seperate the mask into its own single channel patch. If enabled, the RGB image " + "will be saved into the selected output folder whilst the masks will be saved " + "into a sub-folder named `masks`. If not enabled then the mask will be included " + "in the alpha-channel of the RGBA output.", + "datatype": bool, + "group": "output", + }, + "bit_depth": { + "default": "16", + "info": "The bit-depth for the output images:" + "\n\t 8: 8-bit unsigned - Supported by all formats." + "\n\t 16: 16-bit unsigned - Supported by all formats." + "\n\t 32: 32-bit float - Supported by Tiff only.", + "datatype": str, + "choices": ["8", "16", "32"], + "group": "format", + "gui_radio": True, + }, + "format": { + "default": "png", + "info": "File format to save as." + "\n\t png: PNG file format. Transformation matrix is written to the custom iTxt " + "header field 'faceswap'" + "\n\t tiff: TIFF file format. Transformation matrix is written to the " + "'image_description' header field", + "datatype": str, + "choices": ["png", "tiff"], + "group": "format", + "gui_radio": True + }, + "png_compress_level": { + "default": 3, + "info": "ZLIB compression level, 1 gives best speed, 9 gives best compression, 0 gives no " + "compression at all.", + "datatype": int, + "rounding": 1, + "min_max": (0, 9), + "group": "format", + }, + "tiff_compression_method": { + "default": "lzw", + "info": "The compression method to use for Tiff files. Note: For 32bit output, SGILOG " + "compression will always be used regardless of what is selected here.", + "datatype": str, + "choices": ["none", "lzw", "deflate"], + "group": "format", + "gui_radio": True + }, +} diff --git a/plugins/convert/writer/pillow.py b/plugins/convert/writer/pillow.py index ef440f2b6c..a0bf113f62 100644 --- a/plugins/convert/writer/pillow.py +++ b/plugins/convert/writer/pillow.py @@ -1,22 +1,33 @@ #!/usr/bin/env python3 """ Image output writer for faceswap.py converter """ - from io import BytesIO from PIL import Image +import numpy as np + from ._base import Output, logger class Writer(Output): - """ Images output writer using cv2 """ - def __init__(self, output_folder, **kwargs): + """ Images output writer using Pillow + + Parameters + ---------- + output_folder: str + The full path to the output folder where the converted media should be saved + configfile: str, optional + The full path to a custom configuration ini file. If ``None`` is passed + then the file is loaded from the default location. Default: ``None``. + """ + def __init__(self, output_folder: str, **kwargs) -> None: super().__init__(output_folder, **kwargs) - self.check_transparency_format() + self._check_transparency_format() # Correct format namings for writing to byte stream - self.format_dict = dict(jpg="JPEG", jp2="JPEG 2000", tif="TIFF") - self.kwargs = self.get_save_kwargs() + self._format_dict = {"jpg": "JPEG", "jp2": "JPEG 2000", "tif": "TIFF"} + self._separate_mask = self.config["draw_transparent"] and self.config["separate_mask"] + self._kwargs = self._get_save_kwargs() - def check_transparency_format(self): + def _check_transparency_format(self) -> None: """ Make sure that the output format is correct if draw_transparent is selected """ transparent = self.config["draw_transparent"] if not transparent or (transparent and self.config["format"] in ("png", "tif")): @@ -25,10 +36,16 @@ def check_transparency_format(self): "transparency. Changing output format to 'png'") self.config["format"] = "png" - def get_save_kwargs(self): - """ Return the save parameters for the file format """ + def _get_save_kwargs(self) -> dict[str, bool | int | str]: + """ Return the save parameters for the file format + + Returns + ------- + dict + The specific keyword arguments for the selected file format + """ filetype = self.config["format"] - kwargs = dict() + kwargs = {} if filetype in ("gif", "jpg", "png"): kwargs["optimize"] = self.config["optimize"] if filetype == "gif": @@ -40,29 +57,78 @@ def get_save_kwargs(self): logger.debug(kwargs) return kwargs - def write(self, filename, image): - logger.trace("Outputting: (filename: '%s'", filename) - filename = self.output_filename(filename) + def write(self, filename: str, image: list[BytesIO]) -> None: + """ Write out the pre-encoded image to disk. If separate mask has been selected, write out + the encoded mask to a sub-folder in the output directory. + + Parameters + ---------- + filename: str + The full path to write out the image to. + image: list + List of :class:`BytesIO` objects of length 1 (containing just the image to write out) + or length 2 (containing the image and mask to write out) + """ + logger.trace("Outputting: (filename: '%s'", filename) # type:ignore + filenames = self.output_filename(filename, self._separate_mask) try: - with open(filename, "wb") as outfile: - outfile.write(image.read()) - except Exception as err: # pylint: disable=broad-except + for fname, img in zip(filenames, image): + with open(fname, "wb") as outfile: + outfile.write(img.read()) + except Exception as err: # pylint:disable=broad-except logger.error("Failed to save image '%s'. Original Error: %s", filename, err) - def pre_encode(self, image): - """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker """ - logger.trace("Pre-encoding image") - fmt = self.format_dict.get(self.config["format"], None) - fmt = self.config["format"].upper() if fmt is None else fmt + def pre_encode(self, image: np.ndarray, **kwargs) -> list[BytesIO]: + """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker + + Parameters + ---------- + image: :class:`numpy.ndarray` + A 3 or 4 channel BGR swapped frame + + Returns + ------- + list + List of :class:`BytesIO` objects ready for writing. The list will be of length 1 with + image bytes object as the only member unless separate mask has been requested, in which + case it will be length 2 with the image in position 0 and mask in position 1 + """ + logger.trace("Pre-encoding image") # type:ignore + + if self._separate_mask: + encoded_mask = self._encode_image(image[..., -1]) + image = image[..., :3] + + rgb = [2, 1, 0, 3] if image.shape[2] == 4 else [2, 1, 0] + encoded_image = self._encode_image(image[..., rgb]) + + retval = [encoded_image] + + if self._separate_mask: + retval.append(encoded_mask) + + return retval + + def _encode_image(self, image: np.ndarray) -> BytesIO: + """ Encode an image in the correct format as a bytes object for saving + + Parameters + ---------- + image: :class:`np.ndarray` + The single channel mask to encode for saving + + Returns + ------- + :class:`BytesIO` + The image as a bytes object ready for writing to disk + """ + fmt = self._format_dict.get(self.config["format"], self.config["format"].upper()) encoded = BytesIO() - rgb = [2, 1, 0] - if image.shape[2] == 4: - rgb.append(3) - out_image = Image.fromarray(image[..., rgb]) - out_image.save(encoded, fmt, **self.kwargs) + out_image = Image.fromarray(image) + out_image.save(encoded, fmt, **self._kwargs) encoded.seek(0) return encoded - def close(self): - """ Image writer does not need a close method """ + def close(self) -> None: + """ Does nothing as Pillow writer does not need a close method """ return diff --git a/plugins/convert/writer/pillow_defaults.py b/plugins/convert/writer/pillow_defaults.py index d6f17bf936..58bf7e31ea 100755 --- a/plugins/convert/writer/pillow_defaults.py +++ b/plugins/convert/writer/pillow_defaults.py @@ -79,6 +79,21 @@ gui_radio=False, fixed=True, ), + separate_mask=dict( + default=False, + info="Seperate the mask into its own single channel image. This only applies when " + "'draw-transparent' is selected. If enabled, the RGB image will be saved into the " + "selected output folder whilst the masks will be saved into a sub-folder named " + "`masks`. If not enabled then the mask will be included in the alpha-channel of the " + "RGBA output.", + datatype=bool, + rounding=None, + min_max=None, + choices=[], + group="format", + gui_radio=False, + fixed=True, + ), optimize=dict( default=False, info="[gif, jpg and png only] If enabled, indicates that the encoder should make an extra " diff --git a/plugins/extract/__init__.py b/plugins/extract/__init__.py index e69de29bb2..3bffbe70b8 100644 --- a/plugins/extract/__init__.py +++ b/plugins/extract/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python3 +""" Package for Faceswap's extraction pipeline """ +from .extract_media import ExtractMedia +from .pipeline import Extractor diff --git a/plugins/extract/_base.py b/plugins/extract/_base.py index 51c8463942..588394adee 100644 --- a/plugins/extract/_base.py +++ b/plugins/extract/_base.py @@ -2,25 +2,37 @@ """ Base class for Faceswap :mod:`~plugins.extract.detect`, :mod:`~plugins.extract.align` and :mod:`~plugins.extract.mask` Plugins """ +from __future__ import annotations import logging -import os -import sys +import typing as T -from tensorflow.python.framework import errors_impl as tf_errors +from dataclasses import dataclass, field + +import numpy as np +from tensorflow.python.framework import errors_impl as tf_errors # pylint:disable=no-name-in-module # noqa from lib.multithreading import MultiThread from lib.queue_manager import queue_manager from lib.utils import GetModel, FaceswapError from ._config import Config -from .pipeline import ExtractMedia - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - -# TODO CPU mode +from . import ExtractMedia + +if T.TYPE_CHECKING: + from collections.abc import Callable, Generator, Sequence + from queue import Queue + import cv2 + from lib.align import DetectedFace + from lib.model.session import KSession + from .align._base import AlignerBatch + from .detect._base import DetectorBatch + from .mask._base import MaskerBatch + from .recognition._base import RecogBatch + +logger = logging.getLogger(__name__) # TODO Run with warnings mode -def _get_config(plugin_name, configfile=None): +def _get_config(plugin_name: str, configfile: str | None = None) -> dict[str, T.Any]: """ Return the configuration for the requested model Parameters @@ -39,6 +51,54 @@ def _get_config(plugin_name, configfile=None): return Config(plugin_name, configfile=configfile).config_dict +BatchType = T.Union["DetectorBatch", "AlignerBatch", "MaskerBatch", "RecogBatch"] + + +@dataclass +class ExtractorBatch: + """ Dataclass for holding a batch flowing through post Detector plugins. + + The batch size for post Detector plugins is not the same as the overall batch size. + An image may contain 0 or more detected faces, and these need to be split and recombined + to be able to utilize a plugin's internal batch size. + + Plugin types will inherit from this class and add required keys. + + Parameters + ---------- + image: list + List of :class:`numpy.ndarray` containing the original frames + detected_faces: list + List of :class:`~lib.align.DetectedFace` objects + filename: list + List of original frame filenames for the batch + feed: :class:`numpy.ndarray` + Batch of feed images to feed the net with + prediction: :class:`numpy.nd.array` + Batch of predictions. Direct output from the aligner net + data: dict + Any specific data required during the processing phase for a particular plugin + """ + image: list[np.ndarray] = field(default_factory=list) + detected_faces: Sequence[DetectedFace | list[DetectedFace]] = field(default_factory=list) + filename: list[str] = field(default_factory=list) + feed: np.ndarray = np.array([]) + prediction: np.ndarray = np.array([]) + data: list[dict[str, T.Any]] = field(default_factory=list) + + def __repr__(self) -> str: + """ Prettier repr for debug printing """ + data = [{k: (v.shape, v.dtype) if isinstance(v, np.ndarray) else v for k, v in dat.items()} + for dat in self.data] + return (f"{self.__class__.__name__}(" + f"image={[(img.shape, img.dtype) for img in self.image]}, " + f"detected_faces={self.detected_faces}, " + f"filename={self.filename}, " + f"feed={[(f.shape, f.dtype) for f in self.feed]}, " + f"prediction=({self.prediction.shape}, {self.prediction.dtype}), " + f"data={data}") + + class Extractor(): """ Extractor Plugin Object @@ -102,12 +162,16 @@ class Extractor(): plugins.extract.pipeline : The extract pipeline that configures and calls all plugins """ - def __init__(self, git_model_id=None, model_filename=None, exclude_gpus=None, configfile=None, - instance=0): + def __init__(self, + git_model_id: int | None = None, + model_filename: str | list[str] | None = None, + exclude_gpus: list[int] | None = None, + configfile: str | None = None, + instance: int = 0) -> None: logger.debug("Initializing %s: (git_model_id: %s, model_filename: %s, exclude_gpus: %s, " "configfile: %s, instance: %s, )", self.__class__.__name__, git_model_id, model_filename, exclude_gpus, configfile, instance) - + self._is_initialized = False self._instance = instance self._exclude_gpus = exclude_gpus self.config = _get_config(".".join(self.__module__.split(".")[-2:]), configfile=configfile) @@ -118,20 +182,19 @@ def __init__(self, git_model_id=None, model_filename=None, exclude_gpus=None, co be a list of strings """ # << SET THE FOLLOWING IN PLUGINS __init__ IF DIFFERENT FROM DEFAULT >> # - self.name = None - self.input_size = None - self.color_format = "BGR" - self.vram = None - self.vram_warnings = None # Will run at this with warnings - self.vram_per_batch = None + self.name: str | None = None + self.input_size = 0 + self.color_format: T.Literal["BGR", "RGB", "GRAY"] = "BGR" + self.vram = 0 + self.vram_warnings = 0 # Will run at this with warnings + self.vram_per_batch = 0 # << THE FOLLOWING ARE SET IN self.initialize METHOD >> # self.queue_size = 1 """ int: Queue size for all internal queues. Set in :func:`initialize()` """ - self.model = None - """varies: The model for this plugin. - Set in the plugin's :func:`init_model()` method """ + self.model: KSession | cv2.dnn.Net | None = None + """varies: The model for this plugin. Set in the plugin's :func:`init_model()` method """ # For detectors that support batching, this should be set to the calculated batch size # that the amount of available VRAM will support. @@ -139,93 +202,98 @@ def __init__(self, git_model_id=None, model_filename=None, exclude_gpus=None, co """ int: Batchsize for feeding this model. The number of images the model should feed through at once. """ - self._queues = dict() + self._queues: dict[str, Queue] = {} """ dict: in + out queues and internal queues for this plugin, """ - self._threads = [] + self._threads: list[MultiThread] = [] """ list: Internal threads for this plugin """ - self._extract_media = dict() - """ dict: The :class:`plugins.extract.pipeline.ExtractMedia` objects currently being + self._extract_media: dict[str, ExtractMedia] = {} + """ dict: The :class:`~plugins.extract.extract_media.ExtractMedia` objects currently being processed. Stored at input for pairing back up on output of extractor process """ # << THE FOLLOWING PROTECTED ATTRIBUTES ARE SET IN PLUGIN TYPE _base.py >>> # - self._plugin_type = None - """ str: Plugin type. ``detect`` or ``align`` - set in ``._base`` """ + self._plugin_type: T.Literal["align", "detect", "recognition", "mask"] | None = None + """ str: Plugin type. ``detect`, ``align``, ``recognise`` or ``mask`` set in + ``._base`` """ + + # << Objects for splitting frame's detected faces and rejoining them >> + # << for post-detector pliugins >> + self._faces_per_filename: dict[str, int] = {} # Tracking for recompiling batches + self._rollover: ExtractMedia | None = None # batch rollover items + self._output_faces: list[DetectedFace] = [] # Recompiled output faces from plugin logger.debug("Initialized _base %s", self.__class__.__name__) # <<< OVERIDABLE METHODS >>> # - def init_model(self): + def init_model(self) -> None: """ **Override method** Override this method to execute the specific model initialization method """ raise NotImplementedError - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ **Override method** Override this method for specific extractor pre-processing of image Parameters ---------- - batch : dict + batch : :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process - - Notes - ----- - When preparing an input to the model a key ``feed`` must be added - to the :attr:`batch` ``dict`` which contains this input. """ raise NotImplementedError - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ **Override method** Override this method for specific extractor model prediction function Parameters ---------- - batch : dict - Contains the batch that is currently being passed through the plugin process + feed: :class:`numpy.ndarray` + The feed images for the batch Notes ----- - Input for :func:`predict` should have been set in :func:`process_input` with the addition - of a ``feed`` key to the :attr:`batch` ``dict``. + Input for :func:`predict` should have been set in :func:`process_input` - Output from the model should add the key ``prediction`` to the :attr:`batch` ``dict``. + Output from the model should populate the key :attr:`prediction` of the :attr:`batch`. For Detect: - the expected output for the ``prediction`` key of the :attr:`batch` dict should be a + the expected output for the :attr:`prediction` of the :attr:`batch` should be a ``list`` of :attr:`batchsize` of detected face points. These points should be either a ``list``, ``tuple`` or ``numpy.ndarray`` with the first 4 items being the `left`, `top`, `right`, `bottom` points, in that order """ raise NotImplementedError - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ **Override method** Override this method for specific extractor model post predict function Parameters ---------- - batch : dict + batch: :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process Notes ----- For Align: - The key ``landmarks`` must be returned in the :attr:`batch` ``dict`` from this method. - This should be a ``list`` or ``numpy.ndarray`` of :attr:`batchsize` containing a - ``list``, ``tuple`` or ``numpy.ndarray`` of `(x, y)` coordinates of the 68 point + The :attr:`landmarks` must be populated in :attr:`batch` from this method. + This should be a ``list`` or :class:`numpy.ndarray` of :attr:`batchsize` containing a + ``list``, ``tuple`` or :class:`numpy.ndarray` of `(x, y)` coordinates of the 68 point landmarks as calculated from the :attr:`model`. """ raise NotImplementedError - def _predict(self, batch): + def on_completion(self) -> None: + """ Override to perform an action when the extract process has completed. By default, no + action is undertaken """ + return + + def _predict(self, batch: BatchType) -> BatchType: """ **Override method** (at `` level) This method should be overridden at the `` level (IE. @@ -237,12 +305,12 @@ def _predict(self, batch): Parameters ---------- - batch : dict + batch: :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process """ raise NotImplementedError - def _process_input(self, batch): + def _process_input(self, batch: BatchType) -> BatchType: """ **Override method** (at `` level) This method should be overridden at the `` level (IE. @@ -256,17 +324,18 @@ def _process_input(self, batch): Parameters ---------- - batch : dict + batch: :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process Notes ----- - When preparing an input to the model a key ``feed`` must be added - to the :attr:`batch` ``dict`` which contains this input. + When preparing an input to the model a the attribute :attr:`feed` must be added + to the :attr:`batch` which contains this input. """ - return self.process_input(batch) + self.process_input(batch) + return batch - def _process_output(self, batch): + def _process_output(self, batch: BatchType) -> BatchType: """ **Override method** (at `` level) This method should be overridden at the `` level (IE. @@ -280,12 +349,13 @@ def _process_output(self, batch): Parameters ---------- - batch : dict + batch: :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process """ - return self.process_output(batch) + self.process_output(batch) + return batch - def finalize(self, batch): + def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]: """ **Override method** (at `` level) This method should be overridden at the `` level (IE. @@ -297,20 +367,19 @@ def finalize(self, batch): Parameters ---------- - batch : dict + batch: :class:`ExtractorBatch` Contains the batch that is currently being passed through the plugin process - """ raise NotImplementedError - def get_batch(self, queue): + def get_batch(self, queue: Queue) -> tuple[bool, BatchType]: """ **Override method** (at `` level) This method should be overridden at the `` level (IE. :mod:`plugins.extract.detect._base`, :mod:`plugins.extract.align._base` or :mod:`plugins.extract.mask._base`) and should not be overridden within plugins themselves. - Get :class:`~plugins.extract.pipeline.ExtractMedia` items from the queue in batches of + Get :class:`~plugins.extract.extract_media.ExtractMedia` items from the queue in batches of :attr:`batchsize` Parameters @@ -321,7 +390,7 @@ def get_batch(self, queue): raise NotImplementedError # <<< THREADING METHODS >>> # - def start(self): + def start(self) -> None: """ Start all threads Exposed for :mod:`~plugins.extract.pipeline` to start plugin's threads @@ -329,30 +398,63 @@ def start(self): for thread in self._threads: thread.start() - def join(self): + def join(self) -> None: """ Join all threads Exposed for :mod:`~plugins.extract.pipeline` to join plugin's threads """ for thread in self._threads: thread.join() - del thread - def check_and_raise_error(self): + def check_and_raise_error(self) -> None: """ Check all threads for errors Exposed for :mod:`~plugins.extract.pipeline` to check plugin's threads for errors """ for thread in self._threads: - err = thread.check_and_raise_error() - if err is not None: - logger.debug("thread_error_detected") - return True - return False + thread.check_and_raise_error() + + def rollover_collector(self, queue: Queue) -> T.Literal["EOF"] | ExtractMedia: + """ For extractors after the Detectors, the number of detected faces per frame vs extractor + batch size mean that faces will need to be split/re-joined with frames. The rollover + collector can be used to rollover items that don't fit in a batch. + + Collect the item from the :attr:`_rollover` dict or from the queue. Add face count per + frame to self._faces_per_filename for joining batches back up in finalize + + Parameters + ---------- + queue: :class:`queue.Queue` + The input queue to the aligner. Should contain + :class:`~plugins.extract.extract_media.ExtractMedia` objects + + Returns + ------- + :class:`~plugins.extract.extract_media.ExtractMedia` or EOF + The next extract media object, or EOF if pipe has ended + """ + if self._rollover is not None: + logger.trace("Getting from _rollover: (filename: `%s`, faces: %s)", # type:ignore + self._rollover.filename, len(self._rollover.detected_faces)) + item: T.Literal["EOF"] | ExtractMedia = self._rollover + self._rollover = None + else: + next_item = self._get_item(queue) + # Rollover collector should only be used at entry to plugin + assert isinstance(next_item, (ExtractMedia, str)) + item = next_item + if item != "EOF": + logger.trace("Getting from queue: (filename: %s, faces: %s)", # type:ignore + item.filename, len(item.detected_faces)) + self._faces_per_filename[item.filename] = len(item.detected_faces) + return item # <<< PROTECTED ACCESS METHODS >>> # # <<< INIT METHODS >>> # - def _get_model(self, git_model_id, model_filename): + @classmethod + def _get_model(cls, + git_model_id: int | None, + model_filename: str | list[str] | None) -> str | list[str] | None: """ Check if model is available, if not, download and unzip it """ if model_filename is None: logger.debug("No model_filename specified. Returning None") @@ -360,29 +462,30 @@ def _get_model(self, git_model_id, model_filename): if git_model_id is None: logger.debug("No git_model_id specified. Returning None") return None - plugin_path = os.path.join(*self.__module__.split(".")[:-1]) - if os.path.basename(plugin_path) in ("detect", "align", "mask", "recognition"): - base_path = os.path.dirname(os.path.realpath(sys.argv[0])) - cache_path = os.path.join(base_path, plugin_path, ".cache") - else: - cache_path = os.path.join(os.path.dirname(__file__), ".cache") - model = GetModel(model_filename, cache_path, git_model_id) + model = GetModel(model_filename, git_model_id) return model.model_path # <<< PLUGIN INITIALIZATION >>> # - def initialize(self, *args, **kwargs): + def initialize(self, *args, **kwargs) -> None: """ Initialize the extractor plugin Should be called from :mod:`~plugins.extract.pipeline` """ logger.debug("initialize %s: (args: %s, kwargs: %s)", self.__class__.__name__, args, kwargs) + assert self._plugin_type is not None and self.name is not None + if self._is_initialized: + # When batch processing, plugins will be initialized on first job in batch + logger.debug("Plugin already initialized: %s (%s)", + self.name, self._plugin_type.title()) + return + logger.info("Initializing %s (%s)...", self.name, self._plugin_type.title()) self.queue_size = 1 name = self.name.replace(" ", "_").lower() self._add_queues(kwargs["in_queue"], kwargs["out_queue"], - ["predict_{}".format(name), "post_{}".format(name)]) + [f"predict_{name}", f"post_{name}"]) self._compile_threads() try: self.init_model() @@ -398,10 +501,14 @@ def initialize(self, *args, **kwargs): "option to `True`.") raise FaceswapError(msg) from err raise err + self._is_initialized = True logger.info("Initialized %s (%s) with batchsize of %s", self.name, self._plugin_type.title(), self.batchsize) - def _add_queues(self, in_queue, out_queue, queues): + def _add_queues(self, + in_queue: Queue, + out_queue: Queue, + queues: list[str]) -> None: """ Add the queues in_queue and out_queue should be previously created queue manager queues. queues should be a list of queue names """ @@ -409,30 +516,35 @@ def _add_queues(self, in_queue, out_queue, queues): self._queues["out"] = out_queue for q_name in queues: self._queues[q_name] = queue_manager.get_queue( - name="{}{}_{}".format(self._plugin_type, self._instance, q_name), + name=f"{self._plugin_type}{self._instance}_{q_name}", maxsize=self.queue_size) # <<< THREAD METHODS >>> # - def _compile_threads(self): + def _compile_threads(self) -> None: """ Compile the threads into self._threads list """ + assert self.name is not None logger.debug("Compiling %s threads", self._plugin_type) name = self.name.replace(" ", "_").lower() - base_name = "{}_{}".format(self._plugin_type, name) - self._add_thread("{}_input".format(base_name), + base_name = f"{self._plugin_type}_{name}" + self._add_thread(f"{base_name}_input", self._process_input, self._queues["in"], - self._queues["predict_{}".format(name)]) - self._add_thread("{}_predict".format(base_name), + self._queues[f"predict_{name}"]) + self._add_thread(f"{base_name}_predict", self._predict, - self._queues["predict_{}".format(name)], - self._queues["post_{}".format(name)]) - self._add_thread("{}_output".format(base_name), + self._queues[f"predict_{name}"], + self._queues[f"post_{name}"]) + self._add_thread(f"{base_name}_output", self._process_output, - self._queues["post_{}".format(name)], + self._queues[f"post_{name}"], self._queues["out"]) logger.debug("Compiled %s threads: %s", self._plugin_type, self._threads) - def _add_thread(self, name, function, in_queue, out_queue): + def _add_thread(self, + name: str, + function: Callable[[BatchType], BatchType], + in_queue: Queue, + out_queue: Queue) -> None: """ Add a MultiThread thread to self._threads """ logger.debug("Adding thread: (name: %s, function: %s, in_queue: %s, out_queue: %s)", name, function, in_queue, out_queue) @@ -443,24 +555,66 @@ def _add_thread(self, name, function, in_queue, out_queue): out_queue=out_queue)) logger.debug("Added thread: %s", name) - def _thread_process(self, function, in_queue, out_queue): - """ Perform a plugin function in a thread """ - func_name = function.__name__ - logger.debug("threading: (function: '%s')", func_name) + def _obtain_batch_item(self, function: Callable[[BatchType], BatchType], + in_queue: Queue, + out_queue: Queue) -> BatchType | None: + """ Obtain the batch item from the in queue for the current process. + + Parameters + ---------- + function: callable + The current plugin function being run + in_queue: :class:`queue.Queue` + The input queue for the function + out_queue: :class:`queue.Queue` + The output queue from the function + + Returns + ------- + :class:`ExtractorBatch` or ``None`` + The batch, if one exists, or ``None`` if queue is exhausted + """ + batch: T.Literal["EOF"] | BatchType | ExtractMedia + if function.__name__ == "_process_input": # Process input items to batches + exhausted, batch = self.get_batch(in_queue) + if exhausted: + if batch.filename: + # Put the final batch + batch = function(batch) + out_queue.put(batch) + return None + else: + batch = self._get_item(in_queue) + if batch == "EOF": + return None + + # ExtractMedia should only ever be the output of _get_item at the entry to a + # plugin's pipeline (ie in _process_input) + assert not isinstance(batch, ExtractMedia) + return batch + + def _thread_process(self, + function: Callable[[BatchType], BatchType], + in_queue: Queue, + out_queue: Queue) -> None: + """ Perform a plugin function in a thread + + Parameters + ---------- + function: callable + The current plugin function being run + in_queue: :class:`queue.Queue` + The input queue for the function + out_queue: :class:`queue.Queue` + The output queue from the function + """ + logger.debug("threading: (function: '%s')", function.__name__) while True: - if func_name == "_process_input": - # Process input items to batches - exhausted, batch = self.get_batch(in_queue) - if exhausted: - if batch: - # Put the final batch - batch = function(batch) - out_queue.put(batch) - break - else: - batch = self._get_item(in_queue) - if batch == "EOF": - break + batch = self._obtain_batch_item(function, in_queue, out_queue) + if batch is None: + break + if not batch.filename: # Batch not populated. Possible during re-aligns + continue try: batch = function(batch) except tf_errors.UnknownError as err: @@ -475,7 +629,7 @@ def _thread_process(self, function, in_queue, out_queue): "`allow_growth option to `True`.") raise FaceswapError(msg) from err raise err - if func_name == "_process_output": + if function.__name__ == "_process_output": # Process output items to individual items from batch for item in self.finalize(batch): out_queue.put(item) @@ -485,19 +639,14 @@ def _thread_process(self, function, in_queue, out_queue): out_queue.put("EOF") # <<< QUEUE METHODS >>> # - def _get_item(self, queue): + def _get_item(self, queue: Queue) -> T.Literal["EOF"] | ExtractMedia | BatchType: """ Yield one item from a queue """ item = queue.get() if isinstance(item, ExtractMedia): - logger.trace("filename: '%s', image shape: %s, detected_faces: %s, queue: %s, " - "item: %s", + logger.trace("filename: '%s', image shape: %s, detected_faces: %s, " # type:ignore + "queue: %s, item: %s", item.filename, item.image_shape, item.detected_faces, queue, item) self._extract_media[item.filename] = item else: - logger.trace("item: %s, queue: %s", item, queue) + logger.trace("item: %s, queue: %s", item, queue) # type:ignore return item - - @staticmethod - def _dict_lists_to_list_dicts(dictionary): - """ Convert a dictionary of lists to a list of dictionaries """ - return [dict(zip(dictionary, val)) for val in zip(*dictionary.values())] diff --git a/plugins/extract/_config.py b/plugins/extract/_config.py index e2e58643f1..8314cab2f0 100644 --- a/plugins/extract/_config.py +++ b/plugins/extract/_config.py @@ -1,33 +1,140 @@ #!/usr/bin/env python3 """ Default configurations for extract """ +import gettext import logging import os from lib.config import FaceswapConfig -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +# LOCALES +_LANG = gettext.translation("plugins.extract._config", localedir="locales", fallback=True) +_ = _LANG.gettext + +logger = logging.getLogger(__name__) class Config(FaceswapConfig): """ Config File for Extraction """ - def set_defaults(self): + def set_defaults(self) -> None: """ Set the default values for config """ logger.debug("Setting defaults") self.set_globals() self._defaults_from_plugin(os.path.dirname(__file__)) - def set_globals(self): + def set_globals(self) -> None: """ Set the global options for extract """ logger.debug("Setting global config") section = "global" - self.add_section(title=section, info="Options that apply to all extraction plugins") + self.add_section(section, _("Options that apply to all extraction plugins")) + self.add_item( + section=section, + title="allow_growth", + datatype=bool, + default=False, + group=_("settings"), + info=_("Enable the Tensorflow GPU `allow_growth` configuration option. " + "This option prevents Tensorflow from allocating all of the GPU VRAM at launch " + "but can lead to higher VRAM fragmentation and slower performance. Should only " + "be enabled if you are having problems running extraction.")) + self.add_item( + section=section, + title="aligner_min_scale", + datatype=float, + min_max=(0.0, 1.0), + rounding=2, + default=0.07, + group=_("filters"), + info=_("Filters out faces below this size. This is a multiplier of the minimum " + "dimension of the frame (i.e. 1280x720 = 720). If the original face extract " + "box is smaller than the minimum dimension times this multiplier, it is " + "considered a false positive and discarded. Faces which are found to be " + "unusually smaller than the frame tend to be misaligned images, except in " + "extreme long-shots. These can be usually be safely discarded.")) + self.add_item( + section=section, + title="aligner_max_scale", + datatype=float, + min_max=(0.0, 10.0), + rounding=2, + default=2.00, + group=_("filters"), + info=_("Filters out faces above this size. This is a multiplier of the minimum " + "dimension of the frame (i.e. 1280x720 = 720). If the original face extract " + "box is larger than the minimum dimension times this multiplier, it is " + "considered a false positive and discarded. Faces which are found to be " + "unusually larger than the frame tend to be misaligned images except in " + "extreme close-ups. These can be usually be safely discarded.")) + self.add_item( + section=section, + title="aligner_distance", + datatype=float, + min_max=(0.0, 45.0), + rounding=1, + default=22.5, + group=_("filters"), + info=_("Filters out faces who's landmarks are above this distance from an 'average' " + "face. Values above 15 tend to be fairly safe. Values above 10 will remove " + "more false positives, but may also filter out some faces at extreme angles.")) + self.add_item( + section=section, + title="aligner_roll", + datatype=float, + min_max=(0.0, 90.0), + rounding=1, + default=45.0, + group=_("filters"), + info=_("Filters out faces who's calculated roll is greater than zero +/- this value " + "in degrees. Aligned faces should have a roll value close to zero. Values that " + "are a significant distance from 0 degrees tend to be misaligned images. These " + "can usually be safely disgarded.")) + self.add_item( + section=section, + title="aligner_features", + datatype=bool, + default=True, + group=_("filters"), + info=_("Filters out faces where the lowest point of the aligned face's eye or eyebrow " + "is lower than the highest point of the aligned face's mouth. Any faces where " + "this occurs are misaligned and can be safely disgarded.")) + self.add_item( + section=section, + title="filter_refeed", + datatype=bool, + default=True, + group=_("filters"), + info=_("If enabled, and 're-feed' has been selected for extraction, then interim " + "alignments will be filtered prior to averaging the final landmarks. This can " + "help improve the final alignments by removing any obvious misaligns from the " + "interim results, and may also help pick up difficult alignments. If disabled, " + "then all re-feed results will be averaged.")) + self.add_item( + section=section, + title="save_filtered", + datatype=bool, + default=False, + group=_("filters"), + info=_("If enabled, saves any filtered out images into a sub-folder during the " + "extraction process. If disabled, filtered faces are deleted. Note: The faces " + "will always be filtered out of the alignments file, regardless of whether you " + "keep the faces or not.")) + self.add_item( + section=section, + title="realign_refeeds", + datatype=bool, + default=True, + group=_("re-align"), + info=_("If enabled, and 're-align' has been selected for extraction, then all re-feed " + "iterations are re-aligned. If disabled, then only the final averaged output " + "from re-feed will be re-aligned.")) self.add_item( - section=section, title="allow_growth", datatype=bool, default=False, group="settings", - info="[Nvidia Only]. Enable the Tensorflow GPU `allow_growth` configuration option. " - "This option prevents Tensorflow from allocating all of the GPU VRAM at launch " - "but can lead to higher VRAM fragmentation and slower performance. Should only " - "be enabled if you are having problems running extraction.") + section=section, + title="filter_realign", + datatype=bool, + default=True, + group=_("re-align"), + info=_("If enabled, and 're-align' has been selected for extraction, then any " + "alignments which would be filtered out will not be re-aligned.")) diff --git a/plugins/extract/align/_base.py b/plugins/extract/align/_base.py deleted file mode 100644 index 12e3a33b31..0000000000 --- a/plugins/extract/align/_base.py +++ /dev/null @@ -1,387 +0,0 @@ -#!/usr/bin/env python3 -""" Base class for Face Aligner plugins - -All Aligner Plugins should inherit from this class. -See the override methods for which methods are required. - -The plugin will receive a :class:`~plugins.extract.pipeline.ExtractMedia` object. - -For each source item, the plugin must pass a dict to finalize containing: - ->>> {"filename": [], ->>> "landmarks": [list of 68 point face landmarks] ->>> "detected_faces": []} -""" - - -import cv2 -import numpy as np - -from tensorflow.python.framework import errors_impl as tf_errors - -from lib.utils import get_backend, FaceswapError -from plugins.extract._base import Extractor, logger, ExtractMedia - - -class Aligner(Extractor): # pylint:disable=abstract-method - """ Aligner plugin _base Object - - All Aligner plugins must inherit from this class - - Parameters - ---------- - git_model_id: int - The second digit in the github tag that identifies this model. See - https://github.com/deepfakes-models/faceswap-models for more information - model_filename: str - The name of the model file to be loaded - normalize_method: {`None`, 'clahe', 'hist', 'mean'}, optional - Normalize the images fed to the aligner. Default: ``None`` - re_feed: int - The number of times to re-feed a slightly adjusted bounding box into the aligner. - Default: `0` - - Other Parameters - ---------------- - configfile: str, optional - Path to a custom configuration ``ini`` file. Default: Use system configfile - - See Also - -------- - plugins.extract.pipeline : The extraction pipeline for calling plugins - plugins.extract.align : Aligner plugins - plugins.extract._base : Parent class for all extraction plugins - plugins.extract.detect._base : Detector parent class for extraction plugins. - plugins.extract.mask._base : Masker parent class for extraction plugins. - """ - - def __init__(self, git_model_id=None, model_filename=None, - configfile=None, instance=0, normalize_method=None, re_feed=0, **kwargs): - logger.debug("Initializing %s: (normalize_method: %s, re_feed: %s)", - self.__class__.__name__, normalize_method, re_feed) - super().__init__(git_model_id, - model_filename, - configfile=configfile, - instance=instance, - **kwargs) - self._normalize_method = None - self._re_feed = re_feed - self.set_normalize_method(normalize_method) - - self._plugin_type = "align" - self._faces_per_filename = dict() # Tracking for recompiling face batches - self._rollover = None # Items that are rolled over from the previous batch in get_batch - self._output_faces = [] - self._additional_keys = [] - logger.debug("Initialized %s", self.__class__.__name__) - - def set_normalize_method(self, method): - """ Set the normalization method for feeding faces into the aligner. - - Parameters - ---------- - method: {"none", "clahe", "hist", "mean"} - The normalization method to apply to faces prior to feeding into the model - """ - method = None if method is None or method.lower() == "none" else method - self._normalize_method = method - - # << QUEUE METHODS >>> # - def get_batch(self, queue): - """ Get items for inputting into the aligner from the queue in batches - - Items are returned from the ``queue`` in batches of - :attr:`~plugins.extract._base.Extractor.batchsize` - - Items are received as :class:`~plugins.extract.pipeline.ExtractMedia` objects and converted - to ``dict`` for internal processing. - - To ensure consistent batch sizes for aligner the items are split into separate items for - each :class:`~lib.align.DetectedFace` object. - - Remember to put ``'EOF'`` to the out queue after processing - the final batch - - Outputs items in the following format. All lists are of length - :attr:`~plugins.extract._base.Extractor.batchsize`: - - >>> {'filename': [], - >>> 'image': [], - >>> 'detected_faces': [[>> # - def finalize(self, batch): - """ Finalize the output from Aligner - - This should be called as the final task of each `plugin`. - - Pairs the detected faces back up with their original frame before yielding each frame. - - Parameters - ---------- - batch : dict - The final ``dict`` from the `plugin` process. It must contain the `keys`: - ``detected_faces``, ``landmarks``, ``filename`` - - Yields - ------ - :class:`~plugins.extract.pipeline.ExtractMedia` - The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes - and landmarks for the detected faces found in the frame. - """ - - for face, landmarks in zip(batch["detected_faces"], batch["landmarks"]): - if not isinstance(landmarks, np.ndarray): - landmarks = np.array(landmarks) - face.landmarks_xy = landmarks - - logger.trace("Item out: %s", {key: val.shape if isinstance(val, np.ndarray) else val - for key, val in batch.items()}) - - for filename, face in zip(batch["filename"], batch["detected_faces"]): - self._output_faces.append(face) - if len(self._output_faces) != self._faces_per_filename[filename]: - continue - - output = self._extract_media.pop(filename) - output.add_detected_faces(self._output_faces) - self._output_faces = [] - - logger.trace("Final Output: (filename: '%s', image shape: %s, detected_faces: %s, " - "item: %s)", - output.filename, output.image_shape, output.detected_faces, output) - yield output - - # <<< PROTECTED METHODS >>> # - - # << PROCESS_INPUT WRAPPER >> - def _process_input(self, batch): - """ Process the input to the aligner model multiple times based on the user selected - `re-feed` command line option. This adjusts the bounding box for the face to be fed - into the model by a random amount within 0.05 pixels of the detected face's shortest axis. - - References - ---------- - https://studios.disneyresearch.com/2020/06/29/high-resolution-neural-face-swapping-for-visual-effects/ - - Parameters - ---------- - batch: dict - Contains the batch that is currently being passed through the plugin process - - Returns - ------- - dict - The batch with input processed - """ - if not self._additional_keys: - existing_keys = list(batch.keys()) - - original_boxes = np.array([(face.x, face.y, face.w, face.h) - for face in batch["detected_faces"]]) - adjusted_boxes = self._get_adjusted_boxes(original_boxes) - retval = dict() - for bounding_boxes in adjusted_boxes: - for face, box in zip(batch["detected_faces"], bounding_boxes): - face.x, face.y, face.w, face.h = box - - result = self.process_input(batch) - if not self._additional_keys: - self._additional_keys = [key for key in result if key not in existing_keys] - for key in self._additional_keys: - retval.setdefault(key, []).append(batch[key]) - del batch[key] - - # Place the original bounding box back to detected face objects - for face, box in zip(batch["detected_faces"], original_boxes): - face.x, face.y, face.w, face.h = box - - batch.update(retval) - return batch - - def _get_adjusted_boxes(self, original_boxes): - """ Obtain an array of adjusted bounding boxes based on the number of re-feed iterations - that have been selected and the minimum dimension of the original bounding box. - - Parameters - ---------- - original_boxes: :class:`numpy.ndarray` - The original ('x', 'y', 'w', 'h') detected face boxes corresponding to the incoming - detected face objects - - Returns - ------- - :class:`numpy.ndarray` - The original boxes (in position 0) and the randomly adjusted bounding boxes - """ - if self._re_feed == 0: - return original_boxes[None, ...] - beta = 0.05 - max_shift = np.min(original_boxes[..., 2:], axis=1) * beta - rands = np.random.rand(self._re_feed, *original_boxes.shape) * 2 - 1 - new_boxes = np.rint(original_boxes + (rands * max_shift[None, :, None])).astype("int32") - retval = np.concatenate((original_boxes[None, ...], new_boxes)) - logger.trace(retval) - return retval - - # <<< PREDICT WRAPPER >>> # - def _predict(self, batch): - """ Just return the aligner's predict function """ - try: - batch["prediction"] = [self.predict(feed) for feed in batch["feed"]] - return batch - except tf_errors.ResourceExhaustedError as err: - msg = ("You do not have enough GPU memory available to run detection at the " - "selected batch size. You can try a number of things:" - "\n1) Close any other application that is using your GPU (web browsers are " - "particularly bad for this)." - "\n2) Lower the batchsize (the amount of images fed into the model) by " - "editing the plugin settings (GUI: Settings > Configure extract settings, " - "CLI: Edit the file faceswap/config/extract.ini)." - "\n3) Enable 'Single Process' mode.") - raise FaceswapError(msg) from err - except Exception as err: - if get_backend() == "amd": - # pylint:disable=import-outside-toplevel - from lib.plaidml_utils import is_plaidml_error - if (is_plaidml_error(err) and ( - "CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or - "enough memory for the current schedule" in str(err).lower())): - msg = ("You do not have enough GPU memory available to run detection at " - "the selected batch size. You can try a number of things:" - "\n1) Close any other application that is using your GPU (web " - "browsers are particularly bad for this)." - "\n2) Lower the batchsize (the amount of images fed into the " - "model) by editing the plugin settings (GUI: Settings > Configure " - "extract settings, CLI: Edit the file " - "faceswap/config/extract.ini).") - raise FaceswapError(msg) from err - raise - - def _process_output(self, batch): - """ Process the output from the aligner model multiple times based on the user selected - `re-feed amount` configuration option, then average the results for final prediction. - - Parameters - ---------- - batch : dict - Contains the batch that is currently being passed through the plugin process - """ - landmarks = [] - for idx in range(self._re_feed + 1): - subbatch = {key: val - for key, val in batch.items() - if key not in ["feed", "prediction"] + self._additional_keys} - subbatch["prediction"] = batch["prediction"][idx] - for key in self._additional_keys: - subbatch[key] = batch[key][idx] - self.process_output(subbatch) - landmarks.append(subbatch["landmarks"]) - batch["landmarks"] = np.average(landmarks, axis=0) - return batch - - # <<< FACE NORMALIZATION METHODS >>> # - def _normalize_faces(self, faces): - """ Normalizes the face for feeding into model - - The normalization method is dictated by the normalization command line argument - """ - if self._normalize_method is None: - return faces - logger.trace("Normalizing faces") - meth = getattr(self, "_normalize_{}".format(self._normalize_method.lower())) - faces = [meth(face) for face in faces] - logger.trace("Normalized faces") - return faces - - @staticmethod - def _normalize_mean(face): - """ Normalize Face to the Mean """ - face = face / 255.0 - for chan in range(3): - layer = face[:, :, chan] - layer = (layer - layer.min()) / (layer.max() - layer.min()) - face[:, :, chan] = layer - return face * 255.0 - - @staticmethod - def _normalize_hist(face): - """ Equalize the RGB histogram channels """ - for chan in range(3): - face[:, :, chan] = cv2.equalizeHist(face[:, :, chan]) - return face - - @staticmethod - def _normalize_clahe(face): - """ Perform Contrast Limited Adaptive Histogram Equalization """ - clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4, 4)) - for chan in range(3): - face[:, :, chan] = clahe.apply(face[:, :, chan]) - return face diff --git a/plugins/extract/align/_base/__init__.py b/plugins/extract/align/_base/__init__.py new file mode 100644 index 0000000000..6e32deea62 --- /dev/null +++ b/plugins/extract/align/_base/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python3 +""" Base class for Aligner plugins ALL aligners should at least inherit from this class. """ + +from .aligner import Aligner, AlignerBatch, BatchType diff --git a/plugins/extract/align/_base/aligner.py b/plugins/extract/align/_base/aligner.py new file mode 100644 index 0000000000..6746daf631 --- /dev/null +++ b/plugins/extract/align/_base/aligner.py @@ -0,0 +1,834 @@ +#!/usr/bin/env python3 +""" Base class for Face Aligner plugins + +All Aligner Plugins should inherit from this class. +See the override methods for which methods are required. + +The plugin will receive a :class:`~plugins.extract.extract_media.ExtractMedia` object. + +For each source item, the plugin must pass a dict to finalize containing: + +>>> {"filename": [], +>>> "landmarks": [list of 68 point face landmarks] +>>> "detected_faces": []} +""" +from __future__ import annotations +import logging +import typing as T + +from dataclasses import dataclass, field +from time import sleep + +import cv2 +import numpy as np + +from tensorflow.python.framework import errors_impl as tf_errors # pylint:disable=no-name-in-module # noqa + +from lib.align import LandmarkType +from lib.utils import FaceswapError +from plugins.extract import ExtractMedia +from plugins.extract._base import BatchType, ExtractorBatch, Extractor +from .processing import AlignedFilter, ReAlign + +if T.TYPE_CHECKING: + from collections.abc import Generator + from queue import Queue + from lib.align import DetectedFace + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) +_BATCH_IDX: int = 0 + + +def _get_new_batch_id() -> int: + """ Obtain the next available batch index + + Returns + ------- + int + The next available unique batch id + """ + global _BATCH_IDX # pylint:disable=global-statement + _BATCH_IDX += 1 + return _BATCH_IDX + + +@dataclass +class AlignerBatch(ExtractorBatch): + """ Dataclass for holding items flowing through the aligner. + + Inherits from :class:`~plugins.extract._base.ExtractorBatch` + + Parameters + ---------- + batch_id: int + A unique integer for tracking this batch + landmarks: list + List of 68 point :class:`numpy.ndarray` landmark points returned from the aligner + refeeds: list + List of :class:`numpy.ndarrays` for holding each of the feeds that will be put through the + model for each refeed + second_pass: bool, optional + ``True`` if this batch is passing through the aligner for a second time as re-align has + been selected otherwise ``False``. Default: ``False`` + second_pass_masks: :class:`numpy.ndarray`, optional + The masks used to filter out re-feed values for passing to the re-aligner. + """ + batch_id: int = 0 + detected_faces: list[DetectedFace] = field(default_factory=list) + landmarks: np.ndarray = np.array([]) + refeeds: list[np.ndarray] = field(default_factory=list) + second_pass: bool = False + second_pass_masks: np.ndarray = np.array([]) + + def __repr__(self): + """ Prettier repr for debug printing """ + retval = super().__repr__() + retval += (f", batch_id={self.batch_id}, " + f"landmarks=[({self.landmarks.shape}, {self.landmarks.dtype})], " + f"refeeds={[(f.shape, f.dtype) for f in self.refeeds]}, " + f"second_pass={self.second_pass}, " + f"second_pass_masks={self.second_pass_masks})") + return retval + + def __post_init__(self): + """ Make sure that we have been given a non-zero ID """ + assert self.batch_id != 0, ("A batch ID must be specified for Aligner Batches") + + +class Aligner(Extractor): # pylint:disable=abstract-method + """ Aligner plugin _base Object + + All Aligner plugins must inherit from this class + + Parameters + ---------- + git_model_id: int + The second digit in the github tag that identifies this model. See + https://github.com/deepfakes-models/faceswap-models for more information + model_filename: str + The name of the model file to be loaded + normalize_method: {`None`, 'clahe', 'hist', 'mean'}, optional + Normalize the images fed to the aligner. Default: ``None`` + re_feed: int, optional + The number of times to re-feed a slightly adjusted bounding box into the aligner. + Default: `0` + re_align: bool, optional + ``True`` to obtain landmarks by passing the initially aligned face back through the + aligner. Default ``False`` + disable_filter: bool, optional + Disable all aligner filters regardless of config option. Default: ``False`` + Other Parameters + ---------------- + configfile: str, optional + Path to a custom configuration ``ini`` file. Default: Use system configfile + + See Also + -------- + plugins.extract.pipeline : The extraction pipeline for calling plugins + plugins.extract.align : Aligner plugins + plugins.extract._base : Parent class for all extraction plugins + plugins.extract.detect._base : Detector parent class for extraction plugins. + plugins.extract.mask._base : Masker parent class for extraction plugins. + """ + + def __init__(self, + git_model_id: int | None = None, + model_filename: str | None = None, + configfile: str | None = None, + instance: int = 0, + normalize_method: T.Literal["none", "clahe", "hist", "mean"] | None = None, + re_feed: int = 0, + re_align: bool = False, + disable_filter: bool = False, + **kwargs) -> None: + logger.debug("Initializing %s: (normalize_method: %s, re_feed: %s, re_align: %s, " + "disable_filter: %s)", self.__class__.__name__, normalize_method, re_feed, + re_align, disable_filter) + super().__init__(git_model_id, + model_filename, + configfile=configfile, + instance=instance, + **kwargs) + self._plugin_type = "align" + self.realign_centering: CenteringType = "face" # overide for plugin specific centering + + # Override for specific landmark type: + self.landmark_type = LandmarkType.LM_2D_68 + + self._eof_seen = False + self._normalize_method: T.Literal["clahe", "hist", "mean"] | None = None + self._re_feed = re_feed + self._filter = AlignedFilter(feature_filter=self.config["aligner_features"], + min_scale=self.config["aligner_min_scale"], + max_scale=self.config["aligner_max_scale"], + distance=self.config["aligner_distance"], + roll=self.config["aligner_roll"], + save_output=self.config["save_filtered"], + disable=disable_filter) + self._re_align = ReAlign(re_align, + self.config["realign_refeeds"], + self.config["filter_realign"]) + self._needs_refeed_masks: bool = self._re_feed > 0 and ( + self.config["filter_refeed"] or (self._re_align.do_refeeds and + self._re_align.do_filter)) + self.set_normalize_method(normalize_method) + + logger.debug("Initialized %s", self.__class__.__name__) + + def set_normalize_method(self, method: T.Literal["none", "clahe", "hist", "mean"] | None + ) -> None: + """ Set the normalization method for feeding faces into the aligner. + + Parameters + ---------- + method: {"none", "clahe", "hist", "mean"} + The normalization method to apply to faces prior to feeding into the model + """ + method = None if method is None or method.lower() == "none" else method + self._normalize_method = T.cast(T.Literal["clahe", "hist", "mean"] | None, method) + + def initialize(self, *args, **kwargs) -> None: + """ Add a call to add model input size to the re-aligner """ + self._re_align.set_input_size_and_centering(self.input_size, self.realign_centering) + super().initialize(*args, **kwargs) + + def _handle_realigns(self, queue: Queue) -> tuple[bool, AlignerBatch] | None: + """ Handle any items waiting for a second pass through the aligner. + + If EOF has been recieved and items are still being processed through the first pass + then wait for a short time and try again to collect them. + + On EOF return exhausted flag with an empty batch + + Parameters + ---------- + queue : queue.Queue() + The ``queue`` that the plugin will be fed from. + + Returns + ------- + ``None`` or tuple + If items are processed then returns (`bool`, :class:`AlignerBatch`) containing the + exhausted flag and the batch to be processed. If no items are processed returns + ``None`` + """ + if not self._re_align.active: + return None + + exhausted = False + if self._re_align.items_queued: + batch = self._re_align.get_batch() + logger.trace("Re-align batch: %s", batch) # type: ignore[attr-defined] + return exhausted, batch + + if self._eof_seen and self._re_align.items_tracked: + # EOF seen and items still being processed on first pass + logger.debug("Tracked re-align items waiting to be flushed, retrying...") + sleep(0.25) + return self.get_batch(queue) + + if self._eof_seen: + exhausted = True + logger.debug("All items processed. Returning empty batch") + self._filter.output_counts() + self._eof_seen = False # Reset for plugin re-use + return exhausted, AlignerBatch(batch_id=-1) + + return None + + def get_batch(self, queue: Queue) -> tuple[bool, AlignerBatch]: + """ Get items for inputting into the aligner from the queue in batches + + Items are returned from the ``queue`` in batches of + :attr:`~plugins.extract._base.Extractor.batchsize` + + Items are received as :class:`~plugins.extract.extract_media.ExtractMedia` objects and + converted to ``dict`` for internal processing. + + To ensure consistent batch sizes for aligner the items are split into separate items for + each :class:`~lib.align.DetectedFace` object. + + Remember to put ``'EOF'`` to the out queue after processing + the final batch + + Outputs items in the following format. All lists are of length + :attr:`~plugins.extract._base.Extractor.batchsize`: + + >>> {'filename': [], + >>> 'image': [], + >>> 'detected_faces': [[ np.ndarray: + """ Overide for specific plugin processing to convert a batch of face images from UINT8 + (0-255) into the correct format for the plugin's inference + + Parameters + ---------- + faces: :class:`numpy.ndarray` + The batch of faces in UINT8 format + + Returns + ------- + class: `numpy.ndarray` + The batch of faces in the format to feed through the plugin + """ + raise NotImplementedError() + + # <<< FINALIZE METHODS >>> # + def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]: + """ Finalize the output from Aligner + + This should be called as the final task of each `plugin`. + + Pairs the detected faces back up with their original frame before yielding each frame. + + Parameters + ---------- + batch : :class:`AlignerBatch` + The final batch item from the `plugin` process. + + Yields + ------ + :class:`~plugins.extract.extract_media.ExtractMedia` + The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes + and landmarks for the detected faces found in the frame. + """ + assert isinstance(batch, AlignerBatch) + if not batch.second_pass and self._re_align.active: + # Add the batch for second pass re-alignment and return + self._re_align.add_batch(batch) + return + for face, landmarks in zip(batch.detected_faces, batch.landmarks): + if not isinstance(landmarks, np.ndarray): + landmarks = np.array(landmarks) + face.add_landmarks_xy(landmarks) + + logger.trace("Item out: %s", batch) # type: ignore[attr-defined] + + for frame, filename, face in zip(batch.image, batch.filename, batch.detected_faces): + self._output_faces.append(face) + if len(self._output_faces) != self._faces_per_filename[filename]: + continue + + self._output_faces, folders = self._filter(self._output_faces, min(frame.shape[:2])) + + output = self._extract_media.pop(filename) + output.add_detected_faces(self._output_faces) + output.add_sub_folders(folders) + self._output_faces = [] + + logger.trace("Final Output: (filename: '%s', image " # type: ignore[attr-defined] + "shape: %s, detected_faces: %s, item: %s)", output.filename, + output.image_shape, output.detected_faces, output) + yield output + self._re_align.untrack_batch(batch.batch_id) + + def on_completion(self) -> None: + """ Output the filter counts when process has completed """ + self._filter.output_counts() + + # <<< PROTECTED METHODS >>> # + # << PROCESS_INPUT WRAPPER >> + def _get_adjusted_boxes(self, original_boxes: np.ndarray) -> np.ndarray: + """ Obtain an array of adjusted bounding boxes based on the number of re-feed iterations + that have been selected and the minimum dimension of the original bounding box. + + Parameters + ---------- + original_boxes: :class:`numpy.ndarray` + The original ('x', 'y', 'w', 'h') detected face boxes corresponding to the incoming + detected face objects + + Returns + ------- + :class:`numpy.ndarray` + The original boxes (in position 0) and the randomly adjusted bounding boxes + """ + if self._re_feed == 0: + return original_boxes[None, ...] + beta = 0.05 + max_shift = np.min(original_boxes[..., 2:], axis=1) * beta + rands = np.random.rand(self._re_feed, *original_boxes.shape) * 2 - 1 + new_boxes = np.rint(original_boxes + (rands * max_shift[None, :, None])).astype("int32") + retval = np.concatenate((original_boxes[None, ...], new_boxes)) + logger.trace(retval) # type: ignore[attr-defined] + return retval + + def _process_input_first_pass(self, batch: AlignerBatch) -> None: + """ Standard pre-processing for aligners for first pass (if re-align selected) or the + only pass. + + Process the input to the aligner model multiple times based on the user selected + `re-feed` command line option. This adjusts the bounding box for the face to be fed + into the model by a random amount within 0.05 pixels of the detected face's shortest axis. + + References + ---------- + https://studios.disneyresearch.com/2020/06/29/high-resolution-neural-face-swapping-for-visual-effects/ + + Parameters + ---------- + batch: :class:`AlignerBatch` + Contains the batch that is currently being passed through the plugin process + """ + original_boxes = np.array([(face.left, face.top, face.width, face.height) + for face in batch.detected_faces]) + adjusted_boxes = self._get_adjusted_boxes(original_boxes) + + # Put in random re-feed data to the bounding boxes + for bounding_boxes in adjusted_boxes: + for face, box in zip(batch.detected_faces, bounding_boxes): + face.left, face.top, face.width, face.height = box + + self.process_input(batch) + batch.feed = self.faces_to_feed(self._normalize_faces(batch.feed)) + # Move the populated feed into the batch refeed list. It will be overwritten at next + # iteration + batch.refeeds.append(batch.feed) + + # Place the original bounding box back to detected face objects + for face, box in zip(batch.detected_faces, original_boxes): + face.left, face.top, face.width, face.height = box + + def _get_realign_masks(self, batch: AlignerBatch) -> np.ndarray: + """ Obtain the masks required for processing re-aligns + + Parameters + ---------- + batch: :class:`AlignerBatch` + Contains the batch that is currently being passed through the plugin process + + Returns + ------- + :class:`numpy.ndarray` + The filter masks required for masking the re-aligns + """ + if self._re_align.do_refeeds: + retval = batch.second_pass_masks # Masks already calculated during re-feed + elif self._re_align.do_filter: + retval = self._filter.filtered_mask(batch)[None, ...] + else: + retval = np.zeros((batch.landmarks.shape[0], ), dtype="bool")[None, ...] + return retval + + def _process_input_second_pass(self, batch: AlignerBatch) -> None: + """ Process the input for 2nd-pass re-alignment + + Parameters + ---------- + batch: :class:`AlignerBatch` + Contains the batch that is currently being passed through the plugin process + """ + batch.second_pass_masks = self._get_realign_masks(batch) + + if not self._re_align.do_refeeds: + # Expand the dimensions for re-aligns for consistent handling of code + batch.landmarks = batch.landmarks[None, ...] + + refeeds = self._re_align.process_batch(batch) + batch.refeeds = [self.faces_to_feed(self._normalize_faces(faces)) for faces in refeeds] + + def _process_input(self, batch: BatchType) -> AlignerBatch: + """ Perform pre-processing depending on whether this is the first/only pass through the + aligner or the 2nd pass when re-align has been selected + + Parameters + ---------- + batch: :class:`AlignerBatch` + Contains the batch that is currently being passed through the plugin process + + Returns + ------- + :class:`AlignerBatch` + The batch with input processed + """ + assert isinstance(batch, AlignerBatch) + if batch.second_pass: + self._process_input_second_pass(batch) + else: + self._process_input_first_pass(batch) + return batch + + # <<< PREDICT WRAPPER >>> # + def _predict(self, batch: BatchType) -> AlignerBatch: + """ Just return the aligner's predict function + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch to find alignments for + + Returns + ------- + :class:`AlignerBatch` + The batch item with the :attr:`prediction` populated + + Raises + ------ + FaceswapError + If GPU resources are exhausted + """ + assert isinstance(batch, AlignerBatch) + try: + preds = [self.predict(feed) for feed in batch.refeeds] + try: + batch.prediction = np.array(preds) + except ValueError as err: + # If refeed batches are different sizes, Numpy will error, so we need to explicitly + # set the dtype to 'object' rather than let it infer + # numpy error: + # ValueError: setting an array element with a sequence. The requested array has an + # inhomogeneous shape after 1 dimensions. The detected shape was (9,) + + # inhomogeneous part + if "inhomogeneous" in str(err): + logger.trace( # type:ignore[attr-defined] + "Mismatched array sizes, setting dtype to object: %s", + [p.shape for p in preds]) + batch.prediction = np.array(preds, dtype="object") + else: + raise + + return batch + except tf_errors.ResourceExhaustedError as err: + msg = ("You do not have enough GPU memory available to run detection at the " + "selected batch size. You can try a number of things:" + "\n1) Close any other application that is using your GPU (web browsers are " + "particularly bad for this)." + "\n2) Lower the batchsize (the amount of images fed into the model) by " + "editing the plugin settings (GUI: Settings > Configure extract settings, " + "CLI: Edit the file faceswap/config/extract.ini)." + "\n3) Enable 'Single Process' mode.") + raise FaceswapError(msg) from err + + def _process_refeeds(self, batch: AlignerBatch) -> list[AlignerBatch]: + """ Process the output for each selected re-feed + + Parameters + ---------- + batch: :class:`AlignerBatch` + The batch object passing through the aligner + + Returns + ------- + list + List of :class:`AlignerBatch` objects. Each object in the list contains the + results for each selected re-feed + """ + retval: list[AlignerBatch] = [] + if batch.second_pass: + # Re-insert empty sub-patches for re-population in ReAlign for filtered out batches + selected_idx = 0 + for mask in batch.second_pass_masks: + all_filtered = np.all(mask) + if not all_filtered: + feed = batch.refeeds[selected_idx] + pred = batch.prediction[selected_idx] + data = batch.data[selected_idx] if batch.data else {} + selected_idx += 1 + else: # All resuts have been filtered out + feed = pred = np.array([]) + data = {} + + subbatch = AlignerBatch(batch_id=batch.batch_id, + image=batch.image, + detected_faces=batch.detected_faces, + filename=batch.filename, + feed=feed, + prediction=pred, + data=[data], + second_pass=batch.second_pass) + + if not all_filtered: + self.process_output(subbatch) + + retval.append(subbatch) + else: + b_data = batch.data if batch.data else [{}] + for feed, pred, dat in zip(batch.refeeds, batch.prediction, b_data): + subbatch = AlignerBatch(batch_id=batch.batch_id, + image=batch.image, + detected_faces=batch.detected_faces, + filename=batch.filename, + feed=feed, + prediction=pred, + data=[dat], + second_pass=batch.second_pass) + self.process_output(subbatch) + retval.append(subbatch) + return retval + + def _get_refeed_filter_masks(self, + subbatches: list[AlignerBatch], + original_masks: np.ndarray | None = None) -> np.ndarray: + """ Obtain the boolean mask array for masking out failed re-feed results if filter refeed + has been selected + + Parameters + ---------- + subbatches: list + List of sub-batch results for each re-feed performed + original_masks: :class:`numpy.ndarray`, Optional + If passing in the second pass landmarks, these should be the original filter masks so + that we don't calculate the mask again for already filtered faces. Default: ``None`` + + Returns + ------- + :class:`numpy.ndarray` + boolean values for every detected face indicating whether the interim landmarks have + passed the filter test + """ + retval = np.zeros((len(subbatches), subbatches[0].landmarks.shape[0]), dtype="bool") + + if not self._needs_refeed_masks: + return retval + + retval = retval if original_masks is None else original_masks + for subbatch, masks in zip(subbatches, retval): + masks[:] = self._filter.filtered_mask(subbatch, np.flatnonzero(masks)) + return retval + + def _get_mean_landmarks(self, landmarks: np.ndarray, masks: np.ndarray) -> np.ndarray: + """ Obtain the averaged landmarks from the re-fed alignments. If config option + 'filter_refeed' is enabled, then average those results which have not been filtered out + otherwise average all results + + Parameters + ---------- + landmarks: :class:`numpy.ndarray` + The batch of re-fed alignments + masks: :class:`numpy.ndarray` + List of boolean values indicating whether each re-fed alignments passed or failed + the filter test + + Returns + ------- + :class:`numpy.ndarray` + The final averaged landmarks + """ + if any(np.all(masked) for masked in masks.T): + # hacky fix for faces which entirely failed the filter + # We just unmask one value as it is junk anyway and will be discarded on output + for idx, masked in enumerate(masks.T): + if np.all(masked): + masks[0, idx] = False + + masks = np.broadcast_to(np.reshape(masks, (*landmarks.shape[:2], 1, 1)), + landmarks.shape) + return np.ma.array(landmarks, mask=masks).mean(axis=0).data.astype("float32") + + def _process_output_first_pass(self, subbatches: list[AlignerBatch]) -> tuple[np.ndarray, + np.ndarray]: + """ Process the output from the aligner if this is the first or only pass. + + Parameters + ---------- + subbatches: list + List of sub-batch results for each re-feed performed + + Returns + ------- + landmarks: :class:`numpy.ndarray` + If re-align is not selected or if re-align has been selected but only on the final + output (ie: realign_reefeeds is ``False``) then the averaged batch of landmarks for all + re-feeds is returned. + If re-align_refeeds has been selected, then this will output each batch of re-feed + landmarks. + masks: :class:`numpy.ndarray` + Boolean mask corresponding to the re-fed landmarks output indicating any values which + should be filtered out prior to further processing + """ + masks = self._get_refeed_filter_masks(subbatches) + all_landmarks = np.array([sub.landmarks for sub in subbatches]) + + # re-align not selected or not filtering the re-feeds + if not self._re_align.do_refeeds: + retval = self._get_mean_landmarks(all_landmarks, masks) + return retval, masks + + # Re-align selected with filter re-feeds + return all_landmarks, masks + + def _process_output_second_pass(self, + subbatches: list[AlignerBatch], + masks: np.ndarray) -> np.ndarray: + """ Process the output from the aligner if this is the first or only pass. + + Parameters + ---------- + subbatches: list + List of sub-batch results for each re-aligned re-feed performed + masks: :class:`numpy.ndarray` + The original re-feed filter masks from the first pass + """ + self._re_align.process_output(subbatches, masks) + masks = self._get_refeed_filter_masks(subbatches, original_masks=masks) + all_landmarks = np.array([sub.landmarks for sub in subbatches]) + return self._get_mean_landmarks(all_landmarks, masks) + + def _process_output(self, batch: BatchType) -> AlignerBatch: + """ Process the output from the aligner model multiple times based on the user selected + `re-feed amount` configuration option, then average the results for final prediction. + + If the config option 'filter_refeed' is enabled, then mask out any returned alignments + that fail a filter test + + Parameters + ---------- + batch : :class:`AlignerBatch` + Contains the batch that is currently being passed through the plugin process + + Returns + ------- + :class:`AlignerBatch` + The batch item with :attr:`landmarks` populated + """ + assert isinstance(batch, AlignerBatch) + subbatches = self._process_refeeds(batch) + if batch.second_pass: + batch.landmarks = self._process_output_second_pass(subbatches, batch.second_pass_masks) + else: + landmarks, masks = self._process_output_first_pass(subbatches) + batch.landmarks = landmarks + batch.second_pass_masks = masks + return batch + + # <<< FACE NORMALIZATION METHODS >>> # + def _normalize_faces(self, faces: np.ndarray) -> np.ndarray: + """ Normalizes the face for feeding into model + The normalization method is dictated by the normalization command line argument + + Parameters + ---------- + faces: :class:`numpy.ndarray` + The batch of faces to normalize + + Returns + ------- + :class:`numpy.ndarray` + The normalized faces + """ + if self._normalize_method is None: + return faces + logger.trace("Normalizing faces") # type: ignore[attr-defined] + meth = getattr(self, f"_normalize_{self._normalize_method.lower()}") + faces = np.array([meth(face) for face in faces]) + logger.trace("Normalized faces") # type: ignore[attr-defined] + return faces + + @classmethod + def _normalize_mean(cls, face: np.ndarray) -> np.ndarray: + """ Normalize Face to the Mean + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face to normalize + + Returns + ------- + :class:`numpy.ndarray` + The normalized face + """ + face = face / 255.0 + for chan in range(3): + layer = face[:, :, chan] + layer = (layer - layer.min()) / (layer.max() - layer.min()) + face[:, :, chan] = layer + return face * 255.0 + + @classmethod + def _normalize_hist(cls, face: np.ndarray) -> np.ndarray: + """ Equalize the RGB histogram channels + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face to normalize + + Returns + ------- + :class:`numpy.ndarray` + The normalized face + """ + for chan in range(3): + face[:, :, chan] = cv2.equalizeHist(face[:, :, chan]) + return face + + @classmethod + def _normalize_clahe(cls, face: np.ndarray) -> np.ndarray: + """ Perform Contrast Limited Adaptive Histogram Equalization + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face to normalize + + Returns + ------- + :class:`numpy.ndarray` + The normalized face + """ + clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4, 4)) + for chan in range(3): + face[:, :, chan] = clahe.apply(face[:, :, chan]) + return face diff --git a/plugins/extract/align/_base/processing.py b/plugins/extract/align/_base/processing.py new file mode 100644 index 0000000000..efdeec9468 --- /dev/null +++ b/plugins/extract/align/_base/processing.py @@ -0,0 +1,489 @@ +#!/usr/bin/env python3 +""" Processing methods for aligner plugins """ +from __future__ import annotations +import logging +import typing as T + +from threading import Lock + +import numpy as np + +from lib.align import AlignedFace + +if T.TYPE_CHECKING: + from lib.align import DetectedFace + from .aligner import AlignerBatch + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +class AlignedFilter(): + """ Applies filters on the output of the aligner + + Parameters + ---------- + feature_filter: bool + ``True`` to enable filter to check relative position of eyes/eyebrows and mouth. ``False`` + to disable. + min_scale: float + Filters out faces that have been aligned at below this value as a multiplier of the + minimum frame dimension. Set to ``0`` for off. + max_scale: float + Filters out faces that have been aligned at above this value as a multiplier of the + minimum frame dimension. Set to ``0`` for off. + distance: float + Filters out faces that are further than this distance from an "average" face. Set to + ``0`` for off. + roll: float + Filters out faces with a roll value outside of 0 +/- the value given here. Set to ``0`` + for off. + save_output: bool + ``True`` if the filtered faces should be kept as they are being saved. ``False`` if they + should be deleted + disable: bool, Optional + ``True`` to disable the filter regardless of config options. Default: ``False`` + """ + def __init__(self, + feature_filter: bool, + min_scale: float, + max_scale: float, + distance: float, + roll: float, + save_output: bool, + disable: bool = False) -> None: + logger.debug("Initializing %s: (feature_filter: %s, min_scale: %s, max_scale: %s, " + "distance: %s, roll, %s, save_output: %s, disable: %s)", + self.__class__.__name__, feature_filter, min_scale, max_scale, distance, roll, + save_output, disable) + self._features = feature_filter + self._min_scale = min_scale + self._max_scale = max_scale + self._distance = distance / 100. + self._roll = roll + self._save_output = save_output + self._active = not disable and (feature_filter or + max_scale > 0.0 or + min_scale > 0.0 or + distance > 0.0 or + roll > 0.0) + self._counts: dict[str, int] = {"features": 0, + "min_scale": 0, + "max_scale": 0, + "distance": 0, + "roll": 0} + logger.debug("Initialized %s: ", self.__class__.__name__) + + def _scale_test(self, + face: AlignedFace, + minimum_dimension: int) -> T.Literal["min", "max"] | None: + """ Test if a face is below or above the min/max size thresholds. Returns as soon as a test + fails. + + Parameters + ---------- + face: :class:`~lib.aligned.AlignedFace` + The aligned face to test the original size of. + + minimum_dimension: int + The minimum (height, width) of the original frame + + Returns + ------- + "min", "max" or ``None`` + Returns min or max if the face failed the minimum or maximum test respectively. + ``None`` if all tests passed + """ + + if self._min_scale <= 0.0 and self._max_scale <= 0.0: + return None + + roi = face.original_roi.astype("int64") + size = ((roi[1][0] - roi[0][0]) ** 2 + (roi[1][1] - roi[0][1]) ** 2) ** 0.5 + + if self._min_scale > 0.0 and size < minimum_dimension * self._min_scale: + return "min" + + if self._max_scale > 0.0 and size > minimum_dimension * self._max_scale: + return "max" + + return None + + def _handle_filtered(self, + key: str, + face: DetectedFace, + faces: list[DetectedFace], + sub_folders: list[str | None], + sub_folder_index: int) -> None: + """ Add the filtered item to the filter counts. + + If config option `save_filtered` has been enabled then add the face to the output faces + list and update the sub_folder list with the correct name for this face. + + Parameters + ---------- + key: str + The key to use for the filter counts dictionary and the sub_folder name + face: :class:`~lib.align.detected_face.DetectedFace` + The detected face object to be filtered out + faces: list + The list of faces that will be returned from the filter + sub_folders: list + List of sub folder names corresponding to the list of detected face objects + sub_folder_index: int + The index within the sub-folder list that the filtered face belongs to + """ + self._counts[key] += 1 + if not self._save_output: + return + + faces.append(face) + sub_folders[sub_folder_index] = f"_align_filt_{key}" + + def __call__(self, faces: list[DetectedFace], minimum_dimension: int + ) -> tuple[list[DetectedFace], list[str | None]]: + """ Apply the filter to the incoming batch + + Parameters + ---------- + faces: list + List of detected face objects to filter out on size + minimum_dimension: int + The minimum (height, width) of the original frame + + Returns + ------- + detected_faces: list + The filtered list of detected face objects, if saving filtered faces has not been + selected or the full list of detected faces + sub_folders: list + List of ``Nones`` if saving filtered faces has not been selected or list of ``Nones`` + and sub folder names corresponding the filtered face location + """ + sub_folders: list[str | None] = [None for _ in range(len(faces))] + if not self._active: + return faces, sub_folders + + retval: list[DetectedFace] = [] + for idx, face in enumerate(faces): + aligned = AlignedFace(landmarks=face.landmarks_xy, centering="face") + + if self._features and aligned.relative_eye_mouth_position < 0.0: + self._handle_filtered("features", face, retval, sub_folders, idx) + continue + + min_max = self._scale_test(aligned, minimum_dimension) + if min_max in ("min", "max"): + self._handle_filtered(f"{min_max}_scale", face, retval, sub_folders, idx) + continue + + if 0.0 < self._distance < aligned.average_distance: + self._handle_filtered("distance", face, retval, sub_folders, idx) + continue + + if self._roll != 0.0 and not 0.0 < abs(aligned.pose.roll) < self._roll: + self._handle_filtered("roll", face, retval, sub_folders, idx) + continue + + retval.append(face) + return retval, sub_folders + + def filtered_mask(self, + batch: AlignerBatch, + skip: np.ndarray | list[int] | None = None) -> np.ndarray: + """ Obtain a list of boolean values for the given batch indicating whether they pass the + filter test. + + Parameters + ---------- + batch: :class:`AlignerBatch` + The batch of face to obtain masks for + skip: list or :class:`numpy.ndarray`, optional + List or 1D numpy array of indices indicating faces that have already been filter + masked and so should not be filtered again. Values in these index positions will be + returned as ``True`` + + Returns + ------- + :class:`numpy.ndarray` + Boolean mask array corresponding to any of the input DetectedFace objects that passed a + test. ``False`` the face passed the test. ``True`` it failed + """ + skip = [] if skip is None else skip + retval = np.ones((len(batch.detected_faces), ), dtype="bool") + for idx, (landmarks, image) in enumerate(zip(batch.landmarks, batch.image)): + if idx in skip: + continue + face = AlignedFace(landmarks) + if self._features and face.relative_eye_mouth_position < 0.0: + continue + if self._scale_test(face, min(image.shape[:2])) is not None: + continue + if 0.0 < self._distance < face.average_distance: + continue + if self._roll != 0.0 and not 0.0 < abs(face.pose.roll) < self._roll: + continue + retval[idx] = False + return retval + + def output_counts(self): + """ Output the counts of filtered items """ + if not self._active: + return + counts = [f"{key} ({getattr(self, f'_{key}'):.2f}): {count}" + for key, count in self._counts.items() + if count > 0] + if counts: + logger.info("Aligner filtered: (%s)", ", ".join(counts)) + + +class ReAlign(): + """ Holds data and methods for 2nd pass re-aligns + + Parameters + ---------- + active: bool + ``True`` if re-alignment has been requested otherwise ``False`` + do_refeeds: bool + ``True`` if re-feeds should be re-aligned, ``False`` if just the final output of the + re-feeds should be aligned + do_filter: bool + ``True`` if aligner filtered out faces should not be re-aligned. ``False`` if all faces + should be re-aligned + """ + def __init__(self, active: bool, do_refeeds: bool, do_filter: bool) -> None: + logger.debug("Initializing %s: (active: %s, do_refeeds: %s, do_filter: %s)", + self.__class__.__name__, active, do_refeeds, do_filter) + self._active = active + self._do_refeeds = do_refeeds + self._do_filter = do_filter + self._centering: CenteringType = "face" + self._size = 0 + self._tracked_lock = Lock() + self._tracked_batchs: dict[int, + dict[T.Literal["filtered_landmarks"], list[np.ndarray]]] = {} + # TODO. Probably does not need to be a list, just alignerbatch + self._queue_lock = Lock() + self._queued: list[AlignerBatch] = [] + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def active(self) -> bool: + """bool: ``True`` if re_aligns have been selected otherwise ``False``""" + return self._active + + @property + def do_refeeds(self) -> bool: + """bool: ``True`` if re-aligning is active and re-aligning re-feeds has been selected + otherwise ``False``""" + return self._active and self._do_refeeds + + @property + def do_filter(self) -> bool: + """bool: ``True`` if re-aligning is active and faces which failed the aligner filter test + should not be re-aligned otherwise ``False``""" + return self._active and self._do_filter + + @property + def items_queued(self) -> bool: + """bool: ``True`` if re-align is active and items are queued for a 2nd pass otherwise + ``False`` """ + with self._queue_lock: + return self._active and bool(self._queued) + + @property + def items_tracked(self) -> bool: + """bool: ``True`` if items exist in the tracker so still need to be processed """ + with self._tracked_lock: + return bool(self._tracked_batchs) + + def set_input_size_and_centering(self, input_size: int, centering: CenteringType) -> None: + """ Set the input size of the loaded plugin once the model has been loaded + + Parameters + ---------- + input_size: int + The input size, in pixels, of the aligner plugin + centering: ["face", "head" or "legacy"] + The centering to align the image at for re-aligning + """ + logger.debug("input_size: %s, centering: %s", input_size, centering) + self._size = input_size + self._centering = centering + + def track_batch(self, batch_id: int) -> None: + """ Add newly seen batch id from the aligner to the batch tracker, so that we can keep + track of whether there are still batches to be processed when the aligner hits 'EOF' + + Parameters + ---------- + batch_id: int + The batch id to add to batch tracking + """ + if not self._active: + return + logger.trace("Tracking batch id: %s", batch_id) # type: ignore[attr-defined] + with self._tracked_lock: + self._tracked_batchs[batch_id] = {} + + def untrack_batch(self, batch_id: int) -> None: + """ Remove the tracked batch from the tracker once the batch has been fully processed + + Parameters + ---------- + batch_id: int + The batch id to remove from batch tracking + """ + if not self._active: + return + logger.trace("Removing batch id from tracking: %s", batch_id) # type: ignore[attr-defined] + with self._tracked_lock: + del self._tracked_batchs[batch_id] + + def add_batch(self, batch: AlignerBatch) -> None: + """ Add first pass alignments to the queue for picking up for re-alignment, update their + :attr:`second_pass` attribute to ``True`` and clear attributes not required. + + Parameters + ---------- + batch: :class:`AlignerBatch` + aligner batch to perform re-alignment on + """ + with self._queue_lock: + logger.trace("Queueing for second pass: %s", batch) # type: ignore[attr-defined] + batch.second_pass = True + batch.feed = np.array([]) + batch.prediction = np.array([]) + batch.refeeds = [] + batch.data = [] + self._queued.append(batch) + + def get_batch(self) -> AlignerBatch: + """ Retrieve the next batch currently queued for re-alignment + + Returns + ------- + :class:`AlignerBatch` + The next :class:`AlignerBatch` for re-alignment + """ + with self._queue_lock: + retval = self._queued.pop(0) + logger.trace("Retrieving for second pass: %s", # type: ignore[attr-defined] + retval.filename) + return retval + + def process_batch(self, batch: AlignerBatch) -> list[np.ndarray]: + """ Pre process a batch object for re-aligning through the aligner. + + Parameters + ---------- + batch: :class:`AlignerBatch` + aligner batch to perform pre-processing on + + Returns + ------- + list + List of UINT8 aligned faces batch for each selected refeed + """ + logger.trace("Processing batch: %s, landmarks: %s", # type: ignore[attr-defined] + batch.filename, [b.shape for b in batch.landmarks]) + retval: list[np.ndarray] = [] + filtered_landmarks: list[np.ndarray] = [] + for landmarks, masks in zip(batch.landmarks, batch.second_pass_masks): + if not np.all(masks): # At least one face has not already been filtered + aligned_faces = [AlignedFace(lms, + image=image, + size=self._size, + centering=self._centering) + for image, lms, msk in zip(batch.image, landmarks, masks) + if not msk] + faces = np.array([aligned.face for aligned in aligned_faces + if aligned.face is not None]) + retval.append(faces) + batch.data.append({"aligned_faces": aligned_faces}) + + if np.any(masks): + # Track the original landmarks for re-insertion on the other side + filtered_landmarks.append(landmarks[masks]) + + with self._tracked_lock: + self._tracked_batchs[batch.batch_id] = {"filtered_landmarks": filtered_landmarks} + batch.landmarks = np.array([]) # Clear the old landmarks + return retval + + def _transform_to_frame(self, batch: AlignerBatch) -> np.ndarray: + """ Transform the predicted landmarks from the aligned face image back into frame + co-ordinates + + Parameters + ---------- + batch: :class:`AlignerBatch` + An aligner batch containing the aligned faces in the data field and the face + co-ordinate landmarks in the landmarks field + + Returns + ------- + :class:`numpy.ndarray` + The landmarks transformed to frame space + """ + faces: list[AlignedFace] = batch.data[0]["aligned_faces"] + retval = np.array([aligned.transform_points(landmarks, invert=True) + for landmarks, aligned in zip(batch.landmarks, faces)]) + logger.trace("Transformed points: original max: %s, " # type: ignore[attr-defined] + "new max: %s", batch.landmarks.max(), retval.max()) + return retval + + def _re_insert_filtered(self, batch: AlignerBatch, masks: np.ndarray) -> np.ndarray: + """ Re-insert landmarks that were filtered out from the re-align process back into the + landmark results + + Parameters + ---------- + batch: :class:`AlignerBatch` + An aligner batch containing the aligned faces in the data field and the landmarks in + frame space in the landmarks field + masks: np.ndarray + The original filter masks for this batch + + Returns + ------- + :class:`numpy.ndarray` + The full batch of landmarks with filtered out values re-inserted + """ + if not np.any(masks): + logger.trace("No landmarks to re-insert: %s", masks) # type: ignore[attr-defined] + return batch.landmarks + + with self._tracked_lock: + filtered = self._tracked_batchs[batch.batch_id]["filtered_landmarks"].pop(0) + + if np.all(masks): + retval = filtered + else: + retval = np.empty((masks.shape[0], *filtered.shape[1:]), dtype=filtered.dtype) + retval[~masks] = batch.landmarks + retval[masks] = filtered + + logger.trace("Filtered re-inserted: old shape: %s, " # type: ignore[attr-defined] + "new shape: %s)", batch.landmarks.shape, retval.shape) + + return retval + + def process_output(self, subbatches: list[AlignerBatch], batch_masks: np.ndarray) -> None: + """ Process the output from the re-align pass. + + - Transform landmarks from aligned face space to face space + - Re-insert faces that were filtered out from the re-align process back into the + landmarks list + + Parameters + ---------- + subbatches: list + List of sub-batch results for each re-aligned re-feed performed + batch_masks: :class:`numpy.ndarray` + The original re-feed filter masks from the first pass + """ + for batch, masks in zip(subbatches, batch_masks): + if not np.all(masks): + batch.landmarks = self._transform_to_frame(batch) + batch.landmarks = self._re_insert_filtered(batch, masks) diff --git a/plugins/extract/align/cv2_dnn.py b/plugins/extract/align/cv2_dnn.py index d9d1fae091..f646b7cdc2 100644 --- a/plugins/extract/align/cv2_dnn.py +++ b/plugins/extract/align/cv2_dnn.py @@ -23,54 +23,128 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from __future__ import annotations +import logging +import typing as T import cv2 import numpy as np -from ._base import Aligner, logger +from ._base import Aligner, AlignerBatch, BatchType + +if T.TYPE_CHECKING: + from lib.align.detected_face import DetectedFace + +logger = logging.getLogger(__name__) class Align(Aligner): """ Perform transformation to align and get landmarks """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 1 model_filename = "cnn-facial-landmark_v1.pb" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: cv2.dnn.Net + self.model_path: str self.name = "cv2-DNN Aligner" self.input_size = 128 self.color_format = "RGB" self.vram = 0 # Doesn't use GPU self.vram_per_batch = 0 self.batchsize = 1 + self.realign_centering = "legacy" - def init_model(self): + def init_model(self) -> None: """ Initialize CV2 DNN Detector Model""" - self.model = cv2.dnn.readNetFromTensorflow(self.model_path) # pylint: disable=no-member - self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) # pylint: disable=no-member - - def process_input(self, batch): - """ Compile the detected faces for prediction """ - faces, batch["roi"], batch["offsets"] = self.align_image(batch) - faces = self._normalize_faces(faces) - batch["feed"] = np.array(faces, dtype="float32")[..., :3].transpose((0, 3, 1, 2)) - return batch - - def align_image(self, batch): - """ Align the incoming image for prediction """ - logger.trace("Aligning image around center") + self.model = cv2.dnn.readNetFromTensorflow(self.model_path) + self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) + + def faces_to_feed(self, faces: np.ndarray) -> np.ndarray: + """ Convert a batch of face images from UINT8 (0-255) to fp32 (0.0-255.0) + + Parameters + ---------- + faces: :class:`numpy.ndarray` + The batch of faces in UINT8 format + + Returns + ------- + class: `numpy.ndarray` + The batch of faces as fp32 + """ + return faces.astype("float32").transpose((0, 3, 1, 2)) + + def process_input(self, batch: BatchType) -> None: + """ Compile the detected faces for prediction + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch to process input for + + Returns + ------- + :class:`AlignerBatch` + The batch item with the :attr:`feed` populated and any required :attr:`data` added + """ + assert isinstance(batch, AlignerBatch) + lfaces, roi, offsets = self.align_image(batch) + batch.feed = np.array(lfaces)[..., :3] + batch.data.append({"roi": roi, "offsets": offsets}) + + def _get_box_and_offset(self, face: DetectedFace) -> tuple[list[int], int]: + """Obtain the bounding box and offset from a detected face. + + + Parameters + ---------- + face: :class:`~lib.align.DetectedFace` + The detected face object to obtain the bounding box and offset from + + Returns + ------- + box: list + The [left, top, right, bottom] bounding box + offset: int + The offset of the box (difference between half width vs height) + """ + + box = T.cast(list[int], [face.left, + face.top, + face.right, + face.bottom]) + diff_height_width = T.cast(int, face.height) - T.cast(int, face.width) + offset = int(abs(diff_height_width / 2)) + return box, offset + + def align_image(self, batch: AlignerBatch) -> tuple[list[np.ndarray], + list[list[int]], + list[tuple[int, int]]]: + """ Align the incoming image for prediction + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch to align the input for + + Returns + ------- + faces: list + List of feed faces for the aligner + rois: list + List of roi's for the faces + offsets: list + List of offsets for the faces + """ + logger.trace("Aligning image around center") # type:ignore[attr-defined] sizes = (self.input_size, self.input_size) rois = [] faces = [] offsets = [] - for det_face, image in zip(batch["detected_faces"], batch["image"]): - box = (det_face.left, - det_face.top, - det_face.right, - det_face.bottom) - diff_height_width = det_face.h - det_face.w - offset_y = int(abs(diff_height_width / 2)) - box_moved = self.move_box(box, [0, offset_y]) + for det_face, image in zip(batch.detected_faces, batch.image): + box, offset_y = self._get_box_and_offset(det_face) + box_moved = self.move_box(box, (0, offset_y)) # Make box square. roi = self.get_square_box(box_moved) @@ -85,9 +159,24 @@ def align_image(self, batch): offsets.append(offset) return faces, rois, offsets - @staticmethod - def move_box(box, offset): - """Move the box to direction specified by vector offset""" + @classmethod + def move_box(cls, + box: list[int], + offset: tuple[int, int]) -> list[int]: + """Move the box to direction specified by vector offset + + Parameters + ---------- + box: list + The (`left`, `top`, `right`, `bottom`) box positions + offset: tuple + (x, y) offset to move the box + + Returns + ------- + list + The original box shifted by the offset + """ left = box[0] + offset[0] top = box[1] + offset[1] right = box[2] + offset[0] @@ -95,8 +184,19 @@ def move_box(box, offset): return [left, top, right, bottom] @staticmethod - def get_square_box(box): - """Get a square box out of the given box, by expanding it.""" + def get_square_box(box: list[int]) -> list[int]: + """Get a square box out of the given box, by expanding it. + + Parameters + ---------- + box: list + The (`left`, `top`, `right`, `bottom`) box positions + + Returns + ------- + list + The original box but made square + """ left = box[0] top = box[1] right = box[2] @@ -127,15 +227,29 @@ def get_square_box(box): return [left, top, right, bottom] - @staticmethod - def pad_image(box, image): - """Pad image if face-box falls outside of boundaries """ + @classmethod + def pad_image(cls, box: list[int], image: np.ndarray) -> tuple[np.ndarray, tuple[int, int]]: + """Pad image if face-box falls outside of boundaries + + Parameters + ---------- + box: list + The (`left`, `top`, `right`, `bottom`) roi box positions + image: :class:`numpy.ndarray` + The image to be padded + + Returns + ------- + :class:`numpy.ndarray` + The padded image + """ height, width = image.shape[:2] pad_l = 1 - box[0] if box[0] < 0 else 0 pad_t = 1 - box[1] if box[1] < 0 else 0 pad_r = box[2] - width if box[2] > width else 0 pad_b = box[3] - height if box[3] > height else 0 - logger.trace("Padding: (l: %s, t: %s, r: %s, b: %s)", pad_l, pad_t, pad_r, pad_b) + logger.trace("Padding: (l: %s, t: %s, r: %s, b: %s)", # type:ignore[attr-defined] + pad_l, pad_t, pad_r, pad_b) padded_image = cv2.copyMakeBorder(image.copy(), pad_t, pad_b, @@ -144,29 +258,59 @@ def pad_image(box, image): cv2.BORDER_CONSTANT, value=(0, 0, 0)) offsets = (pad_l - pad_r, pad_t - pad_b) - logger.trace("image_shape: %s, Padded shape: %s, box: %s, offsets: %s", + logger.trace("image_shape: %s, Padded shape: %s, box: %s, " # type:ignore[attr-defined] + "offsets: %s", image.shape, padded_image.shape, box, offsets) return padded_image, offsets - def predict(self, batch): - """ Predict the 68 point landmarks """ - logger.trace("Predicting Landmarks") - self.model.setInput(batch) + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Predict the 68 point landmarks + + Parameters + ---------- + feed: :class:`numpy.ndarray` + The batch to feed into the aligner + + Returns + ------- + :class:`numpy.ndarray` + The predictions from the aligner + """ + assert isinstance(self.model, cv2.dnn.Net) + self.model.setInput(feed) retval = self.model.forward() return retval - def process_output(self, batch): - """ Process the output from the model """ + def process_output(self, batch: BatchType) -> None: + """ Process the output from the model + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch from the model with :attr:`predictions` populated + """ + assert isinstance(batch, AlignerBatch) self.get_pts_from_predict(batch) - return batch - @staticmethod - def get_pts_from_predict(batch): - """ Get points from predictor """ - for prediction, roi, offset in zip(batch["prediction"], batch["roi"], batch["offsets"]): - points = np.reshape(prediction, (-1, 2)) - points *= (roi[2] - roi[0]) - points[:, 0] += (roi[0] - offset[0]) - points[:, 1] += (roi[1] - offset[1]) - batch.setdefault("landmarks", []).append(points) - logger.trace("Predicted Landmarks: %s", batch["landmarks"]) + def get_pts_from_predict(self, batch: AlignerBatch): + """ Get points from predictor and populates the :attr:`landmarks` property + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch from the model with :attr:`predictions` populated + """ + landmarks = [] + if batch.second_pass: + batch.landmarks = batch.prediction.reshape(self.batchsize, -1, 2) * self.input_size + else: + for prediction, roi, offset in zip(batch.prediction, + batch.data[0]["roi"], + batch.data[0]["offsets"]): + points = np.reshape(prediction, (-1, 2)) + points *= (roi[2] - roi[0]) + points[:, 0] += (roi[0] - offset[0]) + points[:, 1] += (roi[1] - offset[1]) + landmarks.append(points) + batch.landmarks = np.array(landmarks) + logger.trace("Predicted Landmarks: %s", batch.landmarks) # type:ignore[attr-defined] diff --git a/plugins/extract/align/external.py b/plugins/extract/align/external.py new file mode 100644 index 0000000000..929e9b11c9 --- /dev/null +++ b/plugins/extract/align/external.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +""" Import 68 point landmarks or ROI boxes from a json file """ +import logging +import typing as T +import os +import re + +import numpy as np + +from lib.align import EXTRACT_RATIOS, LandmarkType +from lib.utils import FaceswapError, IMAGE_EXTENSIONS + +from ._base import BatchType, Aligner, AlignerBatch + +logger = logging.getLogger(__name__) + + +class Align(Aligner): + """ Import face detection bounding boxes from an external json file """ + def __init__(self, **kwargs) -> None: + kwargs["normalize_method"] = None # Disable normalization + kwargs["re_feed"] = 0 # Disable re-feed + kwargs["re_align"] = False # Disablle re-align + kwargs["disable_filter"] = True # Disable aligner filters + super().__init__(git_model_id=None, model_filename=None, **kwargs) + + self.name = "External" + self.batchsize = 16 + + self._origin: T.Literal["top-left", + "bottom-left", + "top-right", + "bottom-right"] = self.config["origin"] + + self._re_frame_no: re.Pattern = re.compile(r"\d+$") + self._is_video: bool = False + self._imported: dict[str | int, tuple[int, np.ndarray]] = {} + """dict[str | int, tuple[int, np.ndarray]]: filename as key, value of [number of faces + remaining for the frame, all landmarks in the frame] """ + + self._missing: list[str] = [] + self._roll: dict[T.Literal["bottom-left", "top-right", "bottom-right"], int] = { + "bottom-left": 3, "top-right": 1, "bottom-right": 2} + """dict[Literal["bottom-left", "top-right", "bottom-right"], int]: Amount to roll the + points by for different origins when 4 Point ROI landmarks are provided """ + + centering = self.config["4_point_centering"] + self._adjustment: float = 1. if centering is None else 1. - EXTRACT_RATIOS[centering] + """float: The amount to adjust 4 point ROI landmarks to standardize the points for a + 'head' sized extracted face """ + + def init_model(self) -> None: + """ No initialization to perform """ + logger.debug("No aligner model to initialize") + + def _check_for_video(self, filename: str) -> None: + """ Check a sample filename from the import file for a file extension to set + :attr:`_is_video` + + Parameters + ---------- + filename: str + A sample file name from the imported data + """ + logger.debug("Checking for video from '%s'", filename) + ext = os.path.splitext(filename)[-1] + if ext.lower() not in IMAGE_EXTENSIONS: + self._is_video = True + logger.debug("Set is_video to %s from extension '%s'", self._is_video, ext) + + def _get_key(self, key: str) -> str | int: + """ Obtain the key for the item in the lookup table. If the input are images, the key will + be the image filename. If the input is a video, the key will be the frame number + + Parameters + ---------- + key: str + The initial key value from import data or an import image/frame + + Returns + ------- + str | int + The filename is the input data is images, otherwise the frame number of a video + """ + if not self._is_video: + return key + original_name = os.path.splitext(key)[0] + matches = self._re_frame_no.findall(original_name) + if not matches or len(matches) > 1: + raise FaceswapError(f"Invalid import name: '{key}'. For video files, the key should " + "end with the frame number.") + retval = int(matches[0]) + logger.trace("Obtained frame number %s from key '%s'", # type:ignore[attr-defined] + retval, key) + return retval + + def _import_face(self, face: dict[str, list[int] | list[list[float]]]) -> np.ndarray: + """ Import the landmarks from a single face + + Parameters + ---------- + face: dict[str, list[int] | list[list[float]]] + An import dictionary item for a face + + Returns + ------- + :class:`numpy.ndarray` + The landmark data imported from the json file + + Raises + ------ + FaceSwapError + If the landmarks_2d key does not exist or the landmarks are in an incorrect format + """ + landmarks = face.get("landmarks_2d") + if landmarks is None: + raise FaceswapError("The provided import file is the required key 'landmarks_2d") + if len(landmarks) not in (4, 68): + raise FaceswapError("Imported 'landmarks_2d' should be either 68 facial feature " + "landmarks or 4 ROI corner locations") + retval = np.array(landmarks, dtype="float32") + if retval.shape[-1] != 2: + raise FaceswapError("Imported 'landmarks_2d' should be formatted as a list of (x, y) " + "co-ordinates") + if retval.shape[0] == 4: # Adjust ROI landmarks based on centering selected + center = np.mean(retval, axis=0) + retval = (retval - center) * self._adjustment + center + + return retval + + def import_data(self, data: dict[str, list[dict[str, list[int] | list[list[float]]]]]) -> None: + """ Import the aligner data from the json import file and set to :attr:`_imported` + + Parameters + ---------- + data: dict[str, list[dict[str, list[int] | list[list[float]]]]] + The data to be imported + """ + logger.debug("Data length: %s", len(data)) + self._check_for_video(list(data)[0]) + for key, faces in data.items(): + try: + lms = np.array([self._import_face(face) for face in faces], dtype="float32") + if not np.any(lms): + logger.trace("Skipping frame '%s' with no faces") # type:ignore[attr-defined] + continue + + store_key = self._get_key(key) + self._imported[store_key] = (lms.shape[0], lms) + except FaceswapError as err: + logger.error(str(err)) + msg = f"The imported frame key that failed was '{key}'" + raise FaceswapError(msg) from err + lm_shape = set(v[1].shape[1:] for v in self._imported.values() if v[0] > 0) + if len(lm_shape) > 1: + raise FaceswapError("All external data should have the same number of landmarks. " + f"Found landmarks of shape: {lm_shape}") + if (4, 2) in lm_shape: + self.landmark_type = LandmarkType.LM_2D_4 + + def process_input(self, batch: BatchType) -> None: + """ Put the filenames and original frame dimensions into `batch.feed` so they can be + collected for mapping in `.predict` + + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.AlignerBatch` + The batch to be processed by the plugin + """ + batch.feed = np.array([(self._get_key(os.path.basename(f)), i.shape[:2]) + for f, i in zip(batch.filename, batch.image)], dtype="object") + + def faces_to_feed(self, faces: np.ndarray) -> np.ndarray: + """ No action required for import plugin + + Parameters + ---------- + faces: :class:`numpy.ndarray` + The batch of faces in UINT8 format + + Returns + ------- + class: `numpy.ndarray` + the original batch of faces + """ + return faces + + def _adjust_for_origin(self, landmarks: np.ndarray, frame_dims: tuple[int, int]) -> np.ndarray: + """ Adjust the landmarks to be top-left orientated based on the selected import origin + + Parameters + ---------- + landmarks: :class:`np.ndarray` + The imported facial landmarks box at original (0, 0) origin + frame_dims: tuple[int, int] + The (rows, columns) dimensions of the original frame + + Returns + ------- + :class:`numpy.ndarray` + The adjusted landmarks box for a top-left origin + """ + if not np.any(landmarks) or self._origin == "top-left": + return landmarks + + if LandmarkType.from_shape(landmarks.shape) == LandmarkType.LM_2D_4: + landmarks = np.roll(landmarks, self._roll[self._origin], axis=0) + + if self._origin.startswith("bottom"): + landmarks[:, 1] = frame_dims[0] - landmarks[:, 1] + if self._origin.endswith("right"): + landmarks[:, 0] = frame_dims[1] - landmarks[:, 0] + + return landmarks + + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Pair the input filenames to the import file + + Parameters + ---------- + feed: :class:`numpy.ndarray` + The filenames in the batch to return imported alignments for + + Returns + ------- + :class:`numpy.ndarray` + The predictions for the given filenames + """ + preds = [] + for key, frame_dims in feed: + if key not in self._imported: + self._missing.append(key) + continue + + remaining, all_lms = self._imported[key] + preds.append(self._adjust_for_origin(all_lms[all_lms.shape[0] - remaining], + frame_dims)) + + if remaining == 1: + del self._imported[key] + else: + self._imported[key] = (remaining - 1, all_lms) + + return np.array(preds, dtype="float32") + + def process_output(self, batch: BatchType) -> None: + """ Process the imported data to the landmarks attribute + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch from the model with :attr:`predictions` populated + """ + assert isinstance(batch, AlignerBatch) + batch.landmarks = batch.prediction + logger.trace("Imported landmarks: %s", batch.landmarks) # type:ignore[attr-defined] + + def on_completion(self) -> None: + """ Output information if: + - Imported items were not matched in input data + - Input data was not matched in imported items + """ + super().on_completion() + + if self._missing: + logger.warning("[ALIGN] %s input frames could not be matched in the import file " + "'%s'. Run in verbose mode for a list of frames.", + len(self._missing), self.config["file_name"]) + logger.verbose( # type:ignore[attr-defined] + "[ALIGN] Input frames not in import file: %s", self._missing) + + if self._imported: + logger.warning("[ALIGN] %s items in the import file '%s' could not be matched to any " + "input frames. Run in verbose mode for a list of items.", + len(self._imported), self.config["file_name"]) + logger.verbose( # type:ignore[attr-defined] + "[ALIGN] import file items not in input frames: %s", list(self._imported)) diff --git a/plugins/extract/align/external_defaults.py b/plugins/extract/align/external_defaults.py new file mode 100644 index 0000000000..875abd01d0 --- /dev/null +++ b/plugins/extract/align/external_defaults.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +""" + The default options for the faceswap Import Alignments plugin. + + Defaults files should be named _defaults.py + Any items placed into this file will automatically get added to the relevant config .ini files + within the faceswap/config folder. + + The following variables should be defined: + _HELPTEXT: A string describing what this plugin does + _DEFAULTS: A dictionary containing the options, defaults and meta information. The + dictionary should be defined as: + {: {}} + + should always be lower text. + dictionary requirements are listed below. + + The following keys are expected for the _DEFAULTS dict: + datatype: [required] A python type class. This limits the type of data that can be + provided in the .ini file and ensures that the value is returned in the + correct type to faceswap. Valid data types are: , , + , . + default: [required] The default value for this option. + info: [required] A string describing what this option does. + group: [optional]. A group for grouping options together in the GUI. If not + provided this will not group this option with any others. + choices: [optional] If this option's datatype is of then valid + selections can be defined here. This validates the option and also enables + a combobox / radio option in the GUI. + gui_radio: [optional] If are defined, this indicates that the GUI should use + radio buttons rather than a combobox to display this option. + min_max: [partial] For and data types this is required + otherwise it is ignored. Should be a tuple of min and max accepted values. + This is used for controlling the GUI slider range. Values are not enforced. + rounding: [partial] For and data types this is + required otherwise it is ignored. Used for the GUI slider. For floats, this + is the number of decimal places to display. For ints this is the step size. + fixed: [optional] [train only]. Training configurations are fixed when the model is + created, and then reloaded from the state file. Marking an item as fixed=False + indicates that this value can be changed for existing models, and will override + the value saved in the state file with the updated value in config. If not + provided this will default to True. +""" + + +_HELPTEXT = ( + "Import Aligner options.\n" + "Imports either 68 point 2D landmarks or an aligned bounding box from an external .json file." + ) + + +_DEFAULTS = { + "file_name": { + "default": "import.json", + "info": "The import file should be stored in the same folder as the video (if extracting " + "from a video file) or inside the folder of images (if importing from a folder of images)", + "datatype": str, + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "origin": { + "default": "top-left", + "info": "The origin (0, 0) location of the co-ordinates system used. " + "\n\t top-left: The origin (0, 0) of the canvas is at the top left " + "corner." + "\n\t bottom-left: The origin (0, 0) of the canvas is at the bottom " + "left corner." + "\n\t top-right: The origin (0, 0) of the canvas is at the top right " + "corner." + "\n\t bottom-right: The origin (0, 0) of the canvas is at the bottom " + "right corner.", + "datatype": str, + "choices": ["top-left", "bottom-left", "top-right", "bottom-right"], + "group": "input", + "gui_radio": True + }, + "4_point_centering": { + "default": "head", + "info": "4 point ROI landmarks only. The approximate centering for the location of the " + "corner points to be imported. Default faceswap extracts are generated at 'head' " + "centering, but it is possible to pass in ROI points at a tighter centering. " + "Refer to https://github.com/deepfakes/faceswap/pull/1095 for a visual guide" + "\n\t head: The ROI points represent a loose crop enclosing the whole head." + "\n\t face: The ROI points represent a medium crop enclosing the face." + "\n\t legacy: The ROI points represent a tight crop enclosing the central face " + "area." + "\n\t none: Only required if importing 4 point ROI landmarks back into faceswap " + "having generated them from the 'alignments' tool 'export' job.", + "datatype": str, + "choices": ["head", "face", "legacy", "none"], + "group": "input", + "gui_radio": True + } + +} diff --git a/plugins/extract/align/fan.py b/plugins/extract/align/fan.py index 4f9c9ae55d..a829f3bcac 100644 --- a/plugins/extract/align/fan.py +++ b/plugins/extract/align/fan.py @@ -3,30 +3,43 @@ Code adapted and modified from: https://github.com/1adrianb/face-alignment """ +from __future__ import annotations +import logging +import typing as T + import cv2 import numpy as np from lib.model.session import KSession -from ._base import Aligner, logger +from ._base import Aligner, AlignerBatch, BatchType + +if T.TYPE_CHECKING: + from lib.align import DetectedFace + +logger = logging.getLogger(__name__) class Align(Aligner): """ Perform transformation to align and get landmarks """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 13 model_filename = "face-alignment-network_2d4_keras_v2.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession self.name = "FAN" self.input_size = 256 self.color_format = "RGB" self.vram = 2240 self.vram_warnings = 512 # Will run at this with warnings self.vram_per_batch = 64 - self.batchsize = self.config["batch-size"] + self.realign_centering = "head" + self.batchsize: int = self.config["batch-size"] self.reference_scale = 200. / 195. - def init_model(self): + def init_model(self) -> None: """ Initialize FAN model """ + assert isinstance(self.name, str) + assert isinstance(self.model_path, str) self.model = KSession(self.name, self.model_path, allow_growth=self.config["allow_growth"], @@ -37,70 +50,148 @@ def init_model(self): placeholder = np.zeros(placeholder_shape, dtype="float32") self.model.predict(placeholder) - def process_input(self, batch): - """ Compile the detected faces for prediction """ - logger.debug("Aligning faces around center") - batch["center_scale"] = self.get_center_scale(batch["detected_faces"]) - faces = self.crop(batch) - logger.trace("Aligned image around center") - faces = self._normalize_faces(faces) - batch["feed"] = np.array(faces, dtype="float32")[..., :3] / 255.0 - return batch - - def get_center_scale(self, detected_faces): - """ Get the center and set scale of bounding box """ - logger.debug("Calculating center and scale") + def faces_to_feed(self, faces: np.ndarray) -> np.ndarray: + """ Convert a batch of face images from UINT8 (0-255) to fp32 (0.0-1.0) + + Parameters + ---------- + faces: :class:`numpy.ndarray` + The batch of faces in UINT8 format + + Returns + ------- + class: `numpy.ndarray` + The batch of faces as fp32 in 0.0 to 1.0 range + """ + return faces.astype("float32") / 255. + + def process_input(self, batch: BatchType) -> None: + """ Compile the detected faces for prediction + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch to process input for + """ + assert isinstance(batch, AlignerBatch) + logger.trace("Aligning faces around center") # type:ignore[attr-defined] + center_scale = self.get_center_scale(batch.detected_faces) + batch.feed = np.array(self.crop(batch, center_scale))[..., :3] + batch.data.append({"center_scale": center_scale}) + logger.trace("Aligned image around center") # type:ignore[attr-defined] + + def get_center_scale(self, detected_faces: list[DetectedFace]) -> np.ndarray: + """ Get the center and set scale of bounding box + + Parameters + ---------- + detected_faces: list + List of :class:`~lib.align.DetectedFace` objects for the batch + + Returns + ------- + :class:`numpy.ndarray` + The center and scale of the bounding box + """ + logger.trace("Calculating center and scale") # type:ignore[attr-defined] center_scale = np.empty((len(detected_faces), 68, 3), dtype='float32') for index, face in enumerate(detected_faces): - x_center = (face.left + face.right) / 2.0 - y_center = (face.top + face.bottom) / 2.0 - face.h * 0.12 - scale = (face.w + face.h) * self.reference_scale - center_scale[index, :, 0] = np.full(68, x_center, dtype='float32') - center_scale[index, :, 1] = np.full(68, y_center, dtype='float32') + x_ctr = (T.cast(int, face.left) + face.right) / 2.0 + y_ctr = (T.cast(int, face.top) + face.bottom) / 2.0 - T.cast(int, face.height) * 0.12 + scale = (T.cast(int, face.width) + T.cast(int, face.height)) * self.reference_scale + center_scale[index, :, 0] = np.full(68, x_ctr, dtype='float32') + center_scale[index, :, 1] = np.full(68, y_ctr, dtype='float32') center_scale[index, :, 2] = np.full(68, scale, dtype='float32') - logger.trace("Calculated center and scale: %s", center_scale) + logger.trace("Calculated center and scale: %s", center_scale) # type:ignore[attr-defined] return center_scale - def crop(self, batch): # pylint:disable=too-many-locals - """ Crop image around the center point """ - logger.debug("Cropping images") - sizes = (self.input_size, self.input_size) - batch_shape = batch["center_scale"].shape[:2] + def _crop_image(self, + image: np.ndarray, + top_left: np.ndarray, + bottom_right: np.ndarray) -> np.ndarray: + """ Crop a single image + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to crop + top_left: :class:`numpy.ndarray` + The top left (x, y) point to crop from + bottom_right: :class:`numpy.ndarray` + The bottom right (x, y) point to crop to + + Returns + ------- + :class:`numpy.ndarray` + The cropped image + """ + bottom_right_width, bottom_right_height = bottom_right[0].astype('int32') + top_left_width, top_left_height = top_left[0].astype('int32') + new_dim = (bottom_right_height - top_left_height, + bottom_right_width - top_left_width, + 3 if image.ndim > 2 else 1) + new_img = np.zeros(new_dim, dtype=np.uint8) + + new_x = slice(max(0, -top_left_width), + min(bottom_right_width, image.shape[1]) - top_left_width) + new_y = slice(max(0, -top_left_height), + min(bottom_right_height, image.shape[0]) - top_left_height) + old_x = slice(max(0, top_left_width), min(bottom_right_width, image.shape[1])) + old_y = slice(max(0, top_left_height), min(bottom_right_height, image.shape[0])) + new_img[new_y, new_x] = image[old_y, old_x] + + interp = cv2.INTER_CUBIC if new_dim[0] < self.input_size else cv2.INTER_AREA + return cv2.resize(new_img, + dsize=(self.input_size, self.input_size), + interpolation=interp) + + def crop(self, batch: AlignerBatch, center_scale: np.ndarray) -> list[np.ndarray]: + """ Crop image around the center point + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch to crop the image for + center_scale: :class:`numpy.ndarray` + The center and scale for the bounding box + + Returns + ------- + list + List of cropped images for the batch + """ + logger.trace("Cropping images") # type:ignore[attr-defined] + batch_shape = center_scale.shape[:2] resolutions = np.full(batch_shape, self.input_size, dtype='float32') matrix_ones = np.ones(batch_shape + (3,), dtype='float32') matrix_size = np.full(batch_shape + (3,), self.input_size, dtype='float32') matrix_size[..., 2] = 1.0 - upper_left = self.transform(matrix_ones, batch["center_scale"], resolutions) - bot_right = self.transform(matrix_size, batch["center_scale"], resolutions) + upper_left = self.transform(matrix_ones, center_scale, resolutions) + bot_right = self.transform(matrix_size, center_scale, resolutions) # TODO second pass .. convert to matrix - new_images = [] - for image, top_left, bottom_right in zip(batch["image"], upper_left, bot_right): - height, width = image.shape[:2] - channels = 3 if image.ndim > 2 else 1 - bottom_right_width, bottom_right_height = bottom_right[0].astype('int32') - top_left_width, top_left_height = top_left[0].astype('int32') - new_dim = (bottom_right_height - top_left_height, - bottom_right_width - top_left_width, - channels) - new_img = np.empty(new_dim, dtype=np.uint8) - - new_x = slice(max(0, -top_left_width), min(bottom_right_width, width) - top_left_width) - new_y = slice(max(0, -top_left_height), - min(bottom_right_height, height) - top_left_height) - old_x = slice(max(0, top_left_width), min(bottom_right_width, width)) - old_y = slice(max(0, top_left_height), min(bottom_right_height, height)) - new_img[new_y, new_x] = image[old_y, old_x] - - interp = cv2.INTER_CUBIC if new_dim[0] < self.input_size else cv2.INTER_AREA - new_images.append(cv2.resize(new_img, dsize=sizes, interpolation=interp)) - logger.trace("Cropped images") + new_images = [self._crop_image(image, top_left, bottom_right) + for image, top_left, bottom_right in zip(batch.image, upper_left, bot_right)] + logger.trace("Cropped images") # type:ignore[attr-defined] return new_images - @staticmethod - def transform(points, center_scales, resolutions): - """ Transform Image """ - logger.debug("Transforming Points") + @classmethod + def transform(cls, + points: np.ndarray, + center_scales: np.ndarray, + resolutions: np.ndarray) -> np.ndarray: + """ Transform Image + + Parameters + ---------- + points: :class:`numpy.ndarray` + The points to transform + center_scales: :class:`numpy.ndarray` + The calculated centers and scales for the batch + resolutions: :class:`numpy.ndarray` + The resolutions + """ + logger.trace("Transforming Points") # type:ignore[attr-defined] num_images, num_landmarks = points.shape[:2] transform_matrix = np.eye(3, dtype='float32') transform_matrix = np.repeat(transform_matrix[None, :], num_landmarks, axis=0) @@ -113,45 +204,78 @@ def transform(points, center_scales, resolutions): transform_matrix[:, :, 1, 2] = translations[:, :, 1] # y translation new_points = np.einsum('abij, abj -> abi', transform_matrix, points, optimize='greedy') retval = new_points[:, :, :2].astype('float32') - logger.trace("Transformed Points: %s", retval) + logger.trace("Transformed Points: %s", retval) # type:ignore[attr-defined] return retval - def predict(self, batch): - """ Predict the 68 point landmarks """ - logger.debug("Predicting Landmarks") + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Predict the 68 point landmarks + + Parameters + ---------- + batch: :class:`numpy.ndarray` + The batch to feed into the aligner + + Returns + ------- + :class:`numpy.ndarray` + The predictions from the aligner + """ + logger.trace("Predicting Landmarks") # type:ignore[attr-defined] # TODO Remove lazy transpose and change points from predict to use the correct # order - retval = self.model.predict(batch)[-1].transpose(0, 3, 1, 2) - logger.trace(retval.shape) + retval = self.model.predict(feed)[-1].transpose(0, 3, 1, 2) + logger.trace(retval.shape) # type:ignore[attr-defined] return retval - def process_output(self, batch): - """ Process the output from the model """ + def process_output(self, batch: BatchType) -> None: + """ Process the output from the model + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch from the model with :attr:`predictions` populated + """ + assert isinstance(batch, AlignerBatch) self.get_pts_from_predict(batch) - return batch - def get_pts_from_predict(self, batch): - """ Get points from predictor """ - logger.debug("Obtain points from prediction") - num_images, num_landmarks, height, width = batch["prediction"].shape + def get_pts_from_predict(self, batch: AlignerBatch) -> None: + """ Get points from predictor and populate the :attr:`landmarks` property of the + :class:`AlignerBatch` + + Parameters + ---------- + batch: :class:`AlignerBatch` + The current batch from the model with :attr:`predictions` populated + """ + logger.trace("Obtain points from prediction") # type:ignore[attr-defined] + num_images, num_landmarks = batch.prediction.shape[:2] image_slice = np.repeat(np.arange(num_images)[:, None], num_landmarks, axis=1) landmark_slice = np.repeat(np.arange(num_landmarks)[None, :], num_images, axis=0) resolution = np.full((num_images, num_landmarks), 64, dtype='int32') subpixel_landmarks = np.ones((num_images, num_landmarks, 3), dtype='float32') - flat_indices = batch["prediction"].reshape(num_images, num_landmarks, -1).argmax(-1) - indices = np.array(np.unravel_index(flat_indices, (height, width))) - min_clipped = np.minimum(indices + 1, height - 1) + indices = np.array(np.unravel_index(batch.prediction.reshape(num_images, + num_landmarks, + -1).argmax(-1), + (batch.prediction.shape[2], # height + batch.prediction.shape[3]))) # width + min_clipped = np.minimum(indices + 1, batch.prediction.shape[2] - 1) max_clipped = np.maximum(indices - 1, 0) offsets = [(image_slice, landmark_slice, indices[0], min_clipped[1]), (image_slice, landmark_slice, indices[0], max_clipped[1]), (image_slice, landmark_slice, min_clipped[0], indices[1]), (image_slice, landmark_slice, max_clipped[0], indices[1])] - x_subpixel_shift = batch["prediction"][offsets[0]] - batch["prediction"][offsets[1]] - y_subpixel_shift = batch["prediction"][offsets[2]] - batch["prediction"][offsets[3]] + x_subpixel_shift = batch.prediction[offsets[0]] - batch.prediction[offsets[1]] + y_subpixel_shift = batch.prediction[offsets[2]] - batch.prediction[offsets[3]] # TODO improve rudimentary sub-pixel logic to centroid of 3x3 window algorithm subpixel_landmarks[:, :, 0] = indices[1] + np.sign(x_subpixel_shift) * 0.25 + 0.5 subpixel_landmarks[:, :, 1] = indices[0] + np.sign(y_subpixel_shift) * 0.25 + 0.5 - batch["landmarks"] = self.transform(subpixel_landmarks, batch["center_scale"], resolution) - logger.trace("Obtained points from prediction: %s", batch["landmarks"]) + if batch.second_pass: # Transformation handled by plugin parent for re-aligned faces + batch.landmarks = subpixel_landmarks[..., :2] * 4. + else: + batch.landmarks = self.transform(subpixel_landmarks, + batch.data[0]["center_scale"], + resolution) + logger.trace("Obtained points from prediction: %s", # type:ignore[attr-defined] + batch.landmarks) diff --git a/plugins/extract/align/fan_defaults.py b/plugins/extract/align/fan_defaults.py index 7c6425827c..90d5bc4b3a 100644 --- a/plugins/extract/align/fan_defaults.py +++ b/plugins/extract/align/fan_defaults.py @@ -50,19 +50,19 @@ _DEFAULTS = { - "batch-size": dict( - default=12, - info="The batch size to use. To a point, higher batch sizes equal better performance, " - "but setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can " - "accomodate then this will automatically be lowered." - "\n\tAMD users: A batchsize of 8 requires about 4 GB vram.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ) + "batch-size": { + "default": 12, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered." + "\n\tAMD users: A batchsize of 8 requires about 4 GB vram.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + } } diff --git a/plugins/extract/detect/_base.py b/plugins/extract/detect/_base.py index 97c6effdaa..3c3221d84f 100644 --- a/plugins/extract/detect/_base.py +++ b/plugins/extract/detect/_base.py @@ -4,7 +4,7 @@ All Detector Plugins should inherit from this class. See the override methods for which methods are required. -The plugin will receive a :class:`~plugins.extract.pipeline.ExtractMedia` object. +The plugin will receive a :class:`~plugins.extract.extract_media.ExtractMedia` object. For each source frame, the plugin must pass a dict to finalize containing: @@ -13,17 +13,63 @@ To get a :class:`~lib.align.DetectedFace` object use the function: ->>> face = self.to_detected_face(, , , ) +>>> face = self._to_detected_face(, , , ) """ +from __future__ import annotations +import logging +import typing as T + +from dataclasses import dataclass, field + import cv2 import numpy as np -from tensorflow.python.framework import errors_impl as tf_errors +from tensorflow.python.framework import errors_impl as tf_errors # pylint:disable=no-name-in-module # noqa from lib.align import DetectedFace -from lib.utils import get_backend, FaceswapError +from lib.utils import FaceswapError + +from plugins.extract._base import BatchType, Extractor, ExtractorBatch +from plugins.extract import ExtractMedia + +if T.TYPE_CHECKING: + from collections.abc import Generator + from queue import Queue + +logger = logging.getLogger(__name__) + -from plugins.extract._base import Extractor, logger +@dataclass +class DetectorBatch(ExtractorBatch): + """ Dataclass for holding items flowing through the aligner. + + Inherits from :class:`~plugins.extract._base.ExtractorBatch` + + Parameters + ---------- + rotation_matrix: :class:`numpy.ndarray` + The rotation matrix for any requested rotations + scale: float + The scaling factor to take the input image back to original size + pad: tuple + The amount of padding to apply to the image to feed the network + initial_feed: :class:`numpy.ndarray` + Used to hold the initial :attr:`feed` when rotate images is enabled + """ + detected_faces: list[list["DetectedFace"]] = field(default_factory=list) + rotation_matrix: list[np.ndarray] = field(default_factory=list) + scale: list[float] = field(default_factory=list) + pad: list[tuple[int, int]] = field(default_factory=list) + initial_feed: np.ndarray = np.array([]) + + def __repr__(self): + """ Prettier repr for debug printing """ + retval = super().__repr__() + retval += (f", rotation_matrix={self.rotation_matrix}, " + f"scale={self.scale}, " + f"pad={self.pad}, " + f"initial_feed=({self.initial_feed.shape}, {self.initial_feed.dtype})") + return retval class Detector(Extractor): # pylint:disable=abstract-method @@ -60,8 +106,14 @@ class Detector(Extractor): # pylint:disable=abstract-method plugins.extract.mask._base : Masker parent class for extraction plugins. """ - def __init__(self, git_model_id=None, model_filename=None, - configfile=None, instance=0, rotation=None, min_size=0, **kwargs): + def __init__(self, + git_model_id: int | None = None, + model_filename: str | list[str] | None = None, + configfile: str | None = None, + instance: int = 0, + rotation: str | None = None, + min_size: int = 0, + **kwargs) -> None: logger.debug("Initializing %s: (rotation: %s, min_size: %s)", self.__class__.__name__, rotation, min_size) super().__init__(git_model_id, @@ -77,11 +129,11 @@ def __init__(self, git_model_id=None, model_filename=None, logger.debug("Initialized _base %s", self.__class__.__name__) # <<< QUEUE METHODS >>> # - def get_batch(self, queue): + def get_batch(self, queue: Queue) -> tuple[bool, DetectorBatch]: """ Get items for inputting to the detector plugin in batches - Items are received as :class:`~plugins.extract.pipeline.ExtractMedia` objects and converted - to ``dict`` for internal processing. + Items are received as :class:`~plugins.extract.extract_media.ExtractMedia` objects and + converted to ``dict`` for internal processing. Items are returned from the ``queue`` in batches of :attr:`~plugins.extract._base.Extractor.batchsize` @@ -108,108 +160,161 @@ def get_batch(self, queue): ------- exhausted, bool ``True`` if queue is exhausted, ``False`` if not. - batch, dict - A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`. + batch, :class:`~plugins.extract._base.ExtractorBatch` + The batch object for the current batch """ exhausted = False - batch = dict() + batch = DetectorBatch() for _ in range(self.batchsize): item = self._get_item(queue) if item == "EOF": exhausted = True break - batch.setdefault("filename", []).append(item.filename) + assert isinstance(item, ExtractMedia) + # Put items that are already aligned into the out queue + if item.is_aligned: + self._queues["out"].put(item) + continue + batch.filename.append(item.filename) image, scale, pad = self._compile_detection_image(item) - batch.setdefault("image", []).append(image) - batch.setdefault("scale", []).append(scale) - batch.setdefault("pad", []).append(pad) - - if batch: - batch["image"] = np.array(batch["image"], dtype="float32") - logger.trace("Returning batch: %s", {k: v.shape if isinstance(v, np.ndarray) else v - for k, v in batch.items()}) + batch.image.append(image) + batch.scale.append(scale) + batch.pad.append(pad) + + if batch.filename: + logger.trace("Returning batch: %s", # type: ignore + {k: len(v) if isinstance(v, (list, np.ndarray)) else v + for k, v in batch.__dict__.items()}) else: - logger.trace(item) + logger.trace(item) # type:ignore[attr-defined] + + if not exhausted and not batch.filename: + # This occurs when face filter is fed aligned faces. + # Need to re-run until EOF is hit + return self.get_batch(queue) + return exhausted, batch # <<< FINALIZE METHODS>>> # - def finalize(self, batch): + def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]: """ Finalize the output from Detector This should be called as the final task of each ``plugin``. Parameters ---------- - batch : dict - The final ``dict`` from the `plugin` process. It must contain the keys ``filename``, - ``faces`` + batch : :class:`~plugins.extract._base.ExtractorBatch` + The batch object for the current batch Yields ------ - :class:`~plugins.extract.pipeline.ExtractMedia` + :class:`~plugins.extract.extract_media.ExtractMedia` The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes for the detected faces found in the frame. """ - if not isinstance(batch, dict): - logger.trace("Item out: %s", batch) - return batch + assert isinstance(batch, DetectorBatch) + logger.trace("Item out: %s", # type:ignore[attr-defined] + {k: len(v) if isinstance(v, (list, np.ndarray)) else v + for k, v in batch.__dict__.items()}) - logger.trace("Item out: %s", {k: v.shape if isinstance(v, np.ndarray) else v - for k, v in batch.items()}) - - batch_faces = [[self.to_detected_face(face[0], face[1], face[2], face[3]) + batch_faces = [[self._to_detected_face(face[0], face[1], face[2], face[3]) for face in faces] - for faces in batch["prediction"]] + for faces in batch.prediction] # Rotations - if any(m.any() for m in batch["rotmat"]) and any(batch_faces): + if any(m.any() for m in batch.rotation_matrix) and any(batch_faces): batch_faces = [[self._rotate_face(face, rotmat) if rotmat.any() else face for face in faces] - for faces, rotmat in zip(batch_faces, batch["rotmat"])] + for faces, rotmat in zip(batch_faces, batch.rotation_matrix)] # Remove zero sized faces batch_faces = self._remove_zero_sized_faces(batch_faces) # Scale back out to original frame - batch["detected_faces"] = [[self.to_detected_face((face.left - pad[0]) / scale, - (face.top - pad[1]) / scale, - (face.right - pad[0]) / scale, - (face.bottom - pad[1]) / scale) - for face in faces] - for scale, pad, faces in zip(batch["scale"], - batch["pad"], - batch_faces)] - - if self.min_size > 0 and batch.get("detected_faces", None): - batch["detected_faces"] = self._filter_small_faces(batch["detected_faces"]) - - batch = self._dict_lists_to_list_dicts(batch) - for item in batch: - output = self._extract_media.pop(item["filename"]) - output.add_detected_faces(item["detected_faces"]) - logger.trace("final output: (filename: '%s', image shape: %s, detected_faces: %s, " - "item: %s", output.filename, output.image_shape, output.detected_faces, - output) + batch.detected_faces = [[self._to_detected_face((face.left - pad[0]) / scale, + (face.top - pad[1]) / scale, + (face.right - pad[0]) / scale, + (face.bottom - pad[1]) / scale) + for face in faces + if face.left is not None and face.top is not None] + for scale, pad, faces in zip(batch.scale, + batch.pad, + batch_faces)] + + if self.min_size > 0 and batch.detected_faces: + batch.detected_faces = self._filter_small_faces(batch.detected_faces) + + for idx, filename in enumerate(batch.filename): + output = self._extract_media.pop(filename) + output.add_detected_faces(batch.detected_faces[idx]) + + logger.trace("final output: (filename: '%s', " # type:ignore[attr-defined] + "image shape: %s, detected_faces: %s, item: %s", + output.filename, output.image_shape, output.detected_faces, output) yield output @staticmethod - def to_detected_face(left, top, right, bottom): - """ Return a :class:`~lib.align.DetectedFace` object for the bounding box """ - return DetectedFace(x=int(round(left)), - w=int(round(right - left)), - y=int(round(top)), - h=int(round(bottom - top))) + def _to_detected_face(left: float, top: float, right: float, bottom: float) -> DetectedFace: + """ Convert a bounding box to a detected face object + + Parameters + ---------- + left: float + The left point of the detection bounding box + top: float + The top point of the detection bounding box + right: float + The right point of the detection bounding box + bottom: float + The bottom point of the detection bounding box + + Returns + ------- + class:`~lib.align.DetectedFace` + The detected face object for the given bounding box + """ + return DetectedFace(left=int(round(left)), + width=int(round(right - left)), + top=int(round(top)), + height=int(round(bottom - top))) # <<< PROTECTED ACCESS METHODS >>> # # <<< PREDICT WRAPPER >>> # - def _predict(self, batch): + def _predict(self, batch: BatchType) -> DetectorBatch: """ Wrap models predict function in rotations """ - batch["rotmat"] = [np.array([]) for _ in range(len(batch["feed"]))] - found_faces = [np.array([]) for _ in range(len(batch["feed"]))] + assert isinstance(batch, DetectorBatch) + batch.rotation_matrix = [np.array([]) for _ in range(len(batch.feed))] + found_faces: list[np.ndarray] = [np.array([]) for _ in range(len(batch.feed))] for angle in self.rotation: # Rotate the batch and insert placeholders for already found faces self._rotate_batch(batch, angle) try: - batch = self.predict(batch) + pred = self.predict(batch.feed) + if angle == 0: + batch.prediction = pred + else: + try: + batch.prediction = np.array([b if b.any() else p + for b, p in zip(batch.prediction, pred)]) + except ValueError as err: + # If batches are different sizes after rotation Numpy will error, so we + # need to explicitly set the dtype to 'object' rather than let it infer + # numpy error: + # ValueError: setting an array element with a sequence. The requested array + # has an inhomogeneous shape after 1 dimensions. The detected shape was + # (8,) + inhomogeneous part + if "inhomogeneous" in str(err): + batch.prediction = np.array([b if b.any() else p + for b, p in zip(batch.prediction, pred)], + dtype="object") + logger.trace( # type:ignore[attr-defined] + "Mismatched array sizes, setting dtype to object: %s", + [p.shape for p in batch.prediction]) + else: + raise + + logger.trace("angle: %s, filenames: %s, " # type:ignore[attr-defined] + "prediction: %s", + angle, batch.filename, pred) except tf_errors.ResourceExhaustedError as err: msg = ("You do not have enough GPU memory available to run detection at the " "selected batch size. You can try a number of things:" @@ -220,47 +325,43 @@ def _predict(self, batch): "CLI: Edit the file faceswap/config/extract.ini)." "\n3) Enable 'Single Process' mode.") raise FaceswapError(msg) from err - except Exception as err: - if get_backend() == "amd": - # pylint:disable=import-outside-toplevel - from lib.plaidml_utils import is_plaidml_error - if (is_plaidml_error(err) and ( - "CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or - "enough memory for the current schedule" in str(err).lower())): - msg = ("You do not have enough GPU memory available to run detection at " - "the selected batch size. You can try a number of things:" - "\n1) Close any other application that is using your GPU (web " - "browsers are particularly bad for this)." - "\n2) Lower the batchsize (the amount of images fed into the " - "model) by editing the plugin settings (GUI: Settings > Configure " - "extract settings, CLI: Edit the file " - "faceswap/config/extract.ini).") - raise FaceswapError(msg) from err - raise - - if angle != 0 and any([face.any() for face in batch["prediction"]]): - logger.verbose("found face(s) by rotating image %s degrees", angle) - - found_faces = [face if not found.any() else found - for face, found in zip(batch["prediction"], found_faces)] - - if all([face.any() for face in found_faces]): - logger.trace("Faces found for all images") + + if angle != 0 and any(face.any() for face in batch.prediction): + logger.verbose("found face(s) by rotating image %s " # type:ignore[attr-defined] + "degrees", + angle) + + found_faces = T.cast(list[np.ndarray], ([face if not found.any() else found + for face, found in zip(batch.prediction, + found_faces)])) + if all(face.any() for face in found_faces): + logger.trace("Faces found for all images") # type:ignore[attr-defined] break - batch["prediction"] = found_faces - logger.trace("detect_prediction output: (filenames: %s, prediction: %s, rotmat: %s)", - batch["filename"], batch["prediction"], batch["rotmat"]) + batch.prediction = np.array(found_faces, dtype="object") + logger.trace("detect_prediction output: (filenames: %s, " # type:ignore[attr-defined] + "prediction: %s, rotmat: %s)", + batch.filename, batch.prediction, batch.rotation_matrix) return batch # <<< DETECTION IMAGE COMPILATION METHODS >>> # - def _compile_detection_image(self, item): + def _compile_detection_image(self, item: ExtractMedia + ) -> tuple[np.ndarray, float, tuple[int, int]]: """ Compile the detection image for feeding into the model Parameters ---------- - item: :class:`plugins.extract.pipeline.ExtractMedia` + item: :class:`~plugins.extract.extract_media.ExtractMedia` The input item from the pipeline + + Returns + ------- + image: :class:`numpy.ndarray` + The original image formatted for detection + scale: float + The scaling factor for the image + pad: int + The amount of padding applied to the image """ image = item.get_image_copy(self.color_format) scale = self._set_scale(item.image_size) @@ -268,36 +369,89 @@ def _compile_detection_image(self, item): image = self._scale_image(image, item.image_size, scale) image = self._pad_image(image) - logger.trace("compiled: (images shape: %s, scale: %s, pad: %s)", image.shape, scale, pad) + logger.trace("compiled: (images shape: %s, " # type:ignore[attr-defined] + "scale: %s, pad: %s)", + image.shape, scale, pad) return image, scale, pad - def _set_scale(self, image_size): - """ Set the scale factor for incoming image """ + def _set_scale(self, image_size: tuple[int, int]) -> float: + """ Set the scale factor for incoming image + + Parameters + ---------- + image_size: tuple + The (height, width) of the original image + + Returns + ------- + float + The scaling factor from original image size to model input size + """ scale = self.input_size / max(image_size) - logger.trace("Detector scale: %s", scale) + logger.trace("Detector scale: %s", scale) # type:ignore[attr-defined] return scale - def _set_padding(self, image_size, scale): - """ Set the image padding for non-square images """ + def _set_padding(self, image_size: tuple[int, int], scale: float) -> tuple[int, int]: + """ Set the image padding for non-square images + + Parameters + ---------- + image_size: tuple + The (height, width) of the original image + scale: float + The scaling factor from original image size to model input size + + Returns + ------- + tuple + The amount of padding to apply to the x and y axes + """ pad_left = int(self.input_size - int(image_size[1] * scale)) // 2 pad_top = int(self.input_size - int(image_size[0] * scale)) // 2 return pad_left, pad_top @staticmethod - def _scale_image(image, image_size, scale): - """ Scale the image and optional pad to given size """ + def _scale_image(image: np.ndarray, image_size: tuple[int, int], scale: float) -> np.ndarray: + """ Scale the image and optional pad to given size + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to be scalued + image_size: tuple + The image (height, width) + scale: float + The scaling factor to apply to the image + + Returns + ------- + :class:`numpy.ndarray` + The scaled image + """ interpln = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA if scale != 1.0: dims = (int(image_size[1] * scale), int(image_size[0] * scale)) - logger.trace("Resizing detection image from %s to %s. Scale=%s", + logger.trace("Resizing detection image from %s to %s. " # type:ignore[attr-defined] + "Scale=%s", "x".join(str(i) for i in reversed(image_size)), "x".join(str(i) for i in dims), scale) image = cv2.resize(image, dims, interpolation=interpln) - logger.trace("Resized image shape: %s", image.shape) + logger.trace("Resized image shape: %s", image.shape) # type:ignore[attr-defined] return image - def _pad_image(self, image): - """ Pad a resized image to input size """ + def _pad_image(self, image: np.ndarray) -> np.ndarray: + """ Pad a resized image to input size + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to have padding applied + + Returns + ------- + :class:`numpy.ndarray` + The image with padding applied + """ height, width = image.shape[:2] if width < self.input_size or height < self.input_size: pad_l = (self.input_size - width) // 2 @@ -310,29 +464,54 @@ def _pad_image(self, image): pad_l, pad_r, cv2.BORDER_CONSTANT) - logger.trace("Padded image shape: %s", image.shape) + logger.trace("Padded image shape: %s", image.shape) # type:ignore[attr-defined] return image # <<< FINALIZE METHODS >>> # - def _remove_zero_sized_faces(self, batch_faces): - """ Remove items from batch_faces where detected face is of zero size - or face falls entirely outside of image """ - logger.trace("Input sizes: %s", [len(face) for face in batch_faces]) + def _remove_zero_sized_faces(self, batch_faces: list[list[DetectedFace]] + ) -> list[list[DetectedFace]]: + """ Remove items from batch_faces where detected face is of zero size or face falls + entirely outside of image + + Parameters + ---------- + batch_faces: list + List of detected face objects + + Returns + ------- + list + List of detected face objects with filtered out faces removed + """ + logger.trace("Input sizes: %s", [len(face) for face in batch_faces]) # type: ignore retval = [[face for face in faces - if face.right > 0 and face.left < self.input_size - and face.bottom > 0 and face.top < self.input_size] + if face.right > 0 and face.left is not None and face.left < self.input_size + and face.bottom > 0 and face.top is not None and face.top < self.input_size] for faces in batch_faces] - logger.trace("Output sizes: %s", [len(face) for face in retval]) + logger.trace("Output sizes: %s", [len(face) for face in retval]) # type: ignore return retval - def _filter_small_faces(self, detected_faces): - """ Filter out any faces smaller than the min size threshold """ + def _filter_small_faces(self, detected_faces: list[list[DetectedFace]] + ) -> list[list[DetectedFace]]: + """ Filter out any faces smaller than the min size threshold + + Parameters + ---------- + detected_faces: list + List of detected face objects + + Returns + ------- + list + List of detected face objects with filtered out faces removed + """ retval = [] for faces in detected_faces: this_image = [] for face in faces: - face_size = (face.w ** 2 + face.h ** 2) ** 0.5 + assert face.width is not None and face.height is not None + face_size = (face.width ** 2 + face.height ** 2) ** 0.5 if face_size < self.min_size: logger.debug("Removing detected face: (face_size: %s, min_size: %s", face_size, self.min_size) @@ -343,59 +522,74 @@ def _filter_small_faces(self, detected_faces): # <<< IMAGE ROTATION METHODS >>> # @staticmethod - def _get_rotation_angles(rotation): - """ Set the rotation angles. Includes backwards compatibility for the - 'on' and 'off' options: - - 'on' - increment 90 degrees - - 'off' - disable - - 0 is prepended to the list, as whatever happens, we want to - scan the image in it's upright state """ + def _get_rotation_angles(rotation: str | None) -> list[int]: + """ Set the rotation angles. + + Parameters + ---------- + str + List of requested rotation angles + + Returns + ------- + list + The complete list of rotation angles to apply + """ rotation_angles = [0] - if not rotation or rotation.lower() == "off": + if not rotation: logger.debug("Not setting rotation angles") return rotation_angles - if rotation.lower() == "on": - rotation_angles.extend(range(90, 360, 90)) - else: - passed_angles = [int(angle) - for angle in rotation.split(",") - if int(angle) != 0] - if len(passed_angles) == 1: - rotation_step_size = passed_angles[0] - rotation_angles.extend(range(rotation_step_size, - 360, - rotation_step_size)) - elif len(passed_angles) > 1: - rotation_angles.extend(passed_angles) + passed_angles = [int(angle) + for angle in rotation.split(",") + if int(angle) != 0] + if len(passed_angles) == 1: + rotation_step_size = passed_angles[0] + rotation_angles.extend(range(rotation_step_size, + 360, + rotation_step_size)) + elif len(passed_angles) > 1: + rotation_angles.extend(passed_angles) logger.debug("Rotation Angles: %s", rotation_angles) return rotation_angles - def _rotate_batch(self, batch, angle): + def _rotate_batch(self, batch: DetectorBatch, angle: int) -> None: """ Rotate images in a batch by given angle + if any faces have already been detected for a batch, store the existing rotation - matrix and replace the feed image with a placeholder """ + matrix and replace the feed image with a placeholder + + Parameters + ---------- + batch: :class:`DetectorBatch` + The batch to apply rotation to + angle: int + The amount of degrees to rotate the image by + """ if angle == 0: # Set the initial batch so we always rotate from zero - batch["initial_feed"] = batch["feed"].copy() + batch.initial_feed = batch.feed.copy() return - retval = dict() - for img, faces, rotmat in zip(batch["initial_feed"], batch["prediction"], batch["rotmat"]): + feeds: list[np.ndarray] = [] + rotmats: list[np.ndarray] = [] + for img, faces, rotmat in zip(batch.initial_feed, + batch.prediction, + batch.rotation_matrix): if faces.any(): image = np.zeros_like(img) matrix = rotmat else: image, matrix = self._rotate_image_by_angle(img, angle) - retval.setdefault("feed", []).append(image) - retval.setdefault("rotmat", []).append(matrix) - batch["feed"] = np.array(retval["feed"], dtype="float32") - batch["rotmat"] = retval["rotmat"] + feeds.append(image) + rotmats.append(matrix) + batch.feed = np.array(feeds, dtype="float32") + batch.rotation_matrix = rotmats @staticmethod - def _rotate_face(face, rotation_matrix): + def _rotate_face(face: DetectedFace, rotation_matrix: np.ndarray) -> DetectedFace: """ Rotates the detection bounding box around the given rotation matrix. Parameters @@ -411,7 +605,8 @@ def _rotate_face(face, rotation_matrix): :class:`DetectedFace` The same class with the detection bounding box points rotated by the given matrix. """ - logger.trace("Rotating face: (face: %s, rotation_matrix: %s)", face, rotation_matrix) + logger.trace("Rotating face: (face: %s, rotation_matrix: %s)", # type: ignore + face, rotation_matrix) bounding_box = [[face.left, face.top], [face.right, face.top], [face.right, face.bottom], @@ -424,24 +619,45 @@ def _rotate_face(face, rotation_matrix): rotated = transformed.squeeze() # Bounding box should follow x, y planes, so get min/max for non-90 degree rotations - pt_x = min([pnt[0] for pnt in rotated]) - pt_y = min([pnt[1] for pnt in rotated]) - pt_x1 = max([pnt[0] for pnt in rotated]) - pt_y1 = max([pnt[1] for pnt in rotated]) + pt_x = min(pnt[0] for pnt in rotated) + pt_y = min(pnt[1] for pnt in rotated) + pt_x1 = max(pnt[0] for pnt in rotated) + pt_y1 = max(pnt[1] for pnt in rotated) width = pt_x1 - pt_x height = pt_y1 - pt_y - face.x = int(pt_x) - face.y = int(pt_y) - face.w = int(width) - face.h = int(height) + face.left = int(pt_x) + face.top = int(pt_y) + face.width = int(width) + face.height = int(height) return face - def _rotate_image_by_angle(self, image, angle): + def _rotate_image_by_angle(self, + image: np.ndarray, + angle: int) -> tuple[np.ndarray, np.ndarray]: """ Rotate an image by a given angle. - From: https://stackoverflow.com/questions/22041699 """ - logger.trace("Rotating image: (image: %s, angle: %s)", image.shape, angle) + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to be rotated + angle: int + The angle, in degrees, to rotate the image by + + Returns + ------- + image: :class:`numpy.ndarray` + The rotated image + rotation_matrix: :class:`numpy.ndarray` + The rotation matrix used to rotate the image + + Reference + --------- + https://stackoverflow.com/questions/22041699 + """ + + logger.trace("Rotating image: (image: %s, angle: %s)", # type:ignore[attr-defined] + image.shape, angle) channels_first = image.shape[0] <= 4 if channels_first: image = np.moveaxis(image, 0, 2) @@ -451,7 +667,8 @@ def _rotate_image_by_angle(self, image, angle): rotation_matrix = cv2.getRotationMatrix2D(image_center, -1.*angle, 1.) rotation_matrix[0, 2] += self.input_size / 2 - image_center[0] rotation_matrix[1, 2] += self.input_size / 2 - image_center[1] - logger.trace("Rotated image: (rotation_matrix: %s", rotation_matrix) + logger.trace("Rotated image: (rotation_matrix: %s", # type:ignore[attr-defined] + rotation_matrix) image = cv2.warpAffine(image, rotation_matrix, (self.input_size, self.input_size)) if channels_first: image = np.moveaxis(image, 2, 0) diff --git a/plugins/extract/detect/cv2_dnn.py b/plugins/extract/detect/cv2_dnn.py index 08ce23aa89..9f98918e06 100644 --- a/plugins/extract/detect/cv2_dnn.py +++ b/plugins/extract/detect/cv2_dnn.py @@ -1,14 +1,18 @@ #!/usr/bin/env python3 """ OpenCV DNN Face detection plugin """ +import logging import numpy as np -from ._base import cv2, Detector, logger +from ._base import BatchType, cv2, Detector, DetectorBatch + + +logger = logging.getLogger(__name__) class Detect(Detector): """ CV2 DNN detector for face recognition """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 4 model_filename = ["resnet_ssd_v1.caffemodel", "resnet_ssd_v1.prototxt"] super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) @@ -19,44 +23,45 @@ def __init__(self, **kwargs): self.batchsize = 1 self.confidence = self.config["confidence"] / 100 - def init_model(self): + def init_model(self) -> None: """ Initialize CV2 DNN Detector Model""" - self.model = cv2.dnn.readNetFromCaffe(self.model_path[1], # pylint: disable=no-member + assert isinstance(self.model_path, list) + self.model = cv2.dnn.readNetFromCaffe(self.model_path[1], self.model_path[0]) - self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) # pylint: disable=no-member + self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detection image(s) for prediction """ - batch["feed"] = cv2.dnn.blobFromImages(batch["image"], # pylint: disable=no-member - scalefactor=1.0, - size=(self.input_size, self.input_size), - mean=[104, 117, 123], - swapRB=False, - crop=False) - return batch - - def predict(self, batch): + assert isinstance(batch, DetectorBatch) + batch.feed = cv2.dnn.blobFromImages(batch.image, + scalefactor=1.0, + size=(self.input_size, self.input_size), + mean=[104, 117, 123], + swapRB=False, + crop=False) + + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - self.model.setInput(batch["feed"]) + assert isinstance(self.model, cv2.dnn.Net) + self.model.setInput(feed) predictions = self.model.forward() - batch["prediction"] = self.finalize_predictions(predictions) - return batch + return self.finalize_predictions(predictions) - def finalize_predictions(self, predictions): + def finalize_predictions(self, predictions: np.ndarray) -> np.ndarray: """ Filter faces based on confidence level """ - faces = list() + faces = [] for i in range(predictions.shape[2]): confidence = predictions[0, 0, i, 2] if confidence >= self.confidence: - logger.trace("Accepting due to confidence %s >= %s", + logger.trace("Accepting due to confidence %s >= %s", # type:ignore[attr-defined] confidence, self.confidence) faces.append([(predictions[0, 0, i, 3] * self.input_size), (predictions[0, 0, i, 4] * self.input_size), (predictions[0, 0, i, 5] * self.input_size), (predictions[0, 0, i, 6] * self.input_size)]) - logger.trace("faces: %s", faces) - return [np.array(faces)] + logger.trace("faces: %s", faces) # type:ignore[attr-defined] + return np.array(faces)[None, ...] - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return diff --git a/plugins/extract/detect/cv2_dnn_defaults.py b/plugins/extract/detect/cv2_dnn_defaults.py index 762100a66b..e50c0ecc49 100755 --- a/plugins/extract/detect/cv2_dnn_defaults.py +++ b/plugins/extract/detect/cv2_dnn_defaults.py @@ -50,17 +50,17 @@ ) -_DEFAULTS = dict( - confidence=dict( - default=50, - info="The confidence level at which the detector has succesfully found a face.\nHigher " - "levels will be more discriminating, lower levels will have more false positives.", - datatype=int, - rounding=5, - min_max=(25, 100), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ), -) +_DEFAULTS = { + "confidence": { + "default": 50, + "info": "The confidence level at which the detector has succesfully found a face.\nHigher " + "levels will be more discriminating, lower levels will have more false positives.", + "datatype": int, + "rounding": 5, + "min_max": (25, 100), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, +} diff --git a/plugins/extract/detect/external.py b/plugins/extract/detect/external.py new file mode 100644 index 0000000000..98876e2a95 --- /dev/null +++ b/plugins/extract/detect/external.py @@ -0,0 +1,353 @@ +#!/usr/bin/env python3 +""" Import face detection ROI boxes from a json file """ +from __future__ import annotations + +import logging +import os +import re +import typing as T + +import numpy as np + +from lib.align import AlignedFace +from lib.utils import FaceswapError, IMAGE_EXTENSIONS + +from ._base import Detector + +if T.TYPE_CHECKING: + from lib.align import DetectedFace + from plugins.extract import ExtractMedia + from ._base import BatchType + +logger = logging.getLogger(__name__) + + +class Detect(Detector): + """ Import face detection bounding boxes from an external json file """ + def __init__(self, **kwargs) -> None: + kwargs["rotation"] = None # Disable rotation + kwargs["min_size"] = 0 # Disable min_size + super().__init__(git_model_id=None, model_filename=None, **kwargs) + + self.name = "External" + self.batchsize = 16 + + self._origin: T.Literal["top-left", + "bottom-left", + "top-right", + "bottom-right"] = self.config["origin"] + + self._re_frame_no: re.Pattern = re.compile(r"\d+$") + self._missing: list[str] = [] + self._log_once = True + self._is_video = False + self._imported: dict[str | int, np.ndarray] = {} + """dict[str | int, np.ndarray]: The imported data from external .json file""" + + def init_model(self) -> None: + """ No initialization to perform """ + logger.debug("No detector model to initialize") + + def _compile_detection_image(self, item: ExtractMedia + ) -> tuple[np.ndarray, float, tuple[int, int]]: + """ Override _compile_detection_image method, to obtain the source frame dimensions + + Parameters + ---------- + item: :class:`~plugins.extract.extract_media.ExtractMedia` + The input item from the pipeline + + Returns + ------- + image: :class:`numpy.ndarray` + dummy empty array + scale: float + The scaling factor for the image (1.0) + pad: int + The amount of padding applied to the image (0, 0) + """ + return np.array(item.image_shape[:2], dtype="int64"), 1.0, (0, 0) + + def _check_for_video(self, filename: str) -> None: + """ Check a sample filename from the import file for a file extension to set + :attr:`_is_video` + + Parameters + ---------- + filename: str + A sample file name from the imported data + """ + logger.debug("Checking for video from '%s'", filename) + ext = os.path.splitext(filename)[-1] + if ext.lower() not in IMAGE_EXTENSIONS: + self._is_video = True + logger.debug("Set is_video to %s from extension '%s'", self._is_video, ext) + + def _get_key(self, key: str) -> str | int: + """ Obtain the key for the item in the lookup table. If the input are images, the key will + be the image filename. If the input is a video, the key will be the frame number + + Parameters + ---------- + key: str + The initial key value from import data or an import image/frame + + Returns + ------- + str | int + The filename is the input data is images, otherwise the frame number of a video + """ + if not self._is_video: + return key + original_name = os.path.splitext(key)[0] + matches = self._re_frame_no.findall(original_name) + if not matches or len(matches) > 1: + raise FaceswapError(f"Invalid import name: '{key}'. For video files, the key should " + "end with the frame number.") + retval = int(matches[0]) + logger.trace("Obtained frame number %s from key '%s'", # type:ignore[attr-defined] + retval, key) + return retval + + @classmethod + def _bbox_from_detected(cls, bounding_box: list[int]) -> np.ndarray: + """ Import the detected face roi from a `detected` item in the import file + + Parameters + ---------- + bounding_box: list[int] + a bounding box contained within the import file + + Returns + ------- + :class:`numpy.ndarray` + The "left", "top", "right", "bottom" bounding box for the face + + Raises + ------ + FaceSwapError + If the number of bounding box co-ordinates is incorrect + """ + if len(bounding_box) != 4: + raise FaceswapError("Imported 'detected' bounding boxes should be a list of 4 numbers " + "representing the 'left', 'top', 'right', `bottom` of a face.") + return np.rint(bounding_box) + + def _validate_landmarks(self, landmarks: list[list[float]]) -> np.ndarray: + """ Validate that the there are 4 or 68 landmarks and are a complete list of (x, y) + co-ordinates + + Parameters + ---------- + landmarks: list[float] + The 4 point ROI or 68 point 2D landmarks that are being imported + + Returns + ------- + :class:`numpy.ndarray` + The original landmarks as a numpy array + + Raises + ------ + FaceSwapError + If the landmarks being imported are not correct + """ + if len(landmarks) not in (4, 68): + raise FaceswapError("Imported 'landmarks_2d' should be either 68 facial feature " + "landmarks or 4 ROI corner locations") + retval = np.array(landmarks, dtype="float32") + if retval.shape[-1] != 2: + raise FaceswapError("Imported 'landmarks_2d' should be formatted as a list of (x, y) " + "co-ordinates") + return retval + + def _bbox_from_landmarks2d(self, landmarks: list[list[float]]) -> np.ndarray: + """ Import the detected face roi by estimating from imported landmarks + + Parameters + ---------- + landmarks: list[float] + The 4 point ROI or 68 point 2D landmarks that are being imported + + Returns + ------- + :class:`numpy.ndarray` + The "left", "top", "right", "bottom" bounding box for the face + """ + n_landmarks = self._validate_landmarks(landmarks) + face = AlignedFace(n_landmarks, centering="legacy", coverage_ratio=0.75) + return np.concatenate([np.min(face.original_roi, axis=0), + np.max(face.original_roi, axis=0)]) + + def _import_frame_face(self, + face: dict[str, list[int] | list[list[float]]], + align_origin: T.Literal["top-left", + "bottom-left", + "top-right", + "bottom-right"] | None) -> np.ndarray: + """ Import a detected face ROI from the import file + + Parameters + ---------- + face: dict[str, list[int] | list[list[float]]] + The data that exists within the import file for the frame + align_origin: Literal["top-left", "bottom-left", "top-right", "bottom-right"] | None + The origin of the imported aligner data. Used if the detected ROI is being estimated + from imported aligner data + + Returns + ------- + :class:`numpy.ndarray` + The "left", "top", "right", "bottom" bounding box for the face + + Raises + ------ + FaceSwapError + If the required keys for the bounding boxes are not present for the face + """ + if "detected" in face: + return self._bbox_from_detected(T.cast(list[int], face["detected"])) + if "landmarks_2d" in face: + if self._log_once and align_origin is None: + logger.warning("You are importing Detection data, but have only provided " + "Alignment data. This is most likely incorrect and will lead " + "to poor results") + self._log_once = False + + if self._log_once and align_origin is not None and align_origin != self._origin: + logger.info("Updating Detect origin from Aligner config to '%s'", align_origin) + self._origin = align_origin + self._log_once = False + + return self._bbox_from_landmarks2d(T.cast(list[list[float]], face["landmarks_2d"])) + + raise FaceswapError("The provided import file is missing both of the required keys " + "'detected' and 'landmarks_2d") + + def import_data(self, + data: dict[str, list[dict[str, list[int] | list[list[float]]]]], + align_origin: T.Literal["top-left", + "bottom-left", + "top-right", + "bottom-right"] | None) -> None: + """ Import the detection data from the json import file and set to :attr:`_imported` + + Parameters + ---------- + data: dict[str, list[dict[str, list[int] | list[list[float]]]]] + The data to be imported + align_origin: Literal["top-left", "bottom-left", "top-right", "bottom-right"] | None + The origin of the imported aligner data. Used if the detected ROI is being estimated + from imported aligner data + """ + logger.debug("Data length: %s, align_origin: %s", len(data), align_origin) + self._check_for_video(list(data)[0]) + for key, faces in data.items(): + try: + store_key = self._get_key(key) + self._imported[store_key] = np.array([self._import_frame_face(face, align_origin) + for face in faces], dtype="int32") + except FaceswapError as err: + logger.error(str(err)) + msg = f"The imported frame key that failed was '{key}'" + raise FaceswapError(msg) from err + + def process_input(self, batch: BatchType) -> None: + """ Put the lookup key into `batch.feed` so they can be collected for mapping in `.predict` + + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.DetectorBatch` + The batch to be processed by the plugin + """ + batch.feed = np.array([(self._get_key(os.path.basename(f)), i) + for f, i in zip(batch.filename, batch.image)], dtype="object") + + def _adjust_for_origin(self, box: np.ndarray, frame_dims: tuple[int, int]) -> np.ndarray: + """ Adjust the bounding box to be top-left orientated based on the selected import origin + + Parameters + ---------- + box: :class:`np.ndarray` + The imported bounding box at original (0, 0) origin + frame_dims: tuple[int, int] + The (rows, columns) dimensions of the original frame + + Returns + ------- + :class:`numpy.ndarray` + The adjusted bounding box for a top-left origin + """ + if not np.any(box) or self._origin == "top-left": + return box + if self._origin.startswith("bottom"): + box[:, [1, 3]] = frame_dims[0] - box[:, [1, 3]] + if self._origin.endswith("right"): + box[:, [0, 2]] = frame_dims[1] - box[:, [0, 2]] + + return box + + def predict(self, feed: np.ndarray) -> list[np.ndarray]: # type:ignore[override] + """ Pair the input filenames to the import file + + Parameters + ---------- + feed: :class:`numpy.ndarray` + The filenames with original frame dimensions to obtain the imported bounding boxes for + + Returns + ------- + list[]:class:`numpy.ndarray`] + The bounding boxes for the given filenames + """ + self._missing.extend(f[0] for f in feed if f[0] not in self._imported) + return [self._adjust_for_origin(self._imported.pop(f[0], np.array([], dtype="int32")), + f[1]) + for f in feed] + + def process_output(self, batch: BatchType) -> None: + """ No output processing required for import plugin + + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.DetectorBatch` + The batch to be processed by the plugin + """ + logger.trace("No output processing for import plugin") # type:ignore[attr-defined] + + def _remove_zero_sized_faces(self, batch_faces: list[list[DetectedFace]] + ) -> list[list[DetectedFace]]: + """ Override _remove_zero_sized_faces to just return the faces that have been imported + + Parameters + ---------- + batch_faces: list[list[DetectedFace] + List of detected face objects + + Returns + ------- + list[list[DetectedFace] + Original list of detected face objects + """ + return batch_faces + + def on_completion(self) -> None: + """ Output information if: + - Imported items were not matched in input data + - Input data was not matched in imported items + """ + super().on_completion() + + if self._missing: + logger.warning("[DETECT] %s input frames could not be matched in the import file " + "'%s'. Run in verbose mode for a list of frames.", + len(self._missing), self.config["file_name"]) + logger.verbose( # type:ignore[attr-defined] + "[DETECT] Input frames not in import file: %s", self._missing) + + if self._imported: + logger.warning("[DETECT] %s items in the import file '%s' could not be matched to any " + "input frames. Run in verbose mode for a list of items.", + len(self._imported), self.config["file_name"]) + logger.verbose( # type:ignore[attr-defined] + "[DETECT] import file items not in input frames: %s", list(self._imported)) diff --git a/plugins/extract/detect/external_defaults.py b/plugins/extract/detect/external_defaults.py new file mode 100644 index 0000000000..c444bf419b --- /dev/null +++ b/plugins/extract/detect/external_defaults.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +""" + The default options for the faceswap Import Alignments plugin. + + Defaults files should be named _defaults.py + Any items placed into this file will automatically get added to the relevant config .ini files + within the faceswap/config folder. + + The following variables should be defined: + _HELPTEXT: A string describing what this plugin does + _DEFAULTS: A dictionary containing the options, defaults and meta information. The + dictionary should be defined as: + {: {}} + + should always be lower text. + dictionary requirements are listed below. + + The following keys are expected for the _DEFAULTS dict: + datatype: [required] A python type class. This limits the type of data that can be + provided in the .ini file and ensures that the value is returned in the + correct type to faceswap. Valid data types are: , , + , . + default: [required] The default value for this option. + info: [required] A string describing what this option does. + group: [optional]. A group for grouping options together in the GUI. If not + provided this will not group this option with any others. + choices: [optional] If this option's datatype is of then valid + selections can be defined here. This validates the option and also enables + a combobox / radio option in the GUI. + gui_radio: [optional] If are defined, this indicates that the GUI should use + radio buttons rather than a combobox to display this option. + min_max: [partial] For and data types this is required + otherwise it is ignored. Should be a tuple of min and max accepted values. + This is used for controlling the GUI slider range. Values are not enforced. + rounding: [partial] For and data types this is + required otherwise it is ignored. Used for the GUI slider. For floats, this + is the number of decimal places to display. For ints this is the step size. + fixed: [optional] [train only]. Training configurations are fixed when the model is + created, and then reloaded from the state file. Marking an item as fixed=False + indicates that this value can be changed for existing models, and will override + the value saved in the state file with the updated value in config. If not + provided this will default to True. +""" + + +_HELPTEXT = ( + "Import Detector options.\n" + "Imports a detected face bounding box from an external .json file.\n" + ) + + +_DEFAULTS = { + "file_name": { + "default": "import.json", + "info": "The import file should be stored in the same folder as the video (if extracting " + "from a video file) or inside the folder of images (if importing from a folder of images)", + "datatype": str, + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "origin": { + "default": "top-left", + "info": "The origin (0, 0) location of the co-ordinates system used. " + "\n\t top-left: The origin (0, 0) of the canvas is at the top left " + "corner." + "\n\t bottom-left: The origin (0, 0) of the canvas is at the bottom " + "left corner." + "\n\t top-right: The origin (0, 0) of the canvas is at the top right " + "corner." + "\n\t bottom-right: The origin (0, 0) of the canvas is at the bottom " + "right corner.", + "datatype": str, + "choices": ["top-left", "bottom-left", "top-right", "bottom-right"], + "group": "output", + "gui_radio": True + } +} diff --git a/plugins/extract/detect/mtcnn.py b/plugins/extract/detect/mtcnn.py index ab82c443d5..78859ca249 100644 --- a/plugins/extract/detect/mtcnn.py +++ b/plugins/extract/detect/mtcnn.py @@ -1,38 +1,40 @@ #!/usr/bin/env python3 """ MTCNN Face detection plugin """ - -from __future__ import absolute_import, division, print_function +from __future__ import annotations +import logging +import typing as T import cv2 import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU # noqa:E501 # pylint:disable=import-error + from lib.model.session import KSession -from lib.utils import get_backend -from ._base import Detector, logger +from ._base import BatchType, Detector + +if T.TYPE_CHECKING: + from tensorflow import Tensor -if get_backend() == "amd": - from keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU # noqa pylint:disable=no-name-in-module,import-error +logger = logging.getLogger(__name__) class Detect(Detector): - """ MTCNN detector for face recognition """ - def __init__(self, **kwargs): + """ MTCNN detector for face recognition. """ + def __init__(self, **kwargs) -> None: git_model_id = 2 model_filename = ["mtcnn_det_v2.1.h5", "mtcnn_det_v2.2.h5", "mtcnn_det_v2.3.h5"] super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) self.name = "MTCNN" self.input_size = 640 - self.vram = 320 - self.vram_warnings = 64 # Will run at this with warnings - self.vram_per_batch = 32 + self.vram = 320 if not self.config["cpu"] else 0 + self.vram_warnings = 64 if not self.config["cpu"] else 0 # Will run at this with warnings + self.vram_per_batch = 32 if not self.config["cpu"] else 0 self.batchsize = self.config["batch-size"] - self.kwargs = self.validate_kwargs() + self.kwargs = self._validate_kwargs() self.color_format = "RGB" - def validate_kwargs(self): + def _validate_kwargs(self) -> dict[str, int | float | list[float]]: """ Validate that config options are correct. If not reset to default """ valid = True threshold = [self.config["threshold_1"], @@ -40,7 +42,8 @@ def validate_kwargs(self): self.config["threshold_3"]] kwargs = {"minsize": self.config["minsize"], "threshold": threshold, - "factor": self.config["scalefactor"]} + "factor": self.config["scalefactor"], + "input_size": self.input_size} if kwargs["minsize"] < 10: valid = False @@ -50,36 +53,59 @@ def validate_kwargs(self): valid = False if not valid: - kwargs = {"minsize": 20, # minimum size of face - "threshold": [0.6, 0.7, 0.7], # three steps threshold - "factor": 0.709} # scale factor + kwargs = {} logger.warning("Invalid MTCNN options in config. Running with defaults") + logger.debug("Using mtcnn kwargs: %s", kwargs) return kwargs - def init_model(self): - """ Initialize S3FD Model""" + def init_model(self) -> None: + """ Initialize MTCNN Model. """ + assert isinstance(self.model_path, list) self.model = MTCNN(self.model_path, self.config["allow_growth"], self._exclude_gpus, - **self.kwargs) + self.config["cpu"], + **self.kwargs) # type:ignore + + def process_input(self, batch: BatchType) -> None: + """ Compile the detection image(s) for prediction - def process_input(self, batch): - """ Compile the detection image(s) for prediction """ - batch["feed"] = (batch["image"] - 127.5) / 127.5 - return batch + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.DetectorBatch` + Contains the batch that is currently being passed through the plugin process + """ + batch.feed = (np.array(batch.image, dtype="float32") - 127.5) / 127.5 - def predict(self, batch): - """ Run model to get predictions """ - prediction, points = self.model.detect_faces(batch["feed"]) - logger.trace("filename: %s, prediction: %s, mtcnn_points: %s", - batch["filename"], prediction, points) - batch["prediction"], batch["mtcnn_points"] = prediction, points - return batch + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Run model to get predictions - def process_output(self, batch): - """ Post process the detected faces """ - return batch + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.DetectorBatch` + Contains the batch to pass through the MTCNN model + + Returns + ------- + dict + The batch with the predictions added to the dictionary + """ + assert isinstance(self.model, MTCNN) + prediction, points = self.model.detect_faces(feed) + logger.trace("prediction: %s, mtcnn_points: %s", # type:ignore + prediction, points) + return prediction + + def process_output(self, batch: BatchType) -> None: + """ MTCNN performs no post processing so the original batch is returned + + Parameters + ---------- + batch: :class:`~plugins.extract.detect._base.DetectorBatch` + Contains the batch to apply postprocessing to + """ + return # MTCNN Detector @@ -113,18 +139,57 @@ def process_output(self, batch): class PNet(KSession): - """ Keras P-Net model for MTCNN """ - def __init__(self, model_path, allow_growth, exclude_gpus): + """ Keras P-Net model for MTCNN + + Parameters + ---------- + model_path: str + The path to the keras model file + allow_growth: bool, optional + Enable the Tensorflow GPU allow_growth configuration option. This option prevents + Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation and + slower performance. Default: ``False`` + exclude_gpus: list, optional + A list of indices correlating to connected GPUs that Tensorflow should not use. Pass + ``None`` to not exclude any GPUs. Default: ``None`` + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` + input_size: int + The input size of the model + minsize: int, optional + The minimum size of a face to accept as a detection. Default: `20` + threshold: list, optional + Threshold for P-Net + """ + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None, + cpu_mode: bool, + input_size: int, + min_size: int, + factor: float, + threshold: float) -> None: super().__init__("MTCNN-PNet", model_path, allow_growth=allow_growth, - exclude_gpus=exclude_gpus) + exclude_gpus=exclude_gpus, + cpu_mode=cpu_mode) + self.define_model(self.model_definition) self.load_model_weights() + self._input_size = input_size + self._threshold = threshold + + self._pnet_scales = self._calculate_scales(min_size, factor) + self._pnet_sizes = [(int(input_size * scale), int(input_size * scale)) + for scale in self._pnet_scales] + self._pnet_input: list[np.ndarray] | None = None + @staticmethod - def model_definition(): - """ Keras P-Network for MTCNN """ + def model_definition() -> tuple[list[Tensor], list[Tensor]]: + """ Keras P-Network Definition for MTCNN """ input_ = Input(shape=(None, None, 3)) var_x = Conv2D(10, (3, 3), strides=1, padding='valid', name='conv1')(input_) var_x = PReLU(shared_axes=[1, 2], name='PReLU1')(var_x) @@ -137,20 +202,166 @@ def model_definition(): bbox_regress = Conv2D(4, (1, 1), name='conv4-2')(var_x) return [input_], [classifier, bbox_regress] + def _calculate_scales(self, + minsize: int, + factor: float) -> list[float]: + """ Calculate multi-scale + + Parameters + ---------- + minsize: int + Minimum size for a face to be accepted + factor: float + Scaling factor + + Returns + ------- + list + List of scale floats + """ + factor_count = 0 + var_m = 12.0 / minsize + minl = self._input_size * var_m + # create scale pyramid + scales = [] + while minl >= 12: + scales += [var_m * np.power(factor, factor_count)] + minl = minl * factor + factor_count += 1 + logger.trace(scales) # type:ignore + return scales + + def __call__(self, images: np.ndarray) -> list[np.ndarray]: + """ first stage - fast proposal network (p-net) to obtain face candidates + + Parameters + ---------- + images: :class:`numpy.ndarray` + The batch of images to detect faces in + + Returns + ------- + List + List of face candidates from P-Net + """ + batch_size = images.shape[0] + rectangles: list[list[list[int | float]]] = [[] for _ in range(batch_size)] + scores: list[list[np.ndarray]] = [[] for _ in range(batch_size)] + + if self._pnet_input is None: + self._pnet_input = [np.empty((batch_size, rheight, rwidth, 3), dtype="float32") + for rheight, rwidth in self._pnet_sizes] + + for scale, batch, (rheight, rwidth) in zip(self._pnet_scales, + self._pnet_input, + self._pnet_sizes): + _ = [cv2.resize(images[idx], (rwidth, rheight), dst=batch[idx]) + for idx in range(batch_size)] + cls_prob, roi = self.predict(batch) + cls_prob = cls_prob[..., 1] + out_side = max(cls_prob.shape[1:3]) + cls_prob = np.swapaxes(cls_prob, 1, 2) + roi = np.swapaxes(roi, 1, 3) + for idx in range(batch_size): + # first index 0 = class score, 1 = one hot representation + rect, score = self._detect_face_12net(cls_prob[idx, ...], + roi[idx, ...], + out_side, + 1 / scale) + rectangles[idx].extend(rect) + scores[idx].extend(score) + + return [nms(np.array(rect), np.array(score), 0.7, "iou")[0] # don't output scores + for rect, score in zip(rectangles, scores)] + + def _detect_face_12net(self, + class_probabilities: np.ndarray, + roi: np.ndarray, + size: int, + scale: float) -> tuple[np.ndarray, np.ndarray]: + """ Detect face position and calibrate bounding box on 12net feature map(matrix version) + + Parameters + ---------- + class_probabilities: :class:`numpy.ndarray` + softmax feature map for face classify + roi: :class:`numpy.ndarray` + feature map for regression + size: int + feature map's largest size + scale: float + current input image scale in multi-scales + + Returns + ------- + list + Calibrated face candidates + """ + in_side = 2 * size + 11 + stride = 0. if size == 1 else float(in_side - 12) / (size - 1) + (var_x, var_y) = np.nonzero(class_probabilities >= self._threshold) + boundingbox = np.array([var_x, var_y]).T + + boundingbox = np.concatenate((np.fix((stride * (boundingbox) + 0) * scale), + np.fix((stride * (boundingbox) + 11) * scale)), axis=1) + offset = roi[:4, var_x, var_y].T + boundingbox = boundingbox + offset * 12.0 * scale + rectangles = np.concatenate((boundingbox, + np.array([class_probabilities[var_x, var_y]]).T), axis=1) + rectangles = rect2square(rectangles) + + np.clip(rectangles[..., :4], 0., self._input_size, out=rectangles[..., :4]) + pick = np.where(np.logical_and(rectangles[..., 2] > rectangles[..., 0], + rectangles[..., 3] > rectangles[..., 1]))[0] + rects = rectangles[pick, :4].astype("int") + scores = rectangles[pick, 4] + + return nms(rects, scores, 0.3, "iou") + class RNet(KSession): - """ Keras R-Net model for MTCNN """ - def __init__(self, model_path, allow_growth, exclude_gpus): + """ Keras R-Net model Definition for MTCNN + + Parameters + ---------- + model_path: str + The path to the keras model file + allow_growth: bool, optional + Enable the Tensorflow GPU allow_growth configuration option. This option prevents + Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation and + slower performance. Default: ``False`` + exclude_gpus: list, optional + A list of indices correlating to connected GPUs that Tensorflow should not use. Pass + ``None`` to not exclude any GPUs. Default: ``None`` + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` + input_size: int + The input size of the model + threshold: list, optional + Threshold for R-Net + + """ + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None, + cpu_mode: bool, + input_size: int, + threshold: float) -> None: super().__init__("MTCNN-RNet", model_path, allow_growth=allow_growth, - exclude_gpus=exclude_gpus) + exclude_gpus=exclude_gpus, + cpu_mode=cpu_mode) self.define_model(self.model_definition) self.load_model_weights() + self._input_size = input_size + self._threshold = threshold + @staticmethod - def model_definition(): - """ Keras R-Network for MTCNN """ + def model_definition() -> tuple[list[Tensor], list[Tensor]]: + """ Keras R-Network Definition for MTCNN """ input_ = Input(shape=(24, 24, 3)) var_x = Conv2D(28, (3, 3), strides=1, padding='valid', name='conv1')(input_) var_x = PReLU(shared_axes=[1, 2], name='prelu1')(var_x) @@ -170,19 +381,116 @@ def model_definition(): bbox_regress = Dense(4, name='conv5-2')(var_x) return [input_], [classifier, bbox_regress] + def __call__(self, + images: np.ndarray, + rectangle_batch: list[np.ndarray], + ) -> list[np.ndarray]: + """ second stage - refinement of face candidates with r-net + + Parameters + ---------- + images: :class:`numpy.ndarray` + The batch of images to detect faces in + rectangle_batch: + List of :class:`numpy.ndarray` face candidates from P-Net + + Returns + ------- + List + List of :class:`numpy.ndarray` refined face candidates from R-Net + """ + ret: list[np.ndarray] = [] + for idx, (rectangles, image) in enumerate(zip(rectangle_batch, images)): + if not np.any(rectangles): + ret.append(np.array([])) + continue + + feed_batch = np.empty((rectangles.shape[0], 24, 24, 3), dtype="float32") + + _ = [cv2.resize(image[rect[1]: rect[3], rect[0]: rect[2]], + (24, 24), + dst=feed_batch[idx]) + for idx, rect in enumerate(rectangles)] + + cls_prob, roi_prob = self.predict(feed_batch) + ret.append(self._filter_face_24net(cls_prob, roi_prob, rectangles)) + return ret + + def _filter_face_24net(self, + class_probabilities: np.ndarray, + roi: np.ndarray, + rectangles: np.ndarray, + ) -> np.ndarray: + """ Filter face position and calibrate bounding box on 12net's output + + Parameters + ---------- + class_probabilities: class:`np.ndarray` + Softmax feature map for face classify + roi: :class:`numpy.ndarray` + Feature map for regression + rectangles: list + 12net's predict + + Returns + ------- + list + rectangles in the format [[x, y, x1, y1, score]] + """ + prob = class_probabilities[:, 1] + pick = np.nonzero(prob >= self._threshold) + + bbox = rectangles.T[:4, pick] + scores = np.array([prob[pick]]).T.ravel() + deltas = roi.T[:4, pick] + + dims = np.tile([bbox[2] - bbox[0], bbox[3] - bbox[1]], (2, 1, 1)) + bbox = np.transpose(bbox + deltas * dims).reshape(-1, 4) + bbox = np.clip(rect2square(bbox), 0, self._input_size).astype("int") + return nms(bbox, scores, 0.3, "iou")[0] + class ONet(KSession): - """ Keras O-Net model for MTCNN """ - def __init__(self, model_path, allow_growth, exclude_gpus): + """ Keras O-Net model for MTCNN + + Parameters + ---------- + model_path: str + The path to the keras model file + allow_growth: bool, optional + Enable the Tensorflow GPU allow_growth configuration option. This option prevents + Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation and + slower performance. Default: ``False`` + exclude_gpus: list, optional + A list of indices correlating to connected GPUs that Tensorflow should not use. Pass + ``None`` to not exclude any GPUs. Default: ``None`` + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` + input_size: int + The input size of the model + threshold: list, optional + Threshold for O-Net + """ + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None, + cpu_mode: bool, + input_size: int, + threshold: float) -> None: super().__init__("MTCNN-ONet", model_path, allow_growth=allow_growth, - exclude_gpus=exclude_gpus) + exclude_gpus=exclude_gpus, + cpu_mode=cpu_mode) self.define_model(self.model_definition) self.load_model_weights() + self._input_size = input_size + self._threshold = threshold + @staticmethod - def model_definition(): + def model_definition() -> tuple[list[Tensor], list[Tensor]]: """ Keras O-Network for MTCNN """ input_ = Input(shape=(48, 48, 3)) var_x = Conv2D(32, (3, 3), strides=1, padding='valid', name='conv1')(input_) @@ -206,357 +514,236 @@ def model_definition(): landmark_regress = Dense(10, name='conv6-3')(var_x) return [input_], [classifier, bbox_regress, landmark_regress] + def __call__(self, + images: np.ndarray, + rectangle_batch: list[np.ndarray] + ) -> list[tuple[np.ndarray, np.ndarray]]: + """ Third stage - further refinement and facial landmarks positions with o-net + + Parameters + ---------- + images: :class:`numpy.ndarray` + The batch of images to detect faces in + rectangle_batch: + List of :class:`numpy.ndarray` face candidates from R-Net + + Returns + ------- + List + List of refined final candidates, scores and landmark points from O-Net + """ + ret: list[tuple[np.ndarray, np.ndarray]] = [] + for idx, rectangles in enumerate(rectangle_batch): + if not np.any(rectangles): + ret.append((np.empty((0, 5)), np.empty(0))) + continue + image = images[idx] + feed_batch = np.empty((rectangles.shape[0], 48, 48, 3), dtype="float32") -class MTCNN(): - """ MTCNN Detector for face alignment """ - # TODO Batching for r-net and o-net + _ = [cv2.resize(image[rect[1]: rect[3], rect[0]: rect[2]], + (48, 48), + dst=feed_batch[idx]) + for idx, rect in enumerate(rectangles)] - def __init__(self, model_path, allow_growth, exclude_gpus, minsize, threshold, factor): - """ - minsize: minimum faces' size - threshold: threshold=[th1, th2, th3], th1-3 are three steps threshold - factor: the factor used to create a scaling pyramid of face sizes to - detect in the image. - p-net, r-net, o-net: caffemodel + cls_probs, roi_probs, pts_probs = self.predict(feed_batch) + ret.append(self._filter_face_48net(cls_probs, roi_probs, pts_probs, rectangles)) + return ret + + def _filter_face_48net(self, class_probabilities: np.ndarray, + roi: np.ndarray, + points: np.ndarray, + rectangles: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ Filter face position and calibrate bounding box on 12net's output + + Parameters + ---------- + class_probabilities: :class:`numpy.ndarray` : class_probabilities[1] is face possibility + Array of face probabilities + roi: :class:`numpy.ndarray` + offset + points: :class:`numpy.ndarray` + 5 point face landmark + rectangles: :class:`numpy.ndarray` + 12net's predict, rectangles[i][0:3] is the position, rectangles[i][4] is score + + Returns + ------- + boxes: :class:`numpy.ndarray` + The [l, t, r, b, score] bounding boxes + points: :class:`numpy.ndarray` + The 5 point landmarks """ + prob = class_probabilities[:, 1] + pick = np.nonzero(prob >= self._threshold)[0] + scores = np.array([prob[pick]]).T.ravel() + + bbox = rectangles[pick] + dims = np.array([bbox[..., 2] - bbox[..., 0], bbox[..., 3] - bbox[..., 1]]).T + + pts = np.vstack( + np.hsplit(points[pick], 2)).reshape(2, -1, 5).transpose(1, 2, 0).reshape(-1, 10) + pts = np.tile(dims, (1, 5)) * pts + np.tile(bbox[..., :2], (1, 5)) + + bbox = np.clip(np.floor(bbox + roi[pick] * np.tile(dims, (1, 2))), + 0., + self._input_size) + + indices = np.where( + np.logical_and(bbox[..., 2] > bbox[..., 0], bbox[..., 3] > bbox[..., 1]))[0] + picks = np.concatenate([bbox[indices], pts[indices]], axis=-1) + + results, scores = nms(picks, scores, 0.3, "iom") + return np.concatenate([results[..., :4], scores[..., None]], axis=-1), results[..., 4:].T + + +class MTCNN(): # pylint:disable=too-few-public-methods + """ MTCNN Detector for face alignment + + Parameters + ---------- + model_path: list + List of paths to the 3 MTCNN subnet weights + allow_growth: bool, optional + Enable the Tensorflow GPU allow_growth configuration option. This option prevents + Tensorflow from allocating all of the GPU VRAM, but can lead to higher fragmentation and + slower performance. Default: ``False`` + exclude_gpus: list, optional + A list of indices correlating to connected GPUs that Tensorflow should not use. Pass + ``None`` to not exclude any GPUs. Default: ``None`` + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` + input_size: int, optional + The height, width input size to the model. Default: 640 + minsize: int, optional + The minimum size of a face to accept as a detection. Default: `20` + threshold: list, optional + List of floats for the three steps, Default: `[0.6, 0.7, 0.7]` + factor: float, optional + The factor used to create a scaling pyramid of face sizes to detect in the image. + Default: `0.709` + """ + def __init__(self, + model_path: list[str], + allow_growth: bool, + exclude_gpus: list[int] | None, + cpu_mode: bool, + input_size: int = 640, + minsize: int = 20, + threshold: list[float] | None = None, + factor: float = 0.709) -> None: logger.debug("Initializing: %s: (model_path: '%s', allow_growth: %s, exclude_gpus: %s, " - "minsize: %s, threshold: %s, factor: %s)", self.__class__.__name__, - model_path, allow_growth, exclude_gpus, minsize, threshold, factor) - self.minsize = minsize - self.threshold = threshold - self.factor = factor - - self.pnet = PNet(model_path[0], allow_growth, exclude_gpus) - self.rnet = RNet(model_path[1], allow_growth, exclude_gpus) - self.onet = ONet(model_path[2], allow_growth, exclude_gpus) - self._pnet_scales = None + "input_size: %s, minsize: %s, threshold: %s, factor: %s)", + self.__class__.__name__, model_path, allow_growth, exclude_gpus, + input_size, minsize, threshold, factor) + + threshold = [0.6, 0.7, 0.7] if threshold is None else threshold + self._pnet = PNet(model_path[0], + allow_growth, + exclude_gpus, + cpu_mode, + input_size, + minsize, + factor, + threshold[0]) + self._rnet = RNet(model_path[1], + allow_growth, + exclude_gpus, + cpu_mode, + input_size, + threshold[1]) + self._onet = ONet(model_path[2], + allow_growth, + exclude_gpus, + cpu_mode, + input_size, + threshold[2]) + logger.debug("Initialized: %s", self.__class__.__name__) - def detect_faces(self, batch): + def detect_faces(self, batch: np.ndarray) -> tuple[np.ndarray, tuple[np.ndarray]]: """Detects faces in an image, and returns bounding boxes and points for them. - batch: input batch + + Parameters + ---------- + batch: :class:`numpy.ndarray` + The input batch of images to detect face in + + Returns + ------- + List + list of numpy arrays containing the bounding box and 5 point landmarks + of detected faces """ - origin_h, origin_w = batch.shape[1:3] - rectangles = self.detect_pnet(batch, origin_h, origin_w) - rectangles = self.detect_rnet(batch, rectangles, origin_h, origin_w) - rectangles = self.detect_onet(batch, rectangles, origin_h, origin_w) - ret_boxes = [] - ret_points = [] - for rects in rectangles: - if rects: - total_boxes = np.array([result[:5] for result in rects]) - points = np.array([result[5:] for result in rects]).T - else: - total_boxes = np.empty((0, 9)) - points = np.empty(0) - ret_boxes.append(total_boxes) - ret_points.append(points) - return ret_boxes, ret_points - - def detect_pnet(self, images, height, width): - # pylint: disable=too-many-locals - """ first stage - fast proposal network (p-net) to obtain face candidates """ - if self._pnet_scales is None: - self._pnet_scales = calculate_scales(height, width, self.minsize, self.factor) - rectangles = [[] for _ in range(images.shape[0])] - batch_items = images.shape[0] - for scale in self._pnet_scales: - rwidth, rheight = int(width * scale), int(height * scale) - batch = np.empty((batch_items, rheight, rwidth, 3), dtype="float32") - for idx in range(batch_items): - batch[idx, ...] = cv2.resize(images[idx, ...], (rwidth, rheight)) - output = self.pnet.predict(batch) - cls_prob = output[0][..., 1] - roi = output[1] - out_h, out_w = cls_prob.shape[1:3] - out_side = max(out_h, out_w) - cls_prob = np.swapaxes(cls_prob, 1, 2) - roi = np.swapaxes(roi, 1, 3) - for idx in range(batch_items): - # first index 0 = class score, 1 = one hot representation - rectangle = detect_face_12net(cls_prob[idx, ...], - roi[idx, ...], - out_side, - 1 / scale, - width, - height, - self.threshold[0]) - rectangles[idx].extend(rectangle) - return [nms(x, 0.7, 'iou') for x in rectangles] - - def detect_rnet(self, images, rectangle_batch, height, width): - """ second stage - refinement of face candidates with r-net """ - ret = [] - # TODO: batching - for idx, rectangles in enumerate(rectangle_batch): - if not rectangles: - ret.append([]) - continue - image = images[idx] - crop_number = 0 - predict_24_batch = [] - for rect in rectangles: - crop_img = image[int(rect[1]):int(rect[3]), int(rect[0]):int(rect[2])] - scale_img = cv2.resize(crop_img, (24, 24)) - predict_24_batch.append(scale_img) - crop_number += 1 - predict_24_batch = np.array(predict_24_batch) - output = self.rnet.predict(predict_24_batch, batch_size=128) - cls_prob = output[0] - cls_prob = np.array(cls_prob) - roi_prob = output[1] - roi_prob = np.array(roi_prob) - ret.append(filter_face_24net( - cls_prob, roi_prob, rectangles, width, height, self.threshold[1] - )) - return ret + rectangles = self._pnet(batch) + rectangles = self._rnet(batch, rectangles) - def detect_onet(self, images, rectangle_batch, height, width): - """ third stage - further refinement and facial landmarks positions with o-net """ - ret = [] - # TODO: batching - for idx, rectangles in enumerate(rectangle_batch): - if not rectangles: - ret.append([]) - continue - image = images[idx] - crop_number = 0 - predict_batch = [] - for rect in rectangles: - crop_img = image[int(rect[1]):int(rect[3]), int(rect[0]):int(rect[2])] - scale_img = cv2.resize(crop_img, (48, 48)) - predict_batch.append(scale_img) - crop_number += 1 - predict_batch = np.array(predict_batch) - output = self.onet.predict(predict_batch, batch_size=128) - cls_prob = output[0] - roi_prob = output[1] - pts_prob = output[2] # index - ret.append(filter_face_48net( - cls_prob, - roi_prob, - pts_prob, - rectangles, - width, - height, - self.threshold[2] - )) - return ret + ret_boxes, ret_points = zip(*self._onet(batch, rectangles)) + return np.array(ret_boxes, dtype="object"), ret_points -def detect_face_12net(cls_prob, roi, out_side, scale, width, height, threshold): - # pylint: disable=too-many-locals, too-many-arguments - """ Detect face position and calibrate bounding box on 12net feature map(matrix version) - Input: - cls_prob : softmax feature map for face classify - roi : feature map for regression - out_side : feature map's largest size - scale : current input image scale in multi-scales - width : image's origin width - height : image's origin height - threshold: 0.6 can have 99% recall rate - """ - in_side = 2*out_side+11 - stride = 0 - if out_side != 1: - stride = float(in_side-12)/(out_side-1) - (var_x, var_y) = np.where(cls_prob >= threshold) - boundingbox = np.array([var_x, var_y]).T - bb1 = np.fix((stride * (boundingbox) + 0) * scale) - bb2 = np.fix((stride * (boundingbox) + 11) * scale) - boundingbox = np.concatenate((bb1, bb2), axis=1) - dx_1 = roi[0][var_x, var_y] - dx_2 = roi[1][var_x, var_y] - dx3 = roi[2][var_x, var_y] - dx4 = roi[3][var_x, var_y] - score = np.array([cls_prob[var_x, var_y]]).T - offset = np.array([dx_1, dx_2, dx3, dx4]).T - boundingbox = boundingbox + offset*12.0*scale - rectangles = np.concatenate((boundingbox, score), axis=1) - rectangles = rect2square(rectangles) - pick = [] - for rect in rectangles: - x_1 = int(max(0, rect[0])) - y_1 = int(max(0, rect[1])) - x_2 = int(min(width, rect[2])) - y_2 = int(min(height, rect[3])) - sc_ = rect[4] - if x_2 > x_1 and y_2 > y_1: - pick.append([x_1, y_1, x_2, y_2, sc_]) - return nms(pick, 0.3, "iou") - - -def filter_face_24net(cls_prob, roi, rectangles, width, height, threshold): - # pylint: disable=too-many-locals, too-many-arguments - """ Filter face position and calibrate bounding box on 12net's output - Input: - cls_prob : softmax feature map for face classify - roi_prob : feature map for regression - rectangles: 12net's predict - width : image's origin width - height : image's origin height - threshold : 0.6 can have 97% recall rate - Output: - rectangles: possible face positions - """ - prob = cls_prob[:, 1] - pick = np.where(prob >= threshold) - rectangles = np.array(rectangles) - x_1 = rectangles[pick, 0] - y_1 = rectangles[pick, 1] - x_2 = rectangles[pick, 2] - y_2 = rectangles[pick, 3] - sc_ = np.array([prob[pick]]).T - dx_1 = roi[pick, 0] - dx_2 = roi[pick, 1] - dx3 = roi[pick, 2] - dx4 = roi[pick, 3] - r_width = x_2-x_1 - r_height = y_2-y_1 - x_1 = np.array([(x_1 + dx_1 * r_width)[0]]).T - y_1 = np.array([(y_1 + dx_2 * r_height)[0]]).T - x_2 = np.array([(x_2 + dx3 * r_width)[0]]).T - y_2 = np.array([(y_2 + dx4 * r_height)[0]]).T - rectangles = np.concatenate((x_1, y_1, x_2, y_2, sc_), axis=1) - rectangles = rect2square(rectangles) - pick = [] - for rect in rectangles: - x_1 = int(max(0, rect[0])) - y_1 = int(max(0, rect[1])) - x_2 = int(min(width, rect[2])) - y_2 = int(min(height, rect[3])) - sc_ = rect[4] - if x_2 > x_1 and y_2 > y_1: - pick.append([x_1, y_1, x_2, y_2, sc_]) - return nms(pick, 0.3, 'iou') - - -def filter_face_48net(cls_prob, roi, pts, rectangles, width, height, threshold): - # pylint: disable=too-many-locals, too-many-arguments - """ Filter face position and calibrate bounding box on 12net's output - Input: - cls_prob : cls_prob[1] is face possibility - roi : roi offset - pts : 5 landmark - rectangles: 12net's predict, rectangles[i][0:3] is the position, rectangles[i][4] is score - width : image's origin width - height : image's origin height - threshold : 0.7 can have 94% recall rate on CelebA-database - Output: - rectangles: face positions and landmarks - """ - prob = cls_prob[:, 1] - pick = np.where(prob >= threshold) - rectangles = np.array(rectangles) - x_1 = rectangles[pick, 0] - y_1 = rectangles[pick, 1] - x_2 = rectangles[pick, 2] - y_2 = rectangles[pick, 3] - sc_ = np.array([prob[pick]]).T - dx_1 = roi[pick, 0] - dx_2 = roi[pick, 1] - dx3 = roi[pick, 2] - dx4 = roi[pick, 3] - r_width = x_2-x_1 - r_height = y_2-y_1 - pts0 = np.array([(r_width * pts[pick, 0] + x_1)[0]]).T - pts1 = np.array([(r_height * pts[pick, 5] + y_1)[0]]).T - pts2 = np.array([(r_width * pts[pick, 1] + x_1)[0]]).T - pts3 = np.array([(r_height * pts[pick, 6] + y_1)[0]]).T - pts4 = np.array([(r_width * pts[pick, 2] + x_1)[0]]).T - pts5 = np.array([(r_height * pts[pick, 7] + y_1)[0]]).T - pts6 = np.array([(r_width * pts[pick, 3] + x_1)[0]]).T - pts7 = np.array([(r_height * pts[pick, 8] + y_1)[0]]).T - pts8 = np.array([(r_width * pts[pick, 4] + x_1)[0]]).T - pts9 = np.array([(r_height * pts[pick, 9] + y_1)[0]]).T - x_1 = np.array([(x_1 + dx_1 * r_width)[0]]).T - y_1 = np.array([(y_1 + dx_2 * r_height)[0]]).T - x_2 = np.array([(x_2 + dx3 * r_width)[0]]).T - y_2 = np.array([(y_2 + dx4 * r_height)[0]]).T - rectangles = np.concatenate((x_1, y_1, x_2, y_2, sc_, - pts0, pts1, pts2, pts3, pts4, pts5, pts6, pts7, pts8, pts9), - axis=1) - pick = [] - for rect in rectangles: - x_1 = int(max(0, rect[0])) - y_1 = int(max(0, rect[1])) - x_2 = int(min(width, rect[2])) - y_2 = int(min(height, rect[3])) - if x_2 > x_1 and y_2 > y_1: - pick.append([x_1, y_1, x_2, y_2, - rect[4], rect[5], rect[6], rect[7], rect[8], rect[9], - rect[10], rect[11], rect[12], rect[13], rect[14]]) - return nms(pick, 0.3, 'iom') - - -def nms(rectangles, threshold, method): - # pylint:disable=too-many-locals +def nms(rectangles: np.ndarray, + scores: np.ndarray, + threshold: float, + method: str = "iom") -> tuple[np.ndarray, np.ndarray]: """ apply non-maximum suppression on ROIs in same scale(matrix version) - Input: - rectangles: rectangles[i][0:3] is the position, rectangles[i][4] is score - Output: - rectangles: same as input + + Parameters + ---------- + rectangles: :class:`np.ndarray` + The [b, l, t, r, b] bounding box detection candidates + threshold: float + Threshold for succesful match + method: str, optional + "iom" method or default. Defalt: "iom" + + Returns + ------- + rectangles: :class:`np.ndarray` + The [b, l, t, r, b] bounding boxes + scores :class:`np.ndarray` + The associated scores for the rectangles + """ - if not rectangles: - return rectangles - boxes = np.array(rectangles) - x_1 = boxes[:, 0] - y_1 = boxes[:, 1] - x_2 = boxes[:, 2] - y_2 = boxes[:, 3] - var_s = boxes[:, 4] - area = np.multiply(x_2-x_1+1, y_2-y_1+1) - s_sort = np.array(var_s.argsort()) + if not np.any(rectangles): + return rectangles, scores + bboxes = rectangles[..., :4].T + area = np.multiply(bboxes[2] - bboxes[0] + 1, bboxes[3] - bboxes[1] + 1) + s_sort = scores.argsort() + pick = [] while len(s_sort) > 0: - # s_sort[-1] have highest prob score, s_sort[0:-1]->others - xx_1 = np.maximum(x_1[s_sort[-1]], x_1[s_sort[0:-1]]) - yy_1 = np.maximum(y_1[s_sort[-1]], y_1[s_sort[0:-1]]) - xx_2 = np.minimum(x_2[s_sort[-1]], x_2[s_sort[0:-1]]) - yy_2 = np.minimum(y_2[s_sort[-1]], y_2[s_sort[0:-1]]) - width = np.maximum(0.0, xx_2 - xx_1 + 1) - height = np.maximum(0.0, yy_2 - yy_1 + 1) - inter = width * height - if method == 'iom': + s_bboxes = np.concatenate([ # s_sort[-1] have highest prob score, s_sort[0:-1]->others + np.maximum(bboxes[:2, s_sort[-1], None], bboxes[:2, s_sort[0:-1]]), + np.minimum(bboxes[2:, s_sort[-1], None], bboxes[2:, s_sort[0:-1]])], axis=0) + + inter = (np.maximum(0.0, s_bboxes[2] - s_bboxes[0] + 1) * + np.maximum(0.0, s_bboxes[3] - s_bboxes[1] + 1)) + + if method == "iom": var_o = inter / np.minimum(area[s_sort[-1]], area[s_sort[0:-1]]) else: var_o = inter / (area[s_sort[-1]] + area[s_sort[0:-1]] - inter) + pick.append(s_sort[-1]) s_sort = s_sort[np.where(var_o <= threshold)[0]] - result_rectangle = boxes[pick].tolist() - return result_rectangle - - -def calculate_scales(height, width, minsize, factor): - """ Calculate multi-scale - Input: - height: Original image height - width: Original image width - minsize: Minimum size for a face to be accepted - factor: Scaling factor - Output: - scales : Multi-scale - """ - factor_count = 0 - minl = np.amin([height, width]) - var_m = 12.0 / minsize - minl = minl * var_m - # create scale pyramid - scales = [] - while minl >= 12: - scales += [var_m * np.power(factor, factor_count)] - minl = minl * factor - factor_count += 1 - logger.trace(scales) - return scales - - -def rect2square(rectangles): + + result_rectangle = rectangles[pick] + result_scores = scores[pick] + return result_rectangle, result_scores + + +def rect2square(rectangles: np.ndarray) -> np.ndarray: """ change rectangles into squares (matrix version) - Input: - rectangles: rectangles[i][0:3] is the position, rectangles[i][4] is score - Output: - squares: same as input + + Parameters + ---------- + rectangles: :class:`numpy.ndarray` + [b, x, y, x1, y1] rectangles + + Return + ------ + list + Original rectangle changed to a square """ width = rectangles[:, 2] - rectangles[:, 0] height = rectangles[:, 3] - rectangles[:, 1] diff --git a/plugins/extract/detect/mtcnn_defaults.py b/plugins/extract/detect/mtcnn_defaults.py index 4e73eae737..17396669a1 100755 --- a/plugins/extract/detect/mtcnn_defaults.py +++ b/plugins/extract/detect/mtcnn_defaults.py @@ -51,75 +51,83 @@ _DEFAULTS = { - "minsize": dict( - default=20, - info="The minimum size of a face (in pixels) to be accepted as a positive match.\nLower " - "values use significantly more VRAM and will detect more false positives.", - datatype=int, - rounding=10, - min_max=(20, 1000), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ), - "scalefactor": dict( - default=0.709, - info="The scale factor for the image pyramid.", - datatype=float, - rounding=3, - min_max=(0.1, 0.9), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ), - "batch-size": dict( - default=8, - info="The batch size to use. To a point, higher batch sizes equal better performance, but " - "setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can accomodate " - "then this will automatically be lowered.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ), - "threshold_1": dict( - default=0.6, - info="First stage threshold for face detection. This stage obtains face candidates.", - datatype=float, - rounding=2, - min_max=(0.1, 0.9), - choices=[], - group="threshold", - gui_radio=False, - fixed=True, - ), - "threshold_2": dict( - default=0.7, - info="Second stage threshold for face detection. This stage refines face candidates.", - datatype=float, - rounding=2, - min_max=(0.1, 0.9), - choices=[], - group="threshold", - gui_radio=False, - fixed=True, - ), - "threshold_3": dict( - default=0.7, - info="Third stage threshold for face detection. This stage further refines face " - "candidates.", - datatype=float, - rounding=2, - min_max=(0.1, 0.9), - choices=[], - group="threshold", - gui_radio=False, - fixed=True, - ), + "minsize": { + "default": 20, + "info": "The minimum size of a face (in pixels) to be accepted as a positive match." + "\nLower values use significantly more VRAM and will detect more false positives.", + "datatype": int, + "rounding": 10, + "min_max": (20, 1000), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "scalefactor": { + "default": 0.709, + "info": "The scale factor for the image pyramid.", + "datatype": float, + "rounding": 3, + "min_max": (0.1, 0.9), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "batch-size": { + "default": 8, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "cpu": { + "default": True, + "info": "MTCNN detector still runs fairly quickly on CPU on some setups. " + "Enable CPU mode here to use the CPU for this detector to save some VRAM at a " + "speed cost.", + "datatype": bool, + "group": "settings" + }, + "threshold_1": { + "default": 0.6, + "info": "First stage threshold for face detection. This stage obtains face candidates.", + "datatype": float, + "rounding": 2, + "min_max": (0.1, 0.9), + "choices": [], + "group": "threshold", + "gui_radio": False, + "fixed": True, + }, + "threshold_2": { + "default": 0.7, + "info": "Second stage threshold for face detection. This stage refines face candidates.", + "datatype": float, + "rounding": 2, + "min_max": (0.1, 0.9), + "choices": [], + "group": "threshold", + "gui_radio": False, + "fixed": True, + }, + "threshold_3": { + "default": 0.7, + "info": "Third stage threshold for face detection. This stage further refines face " + "candidates.", + "datatype": float, + "rounding": 2, + "min_max": (0.1, 0.9), + "choices": [], + "group": "threshold", + "gui_radio": False, + "fixed": True, + }, } diff --git a/plugins/extract/detect/s3fd.py b/plugins/extract/detect/s3fd.py index 59205d6c91..89d538b76f 100644 --- a/plugins/extract/detect/s3fd.py +++ b/plugins/extract/detect/s3fd.py @@ -5,29 +5,31 @@ Adapted from S3FD Port in FAN: https://github.com/1adrianb/face-alignment """ +from __future__ import annotations +import logging +import typing as T from scipy.special import logsumexp import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow import keras +from tensorflow.keras import backend as K # pylint:disable=import-error +from tensorflow.keras.layers import ( # pylint:disable=import-error + Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D) + from lib.model.session import KSession -from lib.utils import get_backend -from ._base import Detector, logger +from ._base import BatchType, Detector + +if T.TYPE_CHECKING: + from tensorflow import Tensor -if get_backend() == "amd": - import keras - from keras import backend as K - from keras.layers import Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow import keras - from tensorflow.keras import backend as K # pylint:disable=import-error - from tensorflow.keras.layers import ( # pylint:disable=no-name-in-module,import-error - Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D) +logger = logging.getLogger(__name__) class Detect(Detector): """ S3FD detector for face recognition """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 11 model_filename = "s3fd_keras_v2.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) @@ -38,31 +40,32 @@ def __init__(self, **kwargs): self.vram_per_batch = 208 self.batchsize = self.config["batch-size"] - def init_model(self): + def init_model(self) -> None: """ Initialize S3FD Model""" + assert isinstance(self.model_path, str) confidence = self.config["confidence"] / 100 - model_kwargs = dict(custom_objects=dict(L2Norm=L2Norm, SliceO2K=SliceO2K)) + model_kwargs = {"custom_objects": {"L2Norm": L2Norm, "SliceO2K": SliceO2K}} self.model = S3fd(self.model_path, model_kwargs, self.config["allow_growth"], self._exclude_gpus, confidence) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detection image(s) for prediction """ - batch["feed"] = self.model.prepare_batch(batch["image"]) - return batch + assert isinstance(self.model, S3fd) + batch.feed = self.model.prepare_batch(np.array(batch.image)) - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - predictions = self.model.predict(batch["feed"]) - batch["prediction"] = self.model.finalize_predictions(predictions) - logger.trace("filename: %s, prediction: %s", batch["filename"], batch["prediction"]) - return batch + assert isinstance(self.model, S3fd) + predictions = self.model.predict(feed) + assert isinstance(predictions, list) + return self.model.finalize_predictions(predictions) - def process_output(self, batch): + def process_output(self, batch) -> None: """ Compile found faces for output """ - return batch + return ################################################################################ @@ -78,7 +81,7 @@ class L2Norm(keras.layers.Layer): scale: float, optional The scaling for initial weights. Default: `1.0` """ - def __init__(self, n_channels, scale=1.0, **kwargs): + def __init__(self, n_channels: int, scale: float = 1.0, **kwargs) -> None: super().__init__(**kwargs) self._n_channels = n_channels self._scale = scale @@ -88,7 +91,7 @@ def __init__(self, n_channels, scale=1.0, **kwargs): initializer=keras.initializers.Constant(value=self._scale), dtype="float32") - def call(self, inputs): + def call(self, inputs: Tensor) -> Tensor: # pylint:disable=arguments-differ """ Call the L2 Normalization Layer. Parameters @@ -105,7 +108,7 @@ def call(self, inputs): var_x = inputs / norm * self.w return var_x - def get_config(self): + def get_config(self) -> dict: """ Returns the config of the layer. Returns @@ -121,14 +124,19 @@ def get_config(self): class SliceO2K(keras.layers.Layer): """ Custom Keras Slice layer generated by onnx2keras. """ - def __init__(self, starts, ends, axes=None, steps=None, **kwargs): + def __init__(self, + starts: list[int], + ends: list[int], + axes: list[int] | None = None, + steps: list[int] | None = None, + **kwargs) -> None: self._starts = starts self._ends = ends self._axes = axes self._steps = steps super().__init__(**kwargs) - def _get_slices(self, dimensions): + def _get_slices(self, dimensions: int) -> list[tuple[int, ...]]: """ Obtain slices for the given number of dimensions. Parameters @@ -141,16 +149,12 @@ def _get_slices(self, dimensions): list The slices for the given number of dimensions """ - axes = self._axes - steps = self._steps - if axes is None: - axes = tuple(range(dimensions)) - if steps is None: - steps = (1,) * len(axes) + axes = tuple(range(dimensions)) if self._axes is None else self._axes + steps = (1,) * len(axes) if self._steps is None else self._steps assert len(axes) == len(steps) == len(self._starts) == len(self._ends) return list(zip(axes, self._starts, self._ends, steps)) - def compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape: tuple[int, ...]) -> tuple[int, ...]: """Computes the output shape of the layer. Assumes that the layer will be built to match that input shape provided. @@ -166,25 +170,25 @@ def compute_output_shape(self, input_shape): tuple An output shape tuple. """ - input_shape = list(input_shape) - for a_x, start, end, steps in self._get_slices(len(input_shape)): - size = input_shape[a_x] + in_shape = list(input_shape) + for a_x, start, end, steps in self._get_slices(len(in_shape)): + size = in_shape[a_x] if a_x == 0: raise AttributeError("Can not slice batch axis.") if size is None: if start < 0 or end < 0: raise AttributeError("Negative slices not supported on symbolic axes") logger.warning("Slicing symbolic axis might lead to problems.") - input_shape[a_x] = (end - start) // steps + in_shape[a_x] = (end - start) // steps continue if start < 0: start = size - start if end < 0: end = size - end - input_shape[a_x] = (min(size, end) - start) // steps - return tuple(input_shape) + in_shape[a_x] = (min(size, end) - start) // steps + return tuple(in_shape) - def call(self, inputs, **kwargs): # pylint:disable=unused-argument + def call(self, inputs, **kwargs): # pylint:disable=unused-argument,arguments-differ """This is where the layer's logic lives. Parameters @@ -204,7 +208,7 @@ def call(self, inputs, **kwargs): # pylint:disable=unused-argument retval = inputs[tuple(slices)] return retval - def get_config(self): + def get_config(self) -> dict: """ Returns the config of the layer. Returns @@ -222,7 +226,12 @@ def get_config(self): class S3fd(KSession): """ Keras Network """ - def __init__(self, model_path, model_kwargs, allow_growth, exclude_gpus, confidence): + def __init__(self, + model_path: str, + model_kwargs: dict, + allow_growth: bool, + exclude_gpus: list[int] | None, + confidence: float) -> None: logger.debug("Initializing: %s: (model_path: '%s', model_kwargs: %s, allow_growth: %s, " "exclude_gpus: %s, confidence: %s)", self.__class__.__name__, model_path, model_kwargs, allow_growth, exclude_gpus, confidence) @@ -237,7 +246,7 @@ def __init__(self, model_path, model_kwargs, allow_growth, exclude_gpus, confide self.average_img = np.array([104.0, 117.0, 123.0]) logger.debug("Initialized: %s", self.__class__.__name__) - def model_definition(self): + def model_definition(self) -> tuple[list[Tensor], list[Tensor]]: """ Keras S3FD Model Definition, adapted from FAN pytorch implementation. """ input_ = Input(shape=(640, 640, 3)) var_x = self.conv_block(input_, 64, 1, 2) @@ -306,7 +315,7 @@ def model_definition(self): return [input_], [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6] @classmethod - def conv_block(cls, inputs, filters, idx, recursions): + def conv_block(cls, inputs: Tensor, filters: int, idx: int, recursions: int) -> Tensor: """ First round convolutions with zero padding added. Parameters @@ -338,7 +347,7 @@ def conv_block(cls, inputs, filters, idx, recursions): return var_x @classmethod - def conv_up(cls, inputs, filters, idx): + def conv_up(cls, inputs: Tensor, filters: int, idx: int) -> Tensor: """ Convolution up filter blocks with zero padding added. Parameters @@ -369,7 +378,7 @@ def conv_up(cls, inputs, filters, idx): name=rec_name)(var_x) return var_x - def prepare_batch(self, batch): + def prepare_batch(self, batch: np.ndarray) -> np.ndarray: """ Prepare a batch for prediction. Normalizes the feed images. @@ -387,7 +396,7 @@ def prepare_batch(self, batch): batch = batch - self.average_img return batch - def finalize_predictions(self, bounding_boxes_scales): + def finalize_predictions(self, bounding_boxes_scales: list[np.ndarray]) -> np.ndarray: """ Process the output from the model to obtain faces Parameters @@ -400,11 +409,11 @@ def finalize_predictions(self, bounding_boxes_scales): for img in batch_size: bboxlist = [scale[img:img+1] for scale in bounding_boxes_scales] boxes = self._post_process(bboxlist) - bboxlist = self._nms(boxes, 0.5) - ret.append(bboxlist) - return ret + finallist = self._nms(boxes, 0.5) + ret.append(finallist) + return np.array(ret, dtype="object") - def _post_process(self, bboxlist): + def _post_process(self, bboxlist: list[np.ndarray]) -> np.ndarray: """ Perform post processing on output TODO: do this on the batch. """ @@ -428,12 +437,12 @@ def _post_process(self, bboxlist): return return_numpy @staticmethod - def softmax(inp, axis): + def softmax(inp, axis: int) -> np.ndarray: """Compute softmax values for each sets of scores in x.""" return np.exp(inp - logsumexp(inp, axis=axis, keepdims=True)) @staticmethod - def decode(location, priors): + def decode(location: np.ndarray, priors: np.ndarray) -> np.ndarray: """Decode locations from predictions using priors to undo the encoding we did for offset regression at train time. @@ -457,28 +466,28 @@ def decode(location, priors): return boxes @staticmethod - def _nms(boxes, threshold): + def _nms(boxes: np.ndarray, threshold: float) -> np.ndarray: """ Perform Non-Maximum Suppression """ retained_box_indices = [] areas = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1) ranked_indices = boxes[:, 4].argsort()[::-1] while ranked_indices.size > 0: - best = ranked_indices[0] - rest = ranked_indices[1:] + best_rest = ranked_indices[0], ranked_indices[1:] - max_of_xy = np.maximum(boxes[best, :2], boxes[rest, :2]) - min_of_xy = np.minimum(boxes[best, 2:4], boxes[rest, 2:4]) + max_of_xy = np.maximum(boxes[best_rest[0], :2], boxes[best_rest[1], :2]) + min_of_xy = np.minimum(boxes[best_rest[0], 2:4], boxes[best_rest[1], 2:4]) width_height = np.maximum(0, min_of_xy - max_of_xy + 1) intersection_areas = width_height[:, 0] * width_height[:, 1] - iou = intersection_areas / (areas[best] + areas[rest] - intersection_areas) + iou = intersection_areas / (areas[best_rest[0]] + + areas[best_rest[1]] - intersection_areas) overlapping_boxes = (iou > threshold).nonzero()[0] if len(overlapping_boxes) != 0: overlap_set = ranked_indices[overlapping_boxes + 1] vote = np.average(boxes[overlap_set, :4], axis=0, weights=boxes[overlap_set, 4]) - boxes[best, :4] = vote - retained_box_indices.append(best) + boxes[best_rest[0], :4] = vote + retained_box_indices.append(best_rest[0]) non_overlapping_boxes = (iou <= threshold).nonzero()[0] ranked_indices = ranked_indices[non_overlapping_boxes + 1] diff --git a/plugins/extract/detect/s3fd_defaults.py b/plugins/extract/detect/s3fd_defaults.py index 6c17bba95b..5e219766f4 100755 --- a/plugins/extract/detect/s3fd_defaults.py +++ b/plugins/extract/detect/s3fd_defaults.py @@ -51,32 +51,32 @@ _DEFAULTS = { - "confidence": dict( - default=70, - info="The confidence level at which the detector has succesfully found a face.\n" - "Higher levels will be more discriminating, lower levels will have more false " - "positives.", - datatype=int, - rounding=5, - min_max=(25, 100), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ), - "batch-size": dict( - default=4, - info="The batch size to use. To a point, higher batch sizes equal better performance, " - "but setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can " - "accomodate then this will automatically be lowered." - "\n\tAMD users: A batchsize of 8 requires about 2 GB vram.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ) + "confidence": { + "default": 70, + "info": "The confidence level at which the detector has succesfully found a face.\n" + "Higher levels will be more discriminating, lower levels will have more false " + "positives.", + "datatype": int, + "rounding": 5, + "min_max": (25, 100), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + }, + "batch-size": { + "default": 4, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered." + "\n\tAMD users: A batchsize of 8 requires about 2 GB vram.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + } } diff --git a/plugins/extract/extract_media.py b/plugins/extract/extract_media.py new file mode 100644 index 0000000000..b9d3f84a33 --- /dev/null +++ b/plugins/extract/extract_media.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +""" Object for holding and manipulating media passing through a faceswap extraction pipeline """ +from __future__ import annotations +import logging +import typing as T + +import cv2 + +from lib.logger import parse_class_init + +if T.TYPE_CHECKING: + import numpy as np + from lib.align.alignments import PNGHeaderSourceDict + from lib.align.detected_face import DetectedFace + +logger = logging.getLogger(__name__) + + +class ExtractMedia: + """ An object that passes through the :class:`~plugins.extract.pipeline.Extractor` pipeline. + + Parameters + ---------- + filename: str + The base name of the original frame's filename + image: :class:`numpy.ndarray` + The original frame or a faceswap aligned face image + detected_faces: list, optional + A list of :class:`~lib.align.DetectedFace` objects. Detected faces can be added + later with :func:`add_detected_faces`. Setting ``None`` will default to an empty list. + Default: ``None`` + is_aligned: bool, optional + ``True`` if the :attr:`image` is an aligned faceswap image otherwise ``False``. Used for + face filtering with vggface2. Aligned faceswap images will automatically skip detection, + alignment and masking. Default: ``False`` + """ + + def __init__(self, + filename: str, + image: np.ndarray, + detected_faces: list[DetectedFace] | None = None, + is_aligned: bool = False) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self._filename = filename + self._image: np.ndarray | None = image + self._image_shape = T.cast(tuple[int, int, int], image.shape) + self._detected_faces: list[DetectedFace] = ([] if detected_faces is None + else detected_faces) + self._is_aligned = is_aligned + self._frame_metadata: PNGHeaderSourceDict | None = None + self._sub_folders: list[str | None] = [] + + @property + def filename(self) -> str: + """ str: The base name of the :attr:`image` filename. """ + return self._filename + + @property + def image(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The source frame for this object. """ + assert self._image is not None + return self._image + + @property + def image_shape(self) -> tuple[int, int, int]: + """ tuple: The shape of the stored :attr:`image`. """ + return self._image_shape + + @property + def image_size(self) -> tuple[int, int]: + """ tuple: The (`height`, `width`) of the stored :attr:`image`. """ + return self._image_shape[:2] + + @property + def detected_faces(self) -> list[DetectedFace]: + """list: A list of :class:`~lib.align.DetectedFace` objects in the :attr:`image`. """ + return self._detected_faces + + @property + def is_aligned(self) -> bool: + """ bool. ``True`` if :attr:`image` is an aligned faceswap image otherwise ``False`` """ + return self._is_aligned + + @property + def frame_metadata(self) -> PNGHeaderSourceDict: + """ dict: The frame metadata that has been added from an aligned image. This property + should only be called after :func:`add_frame_metadata` has been called when processing + an aligned face. For all other instances an assertion error will be raised. + + Raises + ------ + AssertionError + If frame metadata has not been populated from an aligned image + """ + assert self._frame_metadata is not None + return self._frame_metadata + + @property + def sub_folders(self) -> list[str | None]: + """ list: The sub_folders that the faces should be output to. Used when binning filter + output is enabled. The list corresponds to the list of detected faces + """ + return self._sub_folders + + def get_image_copy(self, color_format: T.Literal["BGR", "RGB", "GRAY"]) -> np.ndarray: + """ Get a copy of the image in the requested color format. + + Parameters + ---------- + color_format: ['BGR', 'RGB', 'GRAY'] + The requested color format of :attr:`image` + + Returns + ------- + :class:`numpy.ndarray`: + A copy of :attr:`image` in the requested :attr:`color_format` + """ + logger.trace("Requested color format '%s' for frame '%s'", # type:ignore[attr-defined] + color_format, self._filename) + image = getattr(self, f"_image_as_{color_format.lower()}")() + return image + + def add_detected_faces(self, faces: list[DetectedFace]) -> None: + """ Add detected faces to the object. Called at the end of each extraction phase. + + Parameters + ---------- + faces: list + A list of :class:`~lib.align.DetectedFace` objects + """ + logger.trace("Adding detected faces for filename: '%s'. " # type:ignore[attr-defined] + "(faces: %s, lrtb: %s)", self._filename, faces, + [(face.left, face.right, face.top, face.bottom) for face in faces]) + self._detected_faces = faces + + def add_sub_folders(self, folders: list[str | None]) -> None: + """ Add detected faces to the object. Called at the end of each extraction phase. + + Parameters + ---------- + folders: list + A list of str sub folder names or ``None`` if no sub folder is required. Should + correspond to the detected faces list + """ + logger.trace("Adding sub folders for filename: '%s'. " # type:ignore[attr-defined] + "(folders: %s)", self._filename, folders,) + self._sub_folders = folders + + def remove_image(self) -> None: + """ Delete the image and reset :attr:`image` to ``None``. + + Required for multi-phase extraction to avoid the frames stacking RAM. + """ + logger.trace("Removing image for filename: '%s'", # type:ignore[attr-defined] + self._filename) + del self._image + self._image = None + + def set_image(self, image: np.ndarray) -> None: + """ Add the image back into :attr:`image` + + Required for multi-phase extraction adds the image back to this object. + + Parameters + ---------- + image: :class:`numpy.ndarry` + The original frame to be re-applied to for this :attr:`filename` + """ + logger.trace("Reapplying image: (filename: `%s`, " # type:ignore[attr-defined] + "image shape: %s)", self._filename, image.shape) + self._image = image + + def add_frame_metadata(self, metadata: PNGHeaderSourceDict) -> None: + """ Add the source frame metadata from an aligned PNG's header data. + + metadata: dict + The contents of the 'source' field in the PNG header + """ + logger.trace("Adding PNG Source data for '%s': %s", # type:ignore[attr-defined] + self._filename, metadata) + dims = T.cast(tuple[int, int], metadata["source_frame_dims"]) + self._image_shape = (*dims, 3) + self._frame_metadata = metadata + + def _image_as_bgr(self) -> np.ndarray: + """ Get a copy of the source frame in BGR format. + + Returns + ------- + :class:`numpy.ndarray`: + A copy of :attr:`image` in BGR color format """ + return self.image[..., :3].copy() + + def _image_as_rgb(self) -> np.ndarray: + """ Get a copy of the source frame in RGB format. + + Returns + ------- + :class:`numpy.ndarray`: + A copy of :attr:`image` in RGB color format """ + return self.image[..., 2::-1].copy() + + def _image_as_gray(self) -> np.ndarray: + """ Get a copy of the source frame in gray-scale format. + + Returns + ------- + :class:`numpy.ndarray`: + A copy of :attr:`image` in gray-scale color format """ + return cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2GRAY) diff --git a/plugins/extract/mask/_base.py b/plugins/extract/mask/_base.py index 6b982bb882..8b5d71e0ad 100644 --- a/plugins/extract/mask/_base.py +++ b/plugins/extract/mask/_base.py @@ -5,22 +5,52 @@ See the override methods for which methods are required. -The plugin will receive a :class:`~plugins.extract.pipeline.ExtractMedia` object. +The plugin will receive a :class:`~plugins.extract.extract_media.ExtractMedia` object. For each source item, the plugin must pass a dict to finalize containing: >>> {"filename": , >>> "detected_faces": } """ +from __future__ import annotations +import logging +import typing as T + +from dataclasses import dataclass, field import cv2 import numpy as np -from tensorflow.python.framework import errors_impl as tf_errors +from tensorflow.python.framework import errors_impl as tf_errors # pylint:disable=no-name-in-module # noqa + +from lib.align import AlignedFace, LandmarkType, transform_image +from lib.utils import FaceswapError +from plugins.extract import ExtractMedia +from plugins.extract._base import BatchType, ExtractorBatch, Extractor + +if T.TYPE_CHECKING: + from collections.abc import Generator + from queue import Queue + from lib.align import DetectedFace + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +@dataclass +class MaskerBatch(ExtractorBatch): + """ Dataclass for holding items flowing through the aligner. -from lib.align import AlignedFace, transform_image -from lib.utils import get_backend, FaceswapError -from plugins.extract._base import Extractor, ExtractMedia, logger + Inherits from :class:`~plugins.extract._base.ExtractorBatch` + + Parameters + ---------- + roi_masks: list + The region of interest masks for the batch + """ + detected_faces: list[DetectedFace] = field(default_factory=list) + roi_masks: list[np.ndarray] = field(default_factory=list) + feed_faces: list[AlignedFace] = field(default_factory=list) class Masker(Extractor): # pylint:disable=abstract-method @@ -35,9 +65,6 @@ class Masker(Extractor): # pylint:disable=abstract-method https://github.com/deepfakes-models/faceswap-models for more information model_filename: str The name of the model file to be loaded - image_is_aligned: bool, optional - Indicates that the passed in image is an aligned face rather than a frame. - Default: ``False`` Other Parameters ---------------- @@ -53,9 +80,15 @@ class Masker(Extractor): # pylint:disable=abstract-method plugins.extract.align._base : Aligner parent class for extraction plugins. """ - def __init__(self, git_model_id=None, model_filename=None, configfile=None, - instance=0, image_is_aligned=False, **kwargs): - logger.debug("Initializing %s: (configfile: %s, )", self.__class__.__name__, configfile) + _logged_lm_count_once = False + + def __init__(self, + git_model_id: int | None = None, + model_filename: str | None = None, + configfile: str | None = None, + instance: int = 0, + **kwargs) -> None: + logger.debug("Initializing %s: (configfile: %s)", self.__class__.__name__, configfile) super().__init__(git_model_id, model_filename, configfile=configfile, @@ -64,24 +97,41 @@ def __init__(self, git_model_id=None, model_filename=None, configfile=None, self.input_size = 256 # Override for model specific input_size self.coverage_ratio = 1.0 # Override for model specific coverage_ratio + # Override if a specific type of landmark data is required: + self.landmark_type: LandmarkType | None = None + self._plugin_type = "mask" - self._image_is_aligned = image_is_aligned - self._storage_name = self.__module__.split(".")[-1].replace("_", "-") - self._storage_centering = "face" # Centering to store the mask at + self._storage_name = self.__module__.rsplit(".", maxsplit=1)[-1].replace("_", "-") + self._storage_centering: CenteringType = "face" # Centering to store the mask at self._storage_size = 128 # Size to store masks at. Leave this at default - self._faces_per_filename = dict() # Tracking for recompiling face batches - self._rollover = None # Items that are rolled over from the previous batch in get_batch - self._output_faces = [] logger.debug("Initialized %s", self.__class__.__name__) - def get_batch(self, queue): + def _maybe_log_warning(self, face: AlignedFace) -> None: + """ Log a warning, once, if we do not have full facial landmarks + + Parameters + ---------- + face: :class:`~lib.align.aligned_face.AlignedFace` + The aligned face object to test the landmark type for + """ + if face.landmark_type != LandmarkType.LM_2D_4 or self._logged_lm_count_once: + return + + msg = "are likely to be sub-standard" + msg = "can not be be generated" if self.name in ("Components", "Extended") else msg + + logger.warning("Extracted faces do not contain facial landmark data. '%s' masks %s.", + self.name, msg) + self._logged_lm_count_once = True + + def get_batch(self, queue: Queue) -> tuple[bool, MaskerBatch]: """ Get items for inputting into the masker from the queue in batches Items are returned from the ``queue`` in batches of :attr:`~plugins.extract._base.Extractor.batchsize` - Items are received as :class:`~plugins.extract.pipeline.ExtractMedia` objects and converted - to ``dict`` for internal processing. + Items are received as :class:`~plugins.extract.extract_media.ExtractMedia` objects and + converted to ``dict`` for internal processing. To ensure consistent batch sizes for masker the items are split into separate items for each :class:`~lib.align.DetectedFace` object. @@ -104,16 +154,16 @@ def get_batch(self, queue): ------- exhausted, bool ``True`` if queue is exhausted, ``False`` if not - batch, dict - A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`: + batch, :class:`~plugins.extract._base.ExtractorBatch` + The batch object for the current batch """ exhausted = False - batch = dict() + batch = MaskerBatch() idx = 0 while idx < self.batchsize: - item = self._collect_item(queue) + item = self.rollover_collector(queue) if item == "EOF": - logger.trace("EOF received") + logger.trace("EOF received") # type: ignore exhausted = True break # Put frames with no faces into the out queue to keep TQDM consistent @@ -125,7 +175,7 @@ def get_batch(self, queue): image = item.get_image_copy(self.color_format) roi = np.ones((*item.image_size[:2], 1), dtype="float32") - if not self._image_is_aligned: + if not item.is_aligned: # Add the ROI mask to image so we can get the ROI mask with a single warp image = np.concatenate([image, roi], axis=-1) @@ -135,12 +185,14 @@ def get_batch(self, queue): size=self.input_size, coverage_ratio=self.coverage_ratio, dtype="float32", - is_aligned=self._image_is_aligned) + is_aligned=item.is_aligned) + + self._maybe_log_warning(feed_face) - if not self._image_is_aligned: + assert feed_face.face is not None + if not item.is_aligned: # Split roi mask from feed face alpha channel - roi_mask = feed_face.face[..., 3] - feed_face._face = feed_face.face[..., :3] # pylint:disable=protected-access + roi_mask = feed_face.split_mask() else: # We have to do the warp here as AlignedFace did not perform it roi_mask = transform_image(roi, @@ -148,10 +200,10 @@ def get_batch(self, queue): feed_face.size, padding=feed_face.padding) - batch.setdefault("roi_masks", []).append(roi_mask) - batch.setdefault("detected_faces", []).append(face) - batch.setdefault("feed_faces", []).append(feed_face) - batch.setdefault("filename", []).append(item.filename) + batch.roi_masks.append(roi_mask) + batch.detected_faces.append(face) + batch.feed_faces.append(feed_face) + batch.filename.append(item.filename) idx += 1 if idx == self.batchsize: frame_faces = len(item.detected_faces) @@ -159,39 +211,35 @@ def get_batch(self, queue): self._rollover = ExtractMedia( item.filename, item.image, - detected_faces=item.detected_faces[f_idx + 1:]) - logger.trace("Rolled over %s faces of %s to next batch for '%s'", - len(self._rollover.detected_faces), frame_faces, + detected_faces=item.detected_faces[f_idx + 1:], + is_aligned=item.is_aligned) + logger.trace("Rolled over %s faces of %s to next batch " # type:ignore + "for '%s'", len(self._rollover.detected_faces), frame_faces, item.filename) break if batch: - logger.trace("Returning batch: %s", {k: v.shape if isinstance(v, np.ndarray) else v - for k, v in batch.items()}) + logger.trace("Returning batch: %s", # type:ignore + {k: len(v) if isinstance(v, (list, np.ndarray)) else v + for k, v in batch.__dict__.items()}) else: - logger.trace(item) + logger.trace(item) # type:ignore return exhausted, batch - def _collect_item(self, queue): - """ Collect the item from the _rollover dict or from the queue - Add face count per frame to self._faces_per_filename for joining - batches back up in finalize """ - if self._rollover is not None: - logger.trace("Getting from _rollover: (filename: `%s`, faces: %s)", - self._rollover.filename, len(self._rollover.detected_faces)) - item = self._rollover - self._rollover = None - else: - item = self._get_item(queue) - if item != "EOF": - logger.trace("Getting from queue: (filename: %s, faces: %s)", - item.filename, len(item.detected_faces)) - self._faces_per_filename[item.filename] = len(item.detected_faces) - return item - - def _predict(self, batch): + def _predict(self, batch: BatchType) -> MaskerBatch: """ Just return the masker's predict function """ + assert isinstance(batch, MaskerBatch) + assert self.name is not None try: - return self.predict(batch) + # slightly hacky workaround to deal with landmarks based masks: + if self.name.lower() in ("components", "extended"): + feed = np.empty(2, dtype="object") + feed[0] = batch.feed + feed[1] = batch.feed_faces + else: + feed = batch.feed + + batch.prediction = self.predict(feed) + return batch except tf_errors.ResourceExhaustedError as err: msg = ("You do not have enough GPU memory available to run detection at the " "selected batch size. You can try a number of things:" @@ -202,25 +250,8 @@ def _predict(self, batch): "CLI: Edit the file faceswap/config/extract.ini)." "\n3) Enable 'Single Process' mode.") raise FaceswapError(msg) from err - except Exception as err: - if get_backend() == "amd": - # pylint:disable=import-outside-toplevel - from lib.plaidml_utils import is_plaidml_error - if (is_plaidml_error(err) and ( - "CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or - "enough memory for the current schedule" in str(err).lower())): - msg = ("You do not have enough GPU memory available to run detection at " - "the selected batch size. You can try a number of things:" - "\n1) Close any other application that is using your GPU (web " - "browsers are particularly bad for this)." - "\n2) Lower the batchsize (the amount of images fed into the " - "model) by editing the plugin settings (GUI: Settings > Configure " - "extract settings, CLI: Edit the file " - "faceswap/config/extract.ini).") - raise FaceswapError(msg) from err - raise - - def finalize(self, batch): + + def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]: """ Finalize the output from Masker This should be called as the final task of each `plugin`. @@ -235,14 +266,19 @@ def finalize(self, batch): Yields ------ - :class:`~plugins.extract.pipeline.ExtractMedia` + :class:`~plugins.extract.extract_media.ExtractMedia` The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes, landmarks and masks for the detected faces found in the frame. """ - for mask, face, feed_face, roi_mask in zip(batch["prediction"], - batch["detected_faces"], - batch["feed_faces"], - batch["roi_masks"]): + assert isinstance(batch, MaskerBatch) + for mask, face, feed_face, roi_mask in zip(batch.prediction, + batch.detected_faces, + batch.feed_faces, + batch.roi_masks): + if self.name in ("Components", "Extended") and not np.any(mask): + # Components/Extended masks can return empty when called from the manual tool with + # 4 Point ROI landmarks + continue self._crop_out_of_bounds(mask, roi_mask) face.add_mask(self._storage_name, mask, @@ -250,11 +286,12 @@ def finalize(self, batch): feed_face.interpolators[1], storage_size=self._storage_size, storage_centering=self._storage_centering) - del batch["feed_faces"] + del batch.feed - logger.trace("Item out: %s", {key: val.shape if isinstance(val, np.ndarray) else val - for key, val in batch.items()}) - for filename, face in zip(batch["filename"], batch["detected_faces"]): + logger.trace("Item out: %s", # type: ignore + {key: val.shape if isinstance(val, np.ndarray) else val + for key, val in batch.__dict__.items()}) + for filename, face in zip(batch.filename, batch.detected_faces): self._output_faces.append(face) if len(self._output_faces) != self._faces_per_filename[filename]: continue @@ -262,26 +299,27 @@ def finalize(self, batch): output = self._extract_media.pop(filename) output.add_detected_faces(self._output_faces) self._output_faces = [] - logger.trace("Yielding: (filename: '%s', image: %s, detected_faces: %s)", - output.filename, output.image_shape, len(output.detected_faces)) + logger.trace("Yielding: (filename: '%s', image: %s, " # type:ignore + "detected_faces: %s)", output.filename, output.image_shape, + len(output.detected_faces)) yield output # <<< PROTECTED ACCESS METHODS >>> # @classmethod - def _resize(cls, image, target_size): + def _resize(cls, image: np.ndarray, target_size: int) -> np.ndarray: """ resize input and output of mask models appropriately """ height, width, channels = image.shape image_size = max(height, width) scale = target_size / image_size if scale == 1.: return image - method = cv2.INTER_CUBIC if scale > 1. else cv2.INTER_AREA # pylint: disable=no-member + method = cv2.INTER_CUBIC if scale > 1. else cv2.INTER_AREA # pylint:disable=no-member resized = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=method) resized = resized if channels > 1 else resized[..., None] return resized @classmethod - def _crop_out_of_bounds(cls, mask, roi_mask): + def _crop_out_of_bounds(cls, mask: np.ndarray, roi_mask: np.ndarray) -> None: """ Un-mask any area of the predicted mask that falls outside of the original frame. Parameters diff --git a/plugins/extract/mask/bisenet_fp.py b/plugins/extract/mask/bisenet_fp.py index 8136c3078c..cf8a177fe6 100644 --- a/plugins/extract/mask/bisenet_fp.py +++ b/plugins/extract/mask/bisenet_fp.py @@ -4,41 +4,44 @@ Architecture and Pre-Trained Model ported from PyTorch to Keras by TorzDF from https://github.com/zllrunning/face-parsing.PyTorch """ +from __future__ import annotations +import logging +import typing as T + import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import backend as K # pylint:disable=import-error +from tensorflow.keras.layers import ( # pylint:disable=import-error + Activation, Add, BatchNormalization, Concatenate, Conv2D, GlobalAveragePooling2D, Input, + MaxPooling2D, Multiply, Reshape, UpSampling2D, ZeroPadding2D) + from lib.model.session import KSession -from lib.utils import get_backend from plugins.extract._base import _get_config -from ._base import Masker, logger +from ._base import BatchType, Masker, MaskerBatch + +if T.TYPE_CHECKING: + from tensorflow import Tensor -if get_backend() == "amd": - from keras import backend as K - from keras.layers import ( - Activation, Add, BatchNormalization, Concatenate, Conv2D, GlobalAveragePooling2D, Input, - MaxPooling2D, Multiply, Reshape, UpSampling2D, ZeroPadding2D) -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import backend as K # pylint:disable=import-error - from tensorflow.keras.layers import ( # pylint:disable=no-name-in-module,import-error - Activation, Add, BatchNormalization, Concatenate, Conv2D, GlobalAveragePooling2D, Input, - MaxPooling2D, Multiply, Reshape, UpSampling2D, ZeroPadding2D) +logger = logging.getLogger(__name__) class Mask(Masker): """ Neural network to process face image into a segmentation mask of the face """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: self._is_faceswap, version = self._check_weights_selection(kwargs.get("configfile")) git_model_id = 14 model_filename = f"bisnet_face_parsing_v{version}.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession self.name = "BiSeNet - Face Parsing" self.input_size = 512 self.color_format = "RGB" - self.vram = 2304 - self.vram_warnings = 256 - self.vram_per_batch = 64 + self.vram = 2304 if not self.config["cpu"] else 0 + self.vram_warnings = 256 if not self.config["cpu"] else 0 + self.vram_per_batch = 64 if not self.config["cpu"] else 0 self.batchsize = self.config["batch-size"] self._segment_indices = self._get_segment_indices() @@ -46,7 +49,7 @@ def __init__(self, **kwargs): # Separate storage for face and head masks self._storage_name = f"{self._storage_name}_{self._storage_centering}" - def _check_weights_selection(self, configfile): + def _check_weights_selection(self, configfile: str | None) -> tuple[bool, int]: """ Check which weights have been selected. This is required for passing along the correct file name for the corresponding weights @@ -70,7 +73,7 @@ def _check_weights_selection(self, configfile): version = 1 if not is_faceswap else 2 if config.get("include_hair") else 3 return is_faceswap, version - def _get_segment_indices(self): + def _get_segment_indices(self) -> list[int]: """ Obtain the segment indices to include within the face mask area based on user configuration settings. @@ -100,40 +103,40 @@ def _get_segment_indices(self): logger.debug("Selected segment indices: %s", retval) return retval - def init_model(self): + def init_model(self) -> None: """ Initialize the BiSeNet Face Parsing model. """ + assert isinstance(self.model_path, str) lbls = 5 if self._is_faceswap else 19 self.model = BiSeNet(self.model_path, self.config["allow_growth"], self._exclude_gpus, self.input_size, - lbls) + lbls, + self.config["cpu"]) placeholder = np.zeros((self.batchsize, self.input_size, self.input_size, 3), dtype="float32") self.model.predict(placeholder) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ + assert isinstance(batch, MaskerBatch) mean = (0.384, 0.314, 0.279) if self._is_faceswap else (0.485, 0.456, 0.406) std = (0.324, 0.286, 0.275) if self._is_faceswap else (0.229, 0.224, 0.225) - batch["feed"] = ((np.array([feed.face[..., :3] - for feed in batch["feed_faces"]], - dtype="float32") / 255.0) - mean) / std - logger.trace("feed shape: %s", batch["feed"].shape) - return batch + batch.feed = ((np.array([T.cast(np.ndarray, feed.face)[..., :3] + for feed in batch.feed_faces], + dtype="float32") / 255.0) - mean) / std + logger.trace("feed shape: %s", batch.feed.shape) # type:ignore - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - batch["prediction"] = self.model.predict(batch["feed"])[0] - return batch + return self.model.predict(feed)[0] - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - pred = batch["prediction"].argmax(-1).astype("uint8") - batch["prediction"] = np.isin(pred, self._segment_indices).astype("float32") - return batch + pred = batch.prediction.argmax(-1).astype("uint8") + batch.prediction = np.isin(pred, self._segment_indices).astype("float32") # BiSeNet Face-Parsing Model @@ -160,10 +163,10 @@ def process_output(self, batch): # SOFTWARE. -_NAME_TRACKER = set() +_NAME_TRACKER: set[str] = set() -def _get_name(name, start_idx=1): +def _get_name(name: str, start_idx: int = 1) -> str: """ Auto numbering to keep track of layer names. Names are kept the same as the PyTorch original model, to enable easier porting of weights. @@ -217,8 +220,13 @@ class ConvBn(): # pylint:disable=too-few-public-methods The starting index for naming the layers within the block. See :func:`_get_name` for more information. Default: `1` """ - def __init__(self, filters, - kernel_size=3, strides=1, padding=1, activation=True, prefix="", start_idx=1): + def __init__(self, filters: int, + kernel_size: int = 3, + strides: int = 1, + padding: int = 1, + activation: int = True, + prefix: str = "", + start_idx: int = 1) -> None: self._filters = filters self._kernel_size = kernel_size self._strides = strides @@ -227,7 +235,7 @@ def __init__(self, filters, self._prefix = f"{prefix}." if prefix else prefix self._start_idx = start_idx - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the Convolutional Batch Normalization block. Parameters @@ -267,7 +275,7 @@ class ResNet18(): # pylint:disable=too-few-public-methods def __init__(self): self._feature_index = 1 if K.image_data_format() == "channels_first" else -1 - def _basic_block(self, inputs, prefix, filters, strides=1): + def _basic_block(self, inputs: Tensor, prefix: str, filters: int, strides: int = 1) -> Tensor: """ The basic building block for ResNet 18. Parameters @@ -305,7 +313,12 @@ def _basic_block(self, inputs, prefix, filters, strides=1): var_x = Activation("relu", name=f"{prefix}.relu")(var_x) return var_x - def _basic_layer(self, inputs, prefix, filters, num_blocks, strides=1): + def _basic_layer(self, + inputs: Tensor, + prefix: str, + filters: int, + num_blocks: int, + strides: int = 1) -> Tensor: """ The basic layer for ResNet 18. Recursively builds from :func:`_basic_block`. Parameters @@ -332,7 +345,7 @@ def _basic_layer(self, inputs, prefix, filters, num_blocks, strides=1): var_x = self._basic_block(var_x, f"{prefix}.{i + 1}", filters, strides=1) return var_x - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the ResNet 18 block. Parameters @@ -366,10 +379,10 @@ class AttentionRefinementModule(): # pylint:disable=too-few-public-methods The dimensionality of the output space (i.e. the number of output filters in the convolution). """ - def __init__(self, filters): + def __init__(self, filters: int) -> None: self._filters = filters - def __call__(self, inputs, feats): + def __call__(self, inputs: Tensor, feats: int) -> Tensor: """ Call the Attention Refinement block. Parameters @@ -400,7 +413,7 @@ class ContextPath(): # pylint:disable=too-few-public-methods def __init__(self): self._resnet = ResNet18() - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the Context Path block. Parameters @@ -443,10 +456,10 @@ class FeatureFusionModule(): # pylint:disable=too-few-public-methods The dimensionality of the output space (i.e. the number of output filters in the convolution). """ - def __init__(self, filters): + def __init__(self, filters: int) -> None: self._filters = filters - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the Feature Fusion block. Parameters @@ -491,12 +504,12 @@ class BiSeNetOutput(): # pylint:disable=too-few-public-methods label, str, optional The label for this output (for naming). Default: `""` (i.e. empty string, or no label) """ - def __init__(self, filters, num_classes, label=""): + def __init__(self, filters: int, num_classes: int, label: str = "") -> None: self._filters = filters self._num_classes = num_classes self._label = label - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the BiSeNet Output block. Parameters @@ -535,19 +548,28 @@ class BiSeNet(KSession): The input size to the model num_classes: int The number of segmentation classes to create + cpu_mode: bool, optional + ``True`` run the model on CPU. Default: ``False`` """ - def __init__(self, model_path, allow_growth, exclude_gpus, input_size, num_classes): + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None, + input_size: int, + num_classes: int, + cpu_mode: bool) -> None: super().__init__("BiSeNet Face Parsing", model_path, allow_growth=allow_growth, - exclude_gpus=exclude_gpus) + exclude_gpus=exclude_gpus, + cpu_mode=cpu_mode) self._input_size = input_size self._num_classes = num_classes self._cp = ContextPath() self.define_model(self._model_definition) self.load_model_weights() - def _model_definition(self): + def _model_definition(self) -> tuple[Tensor, list[Tensor]]: """ Definition of the VGG Obstructed Model. Returns @@ -558,12 +580,12 @@ def _model_definition(self): """ input_ = Input((self._input_size, self._input_size, 3)) - feat_res8, feat_cp8, feat_cp16 = self._cp(input_) - feat_fuse = FeatureFusionModule(256)([feat_res8, feat_cp8]) + features = self._cp(input_) # res8, cp8, cp16 + feat_fuse = FeatureFusionModule(256)([features[0], features[1]]) feat_out = BiSeNetOutput(256, self._num_classes)(feat_fuse) - feat_out16 = BiSeNetOutput(64, self._num_classes, label="16")(feat_cp8) - feat_out32 = BiSeNetOutput(64, self._num_classes, label="32")(feat_cp16) + feat_out16 = BiSeNetOutput(64, self._num_classes, label="16")(features[1]) + feat_out32 = BiSeNetOutput(64, self._num_classes, label="32")(features[2]) height, width = K.int_shape(input_)[1:3] f_h, f_w = K.int_shape(feat_out)[1:3] diff --git a/plugins/extract/mask/bisenet_fp_defaults.py b/plugins/extract/mask/bisenet_fp_defaults.py index ab556299e5..3b0ae79b92 100644 --- a/plugins/extract/mask/bisenet_fp_defaults.py +++ b/plugins/extract/mask/bisenet_fp_defaults.py @@ -50,50 +50,58 @@ _DEFAULTS = { - "batch-size": dict( - default=8, - info="The batch size to use. To a point, higher batch sizes equal better performance, but " - "setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can accomodate " - "then this will automatically be lowered.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True), - "weights": dict( - default="faceswap", - info="The trained weights to use.\n" - "\n\tfaceswap - Weights trained on wildly varied Faceswap extracted data to better " - "handle varying conditions, obstructions, glasses and multiple targets within a " - "single extracted image." - "\n\toriginal - The original weights trained on the CelebAMask-HQ dataset.", - choices=["faceswap", "original"], - datatype=str, - group="settings", - gui_radio=True, - ), - "include_ears": dict( - default=False, - info="Whether to include ears within the face mask.", - datatype=bool, - group="settings" - ), - "include_hair": dict( - default=False, - info="Whether to include hair within the face mask.", - datatype=bool, - group="settings" - ), - "include_glasses": dict( - default=True, - info="Whether to include glasses within the face mask.\n\tFor 'original' weights " - "excluding glasses will mask out the lenses as well as the frames.\n\tFor 'faceswap' " - "weights, the model has been trained to mask out lenses if eyes cannot be seen (i.e. " - "dark sunglasses) or just the frames if the eyes can be seen. ", - datatype=bool, - group="settings" - ), + "batch-size": { + "default": 8, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True + }, + "cpu": { + "default": False, + "info": "BiseNet mask still runs fairly quickly on CPU on some setups. Enable " + "CPU mode here to use the CPU for this masker to save some VRAM at a speed cost.", + "datatype": bool, + "group": "settings" + }, + "weights": { + "default": "faceswap", + "info": "The trained weights to use.\n" + "\n\tfaceswap - Weights trained on wildly varied Faceswap extracted data to " + "better handle varying conditions, obstructions, glasses and multiple targets " + "within a single extracted image." + "\n\toriginal - The original weights trained on the CelebAMask-HQ dataset.", + "choices": ["faceswap", "original"], + "datatype": str, + "group": "settings", + "gui_radio": True, + }, + "include_ears": { + "default": False, + "info": "Whether to include ears within the face mask.", + "datatype": bool, + "group": "settings" + }, + "include_hair": { + "default": False, + "info": "Whether to include hair within the face mask.", + "datatype": bool, + "group": "settings" + }, + "include_glasses": { + "default": True, + "info": "Whether to include glasses within the face mask.\n\tFor 'original' weights " + "excluding glasses will mask out the lenses as well as the frames.\n\tFor " + "'faceswap' weights, the model has been trained to mask out lenses if eyes cannot " + "be seen (i.e. dark sunglasses) or just the frames if the eyes can be seen.", + "datatype": bool, + "group": "settings" + }, } diff --git a/plugins/extract/mask/components.py b/plugins/extract/mask/components.py index 6f9d1474f0..0a71af4866 100644 --- a/plugins/extract/mask/components.py +++ b/plugins/extract/mask/components.py @@ -1,14 +1,25 @@ #!/usr/bin/env python3 """ Components Mask for faceswap.py """ +from __future__ import annotations +import logging +import typing as T import cv2 import numpy as np -from ._base import Masker, logger + +from lib.align import LandmarkType + +from ._base import BatchType, Masker + +if T.TYPE_CHECKING: + from lib.align.aligned_face import AlignedFace + +logger = logging.getLogger(__name__) class Mask(Masker): """ Perform transformation to align and get landmarks """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = None model_filename = None super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) @@ -17,33 +28,38 @@ def __init__(self, **kwargs): self.vram = 0 # Doesn't use GPU self.vram_per_batch = 0 self.batchsize = 1 + self.landmark_type = LandmarkType.LM_2D_68 - def init_model(self): + def init_model(self) -> None: logger.debug("No mask model to initialize") - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ - batch["feed"] = np.zeros((self.batchsize, self.input_size, self.input_size, 1), - dtype="float32") - return batch + batch.feed = np.zeros((self.batchsize, self.input_size, self.input_size, 1), + dtype="float32") - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - for mask, face in zip(batch["feed"], batch["feed_faces"]): + faces: list[AlignedFace] = feed[1] + feed = feed[0] + for mask, face in zip(feed, faces): + if LandmarkType.from_shape(face.landmarks.shape) != self.landmark_type: + # Called from the manual tool. # TODO This will only work with BS1 + feed = np.zeros_like(feed) + continue parts = self.parse_parts(np.array(face.landmarks)) for item in parts: - item = np.rint(np.concatenate(item)).astype("int32") - hull = cv2.convexHull(item) - cv2.fillConvexPoly(mask, hull, 1.0, lineType=cv2.LINE_AA) - batch["prediction"] = batch["feed"] - return batch + a_item = np.rint(np.concatenate(item)).astype("int32") + hull = cv2.convexHull(a_item) + cv2.fillConvexPoly(mask, hull, [1.0], lineType=cv2.LINE_AA) + return feed - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return @staticmethod - def parse_parts(landmarks): + def parse_parts(landmarks: np.ndarray) -> list[tuple[np.ndarray, ...]]: """ Component face hull mask """ r_jaw = (landmarks[0:9], landmarks[17:18]) l_jaw = (landmarks[8:17], landmarks[26:27]) diff --git a/plugins/extract/mask/custom.py b/plugins/extract/mask/custom.py new file mode 100644 index 0000000000..b1e3328471 --- /dev/null +++ b/plugins/extract/mask/custom.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +""" Components Mask for faceswap.py """ +import logging +import numpy as np +from ._base import BatchType, Masker + +logger = logging.getLogger(__name__) + + +class Mask(Masker): + """ A mask that fills the whole face area with 1s or 0s (depending on user selected settings) + for custom editing. """ + def __init__(self, **kwargs): + git_model_id = None + model_filename = None + super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.input_size = 256 + self.name = "Custom" + self.vram = 0 # Doesn't use GPU + self.vram_per_batch = 0 + self.batchsize = self.config["batch-size"] + self._storage_centering = self.config["centering"] + # Separate storage for face and head masks + self._storage_name = f"{self._storage_name}_{self._storage_centering}" + + def init_model(self) -> None: + logger.debug("No mask model to initialize") + + def process_input(self, batch: BatchType) -> None: + """ Compile the detected faces for prediction """ + batch.feed = np.zeros((self.batchsize, self.input_size, self.input_size, 1), + dtype="float32") + + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Run model to get predictions """ + if self.config["fill"]: + feed[:] = 1.0 + return feed + + def process_output(self, batch: BatchType) -> None: + """ Compile found faces for output """ + return diff --git a/plugins/extract/mask/custom_defaults.py b/plugins/extract/mask/custom_defaults.py new file mode 100644 index 0000000000..9da35416f5 --- /dev/null +++ b/plugins/extract/mask/custom_defaults.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" + The default options for the faceswap BiSeNet Face Parsing plugin. + + Defaults files should be named _defaults.py + Any items placed into this file will automatically get added to the relevant config .ini files + within the faceswap/config folder. + + The following variables should be defined: + _HELPTEXT: A string describing what this plugin does + _DEFAULTS: A dictionary containing the options, defaults and meta information. The + dictionary should be defined as: + {: {}} + + should always be lower text. + dictionary requirements are listed below. + + The following keys are expected for the _DEFAULTS dict: + datatype: [required] A python type class. This limits the type of data that can be + provided in the .ini file and ensures that the value is returned in the + correct type to faceswap. Valid data types are: , , + , . + default: [required] The default value for this option. + info: [required] A string describing what this option does. + group: [optional]. A group for grouping options together in the GUI. If not + provided this will not group this option with any others. + choices: [optional] If this option's datatype is of then valid + selections can be defined here. This validates the option and also enables + a combobox / radio option in the GUI. + gui_radio: [optional] If are defined, this indicates that the GUI should use + radio buttons rather than a combobox to display this option. + min_max: [partial] For and data types this is required + otherwise it is ignored. Should be a tuple of min and max accepted values. + This is used for controlling the GUI slider range. Values are not enforced. + rounding: [partial] For and data types this is + required otherwise it is ignored. Used for the GUI slider. For floats, this + is the number of decimal places to display. For ints this is the step size. + fixed: [optional] [train only]. Training configurations are fixed when the model is + created, and then reloaded from the state file. Marking an item as fixed=False + indicates that this value can be changed for existing models, and will override + the value saved in the state file with the updated value in config. If not + provided this will default to True. +""" + + +_HELPTEXT = ( + "Custom (dummy) Mask options..\n" + "The custom mask just fills a face patch with all 0's (masked out) or all 1's (masked in) for " + "later manual editing. It does not use the GPU for creation." + ) + + +_DEFAULTS = { + "batch-size": { + "default": 8, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "group": "settings" + }, + "centering": { + "default": "face", + "info": "Whether to create a dummy mask with face or head centering.", + "choices": ["face", "head"], + "datatype": str, + "group": "settings", + "gui_radio": True + }, + "fill": { + "default": False, + "info": "Whether the mask should be filled (True) in which case the custom mask will be " + "created with the whole area masked in (i.e. you would need to manually edit out " + "the background) or unfilled (False) in which case you would need to manually " + "edit in the face.", + "datatype": bool, + "group": "settings", + "gui_radio": True, + }, +} diff --git a/plugins/extract/mask/extended.py b/plugins/extract/mask/extended.py index 11df6a5856..d6970cb0e5 100644 --- a/plugins/extract/mask/extended.py +++ b/plugins/extract/mask/extended.py @@ -1,9 +1,20 @@ #!/usr/bin/env python3 """ Extended Mask for faceswap.py """ +from __future__ import annotations +import logging +import typing as T import cv2 import numpy as np -from ._base import Masker, logger + +from lib.align import LandmarkType + +from ._base import BatchType, Masker + +logger = logging.getLogger(__name__) + +if T.TYPE_CHECKING: + from lib.align.aligned_face import AlignedFace class Mask(Masker): @@ -17,34 +28,45 @@ def __init__(self, **kwargs): self.vram = 0 # Doesn't use GPU self.vram_per_batch = 0 self.batchsize = 1 + self.landmark_type = LandmarkType.LM_2D_68 - def init_model(self): + def init_model(self) -> None: logger.debug("No mask model to initialize") - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ - batch["feed"] = np.zeros((self.batchsize, self.input_size, self.input_size, 1), - dtype="float32") - return batch + batch.feed = np.zeros((self.batchsize, self.input_size, self.input_size, 1), + dtype="float32") - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - for mask, face in zip(batch["feed"], batch["feed_faces"]): + faces: list[AlignedFace] = feed[1] + feed = feed[0] + for mask, face in zip(feed, faces): + if LandmarkType.from_shape(face.landmarks.shape) != self.landmark_type: + # Called from the manual tool. # TODO This will only work with BS1 + feed = np.zeros_like(feed) + continue parts = self.parse_parts(np.array(face.landmarks)) for item in parts: - item = np.rint(np.concatenate(item)).astype("int32") - hull = cv2.convexHull(item) - cv2.fillConvexPoly(mask, hull, 1.0, lineType=cv2.LINE_AA) - batch["prediction"] = batch["feed"] - return batch + a_item = np.rint(np.concatenate(item)).astype("int32") + hull = cv2.convexHull(a_item) + cv2.fillConvexPoly(mask, hull, [1.0], lineType=cv2.LINE_AA) + return feed - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return - @staticmethod - def parse_parts(landmarks): - """ Extended face hull mask """ + @classmethod + def _adjust_mask_top(cls, landmarks: np.ndarray) -> None: + """ Adjust the top of the mask to extend above eyebrows + + Parameters + ---------- + landmarks: :class:`numpy.ndarray` + The 68 point landmarks to be adjusted + """ # mid points between the side of face and eye point ml_pnt = (landmarks[36] + landmarks[0]) // 2 mr_pnt = (landmarks[16] + landmarks[45]) // 2 @@ -65,6 +87,10 @@ def parse_parts(landmarks): landmarks[17:22] = top_l + ((top_l - bot_l) // 2) landmarks[22:27] = top_r + ((top_r - bot_r) // 2) + def parse_parts(self, landmarks: np.ndarray) -> list[tuple[np.ndarray, ...]]: + """ Extended face hull mask """ + self._adjust_mask_top(landmarks) + r_jaw = (landmarks[0:9], landmarks[17:18]) l_jaw = (landmarks[8:17], landmarks[26:27]) r_cheek = (landmarks[17:20], landmarks[8:9]) diff --git a/plugins/extract/mask/unet_dfl.py b/plugins/extract/mask/unet_dfl.py index 9dd45613ad..4ca2f3dc07 100644 --- a/plugins/extract/mask/unet_dfl.py +++ b/plugins/extract/mask/unet_dfl.py @@ -12,18 +12,23 @@ Model file sourced from... https://github.com/iperov/DeepFaceLab/blob/master/nnlib/FANSeg_256_full_face.h5 """ +import logging +import typing as T import numpy as np from lib.model.session import KSession -from ._base import Masker, logger +from ._base import BatchType, Masker, MaskerBatch + +logger = logging.getLogger(__name__) class Mask(Masker): """ Neural network to process face image into a segmentation mask of the face """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 6 model_filename = "DFL_256_sigmoid_v1.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession self.name = "U-Net" self.input_size = 256 self.vram = 3424 @@ -32,10 +37,11 @@ def __init__(self, **kwargs): self.batchsize = self.config["batch-size"] self._storage_centering = "legacy" - def init_model(self): + def init_model(self) -> None: + assert self.name is not None and isinstance(self.model_path, str) self.model = KSession(self.name, self.model_path, - model_kwargs=dict(), + model_kwargs={}, allow_growth=self.config["allow_growth"], exclude_gpus=self._exclude_gpus) self.model.load_model() @@ -43,18 +49,19 @@ def init_model(self): dtype="float32") self.model.predict(placeholder) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ - batch["feed"] = np.array([feed.face[..., :3] - for feed in batch["feed_faces"]], dtype="float32") / 255.0 - logger.trace("feed shape: %s", batch["feed"].shape) - return batch + assert isinstance(batch, MaskerBatch) + batch.feed = np.array([T.cast(np.ndarray, feed.face)[..., :3] + for feed in batch.feed_faces], dtype="float32") / 255.0 + logger.trace("feed shape: %s", batch.feed.shape) # type: ignore - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - batch["prediction"] = self.model.predict(batch["feed"]) - return batch + retval = self.model.predict(feed) + assert isinstance(retval, np.ndarray) + return retval - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return diff --git a/plugins/extract/mask/unet_dfl_defaults.py b/plugins/extract/mask/unet_dfl_defaults.py index 1a3fb81890..62514c0188 100644 --- a/plugins/extract/mask/unet_dfl_defaults.py +++ b/plugins/extract/mask/unet_dfl_defaults.py @@ -51,18 +51,18 @@ _DEFAULTS = { - "batch-size": dict( - default=8, - info="The batch size to use. To a point, higher batch sizes equal better performance, but " - "setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can accomodate " - "then this will automatically be lowered.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ) + "batch-size": { + "default": 8, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + } } diff --git a/plugins/extract/mask/vgg_clear.py b/plugins/extract/mask/vgg_clear.py index 2b37f9b118..50165f8015 100644 --- a/plugins/extract/mask/vgg_clear.py +++ b/plugins/extract/mask/vgg_clear.py @@ -1,29 +1,32 @@ #!/usr/bin/env python3 """ VGG Clear face mask plugin. """ +from __future__ import annotations +import logging +import typing as T import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import ( # pylint:disable=import-error + Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, + ZeroPadding2D) + from lib.model.session import KSession -from lib.utils import get_backend -from ._base import Masker, logger +from ._base import BatchType, Masker, MaskerBatch + +if T.TYPE_CHECKING: + from tensorflow import Tensor -if get_backend() == "amd": - from keras.layers import ( - Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, - ZeroPadding2D) -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import ( # pylint:disable=no-name-in-module,import-error - Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, - ZeroPadding2D) +logger = logging.getLogger(__name__) class Mask(Masker): """ Neural network to process face image into a segmentation mask of the face """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 8 model_filename = "Nirkin_300_softmax_v1.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession self.name = "VGG Clear" self.input_size = 300 self.vram = 2944 @@ -31,7 +34,8 @@ def __init__(self, **kwargs): self.vram_per_batch = 400 self.batchsize = self.config["batch-size"] - def init_model(self): + def init_model(self) -> None: + assert isinstance(self.model_path, str) self.model = VGGClear(self.model_path, allow_growth=self.config["allow_growth"], exclude_gpus=self._exclude_gpus) @@ -40,23 +44,23 @@ def init_model(self): dtype="float32") self.model.predict(placeholder) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ - input_ = np.array([feed.face[..., :3] - for feed in batch["feed_faces"]], dtype="float32") - batch["feed"] = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :] - logger.trace("feed shape: %s", batch["feed"].shape) - return batch + assert isinstance(batch, MaskerBatch) + input_ = np.array([T.cast(np.ndarray, feed.face)[..., :3] + for feed in batch.feed_faces], dtype="float32") + batch.feed = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :] + logger.trace("feed shape: %s", batch.feed.shape) # type: ignore - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - predictions = self.model.predict(batch["feed"]) - batch["prediction"] = predictions[..., -1] - return batch + predictions = self.model.predict(feed) + assert isinstance(predictions, np.ndarray) + return predictions[..., -1] - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return class VGGClear(KSession): @@ -87,7 +91,10 @@ class VGGClear(KSession): https://github.com/YuvalNirkin/face_segmentation/releases/download/1.1/face_seg_fcn8s_300_no_aug.zip """ - def __init__(self, model_path, allow_growth, exclude_gpus): + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None): super().__init__("VGG Obstructed", model_path, allow_growth=allow_growth, @@ -96,7 +103,7 @@ def __init__(self, model_path, allow_growth, exclude_gpus): self.load_model_weights() @classmethod - def _model_definition(cls): + def _model_definition(cls) -> tuple[Tensor, Tensor]: """ Definition of the VGG Obstructed Model. Returns @@ -158,13 +165,13 @@ class _ConvBlock(): # pylint:disable=too-few-public-methods iterations: int The number of consecutive Conv2D layers to create """ - def __init__(self, level, filters, iterations): + def __init__(self, level: int, filters: int, iterations: int) -> None: self._name = f"conv{level}_" self._level = level self._filters = filters self._iterator = range(1, iterations + 1) - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the convolutional loop. Parameters @@ -203,12 +210,12 @@ class _ScorePool(): # pylint:disable=too-few-public-methods crop: tuple The amount of 2D cropping to apply. Tuple of `ints` """ - def __init__(self, level, scale, crop): + def __init__(self, level: int, scale: float, crop: tuple[int, int]): self._name = f"_pool{level}" self._cropping = (crop, crop) self._scale = scale - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Score pool block. Parameters diff --git a/plugins/extract/mask/vgg_clear_defaults.py b/plugins/extract/mask/vgg_clear_defaults.py index b9592c5b12..48c5d1f428 100644 --- a/plugins/extract/mask/vgg_clear_defaults.py +++ b/plugins/extract/mask/vgg_clear_defaults.py @@ -50,18 +50,18 @@ _DEFAULTS = { - "batch-size": dict( - default=6, - info="The batch size to use. To a point, higher batch sizes equal better performance, but " - "setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can accomodate " - "then this will automatically be lowered.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ) + "batch-size": { + "default": 6, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + } } diff --git a/plugins/extract/mask/vgg_obstructed.py b/plugins/extract/mask/vgg_obstructed.py index ba8a596e84..a3f543d7e8 100644 --- a/plugins/extract/mask/vgg_obstructed.py +++ b/plugins/extract/mask/vgg_obstructed.py @@ -1,30 +1,32 @@ #!/usr/bin/env python3 """ VGG Obstructed face mask plugin """ +from __future__ import annotations +import logging +import typing as T import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import ( # pylint:disable=import-error + Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, + ZeroPadding2D) from lib.model.session import KSession -from lib.utils import get_backend -from ._base import Masker, logger +from ._base import BatchType, Masker, MaskerBatch -if get_backend() == "amd": - from keras.layers import ( - Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, - ZeroPadding2D) -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import ( # pylint:disable=no-name-in-module,import-error - Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D, - ZeroPadding2D) +if T.TYPE_CHECKING: + from tensorflow import Tensor + +logger = logging.getLogger(__name__) class Mask(Masker): """ Neural network to process face image into a segmentation mask of the face """ - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: git_model_id = 5 model_filename = "Nirkin_500_softmax_v1.h5" super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession self.name = "VGG Obstructed" self.input_size = 500 self.vram = 3936 @@ -32,7 +34,8 @@ def __init__(self, **kwargs): self.vram_per_batch = 304 self.batchsize = self.config["batch-size"] - def init_model(self): + def init_model(self) -> None: + assert isinstance(self.model_path, str) self.model = VGGObstructed(self.model_path, allow_growth=self.config["allow_growth"], exclude_gpus=self._exclude_gpus) @@ -41,22 +44,22 @@ def init_model(self): dtype="float32") self.model.predict(placeholder) - def process_input(self, batch): + def process_input(self, batch: BatchType) -> None: """ Compile the detected faces for prediction """ - input_ = [feed.face[..., :3] for feed in batch["feed_faces"]] - batch["feed"] = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :] - logger.trace("feed shape: %s", batch["feed"].shape) - return batch + assert isinstance(batch, MaskerBatch) + input_ = [T.cast(np.ndarray, feed.face)[..., :3] for feed in batch.feed_faces] + batch.feed = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :] + logger.trace("feed shape: %s", batch.feed.shape) # type:ignore - def predict(self, batch): + def predict(self, feed: np.ndarray) -> np.ndarray: """ Run model to get predictions """ - predictions = self.model.predict(batch["feed"]) - batch["prediction"] = predictions[..., 0] * -1.0 + 1.0 - return batch + predictions = self.model.predict(feed) + assert isinstance(predictions, np.ndarray) + return predictions[..., 0] * -1.0 + 1.0 - def process_output(self, batch): + def process_output(self, batch: BatchType) -> None: """ Compile found faces for output """ - return batch + return class VGGObstructed(KSession): @@ -84,7 +87,10 @@ class VGGObstructed(KSession): Model file sourced from: https://github.com/YuvalNirkin/face_segmentation/releases/download/1.0/face_seg_fcn8s.zip """ - def __init__(self, model_path, allow_growth, exclude_gpus): + def __init__(self, + model_path: str, + allow_growth: bool, + exclude_gpus: list[int] | None) -> None: super().__init__("VGG Obstructed", model_path, allow_growth=allow_growth, @@ -93,7 +99,7 @@ def __init__(self, model_path, allow_growth, exclude_gpus): self.load_model_weights() @classmethod - def _model_definition(cls): + def _model_definition(cls) -> tuple[Tensor, Tensor]: """ Definition of the VGG Obstructed Model. Returns @@ -158,13 +164,13 @@ class _ConvBlock(): # pylint:disable=too-few-public-methods iterations: int The number of consecutive Conv2D layers to create """ - def __init__(self, level, filters, iterations): + def __init__(self, level: int, filters: int, iterations: int) -> None: self._name = f"conv{level}_" self._level = level self._filters = filters self._iterator = range(1, iterations + 1) - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Call the convolutional loop. Parameters @@ -203,12 +209,12 @@ class _ScorePool(): # pylint:disable=too-few-public-methods crop: int The amount of 2D cropping to apply """ - def __init__(self, level, scale, crop): + def __init__(self, level: int, scale: float, crop: int) -> None: self._name = f"_pool{level}" self._cropping = ((crop, crop), (crop, crop)) self._scale = scale - def __call__(self, inputs): + def __call__(self, inputs: Tensor) -> Tensor: """ Score pool block. Parameters diff --git a/plugins/extract/mask/vgg_obstructed_defaults.py b/plugins/extract/mask/vgg_obstructed_defaults.py index a4ca3e28af..7d19354289 100644 --- a/plugins/extract/mask/vgg_obstructed_defaults.py +++ b/plugins/extract/mask/vgg_obstructed_defaults.py @@ -51,18 +51,18 @@ _DEFAULTS = { - "batch-size": dict( - default=2, - info="The batch size to use. To a point, higher batch sizes equal better performance, but " - "setting it too high can harm performance.\n" - "\n\tNvidia users: If the batchsize is set higher than the your GPU can accomodate " - "then this will automatically be lowered.", - datatype=int, - rounding=1, - min_max=(1, 64), - choices=[], - group="settings", - gui_radio=False, - fixed=True, - ) + "batch-size": { + "default": 2, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True, + } } diff --git a/plugins/extract/pipeline.py b/plugins/extract/pipeline.py index a03cc1945c..5a051936bb 100644 --- a/plugins/extract/pipeline.py +++ b/plugins/extract/pipeline.py @@ -7,19 +7,32 @@ This module sets up a pipeline for the extraction workflow, loading detect, align and mask plugins either in parallel or in series, giving easy access to input and output. - - """ - +""" +from __future__ import annotations import logging +import os +import typing as T -import cv2 - +from lib.align import LandmarkType from lib.gpu_stats import GPUStats -from lib.queue_manager import queue_manager, QueueEmpty -from lib.utils import get_backend +from lib.logger import parse_class_init +from lib.queue_manager import EventQueue, queue_manager, QueueEmpty +from lib.serializer import get_serializer +from lib.utils import get_backend, FaceswapError from plugins.plugin_loader import PluginLoader -logger = logging.getLogger(__name__) # pylint:disable=invalid-name +if T.TYPE_CHECKING: + from collections.abc import Generator + from ._base import Extractor as PluginExtractor + from .align._base import Aligner + from .align.external import Align as AlignImport + from .detect._base import Detector + from .detect.external import Detect as DetectImport + from .mask._base import Masker + from .recognition._base import Identity + from . import ExtractMedia + +logger = logging.getLogger(__name__) _INSTANCES = -1 # Tracking for multiple instances of pipeline @@ -39,13 +52,16 @@ class Extractor(): Parameters ---------- - detector: str + detector: str or ``None`` The name of a detector plugin as exists in :mod:`plugins.extract.detect` - aligner: str + aligner: str or ``None`` The name of an aligner plugin as exists in :mod:`plugins.extract.align` - masker: str or list + masker: str or list or ``None`` The name of a masker plugin(s) as exists in :mod:`plugins.extract.mask`. This can be a single masker or a list of multiple maskers + recognition: str or ``None`` + The name of the recognition plugin to use. ``None`` to not do face recognition. + Default: ``None`` configfile: str, optional The path to a custom ``extract.ini`` configfile. If ``None`` then the system :file:`config/extract.ini` file will be used. @@ -61,7 +77,7 @@ class Extractor(): exactly what angles to check. Can also pass in ``'on'`` to increment at 90 degree intervals. Default: ``None`` min_size: int, optional - Used to set the :attr:`plugins.extract.detect.min_size` attribute Filters out faces + Used to set the :attr:`plugins.extract.detect.min_size` attribute. Filters out faces detected below this size. Length, in pixels across the diagonal of the bounding box. Set to ``0`` for off. Default: ``0`` normalize_method: {`None`, 'clahe', 'hist', 'mean'}, optional @@ -70,9 +86,11 @@ class Extractor(): re_feed: int The number of times to re-feed a slightly adjusted bounding box into the aligner. Default: `0` - image_is_aligned: bool, optional - Used to set the :attr:`plugins.extract.mask.image_is_aligned` attribute. Indicates to the - masker that the fed in image is an aligned face rather than a frame. Default: ``False`` + re_align: bool, optional + ``True`` to obtain landmarks by passing the initially aligned face back through the + aligner. Default ``False`` + disable_filter: bool, optional + Disable all aligner filters regardless of config option. Default: ``False`` Attributes ---------- @@ -80,18 +98,26 @@ class Extractor(): The current phase that the pipeline is running. Used in conjunction with :attr:`passes` and :attr:`final_pass` to indicate to the caller which phase is being processed """ - def __init__(self, detector, aligner, masker, configfile=None, multiprocess=False, - exclude_gpus=None, rotate_images=None, min_size=20, normalize_method=None, - re_feed=0, image_is_aligned=False): - logger.debug("Initializing %s: (detector: %s, aligner: %s, masker: %s, configfile: %s, " - "multiprocess: %s, exclude_gpus: %s, rotate_images: %s, min_size: %s, " - "normalize_method: %s, re_feed: %s, image_is_aligned: %s)", - self.__class__.__name__, detector, aligner, masker, configfile, multiprocess, - exclude_gpus, rotate_images, min_size, normalize_method, re_feed, - image_is_aligned) + def __init__(self, + detector: str | None, + aligner: str | None, + masker: str | list[str] | None, + recognition: str | None = None, + configfile: str | None = None, + multiprocess: bool = False, + exclude_gpus: list[int] | None = None, + rotate_images: str | None = None, + min_size: int = 0, + normalize_method: T.Literal["none", "clahe", "hist", "mean"] | None = None, + re_feed: int = 0, + re_align: bool = False, + disable_filter: bool = False) -> None: + logger.debug(parse_class_init(locals())) self._instance = _get_instance() - masker = [masker] if not isinstance(masker, list) else masker - self._flow = self._set_flow(detector, aligner, masker) + maskers = [T.cast(str | None, + masker)] if not isinstance(masker, list) else T.cast(list[str | None], + masker) + self._flow = self._set_flow(detector, aligner, maskers, recognition) self._exclude_gpus = exclude_gpus # We only ever need 1 item in each queue. This is 2 items cached (1 in queue 1 waiting # for queue) at each point. Adding more just stacks RAM with no speed benefit. @@ -99,9 +125,15 @@ def __init__(self, detector, aligner, masker, configfile=None, multiprocess=Fals # TODO Calculate scaling for more plugins than currently exist in _parallel_scaling self._scaling_fallback = 0.4 self._vram_stats = self._get_vram_stats() - self._detect = self._load_detect(detector, rotate_images, min_size, configfile) - self._align = self._load_align(aligner, configfile, normalize_method, re_feed) - self._mask = [self._load_mask(mask, image_is_aligned, configfile) for mask in masker] + self._detect = self._load_detect(detector, aligner, rotate_images, min_size, configfile) + self._align = self._load_align(aligner, + configfile, + normalize_method, + re_feed, + re_align, + disable_filter) + self._recognition = self._load_recognition(recognition, configfile) + self._mask = [self._load_mask(mask, configfile) for mask in maskers] self._is_parallel = self._set_parallel_processing(multiprocess) self._phases = self._set_phases(multiprocess) self._phase_index = 0 @@ -110,7 +142,7 @@ def __init__(self, detector, aligner, masker, configfile=None, multiprocess=Fals logger.debug("Initialized %s", self.__class__.__name__) @property - def input_queue(self): + def input_queue(self) -> EventQueue: """ queue: Return the correct input queue depending on the current phase The input queue is the entry point into the extraction pipeline. An :class:`ExtractMedia` @@ -122,13 +154,13 @@ def input_queue(self): For align/mask (2nd/3rd pass operations) the :attr:`ExtractMedia.detected_faces` should also be populated by calling :func:`ExtractMedia.set_detected_faces`. """ - qname = "extract{}_{}_in".format(self._instance, self._current_phase[0]) + qname = f"extract{self._instance}_{self._current_phase[0]}_in" retval = self._queues[qname] - logger.trace("%s: %s", qname, retval) + logger.trace("%s: %s", qname, retval) # type: ignore return retval @property - def passes(self): + def passes(self) -> int: """ int: Returns the total number of passes the extractor needs to make. This is calculated on several factors (vram available, plugin choice, @@ -146,21 +178,21 @@ def passes(self): >>> extractor.input_queue.put(extract_media) """ retval = len(self._phases) - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def phase_text(self): + def phase_text(self) -> str: """ str: The plugins that are running in the current phase, formatted for info text output. """ plugin_types = set(self._get_plugin_type_and_index(phase)[0] for phase in self._current_phase) retval = ", ".join(plugin_type.title() for plugin_type in list(plugin_types)) - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def final_pass(self): + def final_pass(self) -> bool: """ bool, Return ``True`` if this is the final extractor pass otherwise ``False`` Useful for iterating over the pipeline :attr:`passes` or :func:`detected_faces` and @@ -177,26 +209,45 @@ def final_pass(self): >>> extractor.input_queue.put(extract_media) """ retval = self._phase_index == len(self._phases) - 1 - logger.trace(retval) + logger.trace(retval) # type:ignore[attr-defined] return retval - def set_batchsize(self, plugin_type, batchsize): + @property + def aligner(self) -> Aligner: + """ The currently selected aligner plugin """ + assert self._align is not None + return self._align + + @property + def recognition(self) -> Identity: + """ The currently selected recognition plugin """ + assert self._recognition is not None + return self._recognition + + def reset_phase_index(self) -> None: + """ Reset the current phase index back to 0. Used for when batch processing is used in + extract. """ + self._phase_index = 0 + + def set_batchsize(self, + plugin_type: T.Literal["align", "detect"], + batchsize: int) -> None: """ Set the batch size of a given :attr:`plugin_type` to the given :attr:`batchsize`. This should be set prior to :func:`launch` if the batch size is to be manually overridden Parameters ---------- - plugin_type: {'aligner', 'detector'} + plugin_type: {'align', 'detect'} The plugin_type to be overridden batchsize: int The batch size to use for this plugin type """ logger.debug("Overriding batchsize for plugin_type: %s to: %s", plugin_type, batchsize) - plugin = getattr(self, "_{}".format(plugin_type)) + plugin = getattr(self, f"_{plugin_type}") plugin.batchsize = batchsize - def launch(self): + def launch(self) -> None: """ Launches the plugin(s) This launches the plugins held in the pipeline, and should be called at the beginning @@ -212,7 +263,7 @@ def launch(self): for phase in self._current_phase: self._launch_plugin(phase) - def detected_faces(self): + def detected_faces(self) -> Generator[ExtractMedia, None, None]: """ Generator that returns results, frame by frame from the extraction pipeline This is the exit point for the extraction pipeline and is used to obtain the output @@ -220,7 +271,7 @@ def detected_faces(self): Yields ------ - faces: :class:`ExtractMedia` + faces: :class:`~plugins.extract.extract_media.ExtractMedia` The populated extracted media object. Example @@ -236,8 +287,7 @@ def detected_faces(self): out_queue = self._output_queue while True: try: - if self._check_and_raise_error(): - break + self._check_and_raise_error() faces = out_queue.get(True, 1) if faces == "EOF": break @@ -247,17 +297,92 @@ def detected_faces(self): self._join_threads() if self.final_pass: - # Cleanup queues - for q_name in self._queues: - queue_manager.del_queue(q_name) + for plugin in self._all_plugins: + plugin.on_completion() logger.debug("Detection Complete") else: self._phase_index += 1 logger.debug("Switching to phase: %s", self._current_phase) + def _disable_lm_maskers(self) -> None: + """ Disable any 68 point landmark based maskers if alignment data is not 2D 68 + point landmarks and update the process flow/phases accordingly """ + logger.warning("Alignment data is not 68 point 2D landmarks. Some Faceswap functionality " + "will be unavailable for these faces") + + rem_maskers = [m.name for m in self._mask + if m is not None and m.landmark_type == LandmarkType.LM_2D_68] + self._mask = [m for m in self._mask if m is None or m.name not in rem_maskers] + + self._flow = [ + item for item in self._flow + if not item.startswith("mask") + or item.startswith("mask") and int(item.rsplit("_", maxsplit=1)[-1]) < len(self._mask)] + + self._phases = [[s for s in p if s in self._flow] for p in self._phases + if any(t in p for t in self._flow)] + + for queue in self._queues: + queue_manager.del_queue(queue) + del self._queues + self._queues = self._add_queues() + + logger.warning("The following maskers have been disabled due to unsupported landmarks: %s", + rem_maskers) + + def import_data(self, input_location: str) -> None: + """ Import json data to the detector and/or aligner if 'import' plugin has been selected + + Parameters + ---------- + input_location: str + Full path to the input location for the extract process + """ + assert self._detect is not None + import_plugins: list[DetectImport | AlignImport] = [ + p for p in (self._detect, self.aligner) # type:ignore[misc] + if T.cast(str, p.name).lower() == "external"] + + if not import_plugins: + return + + align_origin = None + assert self.aligner.name is not None + if self.aligner.name.lower() == "external": + align_origin = self.aligner.config["origin"] + + logger.info("Importing external data for %s from json file...", + " and ".join([p.__class__.__name__ for p in import_plugins])) + + folder = input_location + folder = folder if os.path.isdir(folder) else os.path.dirname(folder) + + last_fname = "" + is_68_point = True + for plugin in import_plugins: + plugin_type = plugin.__class__.__name__ + path = os.path.join(folder, plugin.config["file_name"]) + if not os.path.isfile(path): + raise FaceswapError(f"{plugin_type} import file could not be found at '{path}'") + + if path != last_fname: # Different import file for aligner data + last_fname = path + data = get_serializer("json").load(path) + + if plugin_type == "Detect": + plugin.import_data(data, align_origin) # type:ignore[call-arg] + else: + plugin.import_data(data) # type:ignore[call-arg] + is_68_point = plugin.landmark_type == LandmarkType.LM_2D_68 # type:ignore[union-attr] # noqa:E501 # pylint:disable="line-too-long" + + if not is_68_point: + self._disable_lm_maskers() + + logger.info("Imported external data") + # <<< INTERNAL METHODS >>> # @property - def _parallel_scaling(self): + def _parallel_scaling(self) -> dict[int, float]: """ dict: key is number of parallel plugins being loaded, value is the scaling factor that the total base vram for those plugins should be scaled by @@ -277,23 +402,23 @@ def _parallel_scaling(self): 3: 0.55, 4: 0.5, 5: 0.4} - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def _vram_per_phase(self): + def _vram_per_phase(self) -> dict[str, float]: """ dict: The amount of vram required for each phase in :attr:`_flow`. """ - retval = dict() + retval = {} for phase in self._flow: plugin_type, idx = self._get_plugin_type_and_index(phase) - attr = getattr(self, "_{}".format(plugin_type)) + attr = getattr(self, f"_{plugin_type}") attr = attr[idx] if idx is not None else attr retval[phase] = attr.vram - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def _total_vram_required(self): + def _total_vram_required(self) -> float: """ Return vram required for all phases plus the buffer """ vrams = self._vram_per_phase vram_required_count = sum(1 for p in vrams.values() if p > 0) @@ -305,71 +430,89 @@ def _total_vram_required(self): return retval @property - def _current_phase(self): + def _current_phase(self) -> list[str]: """ list: The current phase from :attr:`_phases` that is running through the extractor. """ retval = self._phases[self._phase_index] - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def _final_phase(self): + def _final_phase(self) -> str: """ Return the final phase from the flow list """ retval = self._flow[-1] - logger.trace(retval) + logger.trace(retval) # type: ignore return retval @property - def _output_queue(self): + def _output_queue(self) -> EventQueue: """ Return the correct output queue depending on the current phase """ if self.final_pass: - qname = "extract{}_{}_out".format(self._instance, self._final_phase) + qname = f"extract{self._instance}_{self._final_phase}_out" else: - qname = "extract{}_{}_in".format(self._instance, - self._phases[self._phase_index + 1][0]) + qname = f"extract{self._instance}_{self._phases[self._phase_index + 1][0]}_in" retval = self._queues[qname] - logger.trace("%s: %s", qname, retval) + logger.trace("%s: %s", qname, retval) # type: ignore return retval @property - def _all_plugins(self): + def _all_plugins(self) -> list[PluginExtractor]: """ Return list of all plugin objects in this pipeline """ retval = [] for phase in self._flow: plugin_type, idx = self._get_plugin_type_and_index(phase) - attr = getattr(self, "_{}".format(plugin_type)) + attr = getattr(self, f"_{plugin_type}") attr = attr[idx] if idx is not None else attr retval.append(attr) - logger.trace("All Plugins: %s", retval) + logger.trace("All Plugins: %s", retval) # type: ignore return retval @property - def _active_plugins(self): + def _active_plugins(self) -> list[PluginExtractor]: """ Return the plugins that are currently active based on pass """ retval = [] for phase in self._current_phase: plugin_type, idx = self._get_plugin_type_and_index(phase) - attr = getattr(self, "_{}".format(plugin_type)) + attr = getattr(self, f"_{plugin_type}") retval.append(attr[idx] if idx is not None else attr) - logger.trace("Active plugins: %s", retval) + logger.trace("Active plugins: %s", retval) # type: ignore return retval @staticmethod - def _set_flow(detector, aligner, masker): - """ Set the flow list based on the input plugins """ - logger.debug("detector: %s, aligner: %s, masker: %s", detector, aligner, masker) + def _set_flow(detector: str | None, + aligner: str | None, + masker: list[str | None], + recognition: str | None) -> list[str]: + """ Set the flow list based on the input plugins + + Parameters + ---------- + detector: str or ``None`` + The name of a detector plugin as exists in :mod:`plugins.extract.detect` + aligner: str or ``None + The name of an aligner plugin as exists in :mod:`plugins.extract.align` + masker: str or list or ``None + The name of a masker plugin(s) as exists in :mod:`plugins.extract.mask`. + This can be a single masker or a list of multiple maskers + recognition: str or ``None`` + The name of the recognition plugin to use. ``None`` to not do face recognition. + """ + logger.debug("detector: %s, aligner: %s, masker: %s recognition: %s", + detector, aligner, masker, recognition) retval = [] if detector is not None and detector.lower() != "none": retval.append("detect") if aligner is not None and aligner.lower() != "none": retval.append("align") - retval.extend(["mask_{}".format(idx) + if recognition is not None and recognition.lower() != "none": + retval.append("recognition") + retval.extend([f"mask_{idx}" for idx, mask in enumerate(masker) if mask is not None and mask.lower() != "none"]) logger.debug("flow: %s", retval) return retval @staticmethod - def _get_plugin_type_and_index(flow_phase): + def _get_plugin_type_and_index(flow_phase: str) -> tuple[str, int | None]: """ Obtain the plugin type and index for the plugin for the given flow phase. When multiple plugins for the same phase are allowed (e.g. Mask) this will return @@ -389,20 +532,20 @@ def _get_plugin_type_and_index(flow_phase): The index of this plugin type within the flow, if there are multiple plugins in use otherwise ``None`` if there is only 1 plugin in use for the given phase """ - idx = flow_phase.split("_")[-1] - if idx.isdigit(): - idx = int(idx) + sidx = flow_phase.split("_")[-1] + if sidx.isdigit(): + idx: int | None = int(sidx) plugin_type = "_".join(flow_phase.split("_")[:-1]) else: plugin_type = flow_phase idx = None return plugin_type, idx - def _add_queues(self): + def _add_queues(self) -> dict[str, EventQueue]: """ Add the required processing queues to Queue Manager """ - queues = dict() - tasks = ["extract{}_{}_in".format(self._instance, phase) for phase in self._flow] - tasks.append("extract{}_{}_out".format(self._instance, self._final_phase)) + queues = {} + tasks = [f"extract{self._instance}_{phase}_in" for phase in self._flow] + tasks.append(f"extract{self._instance}_{self._final_phase}_out") for task in tasks: # Limit queue size to avoid stacking ram queue_manager.add_queue(task, maxsize=self._queue_size) @@ -411,7 +554,7 @@ def _add_queues(self): return queues @staticmethod - def _get_vram_stats(): + def _get_vram_stats() -> dict[str, int | str]: """ Obtain statistics on available VRAM and subtract a constant buffer from available vram. Returns @@ -422,14 +565,14 @@ def _get_vram_stats(): vram_buffer = 256 # Leave a buffer for VRAM allocation gpu_stats = GPUStats() stats = gpu_stats.get_card_most_free() - retval = dict(count=gpu_stats.device_count, - device=stats["device"], - vram_free=int(stats["free"] - vram_buffer), - vram_total=int(stats["total"])) + retval: dict[str, int | str] = {"count": gpu_stats.device_count, + "device": stats.device, + "vram_free": int(stats.free - vram_buffer), + "vram_total": int(stats.total)} logger.debug(retval) return retval - def _set_parallel_processing(self, multiprocess): + def _set_parallel_processing(self, multiprocess: bool) -> bool: """ Set whether to run detect, align, and mask together or separately. Parameters @@ -445,21 +588,17 @@ def _set_parallel_processing(self, multiprocess): logger.debug("No GPU detected. Enabling parallel processing.") return True - if get_backend() == "amd": - logger.debug("Parallel processing disabled by amd") - return False - - logger.verbose("%s - %sMB free of %sMB", + logger.verbose("%s - %sMB free of %sMB", # type: ignore self._vram_stats["device"], self._vram_stats["vram_free"], self._vram_stats["vram_total"]) - if self._vram_stats["vram_free"] <= self._total_vram_required: + if T.cast(int, self._vram_stats["vram_free"]) <= self._total_vram_required: logger.warning("Not enough free VRAM for parallel processing. " "Switching to serial") return False return True - def _set_phases(self, multiprocess): + def _set_phases(self, multiprocess: bool) -> list[list[str]]: """ If not enough VRAM is available, then chunk :attr:`_flow` up into phases that will fit into VRAM, otherwise return the single flow. @@ -473,10 +612,9 @@ def _set_phases(self, multiprocess): list: The jobs to be undertaken split into phases that fit into GPU RAM """ - force_single_process = not multiprocess or get_backend() == "amd" - phases = [] - current_phase = [] - available = self._vram_stats["vram_free"] + phases: list[list[str]] = [] + current_phase: list[str] = [] + available = T.cast(int, self._vram_stats["vram_free"]) for phase in self._flow: num_plugins = len([p for p in current_phase if self._vram_per_phase[p] > 0]) num_plugins += 1 if self._vram_per_phase[phase] > 0 else 0 @@ -484,11 +622,11 @@ def _set_phases(self, multiprocess): required = sum(self._vram_per_phase[p] for p in current_phase + [phase]) * scaling logger.debug("Num plugins for phase: %s, scaling: %s, vram required: %s", num_plugins, scaling, required) - if required <= available and not force_single_process: + if required <= available and multiprocess: logger.debug("Required: %s, available: %s. Adding phase '%s' to current phase: %s", required, available, phase, current_phase) current_phase.append(phase) - elif len(current_phase) == 0 or force_single_process: + elif len(current_phase) == 0 or not multiprocess: # Amount of VRAM required to run a single plugin is greater than available. We add # it anyway, and hope it will run with warnings, as the alternative is to not run # at all. @@ -508,105 +646,183 @@ def _set_phases(self, multiprocess): return phases # << INTERNAL PLUGIN HANDLING >> # - def _load_align(self, aligner, configfile, normalize_method, re_feed): - """ Set global arguments and load aligner plugin """ + def _load_align(self, + aligner: str | None, + configfile: str | None, + normalize_method: T.Literal["none", "clahe", "hist", "mean"] | None, + re_feed: int, + re_align: bool, + disable_filter: bool) -> Aligner | None: + """ Set global arguments and load aligner plugin + + Parameters + ---------- + aligner: str + The aligner plugin to load or ``None`` for no aligner + configfile: str + Optional full path to custom config file + normalize_method: str + Optional normalization method to use + re_feed: int + The number of times to adjust the image and re-feed to get an average score + re_align: bool + ``True`` to obtain landmarks by passing the initially aligned face back through the + aligner. + disable_filter: bool + Disable all aligner filters regardless of config option + + Returns + ------- + Aligner plugin if one is specified otherwise ``None`` + """ if aligner is None or aligner.lower() == "none": logger.debug("No aligner selected. Returning None") return None aligner_name = aligner.replace("-", "_").lower() logger.debug("Loading Aligner: '%s'", aligner_name) - aligner = PluginLoader.get_aligner(aligner_name)(exclude_gpus=self._exclude_gpus, - configfile=configfile, - normalize_method=normalize_method, - re_feed=re_feed, - instance=self._instance) - return aligner - - def _load_detect(self, detector, rotation, min_size, configfile): - """ Set global arguments and load detector plugin """ + plugin = PluginLoader.get_aligner(aligner_name)(exclude_gpus=self._exclude_gpus, + configfile=configfile, + normalize_method=normalize_method, + re_feed=re_feed, + re_align=re_align, + disable_filter=disable_filter, + instance=self._instance) + return plugin + + def _load_detect(self, + detector: str | None, + aligner: str | None, + rotation: str | None, + min_size: int, + configfile: str | None) -> Detector | None: + """ Set global arguments and load detector plugin + + Parameters + ---------- + detector: str | None + The name of the face detection plugin to use. ``None`` for no detection + aligner: str | None + The name of the face aligner plugin to use. ``None`` for no aligner + rotation: str | None + The rotation to perform on detection. ``None`` for no rotation + min_size: int + The minimum size of detected faces to accept + configfile: str | None + Full path to a custom config file to use. ``None`` for default config + + Returns + ------- + :class:`~plugins.extract.detect._base.Detector` | None + The face detection plugin to use, or ``None`` if no detection to be performed + """ if detector is None or detector.lower() == "none": logger.debug("No detector selected. Returning None") return None detector_name = detector.replace("-", "_").lower() + + if aligner == "external" and detector_name != "external": + logger.warning("Unsupported '%s' detector selected for 'External' aligner. Switching " + "detector to 'External'", detector_name) + detector_name = aligner + logger.debug("Loading Detector: '%s'", detector_name) - detector = PluginLoader.get_detector(detector_name)(exclude_gpus=self._exclude_gpus, - rotation=rotation, - min_size=min_size, - configfile=configfile, - instance=self._instance) - return detector - - def _load_mask(self, masker, image_is_aligned, configfile): - """ Set global arguments and load masker plugin """ + plugin = PluginLoader.get_detector(detector_name)(exclude_gpus=self._exclude_gpus, + rotation=rotation, + min_size=min_size, + configfile=configfile, + instance=self._instance) + return plugin + + def _load_mask(self, + masker: str | None, + configfile: str | None) -> Masker | None: + """ Set global arguments and load masker plugin + + Parameters + ---------- + masker: str or ``none`` + The name of the masker plugin to use or ``None`` if no masker + configfile: str + Full path to custom config.ini file or ``None`` to use default + + Returns + ------- + :class:`~plugins.extract.mask._base.Masker` or ``None`` + The masker plugin to use or ``None`` if no masker selected + """ if masker is None or masker.lower() == "none": logger.debug("No masker selected. Returning None") return None masker_name = masker.replace("-", "_").lower() logger.debug("Loading Masker: '%s'", masker_name) - masker = PluginLoader.get_masker(masker_name)(exclude_gpus=self._exclude_gpus, - image_is_aligned=image_is_aligned, + plugin = PluginLoader.get_masker(masker_name)(exclude_gpus=self._exclude_gpus, configfile=configfile, instance=self._instance) - return masker - - def _launch_plugin(self, phase): + return plugin + + def _load_recognition(self, + recognition: str | None, + configfile: str | None) -> Identity | None: + """ Set global arguments and load recognition plugin """ + if recognition is None or recognition.lower() == "none": + logger.debug("No recognition selected. Returning None") + return None + recognition_name = recognition.replace("-", "_").lower() + logger.debug("Loading Recognition: '%s'", recognition_name) + plugin = PluginLoader.get_recognition(recognition_name)(exclude_gpus=self._exclude_gpus, + configfile=configfile, + instance=self._instance) + return plugin + + def _launch_plugin(self, phase: str) -> None: """ Launch an extraction plugin """ logger.debug("Launching %s plugin", phase) - in_qname = "extract{}_{}_in".format(self._instance, phase) + in_qname = f"extract{self._instance}_{phase}_in" if phase == self._final_phase: - out_qname = "extract{}_{}_out".format(self._instance, self._final_phase) + out_qname = f"extract{self._instance}_{self._final_phase}_out" else: next_phase = self._flow[self._flow.index(phase) + 1] - out_qname = "extract{}_{}_in".format(self._instance, next_phase) + out_qname = f"extract{self._instance}_{next_phase}_in" logger.debug("in_qname: %s, out_qname: %s", in_qname, out_qname) - kwargs = dict(in_queue=self._queues[in_qname], out_queue=self._queues[out_qname]) + kwargs = {"in_queue": self._queues[in_qname], "out_queue": self._queues[out_qname]} plugin_type, idx = self._get_plugin_type_and_index(phase) - plugin = getattr(self, "_{}".format(plugin_type)) + plugin = getattr(self, f"_{plugin_type}") plugin = plugin[idx] if idx is not None else plugin plugin.initialize(**kwargs) plugin.start() logger.debug("Launched %s plugin", phase) - def _set_extractor_batchsize(self): + def _set_extractor_batchsize(self) -> None: """ Sets the batch size of the requested plugins based on their vram, their vram_per_batch_requirements and the number of plugins being loaded in the current phase. Only adjusts if the the configured batch size requires more vram than is available. Nvidia only. """ - if get_backend() != "nvidia": - logger.debug("Backend is not Nvidia. Not updating batchsize requirements") + backend = get_backend() + if backend not in ("nvidia", "directml", "rocm"): + logger.debug("Not updating batchsize requirements for backend: '%s'", backend) return - if sum([plugin.vram for plugin in self._active_plugins]) == 0: + if sum(plugin.vram for plugin in self._active_plugins) == 0: logger.debug("No plugins use VRAM. Not updating batchsize requirements.") return - batch_required = sum([plugin.vram_per_batch * plugin.batchsize - for plugin in self._active_plugins]) + batch_required = sum(plugin.vram_per_batch * plugin.batchsize + for plugin in self._active_plugins) gpu_plugins = [p for p in self._current_phase if self._vram_per_phase[p] > 0] scaling = self._parallel_scaling.get(len(gpu_plugins), self._scaling_fallback) - plugins_required = sum([self._vram_per_phase[p] for p in gpu_plugins]) * scaling - if plugins_required + batch_required <= self._vram_stats["vram_free"]: + plugins_required = sum(self._vram_per_phase[p] for p in gpu_plugins) * scaling + if plugins_required + batch_required <= T.cast(int, self._vram_stats["vram_free"]): logger.debug("Plugin requirements within threshold: (plugins_required: %sMB, " "vram_free: %sMB)", plugins_required, self._vram_stats["vram_free"]) return # Hacky split across plugins that use vram - available_vram = (self._vram_stats["vram_free"] - plugins_required) // len(gpu_plugins) + available_vram = (T.cast(int, self._vram_stats["vram_free"]) + - plugins_required) // len(gpu_plugins) self._set_plugin_batchsize(gpu_plugins, available_vram) - def set_aligner_normalization_method(self, method): - """ Change the normalization method for faces fed into the aligner. - - Parameters - ---------- - method: {"none", "clahe", "hist", "mean"} - The normalization method to apply to faces prior to feeding into the aligner's model - """ - logger.debug("Setting to: '%s'", method) - self._align.set_normalize_method(method) - - def _set_plugin_batchsize(self, gpu_plugins, available_vram): + def _set_plugin_batchsize(self, gpu_plugins: list[str], available_vram: float) -> None: """ Set the batch size for the given plugin based on given available vram. Do not update plugins which have a vram_per_batch of 0 (CPU plugins) due to zero division error. @@ -645,7 +861,7 @@ def _set_plugin_batchsize(self, gpu_plugins, available_vram): logger.debug("Remaining VRAM to allocate: %sMB", remaining) if batchsizes != requested_batchsizes: - text = ", ".join(["{}: {}".format(plugin.__class__.__name__, batchsize) + text = ", ".join([f"{plugin.__class__.__name__}: {batchsize}" for plugin, batchsize in zip(plugins, batchsizes)]) for plugin, batchsize in zip(plugins, batchsizes): plugin.batchsize = batchsize @@ -656,138 +872,7 @@ def _join_threads(self): for plugin in self._active_plugins: plugin.join() - def _check_and_raise_error(self): + def _check_and_raise_error(self) -> None: """ Check all threads for errors and raise if one occurs """ for plugin in self._active_plugins: - if plugin.check_and_raise_error(): - return True - return False - - -class ExtractMedia(): - """ An object that passes through the :class:`~plugins.extract.pipeline.Extractor` pipeline. - - Parameters - ---------- - filename: str - The base name of the original frame's filename - image: :class:`numpy.ndarray` - The original frame - detected_faces: list, optional - A list of :class:`~lib.align.DetectedFace` objects. Detected faces can be added - later with :func:`add_detected_faces`. Default: ``None`` - """ - - def __init__(self, filename, image, detected_faces=None): - logger.trace("Initializing %s: (filename: '%s', image shape: %s, detected_faces: %s)", - self.__class__.__name__, filename, image.shape, detected_faces) - self._filename = filename - self._image = image - self._image_shape = image.shape - self._detected_faces = detected_faces - - @property - def filename(self): - """ str: The base name of the :attr:`image` filename. """ - return self._filename - - @property - def image(self): - """ :class:`numpy.ndarray`: The source frame for this object. """ - return self._image - - @property - def image_shape(self): - """ tuple: The shape of the stored :attr:`image`. """ - return self._image_shape - - @property - def image_size(self): - """ tuple: The (`height`, `width`) of the stored :attr:`image`. """ - return self._image_shape[:2] - - @property - def detected_faces(self): - """list: A list of :class:`~lib.align.DetectedFace` objects in the - :attr:`image`. """ - return self._detected_faces - - def get_image_copy(self, color_format): - """ Get a copy of the image in the requested color format. - - Parameters - ---------- - color_format: ['BGR', 'RGB', 'GRAY'] - The requested color format of :attr:`image` - - Returns - ------- - :class:`numpy.ndarray`: - A copy of :attr:`image` in the requested :attr:`color_format` - """ - logger.trace("Requested color format '%s' for frame '%s'", color_format, self._filename) - image = getattr(self, "_image_as_{}".format(color_format.lower()))() - return image - - def add_detected_faces(self, faces): - """ Add detected faces to the object. Called at the end of each extraction phase. - - Parameters - ---------- - faces: list - A list of :class:`~lib.align.DetectedFace` objects - """ - logger.trace("Adding detected faces for filename: '%s'. (faces: %s, lrtb: %s)", - self._filename, faces, - [(face.left, face.right, face.top, face.bottom) for face in faces]) - self._detected_faces = faces - - def remove_image(self): - """ Delete the image and reset :attr:`image` to ``None``. - - Required for multi-phase extraction to avoid the frames stacking RAM. - """ - logger.trace("Removing image for filename: '%s'", self._filename) - del self._image - self._image = None - - def set_image(self, image): - """ Add the image back into :attr:`image` - - Required for multi-phase extraction adds the image back to this object. - - Parameters - ---------- - image: :class:`numpy.ndarry` - The original frame to be re-applied to for this :attr:`filename` - """ - logger.trace("Reapplying image: (filename: `%s`, image shape: %s)", - self._filename, image.shape) - self._image = image - - def _image_as_bgr(self): - """ Get a copy of the source frame in BGR format. - - Returns - ------- - :class:`numpy.ndarray`: - A copy of :attr:`image` in BGR color format """ - return self._image[..., :3].copy() - - def _image_as_rgb(self): - """ Get a copy of the source frame in RGB format. - - Returns - ------- - :class:`numpy.ndarray`: - A copy of :attr:`image` in RGB color format """ - return self._image[..., 2::-1].copy() - - def _image_as_gray(self): - """ Get a copy of the source frame in gray-scale format. - - Returns - ------- - :class:`numpy.ndarray`: - A copy of :attr:`image` in gray-scale color format """ - return cv2.cvtColor(self._image.copy(), cv2.COLOR_BGR2GRAY) + plugin.check_and_raise_error() diff --git a/plugins/extract/recognition/_base.py b/plugins/extract/recognition/_base.py new file mode 100644 index 0000000000..61662e623b --- /dev/null +++ b/plugins/extract/recognition/_base.py @@ -0,0 +1,490 @@ +#!/usr/bin/env python3 +""" Base class for Face Recognition plugins + +All Recognition Plugins should inherit from this class. +See the override methods for which methods are required. + +The plugin will receive a :class:`~plugins.extract.extract_media.ExtractMedia` object. + +For each source frame, the plugin must pass a dict to finalize containing: + +>>> {'filename': , +>>> 'detected_faces': >> face = self.to_detected_face(, , , ) +""" +from __future__ import annotations +import logging +import typing as T + +from dataclasses import dataclass, field + +import numpy as np +from tensorflow.python.framework import errors_impl as tf_errors # pylint:disable=no-name-in-module # noqa + +from lib.align import AlignedFace, DetectedFace, LandmarkType +from lib.image import read_image_meta +from lib.utils import FaceswapError +from plugins.extract import ExtractMedia +from plugins.extract._base import BatchType, ExtractorBatch, Extractor + +if T.TYPE_CHECKING: + from collections.abc import Generator + from queue import Queue + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +@dataclass +class RecogBatch(ExtractorBatch): + """ Dataclass for holding items flowing through the aligner. + + Inherits from :class:`~plugins.extract._base.ExtractorBatch` + """ + detected_faces: list[DetectedFace] = field(default_factory=list) + feed_faces: list[AlignedFace] = field(default_factory=list) + + +class Identity(Extractor): # pylint:disable=abstract-method + """ Face Recognition Object + + Parent class for all Recognition plugins + + Parameters + ---------- + git_model_id: int + The second digit in the github tag that identifies this model. See + https://github.com/deepfakes-models/faceswap-models for more information + model_filename: str + The name of the model file to be loaded + + Other Parameters + ---------------- + configfile: str, optional + Path to a custom configuration ``ini`` file. Default: Use system configfile + + See Also + -------- + plugins.extract.pipeline : The extraction pipeline for calling plugins + plugins.extract.detect : Detector plugins + plugins.extract._base : Parent class for all extraction plugins + plugins.extract.align._base : Aligner parent class for extraction plugins. + plugins.extract.mask._base : Masker parent class for extraction plugins. + """ + + _logged_lm_count_once = False + + def __init__(self, + git_model_id: int | None = None, + model_filename: str | None = None, + configfile: str | None = None, + instance: int = 0, + **kwargs): + logger.debug("Initializing %s", self.__class__.__name__) + super().__init__(git_model_id, + model_filename, + configfile=configfile, + instance=instance, + **kwargs) + self.input_size = 256 # Override for model specific input_size + self.centering: CenteringType = "legacy" # Override for model specific centering + self.coverage_ratio = 1.0 # Override for model specific coverage_ratio + + self._plugin_type = "recognition" + self._filter = IdentityFilter(self.config["save_filtered"]) + logger.debug("Initialized _base %s", self.__class__.__name__) + + def _get_detected_from_aligned(self, item: ExtractMedia) -> None: + """ Obtain detected face objects for when loading in aligned faces and a detected face + object does not exist + + Parameters + ---------- + item: :class:`~plugins.extract.extract_media.ExtractMedia` + The extract media to populate the detected face for + """ + detected_face = DetectedFace() + meta = read_image_meta(item.filename).get("itxt", {}).get("alignments") + if meta: + detected_face.from_png_meta(meta) + item.add_detected_faces([detected_face]) + self._faces_per_filename[item.filename] += 1 # Track this added face + logger.debug("Obtained detected face: (filename: %s, detected_face: %s)", + item.filename, item.detected_faces) + + def _maybe_log_warning(self, face: AlignedFace) -> None: + """ Log a warning, once, if we do not have full facial landmarks + + Parameters + ---------- + face: :class:`~lib.align.aligned_face.AlignedFace` + The aligned face object to test the landmark type for + """ + if face.landmark_type != LandmarkType.LM_2D_4 or self._logged_lm_count_once: + return + logger.warning("Extracted faces do not contain facial landmark data. '%s' " + "identity data is likely to be sub-standard.", self.name) + self._logged_lm_count_once = True + + def get_batch(self, queue: Queue) -> tuple[bool, RecogBatch]: + """ Get items for inputting into the recognition from the queue in batches + + Items are returned from the ``queue`` in batches of + :attr:`~plugins.extract._base.Extractor.batchsize` + + Items are received as :class:`~plugins.extract.extract_media.ExtractMedia` objects and + converted to :class:`RecogBatch` for internal processing. + + To ensure consistent batch sizes for masker the items are split into separate items for + each :class:`~lib.align.DetectedFace` object. + + Remember to put ``'EOF'`` to the out queue after processing + the final batch + + Outputs items in the following format. All lists are of length + :attr:`~plugins.extract._base.Extractor.batchsize`: + + >>> {'filename': [], + >>> 'detected_faces': [[ RecogBatch: + """ Just return the recognition's predict function """ + assert isinstance(batch, RecogBatch) + try: + # slightly hacky workaround to deal with landmarks based masks: + batch.prediction = self.predict(batch.feed) + return batch + except tf_errors.ResourceExhaustedError as err: + msg = ("You do not have enough GPU memory available to run recognition at the " + "selected batch size. You can try a number of things:" + "\n1) Close any other application that is using your GPU (web browsers are " + "particularly bad for this)." + "\n2) Lower the batchsize (the amount of images fed into the model) by " + "editing the plugin settings (GUI: Settings > Configure extract settings, " + "CLI: Edit the file faceswap/config/extract.ini)." + "\n3) Enable 'Single Process' mode.") + raise FaceswapError(msg) from err + + def finalize(self, batch: BatchType) -> Generator[ExtractMedia, None, None]: + """ Finalize the output from Masker + + This should be called as the final task of each `plugin`. + + Pairs the detected faces back up with their original frame before yielding each frame. + + Parameters + ---------- + batch : :class:`RecogBatch` + The final batch item from the `plugin` process. + + Yields + ------ + :class:`~plugins.extract.extract_media.ExtractMedia` + The :attr:`DetectedFaces` list will be populated for this class with the bounding + boxes, landmarks and masks for the detected faces found in the frame. + """ + assert isinstance(batch, RecogBatch) + assert isinstance(self.name, str) + for identity, face in zip(batch.prediction, batch.detected_faces): + face.add_identity(self.name.lower(), identity) + del batch.feed + + logger.trace("Item out: %s", # type: ignore + {key: val.shape if isinstance(val, np.ndarray) else val + for key, val in batch.__dict__.items()}) + + for filename, face in zip(batch.filename, batch.detected_faces): + self._output_faces.append(face) + if len(self._output_faces) != self._faces_per_filename[filename]: + continue + + output = self._extract_media.pop(filename) + self._output_faces = self._filter(self._output_faces, output.sub_folders) + + output.add_detected_faces(self._output_faces) + self._output_faces = [] + logger.trace("Yielding: (filename: '%s', image: %s, " # type:ignore + "detected_faces: %s)", output.filename, output.image_shape, + len(output.detected_faces)) + yield output + + def add_identity_filters(self, + filters: np.ndarray, + nfilters: np.ndarray, + threshold: float) -> None: + """ Add identity encodings to filter by identity in the recognition plugin + + Parameters + ---------- + filters: :class:`numpy.ndarray` + The array of filter embeddings to use + nfilters: :class:`numpy.ndarray` + The array of nfilter embeddings to use + threshold: float + The threshold for a positive filter match + """ + logger.debug("Adding identity filters") + self._filter.add_filters(filters, nfilters, threshold) + logger.debug("Added identity filters") + + +class IdentityFilter(): + """ Applies filters on the output of the recognition plugin + + Parameters + ---------- + save_output: bool + ``True`` if the filtered faces should be kept as they are being saved. ``False`` if they + should be deleted + """ + def __init__(self, save_output: bool) -> None: + logger.debug("Initializing %s: (save_output: %s)", self.__class__.__name__, save_output) + self._save_output = save_output + self._filter: np.ndarray | None = None + self._nfilter: np.ndarray | None = None + self._threshold = 0.0 + self._filter_enabled: bool = False + self._nfilter_enabled: bool = False + self._active: bool = False + self._counts = 0 + logger.debug("Initialized %s", self.__class__.__name__) + + def add_filters(self, filters: np.ndarray, nfilters: np.ndarray, threshold) -> None: + """ Add identity encodings to the filter and set whether each filter is enabled + + Parameters + ---------- + filters: :class:`numpy.ndarray` + The array of filter embeddings to use + nfilters: :class:`numpy.ndarray` + The array of nfilter embeddings to use + threshold: float + The threshold for a positive filter match + """ + logger.debug("Adding filters: %s, nfilters: %s, threshold: %s", + filters.shape, nfilters.shape, threshold) + self._filter = filters + self._nfilter = nfilters + self._threshold = threshold + self._filter_enabled = bool(np.any(self._filter)) + self._nfilter_enabled = bool(np.any(self._nfilter)) + self._active = self._filter_enabled or self._nfilter_enabled + logger.debug("filter active: %s, nfilter active: %s, all active: %s", + self._filter_enabled, self._nfilter_enabled, self._active) + + @classmethod + def _find_cosine_similiarity(cls, + source_identities: np.ndarray, + test_identity: np.ndarray) -> np.ndarray: + """ Find the cosine similarity between a source face identity and a test face identity + + Parameters + --------- + source_identities: :class:`numpy.ndarray` + The identity encoding for the source face identities + test_identity: :class:`numpy.ndarray` + The identity encoding for the face identity to test against the sources + + Returns + ------- + :class:`numpy.ndarray`: + The cosine similarity between a face identity and the source identities + """ + s_norm = np.linalg.norm(source_identities, axis=1) + i_norm = np.linalg.norm(test_identity) + retval = source_identities @ test_identity / (s_norm * i_norm) + return retval + + def _get_matches(self, + filter_type: T.Literal["filter", "nfilter"], + identities: np.ndarray) -> np.ndarray: + """ Obtain the average and minimum distances for each face against the source identities + to test against + + Parameters + ---------- + filter_type ["filter", "nfilter"] + The filter type to use for calculating the distance + identities: :class:`numpy.ndarray` + The identity encodings for the current face(s) being checked + + Returns + ------- + :class:`numpy.ndarray` + Boolean array. ``True`` if identity should be filtered otherwise ``False`` + """ + encodings = self._filter if filter_type == "filter" else self._nfilter + assert encodings is not None + distances = np.array([self._find_cosine_similiarity(encodings, identity) + for identity in identities]) + is_match = np.any(distances >= self._threshold, axis=-1) + # Invert for filter (set the `True` match to `False` for should filter) + retval = np.invert(is_match) if filter_type == "filter" else is_match + logger.trace("filter_type: %s, distances shape: %s, is_match: %s, ", # type: ignore + "retval: %s", filter_type, distances.shape, is_match, retval) + return retval + + def _filter_faces(self, + faces: list[DetectedFace], + sub_folders: list[str | None], + should_filter: list[bool]) -> list[DetectedFace]: + """ Filter the detected faces, either removing filtered faces from the list of detected + faces or setting the output subfolder to `"_identity_filt"` for any filtered faces if + saving output is enabled. + + Parameters + ---------- + faces: list + List of detected face objects to filter out on size + sub_folders: list + List of subfolder locations for any faces that have already been filtered when + config option `save_filtered` has been enabled. + should_filter: list + List of 'bool' corresponding to face that have not already been marked for filtering. + ``True`` indicates face should be filtered, ``False`` indicates face should be kept + + Returns + ------- + detected_faces: list + The filtered list of detected face objects, if saving filtered faces has not been + selected or the full list of detected faces + """ + retval: list[DetectedFace] = [] + self._counts += sum(should_filter) + for idx, face in enumerate(faces): + fldr = sub_folders[idx] + if fldr is not None: + # Saving to sub folder is selected and face is already filtered + # so this face was excluded from identity check + retval.append(face) + continue + to_filter = should_filter.pop(0) + if not to_filter or self._save_output: + # Keep the face if not marked as filtered or we are to output to a subfolder + retval.append(face) + if to_filter and self._save_output: + sub_folders[idx] = "_identity_filt" + + return retval + + def __call__(self, + faces: list[DetectedFace], + sub_folders: list[str | None]) -> list[DetectedFace]: + """ Call the identity filter function + + Parameters + ---------- + faces: list + List of detected face objects to filter out on size + sub_folders: list + List of subfolder locations for any faces that have already been filtered when + config option `save_filtered` has been enabled. + + Returns + ------- + detected_faces: list + The filtered list of detected face objects, if saving filtered faces has not been + selected or the full list of detected faces + """ + if not self._active: + return faces + + identities = np.array([face.identity["vggface2"] for face, fldr in zip(faces, sub_folders) + if fldr is None]) + logger.trace("face_count: %s, already_filtered: %s, identity_shape: %s", # type: ignore + len(faces), sum(x is not None for x in sub_folders), identities.shape) + + if not np.any(identities): + logger.trace("All faces already filtered: %s", sub_folders) # type: ignore + return faces + + should_filter: list[np.ndarray] = [] + for f_type in T.get_args(T.Literal["filter", "nfilter"]): + if not getattr(self, f"_{f_type}_enabled"): + continue + should_filter.append(self._get_matches(f_type, identities)) + + # If any of the filter or nfilter evaluate to 'should filter' then filter out face + final_filter: list[bool] = np.array(should_filter).max(axis=0).tolist() + logger.trace("should_filter: %s, final_filter: %s", # type: ignore + should_filter, final_filter) + return self._filter_faces(faces, sub_folders, final_filter) + + def output_counts(self): + """ Output the counts of filtered items """ + if not self._active or not self._counts: + return + logger.info("Identity filtered (%s): %s", self._threshold, self._counts) diff --git a/plugins/extract/recognition/vgg_face2.py b/plugins/extract/recognition/vgg_face2.py new file mode 100644 index 0000000000..acf268bfd4 --- /dev/null +++ b/plugins/extract/recognition/vgg_face2.py @@ -0,0 +1,315 @@ +#!/usr/bin python3 +""" VGG_Face2 inference and sorting """ + +from __future__ import annotations +import logging +import typing as T + +import numpy as np +import psutil +from fastcluster import linkage, linkage_vector + +from lib.model.layers import L2_normalize +from lib.model.session import KSession +from lib.utils import FaceswapError +from ._base import BatchType, RecogBatch, Identity + +if T.TYPE_CHECKING: + from collections.abc import Generator + +logger = logging.getLogger(__name__) + + +class Recognition(Identity): + """ VGG Face feature extraction. + + Extracts feature vectors from faces in order to compare similarity. + + Notes + ----- + Input images should be in BGR Order + + Model exported from: https://github.com/WeidiXie/Keras-VGGFace2-ResNet50 which is based on: + https://www.robots.ox.ac.uk/~vgg/software/vgg_face/ + + + Licensed under Creative Commons Attribution License. + https://creativecommons.org/licenses/by-nc/4.0/ + """ + + def __init__(self, *args, **kwargs) -> None: # pylint:disable=unused-argument + logger.debug("Initializing %s", self.__class__.__name__) + git_model_id = 10 + model_filename = "vggface2_resnet50_v2.h5" + super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) + self.model: KSession + self.name: str = "VGGFace2" + self.input_size = 224 + self.color_format = "BGR" + + self.vram = 2468 if not self.config["cpu"] else 0 + self.vram_warnings = 192 if not self.config["cpu"] else 0 + self.vram_per_batch = 32 if not self.config["cpu"] else 0 + self.batchsize = self.config["batch-size"] + + # Average image provided in https://github.com/ox-vgg/vgg_face2 + self._average_img = np.array([91.4953, 103.8827, 131.0912]) + logger.debug("Initialized %s", self.__class__.__name__) + + # <<< GET MODEL >>> # + def init_model(self) -> None: + """ Initialize VGG Face 2 Model. """ + assert isinstance(self.model_path, str) + model_kwargs = {"custom_objects": {"L2_normalize": L2_normalize}} + self.model = KSession(self.name, + self.model_path, + model_kwargs=model_kwargs, + allow_growth=self.config["allow_growth"], + exclude_gpus=self._exclude_gpus, + cpu_mode=self.config["cpu"]) + self.model.load_model() + + def process_input(self, batch: BatchType) -> None: + """ Compile the detected faces for prediction """ + assert isinstance(batch, RecogBatch) + batch.feed = np.array([T.cast(np.ndarray, feed.face)[..., :3] + for feed in batch.feed_faces], + dtype="float32") - self._average_img + logger.trace("feed shape: %s", batch.feed.shape) # type:ignore + + def predict(self, feed: np.ndarray) -> np.ndarray: + """ Return encodings for given image from vgg_face2. + + Parameters + ---------- + batch: numpy.ndarray + The face to be fed through the predictor. Should be in BGR channel order + + Returns + ------- + numpy.ndarray + The encodings for the face + """ + retval = self.model.predict(feed) + assert isinstance(retval, np.ndarray) + return retval + + def process_output(self, batch: BatchType) -> None: + """ No output processing for vgg_face2 """ + return + + +class Cluster(): # pylint:disable=too-few-public-methods + """ Cluster the outputs from a VGG-Face 2 Model + + Parameters + ---------- + predictions: numpy.ndarray + A stacked matrix of vgg_face2 predictions of the shape (`N`, `D`) where `N` is the + number of observations and `D` are the number of dimensions. NB: The given + :attr:`predictions` will be overwritten to save memory. If you still require the + original values you should take a copy prior to running this method + method: ['single','centroid','median','ward'] + The clustering method to use. + threshold: float, optional + The threshold to start creating bins for. Set to ``None`` to disable binning + """ + + def __init__(self, + predictions: np.ndarray, + method: T.Literal["single", "centroid", "median", "ward"], + threshold: float | None = None) -> None: + logger.debug("Initializing: %s (predictions: %s, method: %s, threshold: %s)", + self.__class__.__name__, predictions.shape, method, threshold) + self._num_predictions = predictions.shape[0] + + self._should_output_bins = threshold is not None + self._threshold = 0.0 if threshold is None else threshold + self._bins: dict[int, int] = {} + self._iterator = self._integer_iterator() + + self._result_linkage = self._do_linkage(predictions, method) + logger.debug("Initialized %s", self.__class__.__name__) + + @classmethod + def _integer_iterator(cls) -> Generator[int, None, None]: + """ Iterator that just yields consecutive integers """ + i = -1 + while True: + i += 1 + yield i + + def _use_vector_linkage(self, dims: int) -> bool: + """ Calculate the RAM that will be required to sort these images and select the appropriate + clustering method. + + From fastcluster documentation: + "While the linkage method requires Θ(N:sup:`2`) memory for clustering of N points, this + [vector] method needs Θ(N D)for N points in RD, which is usually much smaller." + also: + "half the memory can be saved by specifying :attr:`preserve_input`=``False``" + + To avoid under calculating we divide the memory calculation by 1.8 instead of 2 + + Parameters + ---------- + dims: int + The number of dimensions in the vgg_face output + + Returns + ------- + bool: + ``True`` if vector_linkage should be used. ``False`` if linkage should be used + """ + np_float = 24 # bytes size of a numpy float + divider = 1024 * 1024 # bytes to MB + + free_ram = psutil.virtual_memory().available / divider + linkage_required = (((self._num_predictions ** 2) * np_float) / 1.8) / divider + vector_required = ((self._num_predictions * dims) * np_float) / divider + logger.debug("free_ram: %sMB, linkage_required: %sMB, vector_required: %sMB", + int(free_ram), int(linkage_required), int(vector_required)) + + if linkage_required < free_ram: + logger.verbose("Using linkage method") # type:ignore + retval = False + elif vector_required < free_ram: + logger.warning("Not enough RAM to perform linkage clustering. Using vector " + "clustering. This will be significantly slower. Free RAM: %sMB. " + "Required for linkage method: %sMB", + int(free_ram), int(linkage_required)) + retval = True + else: + raise FaceswapError("Not enough RAM available to sort faces. Try reducing " + f"the size of your dataset. Free RAM: {int(free_ram)}MB. " + f"Required RAM: {int(vector_required)}MB") + logger.debug(retval) + return retval + + def _do_linkage(self, + predictions: np.ndarray, + method: T.Literal["single", "centroid", "median", "ward"]) -> np.ndarray: + """ Use FastCluster to perform vector or standard linkage + + Parameters + ---------- + predictions: :class:`numpy.ndarray` + A stacked matrix of vgg_face2 predictions of the shape (`N`, `D`) where `N` is the + number of observations and `D` are the number of dimensions. + method: ['single','centroid','median','ward'] + The clustering method to use. + + Returns + ------- + :class:`numpy.ndarray` + The [`num_predictions`, 4] linkage vector + """ + dims = predictions.shape[-1] + if self._use_vector_linkage(dims): + retval = linkage_vector(predictions, method=method) + else: + retval = linkage(predictions, method=method, preserve_input=False) + logger.debug("Linkage shape: %s", retval.shape) + return retval + + def _process_leaf_node(self, + current_index: int, + current_bin: int) -> list[tuple[int, int]]: + """ Process the output when we have hit a leaf node """ + if not self._should_output_bins: + return [(current_index, 0)] + + if current_bin not in self._bins: + next_val = 0 if not self._bins else max(self._bins.values()) + 1 + self._bins[current_bin] = next_val + return [(current_index, self._bins[current_bin])] + + def _get_bin(self, + tree: np.ndarray, + points: int, + current_index: int, + current_bin: int) -> int: + """ Obtain the bin that we are currently in. + + If we are not currently below the threshold for binning, get a new bin ID from the integer + iterator. + + Parameters + ---------- + tree: numpy.ndarray + A hierarchical tree (dendrogram) + points: int + The number of points given to the clustering process + current_index: int + The position in the tree for the recursive traversal + current_bin int, optional + The ID for the bin we are currently in. Only used when binning is enabled + + Returns + ------- + int + The current bin ID for the node + """ + if tree[current_index - points, 2] >= self._threshold: + current_bin = next(self._iterator) + logger.debug("Creating new bin ID: %s", current_bin) + return current_bin + + def _seriation(self, + tree: np.ndarray, + points: int, + current_index: int, + current_bin: int = 0) -> list[tuple[int, int]]: + """ Seriation method for sorted similarity. + + Seriation computes the order implied by a hierarchical tree (dendrogram). + + Parameters + ---------- + tree: numpy.ndarray + A hierarchical tree (dendrogram) + points: int + The number of points given to the clustering process + current_index: int + The position in the tree for the recursive traversal + current_bin int, optional + The ID for the bin we are currently in. Only used when binning is enabled + + Returns + ------- + list: + The indices in the order implied by the hierarchical tree + """ + if current_index < points: # Output the leaf node + return self._process_leaf_node(current_index, current_bin) + + if self._should_output_bins: + current_bin = self._get_bin(tree, points, current_index, current_bin) + + left = int(tree[current_index-points, 0]) + right = int(tree[current_index-points, 1]) + + serate_left = self._seriation(tree, points, left, current_bin=current_bin) + serate_right = self._seriation(tree, points, right, current_bin=current_bin) + + return serate_left + serate_right # type: ignore + + def __call__(self) -> list[tuple[int, int]]: + """ Process the linkages. + + Transforms a distance matrix into a sorted distance matrix according to the order implied + by the hierarchical tree (dendrogram). + + Returns + ------- + list: + List of indices with the order implied by the hierarchical tree or list of tuples of + (`index`, `bin`) if a binning threshold was provided + """ + logger.info("Sorting face distances. Depending on your dataset this may take some time...") + if self._threshold: + self._threshold = self._result_linkage[:, 2].max() * self._threshold + result_order = self._seriation(self._result_linkage, + self._num_predictions, + self._num_predictions + self._num_predictions - 2) + return result_order diff --git a/plugins/extract/recognition/vgg_face2_defaults.py b/plugins/extract/recognition/vgg_face2_defaults.py new file mode 100644 index 0000000000..67c92783a7 --- /dev/null +++ b/plugins/extract/recognition/vgg_face2_defaults.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +""" + The default options for the faceswap VGG Face2 recognition plugin. + + Defaults files should be named _defaults.py + Any items placed into this file will automatically get added to the relevant config .ini files + within the faceswap/config folder. + + The following variables should be defined: + _HELPTEXT: A string describing what this plugin does + _DEFAULTS: A dictionary containing the options, defaults and meta information. The + dictionary should be defined as: + {: {}} + + should always be lower text. + dictionary requirements are listed below. + + The following keys are expected for the _DEFAULTS dict: + datatype: [required] A python type class. This limits the type of data that can be + provided in the .ini file and ensures that the value is returned in the + correct type to faceswap. Valid data types are: , , + , . + default: [required] The default value for this option. + info: [required] A string describing what this option does. + group: [optional]. A group for grouping options together in the GUI. If not + provided this will not group this option with any others. + choices: [optional] If this option's datatype is of then valid + selections can be defined here. This validates the option and also enables + a combobox / radio option in the GUI. + gui_radio: [optional] If are defined, this indicates that the GUI should use + radio buttons rather than a combobox to display this option. + min_max: [partial] For and data types this is required + otherwise it is ignored. Should be a tuple of min and max accepted values. + This is used for controlling the GUI slider range. Values are not enforced. + rounding: [partial] For and data types this is + required otherwise it is ignored. Used for the GUI slider. For floats, this + is the number of decimal places to display. For ints this is the step size. + fixed: [optional] [train only]. Training configurations are fixed when the model is + created, and then reloaded from the state file. Marking an item as fixed=False + indicates that this value can be changed for existing models, and will override + the value saved in the state file with the updated value in config. If not + provided this will default to True. +""" + + +_HELPTEXT = ( + "VGG Face 2 identity recognition.\n" + "A Keras port of the model trained for VGGFace2: A dataset for recognising faces across pose " + "and age. (https://arxiv.org/abs/1710.08092)" + ) + + +_DEFAULTS = { + "batch-size": { + "default": 16, + "info": "The batch size to use. To a point, higher batch sizes equal better performance, " + "but setting it too high can harm performance.\n" + "\n\tNvidia users: If the batchsize is set higher than the your GPU can " + "accomodate then this will automatically be lowered.", + "datatype": int, + "rounding": 1, + "min_max": (1, 64), + "choices": [], + "group": "settings", + "gui_radio": False, + "fixed": True + }, + "cpu": { + "default": False, + "info": "VGG Face2 still runs fairly quickly on CPU on some setups. Enable " + "CPU mode here to use the CPU for this plugin to save some VRAM at a speed cost.", + "datatype": bool, + "group": "settings" + }, +} diff --git a/plugins/extract/recognition/vgg_face2_keras.py b/plugins/extract/recognition/vgg_face2_keras.py deleted file mode 100644 index 10ca5d7ce9..0000000000 --- a/plugins/extract/recognition/vgg_face2_keras.py +++ /dev/null @@ -1,228 +0,0 @@ -#!/usr/bin python3 -""" VGG_Face2 inference and sorting """ - -import logging -import psutil - -import cv2 -import numpy as np -from fastcluster import linkage, linkage_vector - -from lib.model.layers import L2_normalize -from lib.model.session import KSession -from lib.utils import FaceswapError -from plugins.extract._base import Extractor - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - - -class VGGFace2(Extractor): # pylint:disable=abstract-method - """ VGG Face feature extraction. - - Extracts feature vectors from faces in order to compare similarity. - - Notes - ----- - Input images should be in BGR Order - - Model exported from: https://github.com/WeidiXie/Keras-VGGFace2-ResNet50 which is based on: - https://www.robots.ox.ac.uk/~vgg/software/vgg_face/ - - - Licensed under Creative Commons Attribution License. - https://creativecommons.org/licenses/by-nc/4.0/ - """ - - def __init__(self, *args, **kwargs): # pylint:disable=unused-argument - logger.debug("Initializing %s", self.__class__.__name__) - git_model_id = 10 - model_filename = ["vggface2_resnet50_v2.h5"] - super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) - self._plugin_type = "recognition" - self.name = "VGG_Face2" - self.input_size = 224 - # Average image provided in https://github.com/ox-vgg/vgg_face2 - self._average_img = np.array([91.4953, 103.8827, 131.0912]) - logger.debug("Initialized %s", self.__class__.__name__) - - # <<< GET MODEL >>> # - def init_model(self): - """ Initialize VGG Face 2 Model. """ - model_kwargs = dict(custom_objects={'L2_normalize': L2_normalize}) - self.model = KSession(self.name, - self.model_path, - model_kwargs=model_kwargs, - allow_growth=self.config["allow_growth"], - exclude_gpus=self._exclude_gpus) - self.model.load_model() - - def predict(self, batch): - """ Return encodings for given image from vgg_face2. - - Parameters - ---------- - batch: numpy.ndarray - The face to be fed through the predictor. Should be in BGR channel order - - Returns - ------- - numpy.ndarray - The encodings for the face - """ - face = batch - if face.shape[0] != self.input_size: - face = self._resize_face(face) - face = face[None, :, :, :3] - self._average_img - preds = self.model.predict(face) - return preds[0, :] - - def _resize_face(self, face): - """ Resize incoming face to model_input_size. - - Parameters - ---------- - face: numpy.ndarray - The face to be fed through the predictor. Should be in BGR channel order - - Returns - ------- - numpy.ndarray - The face resized to model input size - """ - sizes = (self.input_size, self.input_size) - interpolation = cv2.INTER_CUBIC if face.shape[0] < self.input_size else cv2.INTER_AREA - face = cv2.resize(face, dsize=sizes, interpolation=interpolation) - return face - - @staticmethod - def find_cosine_similiarity(source_face, test_face): - """ Find the cosine similarity between two faces. - - Parameters - ---------- - source_face: numpy.ndarray - The first face to test against :attr:`test_face` - test_face: numpy.ndarray - The second face to test against :attr:`source_face` - - Returns - ------- - float: - The cosine similarity between the two faces - """ - var_a = np.matmul(np.transpose(source_face), test_face) - var_b = np.sum(np.multiply(source_face, source_face)) - var_c = np.sum(np.multiply(test_face, test_face)) - return 1 - (var_a / (np.sqrt(var_b) * np.sqrt(var_c))) - - def sorted_similarity(self, predictions, method="ward"): - """ Sort a matrix of predictions by similarity. - - Transforms a distance matrix into a sorted distance matrix according to the order implied - by the hierarchical tree (dendrogram). - - Parameters - ---------- - predictions: numpy.ndarray - A stacked matrix of vgg_face2 predictions of the shape (`N`, `D`) where `N` is the - number of observations and `D` are the number of dimensions. NB: The given - :attr:`predictions` will be overwritten to save memory. If you still require the - original values you should take a copy prior to running this method - method: ['single','centroid','median','ward'] - The clustering method to use. - - Returns - ------- - list: - List of indices with the order implied by the hierarchical tree - """ - logger.info("Sorting face distances. Depending on your dataset this may take some time...") - num_predictions, dims = predictions.shape - - kwargs = dict(method=method) - if self._use_vector_linkage(num_predictions, dims): - func = linkage_vector - else: - kwargs["preserve_input"] = False - func = linkage - - result_linkage = func(predictions, **kwargs) - result_order = self._seriation(result_linkage, - num_predictions, - num_predictions + num_predictions - 2) - return result_order - - @staticmethod - def _use_vector_linkage(item_count, dims): - """ Calculate the RAM that will be required to sort these images and select the appropriate - clustering method. - - From fastcluster documentation: - "While the linkage method requires Θ(N:sup:`2`) memory for clustering of N points, this - [vector] method needs Θ(N D)for N points in RD, which is usually much smaller." - also: - "half the memory can be saved by specifying :attr:`preserve_input`=``False``" - - To avoid under calculating we divide the memory calculation by 1.8 instead of 2 - - Parameters - ---------- - item_count: int - The number of images that are to be processed - dims: int - The number of dimensions in the vgg_face output - - Returns - ------- - bool: - ``True`` if vector_linkage should be used. ``False`` if linkage should be used - """ - np_float = 24 # bytes size of a numpy float - divider = 1024 * 1024 # bytes to MB - - free_ram = psutil.virtual_memory().available / divider - linkage_required = (((item_count ** 2) * np_float) / 1.8) / divider - vector_required = ((item_count * dims) * np_float) / divider - logger.debug("free_ram: %sMB, linkage_required: %sMB, vector_required: %sMB", - int(free_ram), int(linkage_required), int(vector_required)) - - if linkage_required < free_ram: - logger.verbose("Using linkage method") - retval = False - elif vector_required < free_ram: - logger.warning("Not enough RAM to perform linkage clustering. Using vector " - "clustering. This will be significantly slower. Free RAM: %sMB. " - "Required for linkage method: %sMB", - int(free_ram), int(linkage_required)) - retval = True - else: - raise FaceswapError("Not enough RAM available to sort faces. Try reducing " - "the size of your dataset. Free RAM: {}MB. " - "Required RAM: {}MB".format(int(free_ram), int(vector_required))) - logger.debug(retval) - return retval - - def _seriation(self, tree, points, current_index): - """ Seriation method for sorted similarity. - - Seriation computes the order implied by a hierarchical tree (dendrogram). - - Parameters - ---------- - tree: numpy.ndarray - A hierarchical tree (dendrogram) - points: int - The number of points given to the clustering process - current_index: int - The position in the tree for the recursive traversal - - Returns - ------- - list: - The indices in the order implied by the hierarchical tree - """ - if current_index < points: - return [current_index] - left = int(tree[current_index-points, 0]) - right = int(tree[current_index-points, 1]) - return self._seriation(tree, points, left) + self._seriation(tree, points, right) diff --git a/plugins/plugin_loader.py b/plugins/plugin_loader.py index 5345216c4c..7d47c20680 100644 --- a/plugins/plugin_loader.py +++ b/plugins/plugin_loader.py @@ -1,11 +1,22 @@ #!/usr/bin/env python3 """ Plugin loader for Faceswap extract, training and convert tasks """ - +from __future__ import annotations import logging import os +import typing as T + from importlib import import_module -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from collections.abc import Callable + from plugins.extract.detect._base import Detector + from plugins.extract.align._base import Aligner + from plugins.extract.mask._base import Masker + from plugins.extract.recognition._base import Identity + from plugins.train.model._base import ModelBase + from plugins.train.trainer._base import TrainerBase + +logger = logging.getLogger(__name__) class PluginLoader(): @@ -21,7 +32,7 @@ class PluginLoader(): >>> aligner = PluginLoader.get_aligner('cv2-dnn') """ @staticmethod - def get_detector(name, disable_logging=False): + def get_detector(name: str, disable_logging: bool = False) -> type[Detector]: """ Return requested detector plugin Parameters @@ -40,7 +51,7 @@ def get_detector(name, disable_logging=False): return PluginLoader._import("extract.detect", name, disable_logging) @staticmethod - def get_aligner(name, disable_logging=False): + def get_aligner(name: str, disable_logging: bool = False) -> type[Aligner]: """ Return requested aligner plugin Parameters @@ -59,7 +70,7 @@ def get_aligner(name, disable_logging=False): return PluginLoader._import("extract.align", name, disable_logging) @staticmethod - def get_masker(name, disable_logging=False): + def get_masker(name: str, disable_logging: bool = False) -> type[Masker]: """ Return requested masker plugin Parameters @@ -78,7 +89,26 @@ def get_masker(name, disable_logging=False): return PluginLoader._import("extract.mask", name, disable_logging) @staticmethod - def get_model(name, disable_logging=False): + def get_recognition(name: str, disable_logging: bool = False) -> type[Identity]: + """ Return requested recognition plugin + + Parameters + ---------- + name: str + The name of the requested reccognition plugin + disable_logging: bool, optional + Whether to disable the INFO log message that the plugin is being imported. + Default: `False` + + Returns + ------- + :class:`plugins.extract.recognition` object: + An extraction recognition plugin + """ + return PluginLoader._import("extract.recognition", name, disable_logging) + + @staticmethod + def get_model(name: str, disable_logging: bool = False) -> type[ModelBase]: """ Return requested training model plugin Parameters @@ -97,7 +127,7 @@ def get_model(name, disable_logging=False): return PluginLoader._import("train.model", name, disable_logging) @staticmethod - def get_trainer(name, disable_logging=False): + def get_trainer(name: str, disable_logging: bool = False) -> type[TrainerBase]: """ Return requested training trainer plugin Parameters @@ -116,7 +146,7 @@ def get_trainer(name, disable_logging=False): return PluginLoader._import("train.trainer", name, disable_logging) @staticmethod - def get_converter(category, name, disable_logging=False): + def get_converter(category: str, name: str, disable_logging: bool = False) -> Callable: """ Return requested converter plugin Converters work slightly differently to other faceswap plugins. They are created to do a @@ -136,10 +166,10 @@ def get_converter(category, name, disable_logging=False): :class:`plugins.convert` object: A converter sub plugin """ - return PluginLoader._import("convert.{}".format(category), name, disable_logging) + return PluginLoader._import(f"convert.{category}", name, disable_logging) @staticmethod - def _import(attr, name, disable_logging): + def _import(attr: str, name: str, disable_logging: bool): """ Import the plugin's module Parameters @@ -164,12 +194,14 @@ def _import(attr, name, disable_logging): return getattr(module, ttl) @staticmethod - def get_available_extractors(extractor_type, add_none=False, extend_plugin=False): + def get_available_extractors(extractor_type: T.Literal["align", "detect", "mask"], + add_none: bool = False, + extend_plugin: bool = False) -> list[str]: """ Return a list of available extractors of the given type Parameters ---------- - extractor_type: {'aligner', 'detector', 'masker'} + extractor_type: {'align', 'detect', 'mask'} The type of extractor to return the plugins for add_none: bool, optional Append "none" to the list of returned plugins. Default: False @@ -194,9 +226,12 @@ def get_available_extractors(extractor_type, add_none=False, extend_plugin=False if not item.name.startswith("_") and not item.name.endswith("defaults.py") and item.name.endswith(".py")] - if extend_plugin and extractor_type == "mask" and "bisenet-fp" in extractors: - extractors.remove("bisenet-fp") - extractors.extend(["bisenet-fp_face", "bisenet-fp_head"]) + extendable = ["bisenet-fp", "custom"] + if extend_plugin and extractor_type == "mask" and any(ext in extendable + for ext in extractors): + for msk in extendable: + extractors.remove(msk) + extractors.extend([f"{msk}_face", f"{msk}_head"]) extractors = sorted(extractors) if add_none: @@ -204,7 +239,7 @@ def get_available_extractors(extractor_type, add_none=False, extend_plugin=False return extractors @staticmethod - def get_available_models(): + def get_available_models() -> list[str]: """ Return a list of available training models Returns @@ -221,7 +256,7 @@ def get_available_models(): return models @staticmethod - def get_default_model(): + def get_default_model() -> str: """ Return the default training model plugin name Returns @@ -234,7 +269,7 @@ def get_default_model(): return 'original' if 'original' in models else models[0] @staticmethod - def get_available_convert_plugins(convert_category, add_none=True): + def get_available_convert_plugins(convert_category: str, add_none: bool = True) -> list[str]: """ Return a list of available converter plugins in the given category Parameters diff --git a/plugins/train/_config.py b/plugins/train/_config.py index 0fd195a1bb..1a354d928c 100644 --- a/plugins/train/_config.py +++ b/plugins/train/_config.py @@ -1,34 +1,120 @@ #!/usr/bin/env python3 """ Default configurations for models """ +import gettext import logging import os from lib.config import FaceswapConfig from plugins.plugin_loader import PluginLoader -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +# LOCALES +_LANG = gettext.translation("plugins.train._config", localedir="locales", fallback=True) +_ = _LANG.gettext -ADDITIONAL_INFO = ("\nNB: Unless specifically stated, values changed here will only take effect " - "when creating a new model.") +logger = logging.getLogger(__name__) + +ADDITIONAL_INFO = _("\nNB: Unless specifically stated, values changed here will only take effect " + "when creating a new model.") + +_LOSS_HELP = { + "ffl": _( + "Focal Frequency Loss. Analyzes the frequency spectrum of the images rather than the " + "images themselves. This loss function can be used on its own, but the original paper " + "found increased benefits when using it as a complementary loss to another spacial loss " + "function (e.g. MSE). Ref: Focal Frequency Loss for Image Reconstruction and Synthesis " + "https://arxiv.org/pdf/2012.12821.pdf NB: This loss does not currently work on AMD " + "cards."), + "flip": _( + "Nvidia FLIP. A perceptual loss measure that approximates the difference perceived by " + "humans as they alternate quickly (or flip) between two images. Used on its own and this " + "loss function creates a distinct grid on the output. However it can be helpful when " + "used as a complimentary loss function. Ref: FLIP: A Difference Evaluator for " + "Alternating Images: " + "https://research.nvidia.com/sites/default/files/node/3260/FLIP_Paper.pdf"), + "gmsd": _( + "Gradient Magnitude Similarity Deviation seeks to match the global standard deviation of " + "the pixel to pixel differences between two images. Similar in approach to SSIM. Ref: " + "Gradient Magnitude Similarity Deviation: An Highly Efficient Perceptual Image Quality " + "Index https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf"), + "l_inf_norm": _( + "The L_inf norm will reduce the largest individual pixel error in an image. As " + "each largest error is minimized sequentially, the overall error is improved. This loss " + "will be extremely focused on outliers."), + "laploss": _( + "Laplacian Pyramid Loss. Attempts to improve results by focussing on edges using " + "Laplacian Pyramids. As this loss function gives priority to edges over other low-" + "frequency information, like color, it should not be used on its own. The original " + "implementation uses this loss as a complimentary function to MSE. " + "Ref: Optimizing the Latent Space of Generative Networks " + "https://arxiv.org/abs/1707.05776"), + "lpips_alex": _( + "LPIPS is a perceptual loss that uses the feature outputs of other pretrained models as a " + "loss metric. Be aware that this loss function will use more VRAM. Used on its own and " + "this loss will create a distinct moire pattern on the output, however it can be helpful " + "as a complimentary loss function. The output of this function is strong, so depending " + "on your chosen primary loss function, you are unlikely going to want to set the weight " + "above about 25%. Ref: The Unreasonable Effectiveness of Deep Features as a Perceptual " + "Metric http://arxiv.org/abs/1801.03924\nThis variant uses the AlexNet backbone. A fairly " + "light and old model which performed best in the paper's original implementation.\nNB: " + "For AMD Users the final linear layer is not implemented."), + "lpips_squeeze": _( + "Same as lpips_alex, but using the SqueezeNet backbone. A more lightweight " + "version of AlexNet.\nNB: For AMD Users the final linear layer is not implemented."), + "lpips_vgg16": _( + "Same as lpips_alex, but using the VGG16 backbone. A more heavyweight model.\n" + "NB: For AMD Users the final linear layer is not implemented."), + "logcosh": _( + "log(cosh(x)) acts similar to MSE for small errors and to MAE for large errors. Like " + "MSE, it is very stable and prevents overshoots when errors are near zero. Like MAE, it " + "is robust to outliers."), + "mae": _( + "Mean absolute error will guide reconstructions of each pixel towards its median value in " + "the training dataset. Robust to outliers but as a median, it can potentially ignore some " + "infrequent image types in the dataset."), + "mse": _( + "Mean squared error will guide reconstructions of each pixel towards its average value in " + "the training dataset. As an avg, it will be susceptible to outliers and typically " + "produces slightly blurrier results. Ref: Multi-Scale Structural Similarity for Image " + "Quality Assessment https://www.cns.nyu.edu/pub/eero/wang03b.pdf"), + "ms_ssim": _( + "Multiscale Structural Similarity Index Metric is similar to SSIM except that it " + "performs the calculations along multiple scales of the input image."), + "smooth_loss": _( + "Smooth_L1 is a modification of the MAE loss to correct two of its disadvantages. " + "This loss has improved stability and guidance for small errors. Ref: A General and " + "Adaptive Robust Loss Function https://arxiv.org/pdf/1701.03077.pdf"), + "ssim": _( + "Structural Similarity Index Metric is a perception-based loss that considers changes in " + "texture, luminance, contrast, and local spatial statistics of an image. Potentially " + "delivers more realistic looking images. Ref: Image Quality Assessment: From Error " + "Visibility to Structural Similarity http://www.cns.nyu.edu/pub/eero/wang03-reprint.pdf"), + "pixel_gradient_diff": _( + "Instead of minimizing the difference between the absolute value of each " + "pixel in two reference images, compute the pixel to pixel spatial difference in each " + "image and then minimize that difference between two images. Allows for large color " + "shifts, but maintains the structure of the image."), + "none": _("Do not use an additional loss function.")} + +_NON_PRIMARY_LOSS = ["flip", "lpips_alex", "lpips_squeeze", "lpips_vgg16", "none"] class Config(FaceswapConfig): """ Config File for Models """ - # pylint: disable=too-many-statements - def set_defaults(self): + # pylint:disable=too-many-statements + def set_defaults(self) -> None: """ Set the default values for config """ logger.debug("Setting defaults") self._set_globals() self._set_loss() self._defaults_from_plugin(os.path.dirname(__file__)) - def _set_globals(self): + def _set_globals(self) -> None: """ Set the global options for training """ logger.debug("Setting global config") section = "global" - self.add_section(title=section, - info="Options that apply to all models" + ADDITIONAL_INFO) + self.add_section(section, + _("Options that apply to all models") + ADDITIONAL_INFO) self.add_item( section=section, title="centering", @@ -37,21 +123,22 @@ def _set_globals(self): default="face", choices=["face", "head", "legacy"], fixed=True, - group="face", - info="How to center the training image. The extracted images are centered on the " - "middle of the skull based on the face's estimated pose. A subsection of these " - "images are used for training. The centering used dictates how this subsection " - "will be cropped from the aligned images." - "\n\tface: Centers the training image on the center of the face, adjusting for " - "pitch and yaw." - "\n\thead: Centers the training image on the center of the head, adjusting for " - "pitch and yaw. NB: You should only select head centering if you intend to " - "include the full head (including hair) in the final swap. This may give mixed " - "results. Additionally, it is only worth choosing head centering if you are " - "training with a mask that includes the hair (e.g. BiSeNet-FP-Head)." - "\n\tlegacy: The 'original' extraction technique. Centers the training image " - "near the tip of the nose with no adjustment. Can result in the edges of the " - "face appearing outside of the training area.") + group=_("face"), + info=_( + "How to center the training image. The extracted images are centered on the " + "middle of the skull based on the face's estimated pose. A subsection of these " + "images are used for training. The centering used dictates how this subsection " + "will be cropped from the aligned images." + "\n\tface: Centers the training image on the center of the face, adjusting for " + "pitch and yaw." + "\n\thead: Centers the training image on the center of the head, adjusting for " + "pitch and yaw. NB: You should only select head centering if you intend to " + "include the full head (including hair) in the final swap. This may give mixed " + "results. Additionally, it is only worth choosing head centering if you are " + "training with a mask that includes the hair (e.g. BiSeNet-FP-Head)." + "\n\tlegacy: The 'original' extraction technique. Centers the training image " + "near the tip of the nose with no adjustment. Can result in the edges of the " + "face appearing outside of the training area.")) self.add_item( section=section, title="coverage", @@ -60,68 +147,71 @@ def _set_globals(self): min_max=(62.5, 100.0), rounding=2, fixed=True, - group="face", - info="How much of the extracted image to train on. A lower coverage will limit the " - "model's scope to a zoomed-in central area while higher amounts can include the " - "entire face. A trade-off exists between lower amounts given more detail " - "versus higher amounts avoiding noticeable swap transitions. For 'Face' " - "centering you will want to leave this above 75%. For Head centering you will " - "most likely want to set this to 100%. Sensible values for 'Legacy' " - "centering are:" - "\n\t62.5% spans from eyebrow to eyebrow." - "\n\t75.0% spans from temple to temple." - "\n\t87.5% spans from ear to ear." - "\n\t100.0% is a mugshot.") - + group=_("face"), + info=_( + "How much of the extracted image to train on. A lower coverage will limit the " + "model's scope to a zoomed-in central area while higher amounts can include the " + "entire face. A trade-off exists between lower amounts given more detail " + "versus higher amounts avoiding noticeable swap transitions. For 'Face' " + "centering you will want to leave this above 75%. For Head centering you will " + "most likely want to set this to 100%. Sensible values for 'Legacy' " + "centering are:" + "\n\t62.5% spans from eyebrow to eyebrow." + "\n\t75.0% spans from temple to temple." + "\n\t87.5% spans from ear to ear." + "\n\t100.0% is a mugshot.")) self.add_item( section=section, title="icnr_init", datatype=bool, default=False, - group="initialization", - info="Use ICNR to tile the default initializer in a repeating pattern. " - "This strategy is designed for pairing with sub-pixel / pixel shuffler " - "to reduce the 'checkerboard effect' in image reconstruction. " - "\n\t https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf") + group=_("initialization"), + info=_( + "Use ICNR to tile the default initializer in a repeating pattern. " + "This strategy is designed for pairing with sub-pixel / pixel shuffler " + "to reduce the 'checkerboard effect' in image reconstruction. " + "\n\t https://arxiv.org/ftp/arxiv/papers/1707/1707.02937.pdf")) self.add_item( section=section, title="conv_aware_init", datatype=bool, default=False, - group="initialization", - info="Use Convolution Aware Initialization for convolutional layers. " - "This can help eradicate the vanishing and exploding gradient problem " - "as well as lead to higher accuracy, lower loss and faster convergence.\nNB:" - "\n\t This can use more VRAM when creating a new model so you may want to " - "lower the batch size for the first run. The batch size can be raised " - "again when reloading the model. " - "\n\t Multi-GPU is not supported for this option, so you should start the model " - "on a single GPU. Once training has started, you can stop training, enable " - "multi-GPU and resume." - "\n\t Building the model will likely take several minutes as the calculations " - "for this initialization technique are expensive. This will only impact starting " - "a new model.") + group=_("initialization"), + info=_( + "Use Convolution Aware Initialization for convolutional layers. " + "This can help eradicate the vanishing and exploding gradient problem " + "as well as lead to higher accuracy, lower loss and faster convergence.\nNB:" + "\n\t This can use more VRAM when creating a new model so you may want to " + "lower the batch size for the first run. The batch size can be raised " + "again when reloading the model. " + "\n\t Multi-GPU is not supported for this option, so you should start the model " + "on a single GPU. Once training has started, you can stop training, enable " + "multi-GPU and resume." + "\n\t Building the model will likely take several minutes as the calculations " + "for this initialization technique are expensive. This will only impact starting " + "a new model.")) self.add_item( section=section, title="optimizer", datatype=str, gui_radio=True, - group="optimizer", + group=_("optimizer"), default="adam", choices=["adabelief", "adam", "nadam", "rms-prop"], - info="The optimizer to use." - "\n\t adabelief - Adapting Stepsizes by the Belief in Observed Gradients. An " - "optimizer with the aim to converge faster, generalize better and remain more " - "stable. (https://arxiv.org/abs/2010.07468). NB: Epsilon for AdaBelief needs to " - "be set to a smaller value than other Optimizers. Generally setting the 'Epsilon " - "Exponent' to around '-16' should work." - "\n\t adam - Adaptive Moment Optimization. A stochastic gradient descent method " - "that is based on adaptive estimation of first-order and second-order moments." - "\n\t nadam - Adaptive Moment Optimization with Nesterov Momentum. Much like " - "Adam but uses a different formula for calculating momentum." - "\n\t rms-prop - Root Mean Square Propagation. Maintains a moving (discounted) " - "average of the square of the gradients. Divides the gradient by the root of " - "this average.") + info=_( + "The optimizer to use." + "\n\t adabelief - Adapting Stepsizes by the Belief in Observed Gradients. An " + "optimizer with the aim to converge faster, generalize better and remain more " + "stable. (https://arxiv.org/abs/2010.07468). NB: Epsilon for AdaBelief needs to " + "be set to a smaller value than other Optimizers. Generally setting the 'Epsilon " + "Exponent' to around '-16' should work." + "\n\t adam - Adaptive Moment Optimization. A stochastic gradient descent method " + "that is based on adaptive estimation of first-order and second-order moments." + "\n\t nadam - Adaptive Moment Optimization with Nesterov Momentum. Much like " + "Adam but uses a different formula for calculating momentum." + "\n\t rms-prop - Root Mean Square Propagation. Maintains a moving (discounted) " + "average of the square of the gradients. Divides the gradient by the root of " + "this average.")) self.add_item( section=section, title="learning_rate", @@ -130,12 +220,13 @@ def _set_globals(self): min_max=(1e-6, 1e-4), rounding=6, fixed=False, - group="optimizer", - info="Learning rate - how fast your network will learn (how large are the " - "modifications to the model weights after one batch of training). Values that " - "are too large might result in model crashes and the inability of the model to " - "find the best solution. Values that are too small might be unable to escape " - "from dead-ends and find the best global minimum.") + group=_("optimizer"), + info=_( + "Learning rate - how fast your network will learn (how large are the " + "modifications to the model weights after one batch of training). Values that " + "are too large might result in model crashes and the inability of the model to " + "find the best solution. Values that are too small might be unable to escape " + "from dead-ends and find the best global minimum.")) self.add_item( section=section, title="epsilon_exponent", @@ -144,67 +235,160 @@ def _set_globals(self): min_max=(-20, 0), rounding=1, fixed=False, - group="optimizer", - info="The epsilon adds a small constant to weight updates to attempt to avoid 'divide " - "by zero' errors. Unless you are using the AdaBelief Optimizer, then Generally " - "this option should be left at default value, For AdaBelief, setting this to " - "around '-16' should work.\n" - "In all instances if you are getting 'NaN' loss values, and have been unable to " - "resolve the issue any other way (for example, increasing batch size, or " - "lowering learning rate), then raising the epsilon can lead to a more stable " - "model. It may, however, come at the cost of slower training and a less accurate " - "final result.\n" - "NB: The value given here is the 'exponent' to the epsilon. For example, " - "choosing '-7' will set the epsilon to 1e-7. Choosing '-3' will set the epsilon " - "to 0.001 (1e-3).") + group=_("optimizer"), + info=_( + "The epsilon adds a small constant to weight updates to attempt to avoid 'divide " + "by zero' errors. Unless you are using the AdaBelief Optimizer, then Generally " + "this option should be left at default value, For AdaBelief, setting this to " + "around '-16' should work.\n" + "In all instances if you are getting 'NaN' loss values, and have been unable to " + "resolve the issue any other way (for example, increasing batch size, or " + "lowering learning rate), then raising the epsilon can lead to a more stable " + "model. It may, however, come at the cost of slower training and a less accurate " + "final result.\n" + "NB: The value given here is the 'exponent' to the epsilon. For example, " + "choosing '-7' will set the epsilon to 1e-7. Choosing '-3' will set the epsilon " + "to 0.001 (1e-3).")) + self.add_item( + section=section, + title="save_optimizer", + datatype=str, + group=_("optimizer"), + default="exit", + fixed=False, + gui_radio=True, + choices=["never", "always", "exit"], + info=_( + "When to save the Optimizer Weights. Saving the optimizer weights is not " + "necessary and will increase the model file size 3x (and by extension the amount " + "of time it takes to save the model). However, it can be useful to save these " + "weights if you want to guarantee that a resumed model carries off exactly from " + "where it left off, rather than spending a few hundred iterations catching up." + "\n\t never - Don't save optimizer weights." + "\n\t always - Save the optimizer weights at every save iteration. Model saving " + "will take longer, due to the increased file size, but you will always have the " + "last saved optimizer state in your model file." + "\n\t exit - Only save the optimizer weights when explicitly terminating a " + "model. This can be when the model is actively stopped or when the target " + "iterations are met. Note: If the training session ends because of another " + "reason (e.g. power outage, Out of Memory Error, NaN detected) then the " + "optimizer weights will NOT be saved.")) + + self.add_item( + section=section, + title="lr_finder_iterations", + datatype=int, + default=1000, + min_max=(100, 10000), + rounding=100, + fixed=True, + group=_("Learning Rate Finder"), + info=_( + "The number of iterations to process to find the optimal learning rate. Higher " + "values will take longer, but will be more accurate.")) + self.add_item( + section=section, + title="lr_finder_mode", + datatype=str, + default="set", + fixed=True, + gui_radio=True, + choices=["set", "graph_and_set", "graph_and_exit"], + group=_("Learning Rate Finder"), + info=_( + "The operation mode for the learning rate finder. Only applicable to new models. " + "For existing models this will always default to 'set'." + "\n\tset - Train with the discovered optimal learning rate." + "\n\tgraph_and_set - Output a graph in the training folder showing the discovered " + "learning rates and train with the optimal learning rate." + "\n\tgraph_and_exit - Output a graph in the training folder with the discovered " + "learning rates and exit.")) + self.add_item( + section=section, + title="lr_finder_strength", + datatype=str, + default="default", + fixed=True, + gui_radio=True, + choices=["default", "aggressive", "extreme"], + group=_("Learning Rate Finder"), + info=_( + "How aggressively to set the Learning Rate. More aggressive can learn faster, but " + "is more likely to lead to exploding gradients." + "\n\tdefault - The default optimal learning rate. A safe choice for nearly all " + "use cases." + "\n\taggressive - Set's a higher learning rate than the default. May learn faster " + "but with a higher chance of exploding gradients." + "\n\textreme - The highest optimal learning rate. A much higher risk of exploding " + "gradients.")) + self.add_item( + section=section, + title="autoclip", + datatype=bool, + default=False, + info=_( + "Apply AutoClipping to the gradients. AutoClip analyzes the " + "gradient weights and adjusts the normalization value dynamically to fit the " + "data. Can help prevent NaNs and improve model optimization at the expense of " + "VRAM. Ref: AutoClip: Adaptive Gradient Clipping for Source Separation Networks " + "https://arxiv.org/abs/2007.14469"), + fixed=False, + gui_radio=True, + group=_("optimizer")) self.add_item( section=section, title="reflect_padding", datatype=bool, default=False, - group="network", - info="Use reflection padding rather than zero padding with convolutions. " - "Each convolution must pad the image boundaries to maintain the proper " - "sizing. More complex padding schemes can reduce artifacts at the " - "border of the image." - "\n\t http://www-cs.engr.ccny.cuny.edu/~wolberg/cs470/hw/hw2_pad.txt") + group=_("network"), + info=_( + "Use reflection padding rather than zero padding with convolutions. " + "Each convolution must pad the image boundaries to maintain the proper " + "sizing. More complex padding schemes can reduce artifacts at the " + "border of the image." + "\n\t http://www-cs.engr.ccny.cuny.edu/~wolberg/cs470/hw/hw2_pad.txt")) self.add_item( section=section, title="allow_growth", datatype=bool, default=False, - group="network", + group=_("network"), fixed=False, - info="[Nvidia Only]. Enable the Tensorflow GPU 'allow_growth' configuration option. " - "This option prevents Tensorflow from allocating all of the GPU VRAM at launch " - "but can lead to higher VRAM fragmentation and slower performance. Should only " - "be enabled if you are receiving errors regarding 'cuDNN fails to initialize' " - "when commencing training.") + info=_( + "Enable the Tensorflow GPU 'allow_growth' configuration option. " + "This option prevents Tensorflow from allocating all of the GPU VRAM at launch " + "but can lead to higher VRAM fragmentation and slower performance. Should only " + "be enabled if you are receiving errors regarding 'cuDNN fails to initialize' " + "when commencing training.")) self.add_item( section=section, title="mixed_precision", datatype=bool, default=False, - group="network", - info="[Nvidia Only], NVIDIA GPUs can run operations in float16 faster than in " - "float32. Mixed precision allows you to use a mix of float16 with float32, to " - "get the performance benefits from float16 and the numeric stability benefits " - "from float32.\n\nWhile mixed precision will run on most Nvidia models, it will " - "only speed up training on more recent GPUs. Those with compute capability 7.0 " - "or higher will see the greatest performance benefit from mixed precision " - "because they have Tensor Cores. Older GPUs offer no math performance benefit " - "for using mixed precision, however memory and bandwidth savings can enable some " - "speedups. Generally RTX GPUs and later will offer the most benefit.") + fixed=False, + group=_("network"), + info=_( + "NVIDIA GPUs can run operations in float16 faster than in " + "float32. Mixed precision allows you to use a mix of float16 with float32, to " + "get the performance benefits from float16 and the numeric stability benefits " + "from float32.\n\nThis is untested on DirectML backend, but will run on most " + "Nvidia models. it will only speed up training on more recent GPUs. Those with " + "compute capability 7.0 or higher will see the greatest performance benefit from " + "mixed precision because they have Tensor Cores. Older GPUs offer no math " + "performance benefit for using mixed precision, however memory and bandwidth " + "savings can enable some speedups. Generally RTX GPUs and later will offer the " + "most benefit.")) self.add_item( section=section, title="nan_protection", datatype=bool, default=True, - group="network", - info="If a 'NaN' is generated in the model, this means that the model has corrupted " - "and the model is likely to start deteriorating from this point on. Enabling NaN " - "protection will stop training immediately in the event of a NaN. The last save " - "will not contain the NaN, so you may still be able to rescue your model.", + group=_("network"), + info=_( + "If a 'NaN' is generated in the model, this means that the model has corrupted " + "and the model is likely to start deteriorating from this point on. Enabling NaN " + "protection will stop training immediately in the event of a NaN. The last save " + "will not contain the NaN, so you may still be able to rescue your model."), fixed=False) self.add_item( section=section, @@ -214,150 +398,198 @@ def _set_globals(self): min_max=(1, 32), rounding=1, fixed=False, - group="convert", - info="[GPU Only]. The number of faces to feed through the model at once when running " - "the Convert process.\n\nNB: Increasing this figure is unlikely to improve " - "convert speed, however, if you are getting Out of Memory errors, then you may " - "want to reduce the batch size.") + group=_("convert"), + info=_( + "[GPU Only]. The number of faces to feed through the model at once when running " + "the Convert process.\n\nNB: Increasing this figure is unlikely to improve " + "convert speed, however, if you are getting Out of Memory errors, then you may " + "want to reduce the batch size.")) - def _set_loss(self): + def _set_loss(self) -> None: + # pylint:disable=line-too-long """ Set the default loss options. Loss Documentation - MAE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine - -learners-should-know-4fb140e9d4b0 - MSE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine - -learners-should-know-4fb140e9d4b0 - LogCosh https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine - -learners-should-know-4fb140e9d4b0 - Smooth L1 https://arxiv.org/pdf/1701.03077.pdf - L_inf_norm https://medium.com/@montjoile/l0-norm-l1-norm-l2-norm-l-infinity - -norm-7a7d18a4f40c - SSIM http://www.cns.nyu.edu/pub/eero/wang03-reprint.pdf - MSSIM https://www.cns.nyu.edu/pub/eero/wang03b.pdf - GMSD https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf - """ + MAE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine-learners-should-know-4fb140e9d4b0 + MSE https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine-learners-should-know-4fb140e9d4b0 + LogCosh https://heartbeat.fritz.ai/5-regression-loss-functions-all-machine-learners-should-know-4fb140e9d4b0 + L_inf_norm https://medium.com/@montjoile/l0-norm-l1-norm-l2-norm-l-infinity-norm-7a7d18a4f40c + """ # noqa + # pylint:enable=line-too-long logger.debug("Setting Loss config") section = "global.loss" - self.add_section(title=section, - info="Loss configuration options\n" - "Loss is the mechanism by which a Neural Network judges how well it " - "thinks that it is recreating a face." + ADDITIONAL_INFO) + self.add_section(section, + _("Loss configuration options\n" + "Loss is the mechanism by which a Neural Network judges how well it " + "thinks that it is recreating a face.") + ADDITIONAL_INFO) self.add_item( section=section, title="loss_function", datatype=str, - group="loss", + group=_("loss"), default="ssim", - choices=["mae", "mse", "logcosh", "smooth_loss", "l_inf_norm", "ssim", "ms_ssim", - "gmsd", "pixel_gradient_diff"], - info="The loss function to use." - "\n\t MAE - Mean absolute error will guide reconstructions of each pixel " - "towards its median value in the training dataset. Robust to outliers but as " - "a median, it can potentially ignore some infrequent image types in the dataset." - "\n\t MSE - Mean squared error will guide reconstructions of each pixel " - "towards its average value in the training dataset. As an avg, it will be " - "susceptible to outliers and typically produces slightly blurrier results." - "\n\t LogCosh - log(cosh(x)) acts similar to MSE for small errors and to " - "MAE for large errors. Like MSE, it is very stable and prevents overshoots " - "when errors are near zero. Like MAE, it is robust to outliers. NB: Due to a bug " - "in PlaidML, this loss does not work on AMD cards." - "\n\t Smooth_L1 --- Modification of the MAE loss to correct two of its " - "disadvantages. This loss has improved stability and guidance for small errors." - "\n\t L_inf_norm --- The L_inf norm will reduce the largest individual pixel " - "error in an image. As each largest error is minimized sequentially, the " - "overall error is improved. This loss will be extremely focused on outliers." - "\n\t SSIM - Structural Similarity Index Metric is a perception-based " - "loss that considers changes in texture, luminance, contrast, and local spatial " - "statistics of an image. Potentially delivers more realistic looking images." - "\n\t MS_SSIM - Multiscale Structural Similarity Index Metric is similar to SSIM " - "except that it performs the calculations along multiple scales of the input " - "image. NB: This loss currently does not work on AMD Cards." - "\n\t GMSD - Gradient Magnitude Similarity Deviation seeks to match " - "the global standard deviation of the pixel to pixel differences between two " - "images. Similar in approach to SSIM. NB: This loss does not currently work on " - "AMD cards." - "\n\t Pixel_Gradient_Difference - Instead of minimizing the difference between " - "the absolute value of each pixel in two reference images, compute the pixel to " - "pixel spatial difference in each image and then minimize that difference " - "between two images. Allows for large color shifts, but maintains the structure " - "of the image.") + fixed=False, + choices=[x for x in sorted(_LOSS_HELP) if x not in _NON_PRIMARY_LOSS], + info=(_("The loss function to use.") + + "\n\n\t" + "\n\n\t".join(f"{k}: {v}" + for k, v in sorted(_LOSS_HELP.items()) + if k not in _NON_PRIMARY_LOSS))) self.add_item( section=section, - title="mask_loss_function", + title="loss_function_2", datatype=str, - group="loss", + group=_("loss"), default="mse", - choices=["mae", "mse"], - info="The loss function to use when learning a mask." - "\n\t MAE - Mean absolute error will guide reconstructions of each pixel " - "towards its median value in the training dataset. Robust to outliers but as " - "a median, it can potentially ignore some infrequent image types in the dataset." - "\n\t MSE - Mean squared error will guide reconstructions of each pixel " - "towards its average value in the training dataset. As an average, it will be " - "susceptible to outliers and typically produces slightly blurrier results.") + fixed=False, + choices=list(sorted(_LOSS_HELP)), + info=(_("The second loss function to use. If using a structural based loss (such as " + "SSIM, MS-SSIM or GMSD) it is common to add an L1 regularization(MAE) or L2 " + "regularization (MSE) function. You can adjust the weighting of this loss " + "function with the loss_weight_2 option.") + + "\n\n\t" + "\n\n\t".join(f"{k}: {v}" for k, v in sorted(_LOSS_HELP.items())))) self.add_item( section=section, - title="l2_reg_term", + title="loss_weight_2", datatype=int, - group="loss", + group=_("loss"), min_max=(0, 400), rounding=1, default=100, - info="The amount of L2 Regularization to apply as a penalty to Structural Similarity " - "loss functions.\n\nNB: You should only adjust this if you know what you are " - "doing!\n\n" - "L2 regularization applies a penalty term to the given Loss function. This " - "penalty will only be applied if SSIM, MS-SSIM or GMSD is selected for the main " - "loss function, otherwise it is ignored." - "\n\nThe value given here is as a percentage weight of the main loss function. " - "For example:" - "\n\t 100 - Will give equal weighting to the main loss and the penalty function. " - "\n\t 25 - Will give the penalty function 1/4 of the weight of the main loss " - "function. " - "\n\t 400 - Will give the penalty function 4x as much importance as the main " - "loss function." - "\n\t 0 - Disables L2 Regularization altogether.") + fixed=False, + info=_( + "The amount of weight to apply to the second loss function.\n\n" + "\n\nThe value given here is as a percentage denoting how much the selected " + "function should contribute to the overall loss cost of the model. For example:" + "\n\t 100 - The loss calculated for the second loss function will be applied at " + "its full amount towards the overall loss score. " + "\n\t 25 - The loss calculated for the second loss function will be reduced by a " + "quarter prior to adding to the overall loss score. " + "\n\t 400 - The loss calculated for the second loss function will be mulitplied " + "4 times prior to adding to the overall loss score. " + "\n\t 0 - Disables the second loss function altogether.")) + self.add_item( + section=section, + title="loss_function_3", + datatype=str, + group=_("loss"), + default="none", + fixed=False, + choices=list(sorted(_LOSS_HELP)), + info=(_("The third loss function to use. You can adjust the weighting of this loss " + "function with the loss_weight_3 option.") + + "\n\n\t" + + "\n\n\t".join(f"{k}: {v}" for k, v in sorted(_LOSS_HELP.items())))) + self.add_item( + section=section, + title="loss_weight_3", + datatype=int, + group=_("loss"), + min_max=(0, 400), + rounding=1, + default=0, + fixed=False, + info=_( + "The amount of weight to apply to the third loss function.\n\n" + "\n\nThe value given here is as a percentage denoting how much the selected " + "function should contribute to the overall loss cost of the model. For example:" + "\n\t 100 - The loss calculated for the third loss function will be applied at " + "its full amount towards the overall loss score. " + "\n\t 25 - The loss calculated for the third loss function will be reduced by a " + "quarter prior to adding to the overall loss score. " + "\n\t 400 - The loss calculated for the third loss function will be mulitplied 4 " + "times prior to adding to the overall loss score. " + "\n\t 0 - Disables the third loss function altogether.")) + self.add_item( + section=section, + title="loss_function_4", + datatype=str, + group=_("loss"), + default="none", + fixed=False, + choices=list(sorted(_LOSS_HELP)), + info=(_("The fourth loss function to use. You can adjust the weighting of this loss " + "function with the loss_weight_3 option.") + + "\n\n\t" + + "\n\n\t".join(f"{k}: {v}" for k, v in sorted(_LOSS_HELP.items())))) + self.add_item( + section=section, + title="loss_weight_4", + datatype=int, + group=_("loss"), + min_max=(0, 400), + rounding=1, + default=0, + fixed=False, + info=_( + "The amount of weight to apply to the fourth loss function.\n\n" + "\n\nThe value given here is as a percentage denoting how much the selected " + "function should contribute to the overall loss cost of the model. For example:" + "\n\t 100 - The loss calculated for the fourth loss function will be applied at " + "its full amount towards the overall loss score. " + "\n\t 25 - The loss calculated for the fourth loss function will be reduced by a " + "quarter prior to adding to the overall loss score. " + "\n\t 400 - The loss calculated for the fourth loss function will be mulitplied " + "4 times prior to adding to the overall loss score. " + "\n\t 0 - Disables the fourth loss function altogether.")) + self.add_item( + section=section, + title="mask_loss_function", + datatype=str, + group=_("loss"), + default="mse", + fixed=False, + choices=["mae", "mse"], + info=_( + "The loss function to use when learning a mask." + "\n\t MAE - Mean absolute error will guide reconstructions of each pixel " + "towards its median value in the training dataset. Robust to outliers but as " + "a median, it can potentially ignore some infrequent image types in the dataset." + "\n\t MSE - Mean squared error will guide reconstructions of each pixel " + "towards its average value in the training dataset. As an average, it will be " + "susceptible to outliers and typically produces slightly blurrier results.")) self.add_item( section=section, title="eye_multiplier", datatype=int, - group="loss", + group=_("loss"), min_max=(1, 40), rounding=1, default=3, fixed=False, - info="The amount of priority to give to the eyes.\n\nThe value given here is as a " - "multiplier of the main loss score. For example:" - "\n\t 1 - The eyes will receive the same priority as the rest of the face. " - "\n\t 10 - The eyes will be given a score 10 times higher than the rest of the " - "face." - "\n\nNB: Penalized Mask Loss must be enable to use this option.") + info=_( + "The amount of priority to give to the eyes.\n\nThe value given here is as a " + "multiplier of the main loss score. For example:" + "\n\t 1 - The eyes will receive the same priority as the rest of the face. " + "\n\t 10 - The eyes will be given a score 10 times higher than the rest of the " + "face." + "\n\nNB: Penalized Mask Loss must be enable to use this option.")) self.add_item( section=section, title="mouth_multiplier", datatype=int, - group="loss", + group=_("loss"), min_max=(1, 40), rounding=1, default=2, fixed=False, - info="The amount of priority to give to the mouth.\n\nThe value given here is as a " - "multiplier of the main loss score. For Example:" - "\n\t 1 - The mouth will receive the same priority as the rest of the face. " - "\n\t 10 - The mouth will be given a score 10 times higher than the rest of the " - "face." - "\n\nNB: Penalized Mask Loss must be enable to use this option.") + info=_( + "The amount of priority to give to the mouth.\n\nThe value given here is as a " + "multiplier of the main loss score. For Example:" + "\n\t 1 - The mouth will receive the same priority as the rest of the face. " + "\n\t 10 - The mouth will be given a score 10 times higher than the rest of the " + "face." + "\n\nNB: Penalized Mask Loss must be enable to use this option.")) self.add_item( section=section, title="penalized_mask_loss", datatype=bool, default=True, - group="loss", - info="Image loss function is weighted by mask presence. For areas of " - "the image without the facial mask, reconstruction errors will be " - "ignored while the masked face area is prioritized. May increase " - "overall quality by focusing attention on the core face area.") + group=_("loss"), + info=_( + "Image loss function is weighted by mask presence. For areas of " + "the image without the facial mask, reconstruction errors will be " + "ignored while the masked face area is prioritized. May increase " + "overall quality by focusing attention on the core face area.")) self.add_item( section=section, title="mask_type", @@ -365,38 +597,54 @@ def _set_loss(self): default="extended", choices=PluginLoader.get_available_extractors("mask", add_none=True, extend_plugin=True), - group="mask", + group=_("mask"), gui_radio=True, - info="The mask to be used for training. If you have selected 'Learn Mask' or " - "'Penalized Mask Loss' you must select a value other than 'none'. The required " - "mask should have been selected as part of the Extract process. If it does not " - "exist in the alignments file then it will be generated prior to training " - "commencing." - "\n\tnone: Don't use a mask." - "\n\tbisenet-fp-face: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked (configurable in mask settings). " - "Use this version of bisenet-fp if your model is trained with 'face' or " - "'legacy' centering." - "\n\tbisenet-fp-head: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked (configurable in mask settings). " - "Use this version of bisenet-fp if your model is trained with 'head' centering." - "\n\tcomponents: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks to create a mask." - "\n\textended: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks and the mask is extended upwards onto the forehead." - "\n\tvgg-clear: Mask designed to provide smart segmentation of mostly frontal " - "faces clear of obstructions. Profile faces and obstructions may result in " - "sub-par performance." - "\n\tvgg-obstructed: Mask designed to provide smart segmentation of mostly " - "frontal faces. The mask model has been specifically trained to recognize " - "some facial obstructions (hands and eyeglasses). Profile faces may result in " - "sub-par performance." - "\n\tunet-dfl: Mask designed to provide smart segmentation of mostly frontal " - "faces. The mask model has been trained by community members and will need " - "testing for further description. Profile faces may result in sub-par " - "performance.") + info=_( + "The mask to be used for training. If you have selected 'Learn Mask' or " + "'Penalized Mask Loss' you must select a value other than 'none'. The required " + "mask should have been selected as part of the Extract process. If it does not " + "exist in the alignments file then it will be generated prior to training " + "commencing." + "\n\tnone: Don't use a mask." + "\n\tbisenet-fp_face: Relatively lightweight NN based mask that provides more " + "refined control over the area to be masked (configurable in mask settings). " + "Use this version of bisenet-fp if your model is trained with 'face' or " + "'legacy' centering." + "\n\tbisenet-fp_head: Relatively lightweight NN based mask that provides more " + "refined control over the area to be masked (configurable in mask settings). " + "Use this version of bisenet-fp if your model is trained with 'head' centering." + "\n\tcomponents: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks to create a mask." + "\n\tcustom_face: Custom user created, face centered mask." + "\n\tcustom_head: Custom user created, head centered mask." + "\n\textended: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks and the mask is extended upwards onto the forehead." + "\n\tvgg-clear: Mask designed to provide smart segmentation of mostly frontal " + "faces clear of obstructions. Profile faces and obstructions may result in " + "sub-par performance." + "\n\tvgg-obstructed: Mask designed to provide smart segmentation of mostly " + "frontal faces. The mask model has been specifically trained to recognize " + "some facial obstructions (hands and eyeglasses). Profile faces may result in " + "sub-par performance." + "\n\tunet-dfl: Mask designed to provide smart segmentation of mostly frontal " + "faces. The mask model has been trained by community members and will need " + "testing for further description. Profile faces may result in sub-par " + "performance.")) + self.add_item( + section=section, + title="mask_dilation", + datatype=float, + min_max=(-5.0, 5.0), + rounding=1, + default=0, + fixed=False, + group=_("mask"), + info=_( + "Dilate or erode the mask. Negative values erode the mask (make it smaller). " + "Positive values dilate the mask (make it larger). The value given is a " + "percentage of the total mask size.")) self.add_item( section=section, title="mask_blur_kernel", @@ -404,12 +652,14 @@ def _set_loss(self): min_max=(0, 9), rounding=1, default=3, - group="mask", - info="Apply gaussian blur to the mask input. This has the effect of smoothing the " - "edges of the mask, which can help with poorly calculated masks and give less " - "of a hard edge to the predicted mask. The size is in pixels (calculated from " - "a 128px mask). Set to 0 to not apply gaussian blur. This value should be odd, " - "if an even number is passed in then it will be rounded to the next odd number.") + fixed=False, + group=_("mask"), + info=_( + "Apply gaussian blur to the mask input. This has the effect of smoothing the " + "edges of the mask, which can help with poorly calculated masks and give less " + "of a hard edge to the predicted mask. The size is in pixels (calculated from " + "a 128px mask). Set to 0 to not apply gaussian blur. This value should be odd, " + "if an even number is passed in then it will be rounded to the next odd number.")) self.add_item( section=section, title="mask_threshold", @@ -417,15 +667,18 @@ def _set_loss(self): default=4, min_max=(0, 50), rounding=1, - group="mask", - info="Sets pixels that are near white to white and near black to black. Set to 0 for " - "off.") + fixed=False, + group=_("mask"), + info=_( + "Sets pixels that are near white to white and near black to black. Set to 0 for " + "off.")) self.add_item( section=section, title="learn_mask", datatype=bool, default=False, - group="mask", - info="Dedicate a portion of the model to learning how to duplicate the input " - "mask. Increases VRAM usage in exchange for learning a quick ability to try " - "to replicate more complex mask models.") + group=_("mask"), + info=_( + "Dedicate a portion of the model to learning how to duplicate the input " + "mask. Increases VRAM usage in exchange for learning a quick ability to try " + "to replicate more complex mask models.")) diff --git a/plugins/train/model/_base.py b/plugins/train/model/_base.py deleted file mode 100644 index 5c327d80bc..0000000000 --- a/plugins/train/model/_base.py +++ /dev/null @@ -1,1737 +0,0 @@ -#!/usr/bin/env python3 -""" -Base class for Models. ALL Models should at least inherit from this class. - -See :mod:`~plugins.train.model.original` for an annotated example for how to create model plugins. -""" -import logging -import os -import platform -import sys -import time - -from collections import OrderedDict -from contextlib import nullcontext - -import numpy as np -import tensorflow as tf - -from lib.serializer import get_serializer -from lib.model.backup_restore import Backup -from lib.model import losses, optimizers -from lib.model.nn_blocks import set_config as set_nnblock_config -from lib.utils import get_backend, FaceswapError -from plugins.train._config import Config - -if get_backend() == "amd": - from keras import losses as k_losses - from keras import backend as K - from keras.layers import Input - from keras.models import load_model, Model as KModel -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import losses as k_losses # pylint:disable=import-error - from tensorflow.keras import backend as K # pylint:disable=import-error - from tensorflow.keras.layers import Input # pylint:disable=import-error,no-name-in-module - from tensorflow.keras.models import load_model, Model as KModel # noqa pylint:disable=import-error,no-name-in-module - - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name -_CONFIG = None - - -def KerasModel(inputs, outputs, name): # pylint:disable=invalid-name - """ wrapper for :class:`keras.models.Model`. - - There are some minor foibles between Keras 2.2 and the Tensorflow version of Keras, so this - catches potential issues and fixes prior to returning the requested model. - - All models created within plugins should use this method, and should not call keras directly - for a model. - - Parameters - ---------- - inputs: a keras.Input object or list of keras.Input objects. - The input(s) of the model - outputs: keras objects - The output(s) of the model. - name: str - The name of the model. - - Returns - ------- - :class:`keras.models.Model` - A Keras Model - """ - if get_backend() == "amd": - logger.debug("Flattening inputs (%s) and outputs (%s) for AMD", inputs, outputs) - inputs = np.array(inputs).flatten().tolist() - outputs = np.array(outputs).flatten().tolist() - logger.debug("Flattened inputs (%s) and outputs (%s)", inputs, outputs) - return KModel(inputs, outputs, name=name) - - -def _get_all_sub_models(model, models=None): - """ For a given model, return all sub-models that occur (recursively) as children. - - Parameters - ---------- - model: :class:`keras.models.Model` - A Keras model to scan for sub models - models: `None` - Do not provide this parameter. It is used for recursion - - Returns - ------- - list - A list of all :class:`keras.models.Model`s found within the given model. The provided - model will always be returned in the first position - """ - if models is None: - models = [model] - else: - models.append(model) - for layer in model.layers: - if isinstance(layer, KModel): - _get_all_sub_models(layer, models=models) - return models - - -class ModelBase(): - """ Base class that all model plugins should inherit from. - - Parameters - ---------- - model_dir: str - The full path to the model save location - arguments: :class:`argparse.Namespace` - The arguments that were passed to the train or convert process as generated from - Faceswap's command line arguments - predict: bool, optional - ``True`` if the model is being loaded for inference, ``False`` if the model is being loaded - for training. Default: ``False`` - - Attributes - ---------- - input_shape: tuple or list - A `tuple` of `ints` defining the shape of the faces that the model takes as input. This - should be overridden by model plugins in their :func:`__init__` function. If the input size - is the same for both sides of the model, then this can be a single 3 dimensional `tuple`. - If the inputs have different sizes for `"A"` and `"B"` this should be a `list` of 2 3 - dimensional shape `tuples`, 1 for each side respectively. - trainer: str - Currently there is only one trainer available (`"original"`), so at present this attribute - can be ignored. If/when more trainers are added, then this attribute should be overridden - with the trainer name that a model requires in the model plugin's - :func:`__init__` function. - """ - def __init__(self, model_dir, arguments, predict=False): - logger.debug("Initializing ModelBase (%s): (model_dir: '%s', arguments: %s, predict: %s)", - self.__class__.__name__, model_dir, arguments, predict) - - self.input_shape = None # Must be set within the plugin after initializing - self.trainer = "original" # Override for plugin specific trainer - self.color_order = "bgr" # Override for plugin specific image color channel order - - self._args = arguments - self._is_predict = predict - self._model = None - - self._configfile = arguments.configfile if hasattr(arguments, "configfile") else None - self._load_config() - - if self.config["penalized_mask_loss"] and self.config["mask_type"] is None: - raise FaceswapError("Penalized Mask Loss has been selected but you have not chosen a " - "Mask to use. Please select a mask or disable Penalized Mask " - "Loss.") - - self._io = _IO(self, model_dir, self._is_predict) - self._check_multiple_models() - - self._state = State(model_dir, - self.name, - self._config_changeable_items, - False if self._is_predict else self._args.no_logs) - self._settings = _Settings(self._args, - self.config["mixed_precision"], - self.config["allow_growth"], - self._is_predict) - self._loss = _Loss() - - logger.debug("Initialized ModelBase (%s)", self.__class__.__name__) - - @property - def model(self): - """:class:`Keras.models.Model`: The compiled model for this plugin. """ - return self._model - - @property - def command_line_arguments(self): - """ :class:`argparse.Namespace`: The command line arguments passed to the model plugin from - either the train or convert script """ - return self._args - - @property - def coverage_ratio(self): - """ float: The ratio of the training image to crop out and train on as defined in user - configuration options. - - NB: The coverage ratio is a raw float, but will be applied to integer pixel images. - - To ensure consistent rounding and guaranteed even image size, the calculation for coverage - should always be: :math:`(original_size * coverage_ratio // 2) * 2` - """ - return self.config.get("coverage", 62.5) / 100 - - @property - def model_dir(self): - """str: The full path to the model folder location. """ - return self._io._model_dir # pylint:disable=protected-access - - @property - def config(self): - """ dict: The configuration dictionary for current plugin, as set by the user's - configuration settings. """ - global _CONFIG # pylint: disable=global-statement - if not _CONFIG: - model_name = self._config_section - logger.debug("Loading config for: %s", model_name) - _CONFIG = Config(model_name, configfile=self._configfile).config_dict - return _CONFIG - - @property - def name(self): - """ str: The name of this model based on the plugin name. """ - basename = os.path.basename(sys.modules[self.__module__].__file__) - return os.path.splitext(basename)[0].lower() - - @property - def model_name(self): - """ str: The name of the keras model. Generally this will be the same as :attr:`name` - but some plugins will override this when they contain multiple architectures """ - return self.name - - @property - def output_shapes(self): - """ list: A list of list of shape tuples for the outputs of the model with the batch - dimension removed. The outer list contains 2 sub-lists (one for each side "a" and "b"). - The inner sub-lists contain the output shapes for that side. """ - shapes = [tuple(K.int_shape(output)[-3:]) for output in self._model.outputs] - return [shapes[:len(shapes) // 2], shapes[len(shapes) // 2:]] - - @property - def iterations(self): - """ int: The total number of iterations that the model has trained. """ - return self._state.iterations - - # Private properties - @property - def _config_section(self): - """ str: The section name for the current plugin for loading configuration options from the - config file. """ - return ".".join(self.__module__.split(".")[-2:]) - - @property - def _config_changeable_items(self): - """ dict: The configuration options that can be updated after the model has already been - created. """ - return Config(self._config_section, configfile=self._configfile).changeable_items - - @property - def state(self): - """:class:`State`: The state settings for the current plugin. """ - return self._state - - def _load_config(self): - """ Load the global config for reference in :attr:`config` and set the faceswap blocks - configuration options in `lib.model.nn_blocks` """ - global _CONFIG # pylint: disable=global-statement - if not _CONFIG: - model_name = self._config_section - logger.debug("Loading config for: %s", model_name) - _CONFIG = Config(model_name, configfile=self._configfile).config_dict - - nn_block_keys = ['icnr_init', 'conv_aware_init', 'reflect_padding'] - set_nnblock_config({key: _CONFIG.pop(key) - for key in nn_block_keys}) - - def _check_multiple_models(self): - """ Check whether multiple models exist in the model folder, and that no models exist that - were trained with a different plugin than the requested plugin. - - Raises - ------ - FaceswapError - If multiple model files, or models for a different plugin from that requested exists - within the model folder - """ - multiple_models = self._io.multiple_models_in_folder - if multiple_models is None: - logger.debug("Contents of model folder are valid") - return - - if len(multiple_models) == 1: - msg = (f"You have requested to train with the '{self.name}' plugin, but a model file " - f"for the '{multiple_models[0]}' plugin already exists in the folder " - f"'{self.model_dir}'.\nPlease select a different model folder.") - else: - ptypes = "', '".join(multiple_models) - msg = (f"There are multiple plugin types ('{ptypes}') stored in the model folder '" - f"{self.model_dir}'. This is not supported.\nPlease split the model files into " - "their own folders before proceeding") - raise FaceswapError(msg) - - def build(self): - """ Build the model and assign to :attr:`model`. - - Within the defined strategy scope, either builds the model from scratch or loads an - existing model if one exists. - - If running inference, then the model is built only for the required side to perform the - swap function, otherwise the model is then compiled with the optimizer and chosen - loss function(s). - - Finally, a model summary is outputted to the logger at verbose level. - """ - self._update_legacy_models() - is_summary = hasattr(self._args, "summary") and self._args.summary - with self._settings.strategy_scope(): - if self._io.model_exists: - model = self._io._load() # pylint:disable=protected-access - if self._is_predict: - inference = _Inference(model, self._args.swap_model) - self._model = inference.model - else: - self._model = model - else: - self._validate_input_shape() - inputs = self._get_inputs() - self._model = self.build_model(inputs) - if not is_summary and not self._is_predict: - self._compile_model() - self._output_summary() - - def _update_legacy_models(self): - """ Load weights from legacy split models into new unified model, archiving old model files - to a new folder. """ - if self._legacy_mapping() is None: - return - if not all(os.path.isfile(os.path.join(self.model_dir, fname)) - for fname in self._legacy_mapping()): - return - archive_dir = f"{self.model_dir}_TF1_Archived" - if os.path.exists(archive_dir): - raise FaceswapError("We need to update your model files for use with Tensorflow 2.x, " - "but the archive folder already exists. Please remove the " - f"following folder to continue: '{archive_dir}'") - - logger.info("Updating legacy models for Tensorflow 2.x") - logger.info("Your Tensorflow 1.x models will be archived in the following location: '%s'", - archive_dir) - os.rename(self.model_dir, archive_dir) - os.mkdir(self.model_dir) - new_model = self.build_model(self._get_inputs()) - for model_name, layer_name in self._legacy_mapping().items(): - old_model = load_model(os.path.join(archive_dir, model_name), compile=False) - layer = [layer for layer in new_model.layers if layer.name == layer_name] - if not layer: - logger.warning("Skipping legacy weights from '%s'...", model_name) - continue - layer = layer[0] - logger.info("Updating legacy weights from '%s'...", model_name) - layer.set_weights(old_model.get_weights()) - filename = self._io._filename # pylint:disable=protected-access - logger.info("Saving Tensorflow 2.x model to '%s'", filename) - new_model.save(filename) - # Penalized Loss and Learn Mask used to be disabled automatically if a mask wasn't - # selected, so disable it if enabled, but mask_type is None - if self.config["mask_type"] is None: - self.config["penalized_mask_loss"] = False - self.config["learn_mask"] = False - self.config["eye_multiplier"] = 1 - self.config["mouth_multiplier"] = 1 - self._state.save() - - def _validate_input_shape(self): - """ Validate that the input shape is either a single shape tuple of 3 dimensions or - a list of 2 shape tuples of 3 dimensions. """ - assert len(self.input_shape) in (2, 3), "Input shape should either be a single 3 " \ - "dimensional shape tuple for use in both sides of the model, or a list of 2 3 " \ - "dimensional shape tuples for use in the 'A' and 'B' sides of the model" - if len(self.input_shape) == 2: - assert [len(shape) == 3 for shape in self.input_shape], "All input shapes should " \ - "have 3 dimensions" - - def _get_inputs(self): - """ Obtain the standardized inputs for the model. - - The inputs will be returned for the "A" and "B" sides in the shape as defined by - :attr:`input_shape`. - - Returns - ------- - list - A list of :class:`keras.layers.Input` tensors. This will be a list of 2 tensors (one - for each side) each of shapes :attr:`input_shape`. - """ - logger.debug("Getting inputs") - if len(self.input_shape) == 3: - input_shapes = [self.input_shape, self.input_shape] - else: - input_shapes = self.input_shape - inputs = [Input(shape=shape, name=f"face_in_{side}") - for side, shape in zip(("a", "b"), input_shapes)] - logger.debug("inputs: %s", inputs) - return inputs - - def build_model(self, inputs): - """ Override for Model Specific autoencoder builds. - - Parameters - ---------- - inputs: list - A list of :class:`keras.layers.Input` tensors. This will be a list of 2 tensors (one - for each side) each of shapes :attr:`input_shape`. - """ - raise NotImplementedError - - def _output_summary(self): - """ Output the summary of the model and all sub-models to the verbose logger. """ - if hasattr(self._args, "summary") and self._args.summary: - print_fn = None # Print straight to stdout - else: - # print to logger - print_fn = lambda x: logger.verbose("%s", x) # noqa - for idx, model in enumerate(_get_all_sub_models(self._model)): - if idx == 0: - parent = model - continue - model.summary(line_length=100, print_fn=print_fn) - parent.summary(line_length=100, print_fn=print_fn) - - def save(self): - """ Save the model to disk. - - Saves the serialized model, with weights, to the folder location specified when - initializing the plugin. If loss has dropped on both sides of the model, then - a backup is taken. - """ - self._io._save() # pylint:disable=protected-access - - def snapshot(self): - """ Creates a snapshot of the model folder to the models parent folder, with the number - of iterations completed appended to the end of the model name. """ - self._io._snapshot() # pylint:disable=protected-access - - def _compile_model(self): - """ Compile the model to include the Optimizer and Loss Function(s). """ - logger.debug("Compiling Model") - - optimizer = _Optimizer(self.config["optimizer"], - self.config["learning_rate"], - self.config.get("clipnorm", False), - 10 ** int(self.config["epsilon_exponent"]), - self._args).optimizer - if self._settings.use_mixed_precision: - optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer) - - if get_backend() == "amd": - self._rewrite_plaid_outputs() - - weights = _Weights(self) - weights.load(self._io.model_exists) - weights.freeze() - - self._loss.configure(self._model) - self._model.compile(optimizer=optimizer, loss=self._loss.functions) - self._state.add_session_loss_names(self._loss.names) - logger.debug("Compiled Model: %s", self._model) - - def _rewrite_plaid_outputs(self): - """ Rewrite the output names for models using the PlaidML (Keras 2.2.4) backend - - Keras 2.2.4 duplicates model output names if any of the models have multiple outputs - so we need to rename the outputs so we can successfully map the loss dictionaries. - - This is a bit of a hack, but it does work. - """ - # TODO Remove this rewrite code if PlaidML updates to a version of Keras where this is - # no longer necessary - if len(self._model.output_names) == len(set(self._model.output_names)): - logger.debug("Output names are unique, not rewriting: %s", self._model.output_names) - return - seen = {name: 0 for name in set(self._model.output_names)} - new_names = [] - for name in self._model.output_names: - new_names.append(f"{name}_{seen[name]}") - seen[name] += 1 - logger.debug("Output names rewritten: (old: %s, new: %s)", - self._model.output_names, new_names) - self._model.output_names = new_names - - def _legacy_mapping(self): # pylint:disable=no-self-use - """ The mapping of separate model files to single model layers for transferring of legacy - weights. - - Returns - ------- - dict or ``None`` - Dictionary of original H5 filenames for legacy models mapped to new layer names or - ``None`` if the model did not exist in Faceswap prior to Tensorflow 2 - """ - return None - - def add_history(self, loss): - """ Add the current iteration's loss history to :attr:`_io.history`. - - Called from the trainer after each iteration, for tracking loss drop over time between - save iterations. - - Parameters - ---------- - loss: list - The loss values for the A and B side for the current iteration. This should be the - collated loss values for each side. - """ - self._io.history[0].append(loss[0]) - self._io.history[1].append(loss[1]) - - -class _IO(): - """ Model saving and loading functions. - - Handles the loading and saving of the plugin model from disk as well as the model backup and - snapshot functions. - - Parameters - ---------- - plugin: :class:`Model` - The parent plugin class that owns the IO functions. - model_dir: str - The full path to the model save location - is_predict: bool - ``True`` if the model is being loaded for inference. ``False`` if the model is being loaded - for training. - """ - def __init__(self, plugin, model_dir, is_predict): - self._plugin = plugin - self._is_predict = is_predict - self._model_dir = model_dir - self._history = [[], []] # Loss histories per save iteration - self._backup = Backup(self._model_dir, self._plugin.name) - - @property - def _filename(self): - """str: The filename for this model.""" - return os.path.join(self._model_dir, f"{self._plugin.name}.h5") - - @property - def model_exists(self): - """ bool: ``True`` if a model of the type being loaded exists within the model folder - location otherwise ``False``. - """ - return os.path.isfile(self._filename) - - @property - def history(self): - """ list: list of loss histories per side for the current save iteration. """ - return self._history - - @property - def multiple_models_in_folder(self): - """ :list: or ``None`` If there are multiple model types in the requested folder, or model - types that don't correspond to the requested plugin type, then returns the list of plugin - names that exist in the folder, otherwise returns ``None`` """ - plugins = [fname.replace(".h5", "") - for fname in os.listdir(self._model_dir) - if fname.endswith(".h5")] - test_names = plugins + [self._plugin.name] - test = False if not test_names else os.path.commonprefix(test_names) == "" - retval = None if not test else plugins - logger.debug("plugin name: %s, plugins: %s, test result: %s, retval: %s", - self._plugin.name, plugins, test, retval) - return retval - - def _load(self): - """ Loads the model from disk - - If the predict function is to be called and the model cannot be found in the model folder - then an error is logged and the process exits. - - When loading the model, the plugin model folder is scanned for custom layers which are - added to Keras' custom objects. - - Returns - ------- - :class:`keras.models.Model` - The saved model loaded from disk - """ - logger.debug("Loading model: %s", self._filename) - if self._is_predict and not self.model_exists: - logger.error("Model could not be found in folder '%s'. Exiting", self._model_dir) - sys.exit(1) - - try: - model = load_model(self._filename, compile=False) - except RuntimeError as err: - if "unable to get link info" in str(err).lower(): - msg = (f"Unable to load the model from '{self._filename}'. This may be a " - "temporary error but most likely means that your model has corrupted.\n" - "You can try to load the model again but if the problem persists you " - "should use the Restore Tool to restore your model from backup.\n" - f"Original error: {str(err)}") - raise FaceswapError(msg) from err - raise err - except KeyError as err: - if "unable to open object" in str(err).lower(): - msg = (f"Unable to load the model from '{self._filename}'. This may be a " - "temporary error but most likely means that your model has corrupted.\n" - "You can try to load the model again but if the problem persists you " - "should use the Restore Tool to restore your model from backup.\n" - f"Original error: {str(err)}") - raise FaceswapError(msg) from err - raise err - - logger.info("Loaded model from disk: '%s'", self._filename) - return model - - def _save(self): - """ Backup and save the model and state file. - - Notes - ----- - The backup function actually backups the model from the previous save iteration rather than - the current save iteration. This is not a bug, but protection against long save times, as - models can get quite large, so renaming the current model file rather than copying it can - save substantial amount of time. - """ - logger.debug("Backing up and saving models") - print("") # Insert a new line to avoid spamming the same row as loss output - save_averages = self._get_save_averages() - if save_averages and self._should_backup(save_averages): - self._backup.backup_model(self._filename) - # pylint:disable=protected-access - self._backup.backup_model(self._plugin.state._filename) - - self._plugin.model.save(self._filename, include_optimizer=False) - self._plugin.state.save() - - msg = "[Saved models]" - if save_averages: - lossmsg = [f"face_{side}: {avg:.5f}" - for side, avg in zip(("a", "b"), save_averages)] - msg += f" - Average loss since last save: {', '.join(lossmsg)}" - logger.info(msg) - - def _get_save_averages(self): - """ Return the average loss since the last save iteration and reset historical loss """ - logger.debug("Getting save averages") - if not all(loss for loss in self._history): - logger.debug("No loss in history") - retval = [] - else: - retval = [sum(loss) / len(loss) for loss in self._history] - self._history = [[], []] # Reset historical loss - logger.debug("Average losses since last save: %s", retval) - return retval - - def _should_backup(self, save_averages): - """ Check whether the loss averages for this save iteration is the lowest that has been - seen. - - This protects against model corruption by only backing up the model if both sides have - seen a total fall in loss. - - Notes - ----- - This is by no means a perfect system. If the model corrupts at an iteration close - to a save iteration, then the averages may still be pushed lower than a previous - save average, resulting in backing up a corrupted model. - - Parameters - ---------- - save_averages: list - The average loss for each side for this save iteration - """ - backup = True - for side, loss in zip(("a", "b"), save_averages): - if not self._plugin.state.lowest_avg_loss.get(side, None): - logger.debug("Set initial save iteration loss average for '%s': %s", side, loss) - self._plugin.state.lowest_avg_loss[side] = loss - continue - backup = loss < self._plugin.state.lowest_avg_loss[side] if backup else backup - - if backup: # Update lowest loss values to the state file - # pylint:disable=unnecessary-comprehension - old_avgs = {key: val for key, val in self._plugin.state.lowest_avg_loss.items()} - self._plugin.state.lowest_avg_loss["a"] = save_averages[0] - self._plugin.state.lowest_avg_loss["b"] = save_averages[1] - logger.debug("Updated lowest historical save iteration averages from: %s to: %s", - old_avgs, self._plugin.state.lowest_avg_loss) - - logger.debug("Should backup: %s", backup) - return backup - - def _snapshot(self): - """ Perform a model snapshot. - - Notes - ----- - Snapshot function is called 1 iteration after the model was saved, so that it is built from - the latest save, hence iteration being reduced by 1. - """ - logger.debug("Performing snapshot. Iterations: %s", self._plugin.iterations) - self._backup.snapshot_models(self._plugin.iterations - 1) - logger.debug("Performed snapshot") - - -class _Settings(): - """ Tensorflow core training settings. - - Sets backend tensorflow settings prior to launching the model. - - Tensorflow 2 uses distribution strategies for multi-GPU/system training. These are context - managers. To enable the code to be more readable, we handle strategies the same way for Nvidia - and AMD backends. PlaidML does not support strategies, but we need to still create a context - manager so that we don't need branching logic. - - Parameters - ---------- - arguments: :class:`argparse.Namespace` - The arguments that were passed to the train or convert process as generated from - Faceswap's command line arguments - mixed_precision: bool - ``True`` if Mixed Precision training should be used otherwise ``False`` - allow_growth: bool - ``True`` if the Tensorflow allow_growth parameter should be set otherwise ``False`` - is_predict: bool, optional - ``True`` if the model is being loaded for inference, ``False`` if the model is being loaded - for training. Default: ``False`` - """ - def __init__(self, arguments, mixed_precision, allow_growth, is_predict): - logger.debug("Initializing %s: (arguments: %s, mixed_precision: %s, allow_growth: %s, " - "is_predict: %s)", self.__class__.__name__, arguments, mixed_precision, - allow_growth, is_predict) - self._set_tf_settings(allow_growth, arguments.exclude_gpus) - - use_mixed_precision = not is_predict and mixed_precision and get_backend() == "nvidia" - self._use_mixed_precision = self._set_keras_mixed_precision(use_mixed_precision, - bool(arguments.exclude_gpus)) - - distributed = False if not hasattr(arguments, "distributed") else arguments.distributed - self._strategy = self._get_strategy(distributed) - logger.debug("Initialized %s", self.__class__.__name__) - - @property - def use_strategy(self): - """ bool: ``True`` if a distribution strategy is to be used otherwise ``False``. """ - return self._strategy is not None - - @property - def use_mixed_precision(self): - """ bool: ``True`` if mixed precision training has been enabled, otherwise ``False``. """ - return self._use_mixed_precision - - @classmethod - def _set_tf_settings(cls, allow_growth, exclude_devices): - """ Specify Devices to place operations on and Allow TensorFlow to manage VRAM growth. - - Enables the Tensorflow allow_growth option if requested in the command line arguments - - Parameters - ---------- - allow_growth: bool - ``True`` if the Tensorflow allow_growth parameter should be set otherwise ``False`` - exclude_devices: list or ``None`` - List of GPU device indices that should not be made available to Tensorflow. Pass - ``None`` if all devices should be made available - """ - if get_backend() == "amd": - return # No settings for AMD - if get_backend() == "cpu": - logger.verbose("Hiding GPUs from Tensorflow") - tf.config.set_visible_devices([], "GPU") - return - - if not exclude_devices and not allow_growth: - logger.debug("Not setting any specific Tensorflow settings") - return - - gpus = tf.config.list_physical_devices('GPU') - if exclude_devices: - gpus = [gpu for idx, gpu in enumerate(gpus) if idx not in exclude_devices] - logger.debug("Filtering devices to: %s", gpus) - tf.config.set_visible_devices(gpus, "GPU") - - if allow_growth: - logger.debug("Setting Tensorflow 'allow_growth' option") - for gpu in gpus: - logger.info("Setting allow growth for GPU: %s", gpu) - tf.config.experimental.set_memory_growth(gpu, True) - logger.debug("Set Tensorflow 'allow_growth' option") - - @classmethod - def _set_keras_mixed_precision(cls, use_mixed_precision, exclude_gpus): - """ Enable the Keras experimental Mixed Precision API. - - Enables the Keras experimental Mixed Precision API if requested in the user configuration - file. - - Parameters - ---------- - use_mixed_precision: bool - ``True`` if experimental mixed precision support should be enabled for Nvidia GPUs - otherwise ``False``. - exclude_gpus: bool - ``True`` If connected GPUs are being excluded otherwise ``False``. - """ - logger.debug("use_mixed_precision: %s, exclude_gpus: %s", - use_mixed_precision, exclude_gpus) - if not use_mixed_precision: - logger.debug("Not enabling 'mixed_precision' (backend: %s, use_mixed_precision: %s)", - get_backend(), use_mixed_precision) - return False - logger.info("Enabling Mixed Precision Training.") - - policy = tf.keras.mixed_precision.Policy('mixed_float16') - tf.keras.mixed_precision.set_global_policy(policy) - logger.debug("Enabled mixed precision. (Compute dtype: %s, variable_dtype: %s)", - policy.compute_dtype, policy.variable_dtype) - return True - - @classmethod - def _get_strategy(cls, distributed): - """ If we are running on Nvidia backend and the strategy is not `"default"` then return - the correct tensorflow distribution strategy, otherwise return ``None``. - - Notes - ----- - By default Tensorflow defaults mirrored strategy to use the Nvidia NCCL method for - reductions, however this is only available in Linux, so the method used falls back to - `Hierarchical Copy All Reduce` if the OS is not Linux. - - Parameters - ---------- - distributed: bool - ``True`` if Tensorflow mirrored strategy should be used for multiple GPU training. - ``False`` if the default strategy should be used. - - Returns - ------- - :class:`tensorflow.python.distribute.Strategy` or `None` - The request Tensorflow Strategy if the backend is Nvidia and the strategy is not - `"Default"` otherwise ``None`` - """ - if get_backend() != "nvidia": - retval = None - elif distributed: - if platform.system().lower() == "linux": - cross_device_ops = tf.distribute.NcclAllReduce() - else: - cross_device_ops = tf.distribute.HierarchicalCopyAllReduce() - logger.debug("cross_device_ops: %s", cross_device_ops) - retval = tf.distribute.MirroredStrategy(cross_device_ops=cross_device_ops) - else: - retval = tf.distribute.get_strategy() - logger.debug("Using strategy: %s", retval) - return retval - - def strategy_scope(self): - """ Return the strategy scope if we have set a strategy, otherwise return a null - context. - - Returns - ------- - :func:`tensorflow.python.distribute.Strategy.scope` or :func:`contextlib.nullcontext` - The tensorflow strategy scope if a strategy is valid in the current scenario. A null - context manager if the strategy is not valid in the current scenario - """ - retval = nullcontext() if self._strategy is None else self._strategy.scope() - logger.debug("Using strategy scope: %s", retval) - return retval - - -class _Weights(): - """ Handling of freezing and loading model weights - - Parameters - ---------- - plugin: :class:`Model` - The parent plugin class that owns the IO functions. - """ - def __init__(self, plugin): - logger.debug("Initializing %s: (plugin: %s)", self.__class__.__name__, plugin) - self._model = plugin.model - self._name = plugin.model_name - self._do_freeze = plugin._args.freeze_weights - self._weights_file = self._check_weights_file(plugin._args.load_weights) - - freeze_layers = plugin.config.get("freeze_layers") # Standardized config for freezing - load_layers = plugin.config.get("load_layers") # Standardized config for loading - self._freeze_layers = freeze_layers if freeze_layers else ["encoder"] # No plugin config - self._load_layers = load_layers if load_layers else ["encoder"] # No plugin config - logger.debug("Initialized %s", self.__class__.__name__) - - @classmethod - def _check_weights_file(cls, weights_file): - """ Validate that we have a valid path to a .h5 file. - - Parameters - ---------- - weights_file: str - The full path to a weights file - - Returns - ------- - str - The full path to a weights file - """ - if not weights_file: - logger.debug("No weights file selected.") - return None - - msg = "" - if not os.path.exists(weights_file): - msg = f"Load weights selected, but the path '{weights_file}' does not exist." - elif not os.path.splitext(weights_file)[-1].lower() == ".h5": - msg = (f"Load weights selected, but the path '{weights_file}' is not a valid Keras " - f"model (.h5) file.") - - if msg: - msg += " Please check and try again." - raise FaceswapError(msg) - - logger.verbose("Using weights file: %s", weights_file) - return weights_file - - def freeze(self): - """ If freeze has been selected in the cli arguments, then freeze those models indicated - in the plugin's configuration. """ - # Blanket unfreeze layers, as checking the value of :attr:`layer.trainable` appears to - # return ``True`` even when the weights have been frozen - for layer in _get_all_sub_models(self._model): - layer.trainable = True - - if not self._do_freeze: - logger.debug("Freeze weights deselected. Not freezing") - return - - for layer in _get_all_sub_models(self._model): - if layer.name in self._freeze_layers: - logger.info("Freezing weights for '%s' in model '%s'", layer.name, self._name) - layer.trainable = False - self._freeze_layers.remove(layer.name) - if self._freeze_layers: - logger.warning("The following layers were set to be frozen but do not exist in the " - "model: %s", self._freeze_layers) - - def load(self, model_exists): - """ Load weights for newly created models, or output warning for pre-existing models. - - Parameters - ---------- - model_exists: bool - ``True`` if a model pre-exists and is being resumed, ``False`` if this is a new model - """ - if not self._weights_file: - logger.debug("No weights file provided. Not loading weights.") - return - if model_exists and self._weights_file: - logger.warning("Ignoring weights file '%s' as this model is resuming.", - self._weights_file) - return - - weights_models = self._get_weights_model() - all_models = _get_all_sub_models(self._model) - - for model_name in self._load_layers: - sub_model = next((lyr for lyr in all_models if lyr.name == model_name), None) - sub_weights = next((lyr for lyr in weights_models if lyr.name == model_name), None) - - if not sub_model or not sub_weights: - msg = f"Skipping layer {model_name} as not in " - msg += "current_model." if not sub_model else f"weights '{self._weights_file}.'" - logger.warning(msg) - continue - - logger.info("Loading weights for layer '%s'", model_name) - skipped_ops = 0 - loaded_ops = 0 - for layer in sub_model.layers: - success = self._load_layer_weights(layer, sub_weights, model_name) - if success == 0: - skipped_ops += 1 - elif success == 1: - loaded_ops += 1 - - del weights_models - - if loaded_ops == 0: - raise FaceswapError(f"No weights were succesfully loaded from your weights file: " - f"'{self._weights_file}'. Please check and try again.") - if skipped_ops > 0: - logger.warning("%s weight(s) were unable to be loaded for your model. This is most " - "likely because the weights you are trying to load were trained with " - "different settings than you have set for your current model.", - skipped_ops) - - def _get_weights_model(self): - """ Obtain a list of all sub-models contained within the weights model. - - Returns - ------- - list - List of all models contained within the .h5 file - - Raises - ------ - FaceswapError - In the event of a failure to load the weights, or the weights belonging to a different - model - """ - retval = _get_all_sub_models(load_model(self._weights_file, compile=False)) - if not retval: - raise FaceswapError(f"Error loading weights file {self._weights_file}.") - - if retval[0].name != self._name: - raise FaceswapError(f"You are attempting to load weights from a '{retval[0].name}' " - f"model into a '{self._name}' model. This is not supported.") - return retval - - def _load_layer_weights(self, layer, sub_weights, model_name): - """ Load the weights for a single layer. - - Parameters - ---------- - layer: :class:`keras.layers.Layer` - The layer to set the weights for - sub_weights: list - The list of layers in the weights model to load weights from - model_name: str - The name of the current sub-model that is having it's weights loaded - - Returns - ------- - int - `-1` if the layer has no weights to load. `0` if weights loading was unsuccessful. `1` - if weights loading was successful - """ - old_weights = layer.get_weights() - if not old_weights: - logger.debug("Skipping layer without weights: %s", layer.name) - return -1 - - layer_weights = next((lyr for lyr in sub_weights.layers if lyr.name == layer.name), None) - if not layer_weights: - logger.warning("The weights file '%s' for layer '%s' does not contain weights for " - "'%s'. Skipping", self._weights_file, model_name, layer.name) - return 0 - - new_weights = layer_weights.get_weights() - if old_weights[0].shape != new_weights[0].shape: - logger.warning("The weights for layer '%s' are of incompatible shapes. Skipping.", - layer.name) - return 0 - logger.verbose("Setting weights for '%s'", layer.name) - layer.set_weights(layer_weights.get_weights()) - return 1 - - -class _Optimizer(): # pylint:disable=too-few-public-methods - """ Obtain the selected optimizer with the appropriate keyword arguments. - - Parameters - ---------- - optimizer: str - The selected optimizer name for the plugin - learning_rate: float - The selected learning rate to use - clipnorm: bool - Whether to clip gradients to avoid exploding/vanishing gradients - epsilon: float - The value to use for the epsilon of the optimizer - arguments: :class:`argparse.Namespace` - The arguments that were passed to the train or convert process as generated from - Faceswap's command line arguments - """ - def __init__(self, optimizer, learning_rate, clipnorm, epsilon, arguments): - logger.debug("Initializing %s: (optimizer: %s, learning_rate: %s, clipnorm: %s, " - "epsilon: %s, arguments: %s)", self.__class__.__name__, - optimizer, learning_rate, clipnorm, epsilon, arguments) - valid_optimizers = {"adabelief": (optimizers.AdaBelief, - dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)), - "adam": (optimizers.Adam, - dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)), - "nadam": (optimizers.Nadam, - dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)), - "rms-prop": (optimizers.RMSprop, dict(epsilon=epsilon))} - self._optimizer, self._kwargs = valid_optimizers[optimizer] - - self._configure(learning_rate, clipnorm, arguments) - logger.verbose("Using %s optimizer", optimizer.title()) - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def optimizer(self): - """ :class:`keras.optimizers.Optimizer`: The requested optimizer. """ - return self._optimizer(**self._kwargs) - - def _configure(self, learning_rate, clipnorm, arguments): - """ Configure the optimizer based on user settings. - - Parameters - ---------- - learning_rate: float - The selected learning rate to use - clipnorm: bool - Whether to clip gradients to avoid exploding/vanishing gradients - arguments: :class:`argparse.Namespace` - The arguments that were passed to the train or convert process as generated from - Faceswap's command line arguments - - Notes - ----- - Clip-norm is ballooning VRAM usage, which is not expected behavior and may be a bug in - Keras/Tensorflow. - - PlaidML has a bug regarding the clip-norm parameter See: - https://github.com/plaidml/plaidml/issues/228. We workaround by simply not adding this - parameter for AMD backend users. - """ - lr_key = "lr" if get_backend() == "amd" else "learning_rate" - self._kwargs[lr_key] = learning_rate - - if clipnorm and (arguments.distributed or _CONFIG["mixed_precision"]): - logger.warning("Clipnorm has been selected, but is unsupported when using distributed " - "or mixed_precision training, so has been disabled. If you wish to " - "enable clipnorm, then you must disable these options.") - clipnorm = False - if clipnorm and get_backend() == "amd": - # TODO add clipnorm in for plaidML when it is fixed upstream. Still not fixed in - # release 0.7.0. - logger.warning("Due to a bug in plaidML, clipnorm cannot be used on AMD backends so " - "has been disabled") - clipnorm = False - if clipnorm: - self._kwargs["clipnorm"] = 1.0 - - logger.debug("optimizer kwargs: %s", self._kwargs) - - -class _Loss(): - """ Holds loss names and functions for an Autoencoder. """ - def __init__(self): - logger.debug("Initializing %s", self.__class__.__name__) - self._loss_dict = dict(mae=k_losses.mean_absolute_error, - mse=k_losses.mean_squared_error, - logcosh=k_losses.logcosh, - smooth_loss=losses.GeneralizedLoss(), - l_inf_norm=losses.LInfNorm(), - ssim=losses.DSSIMObjective(), - ms_ssim=losses.MSSSIMLoss(), - gmsd=losses.GMSDLoss(), - pixel_gradient_diff=losses.GradientLoss()) - self._uses_l2_reg = ["ssim", "ms_ssim", "gmsd"] - self._inputs = None - self._names = [] - self._funcs = {} - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def names(self): - """ list: The list of loss names for the model. """ - return self._names - - @property - def functions(self): - """ dict: The loss functions that apply to each model output. """ - return self._funcs - - @property - def _config(self): - """ :dict: The configuration options for this plugin """ - return _CONFIG - - @property - def _mask_inputs(self): - """ list: The list of input tensors to the model that contain the mask. Returns ``None`` - if there is no mask input to the model. """ - mask_inputs = [inp for inp in self._inputs if inp.name.startswith("mask")] - return None if not mask_inputs else mask_inputs - - @property - def _mask_shapes(self): - """ list: The list of shape tuples for the mask input tensors for the model. Returns - ``None`` if there is no mask input. """ - if self._mask_inputs is None: - return None - return [K.int_shape(mask_input) for mask_input in self._mask_inputs] - - def configure(self, model): - """ Configure the loss functions for the given inputs and outputs. - - Parameters - ---------- - model: :class:`keras.models.Model` - The model that is to be trained - """ - self._inputs = model.inputs - self._set_loss_names(model.outputs) - self._set_loss_functions(model.output_names) - self._names.insert(0, "total") - - def _set_loss_names(self, outputs): - """ Name the losses based on model output. - - This is used for correct naming in the state file, for display purposes only. - - Adds the loss names to :attr:`names` - - Notes - ----- - TODO Currently there is an issue in Tensorflow that wraps all outputs in an Identity layer - when running in Eager Execution mode, which means we cannot use the name of the output - layers to name the losses (https://github.com/tensorflow/tensorflow/issues/32180). - With this in mind, losses are named based on their shapes - - Parameters - ---------- - outputs: list - A list of output tensors from the model plugin - """ - # TODO Use output names if/when these are fixed upstream - split_outputs = [outputs[:len(outputs) // 2], outputs[len(outputs) // 2:]] - for side, side_output in zip(("a", "b"), split_outputs): - output_names = [output.name for output in side_output] - output_shapes = [K.int_shape(output)[1:] for output in side_output] - output_types = ["mask" if shape[-1] == 1 else "face" for shape in output_shapes] - logger.debug("side: %s, output names: %s, output_shapes: %s, output_types: %s", - side, output_names, output_shapes, output_types) - for idx, name in enumerate(output_types): - suffix = "" if output_types.count(name) == 1 else f"_{idx}" - self._names.append(f"{name}_{side}{suffix}") - logger.debug(self._names) - - def _set_loss_functions(self, output_names): - """ Set the loss functions and their associated weights. - - Adds the loss functions to the :attr:`functions` dictionary. - - Parameters - ---------- - output_names: list - The output names from the model - """ - mask_channels = self._get_mask_channels() - face_loss = self._loss_dict[self._config["loss_function"]] - - for name, output_name in zip(self._names, output_names): - if name.startswith("mask"): - loss_func = self._loss_dict[self._config["mask_loss_function"]] - else: - loss_func = losses.LossWrapper() - loss_func.add_loss(face_loss, mask_channel=mask_channels[0]) - self._add_l2_regularization_term(loss_func, mask_channels[0]) - - channel_idx = 1 - for multiplier in ("eye_multiplier", "mouth_multiplier"): - mask_channel = mask_channels[channel_idx] - if self._config[multiplier] > 1: - loss_func.add_loss(face_loss, - weight=self._config[multiplier] * 1.0, - mask_channel=mask_channel) - self._add_l2_regularization_term(loss_func, mask_channel) - channel_idx += 1 - - logger.debug("%s: (output_name: '%s', function: %s)", name, output_name, loss_func) - self._funcs[output_name] = loss_func - logger.debug("functions: %s", self._funcs) - - def _add_l2_regularization_term(self, loss_wrapper, mask_channel): - """ Check if an L2 Regularization term should be added and add to the loss function - wrapper. - - Parameters - ---------- - loss_wrapper: :class:`lib.model.losses.LossWrapper` - The wrapper loss function that holds the face losses - mask_channel: int - The channel that holds the mask in `y_true`, if a mask is used for the loss. - `-1` if the input is not masked - """ - if self._config["loss_function"] in self._uses_l2_reg and self._config["l2_reg_term"] > 0: - logger.debug("Adding L2 Regularization for Structural Loss") - loss_wrapper.add_loss(self._loss_dict["mse"], - weight=self._config["l2_reg_term"] / 100.0, - mask_channel=mask_channel) - - def _get_mask_channels(self): - """ Obtain the channels from the face targets that the masks reside in from the training - data generator. - - Returns - ------- - list: - A list of channel indices that contain the mask for the corresponding config item - """ - eye_multiplier = self._config["eye_multiplier"] - mouth_multiplier = self._config["mouth_multiplier"] - if not self._config["penalized_mask_loss"] and (eye_multiplier > 1 or - mouth_multiplier > 1): - logger.warning("You have selected eye/mouth loss multipliers greater than 1x, but " - "Penalized Mask Loss is disabled. Disabling all multipliers.") - eye_multiplier = 1 - mouth_multiplier = 1 - uses_masks = (self._config["penalized_mask_loss"], - eye_multiplier > 1, - mouth_multiplier > 1) - mask_channels = [-1 for _ in range(len(uses_masks))] - current_channel = 3 - for idx, mask_required in enumerate(uses_masks): - if mask_required: - mask_channels[idx] = current_channel - current_channel += 1 - logger.debug("uses_masks: %s, mask_channels: %s", uses_masks, mask_channels) - return mask_channels - - -class State(): - """ Holds state information relating to the plugin's saved model. - - Parameters - ---------- - model_dir: str - The full path to the model save location - model_name: str - The name of the model plugin - config_changeable_items: dict - Configuration options that can be altered when resuming a model, and their current values - no_logs: bool - ``True`` if Tensorboard logs should not be generated, otherwise ``False`` - """ - def __init__(self, model_dir, model_name, config_changeable_items, no_logs): - logger.debug("Initializing %s: (model_dir: '%s', model_name: '%s', " - "config_changeable_items: '%s', no_logs: %s", self.__class__.__name__, - model_dir, model_name, config_changeable_items, no_logs) - self._serializer = get_serializer("json") - filename = f"{model_name}_state.{self._serializer.file_extension}" - self._filename = os.path.join(model_dir, filename) - self._name = model_name - self._iterations = 0 - self._sessions = {} - self._lowest_avg_loss = {} - self._config = {} - self._load(config_changeable_items) - self._session_id = self._new_session_id() - self._create_new_session(no_logs, config_changeable_items) - logger.debug("Initialized %s:", self.__class__.__name__) - - @property - def loss_names(self): - """ list: The loss names for the current session """ - return self._sessions[self._session_id]["loss_names"] - - @property - def current_session(self): - """ dict: The state dictionary for the current :attr:`session_id`. """ - return self._sessions[self._session_id] - - @property - def iterations(self): - """ int: The total number of iterations that the model has trained. """ - return self._iterations - - @property - def lowest_avg_loss(self): - """dict: The lowest average save interval loss seen for each side. """ - return self._lowest_avg_loss - - @property - def session_id(self): - """ int: The current training session id. """ - return self._session_id - - def _new_session_id(self): - """ Generate a new session id. Returns 1 if this is a new model, or the last session id + 1 - if it is a pre-existing model. - - Returns - ------- - int - The newly generated session id - """ - if not self._sessions: - session_id = 1 - else: - session_id = max(int(key) for key in self._sessions.keys()) + 1 - logger.debug(session_id) - return session_id - - def _create_new_session(self, no_logs, config_changeable_items): - """ Initialize a new session, creating the dictionary entry for the session in - :attr:`_sessions`. - - Parameters - ---------- - no_logs: bool - ``True`` if Tensorboard logs should not be generated, otherwise ``False`` - config_changeable_items: dict - Configuration options that can be altered when resuming a model, and their current - values - """ - logger.debug("Creating new session. id: %s", self._session_id) - self._sessions[self._session_id] = dict(timestamp=time.time(), - no_logs=no_logs, - loss_names=[], - batchsize=0, - iterations=0, - config=config_changeable_items) - - def add_session_loss_names(self, loss_names): - """ Add the session loss names to the sessions dictionary. - - The loss names are used for Tensorboard logging - - Parameters - ---------- - loss_names: list - The list of loss names for this session. - """ - logger.debug("Adding session loss_names: %s", loss_names) - self._sessions[self._session_id]["loss_names"] = loss_names - - def add_session_batchsize(self, batch_size): - """ Add the session batch size to the sessions dictionary. - - Parameters - ---------- - batch_size: int - The batch size for the current training session - """ - logger.debug("Adding session batch size: %s", batch_size) - self._sessions[self._session_id]["batchsize"] = batch_size - - def increment_iterations(self): - """ Increment :attr:`iterations` and session iterations by 1. """ - self._iterations += 1 - self._sessions[self._session_id]["iterations"] += 1 - - def _load(self, config_changeable_items): - """ Load a state file and set the serialized values to the class instance. - - Updates the model's config with the values stored in the state file. - - Parameters - ---------- - config_changeable_items: dict - Configuration options that can be altered when resuming a model, and their current - values - """ - logger.debug("Loading State") - if not os.path.exists(self._filename): - logger.info("No existing state file found. Generating.") - return - state = self._serializer.load(self._filename) - self._name = state.get("name", self._name) - self._sessions = state.get("sessions", {}) - self._lowest_avg_loss = state.get("lowest_avg_loss", {}) - self._iterations = state.get("iterations", 0) - self._config = state.get("config", {}) - logger.debug("Loaded state: %s", state) - self._replace_config(config_changeable_items) - - def save(self): - """ Save the state values to the serialized state file. """ - logger.debug("Saving State") - state = {"name": self._name, - "sessions": self._sessions, - "lowest_avg_loss": self._lowest_avg_loss, - "iterations": self._iterations, - "config": _CONFIG} - self._serializer.save(self._filename, state) - logger.debug("Saved State") - - def _replace_config(self, config_changeable_items): - """ Replace the loaded config with the one contained within the state file. - - Check for any `fixed`=``False`` parameter changes and log info changes. - - Update any legacy config items to their current versions. - - Parameters - ---------- - config_changeable_items: dict - Configuration options that can be altered when resuming a model, and their current - values - """ - global _CONFIG # pylint: disable=global-statement - legacy_update = self._update_legacy_config() - # Add any new items to state config for legacy purposes and set sensible defaults for - # any values that may have been changed in the config file which could be detrimental. - legacy_defaults = dict(centering="legacy", - mask_loss_function="mse", - l2_reg_term=100, - optimizer="adam", - mixed_precision=False) - for key, val in _CONFIG.items(): - if key not in self._config.keys(): - setting = legacy_defaults.get(key, val) - logger.info("Adding new config item to state file: '%s': '%s'", key, setting) - self._config[key] = setting - self._update_changed_config_items(config_changeable_items) - logger.debug("Replacing config. Old config: %s", _CONFIG) - _CONFIG = self._config - if legacy_update: - self.save() - logger.debug("Replaced config. New config: %s", _CONFIG) - logger.info("Using configuration saved in state file") - - def _update_legacy_config(self): - """ Legacy updates for new config additions. - - When new config items are added to the Faceswap code, existing model state files need to be - updated to handle these new items. - - Current existing legacy update items: - - * loss - If old `dssim_loss` is ``true`` set new `loss_function` to `ssim` otherwise - set it to `mae`. Remove old `dssim_loss` item - - * masks - If `learn_mask` does not exist then it is set to ``True`` if `mask_type` is - not ``None`` otherwise it is set to ``False``. - - * masks type - Replace removed masks 'dfl_full' and 'facehull' with `components` mask - - Returns - ------- - bool - ``True`` if legacy items exist and state file has been updated, otherwise ``False`` - """ - logger.debug("Checking for legacy state file update") - priors = ["dssim_loss", "mask_type", "mask_type"] - new_items = ["loss_function", "learn_mask", "mask_type"] - updated = False - for old, new in zip(priors, new_items): - if old not in self._config: - logger.debug("Legacy item '%s' not in config. Skipping update", old) - continue - - # dssim_loss > loss_function - if old == "dssim_loss": - self._config[new] = "ssim" if self._config[old] else "mae" - del self._config[old] - updated = True - logger.info("Updated config from legacy dssim format. New config loss " - "function: '%s'", self._config[new]) - continue - - # Add learn mask option and set to True if model has "penalized_mask_loss" specified - if old == "mask_type" and new == "learn_mask" and new not in self._config: - self._config[new] = self._config["mask_type"] is not None - updated = True - logger.info("Added new 'learn_mask' config item for this model. Value set to: %s", - self._config[new]) - continue - - # Replace removed masks with most similar equivalent - if old == "mask_type" and new == "mask_type" and self._config[old] in ("facehull", - "dfl_full"): - old_mask = self._config[old] - self._config[new] = "components" - updated = True - logger.info("Updated 'mask_type' from '%s' to '%s' for this model", - old_mask, self._config[new]) - - logger.debug("State file updated for legacy config: %s", updated) - return updated - - def _update_changed_config_items(self, config_changeable_items): - """ Update any parameters which are not fixed and have been changed. - - Parameters - ---------- - config_changeable_items: dict - Configuration options that can be altered when resuming a model, and their current - values - """ - if not config_changeable_items: - logger.debug("No changeable parameters have been updated") - return - for key, val in config_changeable_items.items(): - old_val = self._config[key] - if old_val == val: - continue - self._config[key] = val - logger.info("Config item: '%s' has been updated from '%s' to '%s'", key, old_val, val) - - -class _Inference(): # pylint:disable=too-few-public-methods - """ Calculates required layers and compiles a saved model for inference. - - Parameters - ---------- - saved_model: :class:`keras.models.Model` - The saved trained Faceswap model - switch_sides: bool - ``True`` if the swap should be performed "B" > "A" ``False`` if the swap should be - "A" > "B" - """ - def __init__(self, saved_model, switch_sides): - logger.debug("Initializing: %s (saved_model: %s, switch_sides: %s)", - self.__class__.__name__, saved_model, switch_sides) - self._config = saved_model.get_config() - - self._input_idx = 1 if switch_sides else 0 - self._output_idx = 0 if switch_sides else 1 - - self._input_names = [inp[0] for inp in self._config["input_layers"]] - self._model = self._make_inference_model(saved_model) - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def model(self): - """ :class:`keras.models.Model`: The Faceswap model, compiled for inference. """ - return self._model - - def _get_nodes(self, nodes): - """ Given in input list of nodes from a :attr:`keras.models.Model.get_config` dictionary, - filters the layer name(s) and output index of the node, splitting to the correct output - index in the event of multiple inputs. - - Parameters - ---------- - nodes: list - A node entry from the :attr:`keras.models.Model.get_config` dictionary - - Returns - ------- - list - The (node name, output index) for each node passed in - """ - nodes = np.array(nodes, dtype="object")[..., :3] - num_layers = nodes.shape[0] - nodes = nodes[self._output_idx] if num_layers == 2 else nodes[0] - retval = [(node[0], node[2]) for node in nodes] - return retval - - def _make_inference_model(self, saved_model): - """ Extract the sub-models from the saved model that are required for inference. - - Parameters - ---------- - saved_model: :class:`keras.models.Model` - The saved trained Faceswap model - - Returns - ------- - :class:`keras.models.Model` - The model compiled for inference - """ - logger.debug("Compiling inference model. saved_model: %s", saved_model) - struct = self._get_filtered_structure() - model_inputs = self._get_inputs(saved_model.inputs) - compiled_layers = {} - for layer in saved_model.layers: - if layer.name not in struct: - logger.debug("Skipping unused layer: '%s'", layer.name) - continue - inbound = struct[layer.name] - logger.debug("Processing layer '%s': (layer: %s, inbound_nodes: %s)", - layer.name, layer, inbound) - if not inbound: - model = model_inputs - logger.debug("Adding model inputs %s: %s", layer.name, model) - else: - layer_inputs = [] - for inp in inbound: - inbound_layer = compiled_layers[inp[0]] - if isinstance(inbound_layer, list) and len(inbound_layer) > 1: - # Multi output inputs - inbound_output_idx = inp[1] - next_input = inbound_layer[inbound_output_idx] - logger.debug("Selecting output index %s from multi output inbound layer: " - "%s (using: %s)", inbound_output_idx, inbound_layer, - next_input) - else: - next_input = inbound_layer - - if get_backend() == "amd" and isinstance(next_input, list): - # tensorflow.keras and keras 2.2 behave differently for layer inputs - layer_inputs.extend(next_input) - else: - layer_inputs.append(next_input) - - logger.debug("Compiling layer '%s': layer inputs: %s", layer.name, layer_inputs) - model = layer(layer_inputs) - compiled_layers[layer.name] = model - retval = KerasModel(model_inputs, model, name=f"{saved_model.name}_inference") - logger.debug("Compiled inference model '%s': %s", retval.name, retval) - return retval - - def _get_filtered_structure(self): - """ Obtain the structure of the inference model. - - This parses the model config (in reverse) to obtain the required layers for an inference - model. - - Returns - ------- - :class:`collections.OrderedDict` - The layer name as key with the input name and output index as value. - """ - # Filter output layer - out = np.array(self._config["output_layers"], dtype="object") - if out.ndim == 2: - out = np.expand_dims(out, axis=1) # Needs to be expanded for _get_nodes - outputs = self._get_nodes(out) - - # Iterate backwards from the required output to get the reversed model structure - current_layers = [outputs[0]] - next_layers = [] - struct = OrderedDict() - drop_input = self._input_names[abs(self._input_idx - 1)] - switch_input = self._input_names[self._input_idx] - while True: - layer_info = current_layers.pop(0) - current_layer = next(lyr for lyr in self._config["layers"] - if lyr["name"] == layer_info[0]) - inbound = current_layer["inbound_nodes"] - - if not inbound: - break - - inbound_info = self._get_nodes(inbound) - - if any(inb[0] == drop_input for inb in inbound_info): # Switch inputs - inbound_info = [(switch_input if inb[0] == drop_input else inb[0], inb[1]) - for inb in inbound_info] - struct[layer_info[0]] = inbound_info - next_layers.extend(inbound_info) - - if not current_layers: - current_layers = next_layers - next_layers = [] - - struct[switch_input] = [] # Add the input layer - logger.debug("Model structure: %s", struct) - return struct - - def _get_inputs(self, inputs): - """ Obtain the inputs for the requested swap direction. - - Parameters - ---------- - inputs: list - The full list of input tensors to the saved faceswap training model - - Returns - ------- - list - List of input tensors to feed the model for the requested swap direction - """ - input_split = len(inputs) // 2 - start_idx = input_split * self._input_idx - retval = inputs[start_idx: start_idx + input_split] - logger.debug("model inputs: %s, input_split: %s, start_idx: %s, inference_inputs: %s", - inputs, input_split, start_idx, retval) - return retval diff --git a/plugins/train/model/_base/__init__.py b/plugins/train/model/_base/__init__.py new file mode 100644 index 0000000000..c26c15103e --- /dev/null +++ b/plugins/train/model/_base/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python3 +""" Base class for Models plugins ALL Models should at least inherit from this class. """ + +from .model import get_all_sub_models, ModelBase diff --git a/plugins/train/model/_base/io.py b/plugins/train/model/_base/io.py new file mode 100644 index 0000000000..c52bf53094 --- /dev/null +++ b/plugins/train/model/_base/io.py @@ -0,0 +1,470 @@ +#!/usr/bin/env python3 +""" +IO handling for the model base plugin. + +The objects in this module should not be called directly, but are called from +:class:`~plugins.train.model._base.ModelBase` + +This module handles: + - The loading, saving and backing up of keras models to and from disk. + - The loading and freezing of weights for model plugins. +""" +from __future__ import annotations +import logging +import os +import sys +import typing as T + +import tensorflow as tf + +from lib.model.backup_restore import Backup +from lib.utils import FaceswapError + +if T.TYPE_CHECKING: + from .model import ModelBase + +kmodels = tf.keras.models +logger = logging.getLogger(__name__) + + +def get_all_sub_models( + model: tf.keras.models.Model, + models: list[tf.keras.models.Model] | None = None) -> list[tf.keras.models.Model]: + """ For a given model, return all sub-models that occur (recursively) as children. + + Parameters + ---------- + model: :class:`tensorflow.keras.models.Model` + A Keras model to scan for sub models + models: `None` + Do not provide this parameter. It is used for recursion + + Returns + ------- + list + A list of all :class:`tensorflow.keras.models.Model` objects found within the given model. + The provided model will always be returned in the first position + """ + if models is None: + models = [model] + else: + models.append(model) + for layer in model.layers: + if isinstance(layer, kmodels.Model): + get_all_sub_models(layer, models=models) + return models + + +class IO(): + """ Model saving and loading functions. + + Handles the loading and saving of the plugin model from disk as well as the model backup and + snapshot functions. + + Parameters + ---------- + plugin: :class:`Model` + The parent plugin class that owns the IO functions. + model_dir: str + The full path to the model save location + is_predict: bool + ``True`` if the model is being loaded for inference. ``False`` if the model is being loaded + for training. + save_optimizer: ["never", "always", "exit"] + When to save the optimizer weights. `"never"` never saves the optimizer weights. `"always"` + always saves the optimizer weights. `"exit"` only saves the optimizer weights on an exit + request. + """ + def __init__(self, + plugin: ModelBase, + model_dir: str, + is_predict: bool, + save_optimizer: T.Literal["never", "always", "exit"]) -> None: + self._plugin = plugin + self._is_predict = is_predict + self._model_dir = model_dir + self._save_optimizer = save_optimizer + self._history: list[list[float]] = [[], []] # Loss histories per save iteration + self._backup = Backup(self._model_dir, self._plugin.name) + + @property + def model_dir(self) -> str: + """ str: The full path to the model folder """ + return self._model_dir + + @property + def filename(self) -> str: + """str: The filename for this model.""" + return os.path.join(self._model_dir, f"{self._plugin.name}.h5") + + @property + def model_exists(self) -> bool: + """ bool: ``True`` if a model of the type being loaded exists within the model folder + location otherwise ``False``. + """ + return os.path.isfile(self.filename) + + @property + def history(self) -> list[list[float]]: + """ list: list of loss histories per side for the current save iteration. """ + return self._history + + @property + def multiple_models_in_folder(self) -> list[str] | None: + """ :list: or ``None`` If there are multiple model types in the requested folder, or model + types that don't correspond to the requested plugin type, then returns the list of plugin + names that exist in the folder, otherwise returns ``None`` """ + plugins = [fname.replace(".h5", "") + for fname in os.listdir(self._model_dir) + if fname.endswith(".h5")] + test_names = plugins + [self._plugin.name] + test = False if not test_names else os.path.commonprefix(test_names) == "" + retval = None if not test else plugins + logger.debug("plugin name: %s, plugins: %s, test result: %s, retval: %s", + self._plugin.name, plugins, test, retval) + return retval + + def load(self) -> tf.keras.models.Model: + """ Loads the model from disk + + If the predict function is to be called and the model cannot be found in the model folder + then an error is logged and the process exits. + + When loading the model, the plugin model folder is scanned for custom layers which are + added to Keras' custom objects. + + Returns + ------- + :class:`tensorflow.keras.models.Model` + The saved model loaded from disk + """ + logger.debug("Loading model: %s", self.filename) + if self._is_predict and not self.model_exists: + logger.error("Model could not be found in folder '%s'. Exiting", self._model_dir) + sys.exit(1) + + try: + model = kmodels.load_model(self.filename, compile=False) + except RuntimeError as err: + if "unable to get link info" in str(err).lower(): + msg = (f"Unable to load the model from '{self.filename}'. This may be a " + "temporary error but most likely means that your model has corrupted.\n" + "You can try to load the model again but if the problem persists you " + "should use the Restore Tool to restore your model from backup.\n" + f"Original error: {str(err)}") + raise FaceswapError(msg) from err + raise err + except KeyError as err: + if "unable to open object" in str(err).lower(): + msg = (f"Unable to load the model from '{self.filename}'. This may be a " + "temporary error but most likely means that your model has corrupted.\n" + "You can try to load the model again but if the problem persists you " + "should use the Restore Tool to restore your model from backup.\n" + f"Original error: {str(err)}") + raise FaceswapError(msg) from err + raise err + + logger.info("Loaded model from disk: '%s'", self.filename) + return model + + def save(self, + is_exit: bool = False, + force_save_optimizer: bool = False) -> None: + """ Backup and save the model and state file. + + Parameters + ---------- + is_exit: bool, optional + ``True`` if the save request has come from an exit process request otherwise ``False``. + Default: ``False`` + force_save_optimizer: bool, optional + ``True`` to force saving the optimizer weights with the model, otherwise ``False``. + Default:``False`` + + Notes + ----- + The backup function actually backups the model from the previous save iteration rather than + the current save iteration. This is not a bug, but protection against long save times, as + models can get quite large, so renaming the current model file rather than copying it can + save substantial amount of time. + """ + logger.debug("Backing up and saving models") + print("") # Insert a new line to avoid spamming the same row as loss output + save_averages = self._get_save_averages() + if save_averages and self._should_backup(save_averages): + self._backup.backup_model(self.filename) + self._backup.backup_model(self._plugin.state.filename) + + include_optimizer = (force_save_optimizer or + self._save_optimizer == "always" or + (self._save_optimizer == "exit" and is_exit)) + + try: + self._plugin.model.save(self.filename, include_optimizer=include_optimizer) + except ValueError as err: + if include_optimizer and "name already exists" in str(err): + logger.warning("Due to a bug in older versions of Tensorflow, optimizer state " + "cannot be saved for this model.") + self._plugin.model.save(self.filename, include_optimizer=False) + else: + raise + + self._plugin.state.save() + + msg = "[Saved optimizer state for Snapshot]" if force_save_optimizer else "[Saved model]" + if save_averages: + lossmsg = [f"face_{side}: {avg:.5f}" + for side, avg in zip(("a", "b"), save_averages)] + msg += f" - Average loss since last save: {', '.join(lossmsg)}" + logger.info(msg) + + def _get_save_averages(self) -> list[float]: + """ Return the average loss since the last save iteration and reset historical loss """ + logger.debug("Getting save averages") + if not all(loss for loss in self._history): + logger.debug("No loss in history") + retval = [] + else: + retval = [sum(loss) / len(loss) for loss in self._history] + self._history = [[], []] # Reset historical loss + logger.debug("Average losses since last save: %s", retval) + return retval + + def _should_backup(self, save_averages: list[float]) -> bool: + """ Check whether the loss averages for this save iteration is the lowest that has been + seen. + + This protects against model corruption by only backing up the model if both sides have + seen a total fall in loss. + + Notes + ----- + This is by no means a perfect system. If the model corrupts at an iteration close + to a save iteration, then the averages may still be pushed lower than a previous + save average, resulting in backing up a corrupted model. + + Parameters + ---------- + save_averages: list + The average loss for each side for this save iteration + """ + backup = True + for side, loss in zip(("a", "b"), save_averages): + if not self._plugin.state.lowest_avg_loss.get(side, None): + logger.debug("Set initial save iteration loss average for '%s': %s", side, loss) + self._plugin.state.lowest_avg_loss[side] = loss + continue + backup = loss < self._plugin.state.lowest_avg_loss[side] if backup else backup + + if backup: # Update lowest loss values to the state file + # pylint:disable=unnecessary-comprehension + old_avgs = {key: val for key, val in self._plugin.state.lowest_avg_loss.items()} + self._plugin.state.lowest_avg_loss["a"] = save_averages[0] + self._plugin.state.lowest_avg_loss["b"] = save_averages[1] + logger.debug("Updated lowest historical save iteration averages from: %s to: %s", + old_avgs, self._plugin.state.lowest_avg_loss) + + logger.debug("Should backup: %s", backup) + return backup + + def snapshot(self) -> None: + """ Perform a model snapshot. + + Notes + ----- + Snapshot function is called 1 iteration after the model was saved, so that it is built from + the latest save, hence iteration being reduced by 1. + """ + logger.debug("Performing snapshot. Iterations: %s", self._plugin.iterations) + self._backup.snapshot_models(self._plugin.iterations - 1) + logger.debug("Performed snapshot") + + +class Weights(): + """ Handling of freezing and loading model weights + + Parameters + ---------- + plugin: :class:`Model` + The parent plugin class that owns the IO functions. + """ + def __init__(self, plugin: ModelBase) -> None: + logger.debug("Initializing %s: (plugin: %s)", self.__class__.__name__, plugin) + self._model = plugin.model + self._name = plugin.model_name + self._do_freeze = plugin._args.freeze_weights + self._weights_file = self._check_weights_file(plugin._args.load_weights) + + freeze_layers = plugin.config.get("freeze_layers") # Standardized config for freezing + load_layers = plugin.config.get("load_layers") # Standardized config for loading + self._freeze_layers = freeze_layers if freeze_layers else ["encoder"] # No plugin config + self._load_layers = load_layers if load_layers else ["encoder"] # No plugin config + logger.debug("Initialized %s", self.__class__.__name__) + + @classmethod + def _check_weights_file(cls, weights_file: str) -> str | None: + """ Validate that we have a valid path to a .h5 file. + + Parameters + ---------- + weights_file: str + The full path to a weights file + + Returns + ------- + str + The full path to a weights file + """ + if not weights_file: + logger.debug("No weights file selected.") + return None + + msg = "" + if not os.path.exists(weights_file): + msg = f"Load weights selected, but the path '{weights_file}' does not exist." + elif not os.path.splitext(weights_file)[-1].lower() == ".h5": + msg = (f"Load weights selected, but the path '{weights_file}' is not a valid Keras " + f"model (.h5) file.") + + if msg: + msg += " Please check and try again." + raise FaceswapError(msg) + + logger.verbose("Using weights file: %s", weights_file) # type:ignore + return weights_file + + def freeze(self) -> None: + """ If freeze has been selected in the cli arguments, then freeze those models indicated + in the plugin's configuration. """ + # Blanket unfreeze layers, as checking the value of :attr:`layer.trainable` appears to + # return ``True`` even when the weights have been frozen + for layer in get_all_sub_models(self._model): + layer.trainable = True + + if not self._do_freeze: + logger.debug("Freeze weights deselected. Not freezing") + return + + for layer in get_all_sub_models(self._model): + if layer.name in self._freeze_layers: + logger.info("Freezing weights for '%s' in model '%s'", layer.name, self._name) + layer.trainable = False + self._freeze_layers.remove(layer.name) + if self._freeze_layers: + logger.warning("The following layers were set to be frozen but do not exist in the " + "model: %s", self._freeze_layers) + + def load(self, model_exists: bool) -> None: + """ Load weights for newly created models, or output warning for pre-existing models. + + Parameters + ---------- + model_exists: bool + ``True`` if a model pre-exists and is being resumed, ``False`` if this is a new model + """ + if not self._weights_file: + logger.debug("No weights file provided. Not loading weights.") + return + if model_exists and self._weights_file: + logger.warning("Ignoring weights file '%s' as this model is resuming.", + self._weights_file) + return + + weights_models = self._get_weights_model() + all_models = get_all_sub_models(self._model) + + for model_name in self._load_layers: + sub_model = next((lyr for lyr in all_models if lyr.name == model_name), None) + sub_weights = next((lyr for lyr in weights_models if lyr.name == model_name), None) + + if not sub_model or not sub_weights: + msg = f"Skipping layer {model_name} as not in " + msg += "current_model." if not sub_model else f"weights '{self._weights_file}.'" + logger.warning(msg) + continue + + logger.info("Loading weights for layer '%s'", model_name) + skipped_ops = 0 + loaded_ops = 0 + for layer in sub_model.layers: + success = self._load_layer_weights(layer, sub_weights, model_name) + if success == 0: + skipped_ops += 1 + elif success == 1: + loaded_ops += 1 + + del weights_models + + if loaded_ops == 0: + raise FaceswapError(f"No weights were succesfully loaded from your weights file: " + f"'{self._weights_file}'. Please check and try again.") + if skipped_ops > 0: + logger.warning("%s weight(s) were unable to be loaded for your model. This is most " + "likely because the weights you are trying to load were trained with " + "different settings than you have set for your current model.", + skipped_ops) + + def _get_weights_model(self) -> list[tf.keras.models.Model]: + """ Obtain a list of all sub-models contained within the weights model. + + Returns + ------- + list + List of all models contained within the .h5 file + + Raises + ------ + FaceswapError + In the event of a failure to load the weights, or the weights belonging to a different + model + """ + retval = get_all_sub_models(kmodels.load_model(self._weights_file, compile=False)) + if not retval: + raise FaceswapError(f"Error loading weights file {self._weights_file}.") + + if retval[0].name != self._name: + raise FaceswapError(f"You are attempting to load weights from a '{retval[0].name}' " + f"model into a '{self._name}' model. This is not supported.") + return retval + + def _load_layer_weights(self, + layer: tf.keras.layers.Layer, + sub_weights: tf.keras.layers.Layer, + model_name: str) -> T.Literal[-1, 0, 1]: + """ Load the weights for a single layer. + + Parameters + ---------- + layer: :class:`tensorflow.keras.layers.Layer` + The layer to set the weights for + sub_weights: list + The list of layers in the weights model to load weights from + model_name: str + The name of the current sub-model that is having it's weights loaded + + Returns + ------- + int + `-1` if the layer has no weights to load. `0` if weights loading was unsuccessful. `1` + if weights loading was successful + """ + old_weights = layer.get_weights() + if not old_weights: + logger.debug("Skipping layer without weights: %s", layer.name) + return -1 + + layer_weights = next((lyr for lyr in sub_weights.layers + if lyr.name == layer.name), None) + if not layer_weights: + logger.warning("The weights file '%s' for layer '%s' does not contain weights for " + "'%s'. Skipping", self._weights_file, model_name, layer.name) + return 0 + + new_weights = layer_weights.get_weights() + if old_weights[0].shape != new_weights[0].shape: + logger.warning("The weights for layer '%s' are of incompatible shapes. Skipping.", + layer.name) + return 0 + logger.verbose("Setting weights for '%s'", layer.name) # type:ignore + layer.set_weights(layer_weights.get_weights()) + return 1 diff --git a/plugins/train/model/_base/model.py b/plugins/train/model/_base/model.py new file mode 100644 index 0000000000..5bc1161b17 --- /dev/null +++ b/plugins/train/model/_base/model.py @@ -0,0 +1,948 @@ +#!/usr/bin/env python3 +""" +Base class for Models. ALL Models should at least inherit from this class. + +See :mod:`~plugins.train.model.original` for an annotated example for how to create model plugins. +""" +from __future__ import annotations +import logging +import os +import sys +import time +import typing as T + +from collections import OrderedDict + +import numpy as np +import tensorflow as tf + +from lib.serializer import get_serializer +from lib.model.nn_blocks import set_config as set_nnblock_config +from lib.utils import FaceswapError +from plugins.train._config import Config + +from .io import IO, get_all_sub_models, Weights +from .settings import Loss, Optimizer, Settings + +if T.TYPE_CHECKING: + import argparse + from lib.config import ConfigValueType + +keras = tf.keras +K = tf.keras.backend + + +logger = logging.getLogger(__name__) +_CONFIG: dict[str, ConfigValueType] = {} + + +class ModelBase(): + """ Base class that all model plugins should inherit from. + + Parameters + ---------- + model_dir: str + The full path to the model save location + arguments: :class:`argparse.Namespace` + The arguments that were passed to the train or convert process as generated from + Faceswap's command line arguments + predict: bool, optional + ``True`` if the model is being loaded for inference, ``False`` if the model is being loaded + for training. Default: ``False`` + + Attributes + ---------- + input_shape: tuple or list + A `tuple` of `ints` defining the shape of the faces that the model takes as input. This + should be overridden by model plugins in their :func:`__init__` function. If the input size + is the same for both sides of the model, then this can be a single 3 dimensional `tuple`. + If the inputs have different sizes for `"A"` and `"B"` this should be a `list` of 2 3 + dimensional shape `tuples`, 1 for each side respectively. + trainer: str + Currently there is only one trainer available (`"original"`), so at present this attribute + can be ignored. If/when more trainers are added, then this attribute should be overridden + with the trainer name that a model requires in the model plugin's + :func:`__init__` function. + """ + def __init__(self, + model_dir: str, + arguments: argparse.Namespace, + predict: bool = False) -> None: + logger.debug("Initializing ModelBase (%s): (model_dir: '%s', arguments: %s, predict: %s)", + self.__class__.__name__, model_dir, arguments, predict) + + # Input shape must be set within the plugin after initializing + self.input_shape: tuple[int, ...] = () + self.trainer = "original" # Override for plugin specific trainer + self.color_order: T.Literal["bgr", "rgb"] = "bgr" # Override for image color channel order + + self._args = arguments + self._is_predict = predict + self._model: tf.keras.models.Model | None = None + + self._configfile = arguments.configfile if hasattr(arguments, "configfile") else None + self._load_config() + + if self.config["penalized_mask_loss"] and self.config["mask_type"] is None: + raise FaceswapError("Penalized Mask Loss has been selected but you have not chosen a " + "Mask to use. Please select a mask or disable Penalized Mask " + "Loss.") + + if self.config["learn_mask"] and self.config["mask_type"] is None: + raise FaceswapError("'Learn Mask' has been selected but you have not chosen a Mask to " + "use. Please select a mask or disable 'Learn Mask'.") + + self._mixed_precision = self.config["mixed_precision"] + self._io = IO(self, model_dir, self._is_predict, self.config["save_optimizer"]) + self._check_multiple_models() + + self._state = State(model_dir, + self.name, + self._config_changeable_items, + False if self._is_predict else self._args.no_logs) + self._settings = Settings(self._args, + self._mixed_precision, + self.config["allow_growth"], + self._is_predict) + self._loss = Loss(self.config, self.color_order) + + logger.debug("Initialized ModelBase (%s)", self.__class__.__name__) + + @property + def model(self) -> tf.keras.models.Model: + """:class:`Keras.models.Model`: The compiled model for this plugin. """ + return self._model + + @property + def command_line_arguments(self) -> argparse.Namespace: + """ :class:`argparse.Namespace`: The command line arguments passed to the model plugin from + either the train or convert script """ + return self._args + + @property + def coverage_ratio(self) -> float: + """ float: The ratio of the training image to crop out and train on as defined in user + configuration options. + + NB: The coverage ratio is a raw float, but will be applied to integer pixel images. + + To ensure consistent rounding and guaranteed even image size, the calculation for coverage + should always be: :math:`(original_size * coverage_ratio // 2) * 2` + """ + return self.config.get("coverage", 62.5) / 100 + + @property + def io(self) -> IO: # pylint:disable=invalid-name + """ :class:`~plugins.train.model.io.IO`: Input/Output operations for the model """ + return self._io + + @property + def config(self) -> dict: + """ dict: The configuration dictionary for current plugin, as set by the user's + configuration settings. """ + global _CONFIG # pylint:disable=global-statement + if not _CONFIG: + model_name = self._config_section + logger.debug("Loading config for: %s", model_name) + _CONFIG = Config(model_name, configfile=self._configfile).config_dict + return _CONFIG + + @property + def name(self) -> str: + """ str: The name of this model based on the plugin name. """ + _name = sys.modules[self.__module__].__file__ + assert isinstance(_name, str) + return os.path.splitext(os.path.basename(_name))[0].lower() + + @property + def model_name(self) -> str: + """ str: The name of the keras model. Generally this will be the same as :attr:`name` + but some plugins will override this when they contain multiple architectures """ + return self.name + + @property + def input_shapes(self) -> list[tuple[None, int, int, int]]: + """ list: A flattened list corresponding to all of the inputs to the model. """ + shapes = [T.cast(tuple[None, int, int, int], K.int_shape(inputs)) + for inputs in self.model.inputs] + return shapes + + @property + def output_shapes(self) -> list[tuple[None, int, int, int]]: + """ list: A flattened list corresponding to all of the outputs of the model. """ + shapes = [T.cast(tuple[None, int, int, int], K.int_shape(output)) + for output in self.model.outputs] + return shapes + + @property + def iterations(self) -> int: + """ int: The total number of iterations that the model has trained. """ + return self._state.iterations + + # Private properties + @property + def _config_section(self) -> str: + """ str: The section name for the current plugin for loading configuration options from the + config file. """ + return ".".join(self.__module__.split(".")[-2:]) + + @property + def _config_changeable_items(self) -> dict: + """ dict: The configuration options that can be updated after the model has already been + created. """ + return Config(self._config_section, configfile=self._configfile).changeable_items + + @property + def state(self) -> "State": + """:class:`State`: The state settings for the current plugin. """ + return self._state + + def _load_config(self) -> None: + """ Load the global config for reference in :attr:`config` and set the faceswap blocks + configuration options in `lib.model.nn_blocks` """ + global _CONFIG # pylint:disable=global-statement + if not _CONFIG: + model_name = self._config_section + logger.debug("Loading config for: %s", model_name) + _CONFIG = Config(model_name, configfile=self._configfile).config_dict + + nn_block_keys = ['icnr_init', 'conv_aware_init', 'reflect_padding'] + set_nnblock_config({key: _CONFIG.pop(key) + for key in nn_block_keys}) + + def _check_multiple_models(self) -> None: + """ Check whether multiple models exist in the model folder, and that no models exist that + were trained with a different plugin than the requested plugin. + + Raises + ------ + FaceswapError + If multiple model files, or models for a different plugin from that requested exists + within the model folder + """ + multiple_models = self._io.multiple_models_in_folder + if multiple_models is None: + logger.debug("Contents of model folder are valid") + return + + if len(multiple_models) == 1: + msg = (f"You have requested to train with the '{self.name}' plugin, but a model file " + f"for the '{multiple_models[0]}' plugin already exists in the folder " + f"'{self.io.model_dir}'.\nPlease select a different model folder.") + else: + ptypes = "', '".join(multiple_models) + msg = (f"There are multiple plugin types ('{ptypes}') stored in the model folder '" + f"{self.io.model_dir}'. This is not supported.\nPlease split the model files " + "into their own folders before proceeding") + raise FaceswapError(msg) + + def build(self) -> None: + """ Build the model and assign to :attr:`model`. + + Within the defined strategy scope, either builds the model from scratch or loads an + existing model if one exists. + + If running inference, then the model is built only for the required side to perform the + swap function, otherwise the model is then compiled with the optimizer and chosen + loss function(s). + + Finally, a model summary is outputted to the logger at verbose level. + """ + self._update_legacy_models() + is_summary = hasattr(self._args, "summary") and self._args.summary + with self._settings.strategy_scope(): + if self._io.model_exists: + model = self.io.load() + if self._is_predict: + inference = _Inference(model, self._args.swap_model) + self._model = inference.model + else: + self._model = model + else: + self._validate_input_shape() + inputs = self._get_inputs() + if not self._settings.use_mixed_precision and not is_summary: + # Store layer names which can be switched to mixed precision + model, mp_layers = self._settings.get_mixed_precision_layers(self.build_model, + inputs) + self._state.add_mixed_precision_layers(mp_layers) + self._model = model + else: + self._model = self.build_model(inputs) + if not is_summary and not self._is_predict: + self._compile_model() + self._output_summary() + + def _update_legacy_models(self) -> None: + """ Load weights from legacy split models into new unified model, archiving old model files + to a new folder. """ + legacy_mapping = self._legacy_mapping() # pylint:disable=assignment-from-none + if legacy_mapping is None: + return + + if not all(os.path.isfile(os.path.join(self.io.model_dir, fname)) + for fname in legacy_mapping): + return + archive_dir = f"{self.io.model_dir}_TF1_Archived" + if os.path.exists(archive_dir): + raise FaceswapError("We need to update your model files for use with Tensorflow 2.x, " + "but the archive folder already exists. Please remove the " + f"following folder to continue: '{archive_dir}'") + + logger.info("Updating legacy models for Tensorflow 2.x") + logger.info("Your Tensorflow 1.x models will be archived in the following location: '%s'", + archive_dir) + os.rename(self.io.model_dir, archive_dir) + os.mkdir(self.io.model_dir) + new_model = self.build_model(self._get_inputs()) + for model_name, layer_name in legacy_mapping.items(): + old_model: tf.keras.models.Model = keras.models.load_model( + os.path.join(archive_dir, model_name), + compile=False) + layer = [layer for layer in new_model.layers if layer.name == layer_name] + if not layer: + logger.warning("Skipping legacy weights from '%s'...", model_name) + continue + klayer: tf.keras.layers.Layer = layer[0] + logger.info("Updating legacy weights from '%s'...", model_name) + klayer.set_weights(old_model.get_weights()) + filename = self._io.filename + logger.info("Saving Tensorflow 2.x model to '%s'", filename) + new_model.save(filename) + # Penalized Loss and Learn Mask used to be disabled automatically if a mask wasn't + # selected, so disable it if enabled, but mask_type is None + if self.config["mask_type"] is None: + self.config["penalized_mask_loss"] = False + self.config["learn_mask"] = False + self.config["eye_multiplier"] = 1 + self.config["mouth_multiplier"] = 1 + self._state.save() + + def _validate_input_shape(self) -> None: + """ Validate that the input shape is either a single shape tuple of 3 dimensions or + a list of 2 shape tuples of 3 dimensions. """ + assert len(self.input_shape) == 3, "Input shape should be a 3 dimensional shape tuple" + + def _get_inputs(self) -> list[tf.keras.layers.Input]: + """ Obtain the standardized inputs for the model. + + The inputs will be returned for the "A" and "B" sides in the shape as defined by + :attr:`input_shape`. + + Returns + ------- + list + A list of :class:`keras.layers.Input` tensors. This will be a list of 2 tensors (one + for each side) each of shapes :attr:`input_shape`. + """ + logger.debug("Getting inputs") + input_shapes = [self.input_shape, self.input_shape] + inputs = [keras.layers.Input(shape=shape, name=f"face_in_{side}") + for side, shape in zip(("a", "b"), input_shapes)] + logger.debug("inputs: %s", inputs) + return inputs + + def build_model(self, inputs: list[tf.keras.layers.Input]) -> tf.keras.models.Model: + """ Override for Model Specific autoencoder builds. + + Parameters + ---------- + inputs: list + A list of :class:`keras.layers.Input` tensors. This will be a list of 2 tensors (one + for each side) each of shapes :attr:`input_shape`. + + Returns + ------- + :class:`keras.models.Model` + See Keras documentation for the correct structure, but note that parameter :attr:`name` + is a required rather than an optional argument in Faceswap. You should assign this to + the attribute ``self.name`` that is automatically generated from the plugin's filename. + """ + raise NotImplementedError + + def _output_summary(self) -> None: + """ Output the summary of the model and all sub-models to the verbose logger. """ + if hasattr(self._args, "summary") and self._args.summary: + print_fn = None # Print straight to stdout + else: + # print to logger + print_fn = lambda x: logger.verbose("%s", x) #type:ignore[attr-defined] # noqa[E731] # pylint:disable=C3001 + for idx, model in enumerate(get_all_sub_models(self.model)): + if idx == 0: + parent = model + continue + model.summary(line_length=100, print_fn=print_fn) + parent.summary(line_length=100, print_fn=print_fn) + + def _compile_model(self) -> None: + """ Compile the model to include the Optimizer and Loss Function(s). """ + logger.debug("Compiling Model") + + if self.state.model_needs_rebuild: + self._model = self._settings.check_model_precision(self._model, self._state) + + optimizer = Optimizer(self.config["optimizer"], + self.config["learning_rate"], + self.config["autoclip"], + 10 ** int(self.config["epsilon_exponent"])).optimizer + if self._settings.use_mixed_precision: + optimizer = self._settings.loss_scale_optimizer(optimizer) + + weights = Weights(self) + weights.load(self._io.model_exists) + weights.freeze() + + self._loss.configure(self.model) + self.model.compile(optimizer=optimizer, loss=self._loss.functions) + self._state.add_session_loss_names(self._loss.names) + logger.debug("Compiled Model: %s", self.model) + + def _legacy_mapping(self) -> dict | None: + """ The mapping of separate model files to single model layers for transferring of legacy + weights. + + Returns + ------- + dict or ``None`` + Dictionary of original H5 filenames for legacy models mapped to new layer names or + ``None`` if the model did not exist in Faceswap prior to Tensorflow 2 + """ + return None + + def add_history(self, loss: list[float]) -> None: + """ Add the current iteration's loss history to :attr:`_io.history`. + + Called from the trainer after each iteration, for tracking loss drop over time between + save iterations. + + Parameters + ---------- + loss: list + The loss values for the A and B side for the current iteration. This should be the + collated loss values for each side. + """ + self._io.history[0].append(loss[0]) + self._io.history[1].append(loss[1]) + + +class State(): + """ Holds state information relating to the plugin's saved model. + + Parameters + ---------- + model_dir: str + The full path to the model save location + model_name: str + The name of the model plugin + config_changeable_items: dict + Configuration options that can be altered when resuming a model, and their current values + no_logs: bool + ``True`` if Tensorboard logs should not be generated, otherwise ``False`` + """ + def __init__(self, + model_dir: str, + model_name: str, + config_changeable_items: dict, + no_logs: bool) -> None: + logger.debug("Initializing %s: (model_dir: '%s', model_name: '%s', " + "config_changeable_items: '%s', no_logs: %s", self.__class__.__name__, + model_dir, model_name, config_changeable_items, no_logs) + self._serializer = get_serializer("json") + filename = f"{model_name}_state.{self._serializer.file_extension}" + self._filename = os.path.join(model_dir, filename) + self._name = model_name + self._iterations = 0 + self._mixed_precision_layers: list[str] = [] + self._rebuild_model = False + self._sessions: dict[int, dict] = {} + self._lowest_avg_loss: dict[str, float] = {} + self._config: dict[str, ConfigValueType] = {} + self._load(config_changeable_items) + self._session_id = self._new_session_id() + self._create_new_session(no_logs, config_changeable_items) + logger.debug("Initialized %s:", self.__class__.__name__) + + @property + def filename(self) -> str: + """ str: Full path to the state filename """ + return self._filename + + @property + def loss_names(self) -> list[str]: + """ list: The loss names for the current session """ + return self._sessions[self._session_id]["loss_names"] + + @property + def current_session(self) -> dict: + """ dict: The state dictionary for the current :attr:`session_id`. """ + return self._sessions[self._session_id] + + @property + def iterations(self) -> int: + """ int: The total number of iterations that the model has trained. """ + return self._iterations + + @property + def lowest_avg_loss(self) -> dict: + """dict: The lowest average save interval loss seen for each side. """ + return self._lowest_avg_loss + + @property + def session_id(self) -> int: + """ int: The current training session id. """ + return self._session_id + + @property + def sessions(self) -> dict[int, dict[str, T.Any]]: + """ dict[int, dict[str, Any]]: The session information for each session in the state + file """ + return {int(k): v for k, v in self._sessions.items()} + + @property + def mixed_precision_layers(self) -> list[str]: + """list: Layers that can be switched between mixed-float16 and float32. """ + return self._mixed_precision_layers + + @property + def model_needs_rebuild(self) -> bool: + """bool: ``True`` if mixed precision policy has changed so model needs to be rebuilt + otherwise ``False`` """ + return self._rebuild_model + + def _new_session_id(self) -> int: + """ Generate a new session id. Returns 1 if this is a new model, or the last session id + 1 + if it is a pre-existing model. + + Returns + ------- + int + The newly generated session id + """ + if not self._sessions: + session_id = 1 + else: + session_id = max(int(key) for key in self._sessions.keys()) + 1 + logger.debug(session_id) + return session_id + + def _create_new_session(self, no_logs: bool, config_changeable_items: dict) -> None: + """ Initialize a new session, creating the dictionary entry for the session in + :attr:`_sessions`. + + Parameters + ---------- + no_logs: bool + ``True`` if Tensorboard logs should not be generated, otherwise ``False`` + config_changeable_items: dict + Configuration options that can be altered when resuming a model, and their current + values + """ + logger.debug("Creating new session. id: %s", self._session_id) + self._sessions[self._session_id] = {"timestamp": time.time(), + "no_logs": no_logs, + "loss_names": [], + "batchsize": 0, + "iterations": 0, + "config": config_changeable_items} + + def update_session_config(self, key: str, value: T.Any) -> None: + """ Update a configuration item of the currently loaded session. + + Parameters + ---------- + key: str + The configuration item to update for the current session + value: any + The value to update to + """ + old_val = self.current_session["config"][key] + assert isinstance(value, type(old_val)) + logger.debug("Updating configuration item '%s' from '%s' to '%s'", key, old_val, value) + self.current_session["config"][key] = value + + def add_session_loss_names(self, loss_names: list[str]) -> None: + """ Add the session loss names to the sessions dictionary. + + The loss names are used for Tensorboard logging + + Parameters + ---------- + loss_names: list + The list of loss names for this session. + """ + logger.debug("Adding session loss_names: %s", loss_names) + self._sessions[self._session_id]["loss_names"] = loss_names + + def add_session_batchsize(self, batch_size: int) -> None: + """ Add the session batch size to the sessions dictionary. + + Parameters + ---------- + batch_size: int + The batch size for the current training session + """ + logger.debug("Adding session batch size: %s", batch_size) + self._sessions[self._session_id]["batchsize"] = batch_size + + def increment_iterations(self) -> None: + """ Increment :attr:`iterations` and session iterations by 1. """ + self._iterations += 1 + self._sessions[self._session_id]["iterations"] += 1 + + def add_mixed_precision_layers(self, layers: list[str]) -> None: + """ Add the list of model's layers that are compatible for mixed precision to the + state dictionary """ + logger.debug("Storing mixed precision layers: %s", layers) + self._mixed_precision_layers = layers + + def _load(self, config_changeable_items: dict) -> None: + """ Load a state file and set the serialized values to the class instance. + + Updates the model's config with the values stored in the state file. + + Parameters + ---------- + config_changeable_items: dict + Configuration options that can be altered when resuming a model, and their current + values + """ + logger.debug("Loading State") + if not os.path.exists(self._filename): + logger.info("No existing state file found. Generating.") + return + state = self._serializer.load(self._filename) + self._name = state.get("name", self._name) + self._sessions = state.get("sessions", {}) + self._lowest_avg_loss = state.get("lowest_avg_loss", {}) + self._iterations = state.get("iterations", 0) + self._mixed_precision_layers = state.get("mixed_precision_layers", []) + self._config = state.get("config", {}) + logger.debug("Loaded state: %s", state) + self._replace_config(config_changeable_items) + + def save(self) -> None: + """ Save the state values to the serialized state file. """ + logger.debug("Saving State") + state = {"name": self._name, + "sessions": self._sessions, + "lowest_avg_loss": self._lowest_avg_loss, + "iterations": self._iterations, + "mixed_precision_layers": self._mixed_precision_layers, + "config": _CONFIG} + self._serializer.save(self._filename, state) + logger.debug("Saved State") + + def _replace_config(self, config_changeable_items) -> None: + """ Replace the loaded config with the one contained within the state file. + + Check for any `fixed`=``False`` parameter changes and log info changes. + + Update any legacy config items to their current versions. + + Parameters + ---------- + config_changeable_items: dict + Configuration options that can be altered when resuming a model, and their current + values + """ + global _CONFIG # pylint:disable=global-statement + if _CONFIG is None: + return + legacy_update = self._update_legacy_config() + # Add any new items to state config for legacy purposes where the new default may be + # detrimental to an existing model. + legacy_defaults: dict[str, str | int | bool] = {"centering": "legacy", + "mask_loss_function": "mse", + "l2_reg_term": 100, + "optimizer": "adam", + "mixed_precision": False} + for key, val in _CONFIG.items(): + if key not in self._config.keys(): + setting: ConfigValueType = legacy_defaults.get(key, val) + logger.info("Adding new config item to state file: '%s': '%s'", key, setting) + self._config[key] = setting + self._update_changed_config_items(config_changeable_items) + logger.debug("Replacing config. Old config: %s", _CONFIG) + _CONFIG = self._config + if legacy_update: + self.save() + logger.debug("Replaced config. New config: %s", _CONFIG) + logger.info("Using configuration saved in state file") + + def _update_legacy_config(self) -> bool: + """ Legacy updates for new config additions. + + When new config items are added to the Faceswap code, existing model state files need to be + updated to handle these new items. + + Current existing legacy update items: + + * loss - If old `dssim_loss` is ``true`` set new `loss_function` to `ssim` otherwise + set it to `mae`. Remove old `dssim_loss` item + + * l2_reg_term - If this exists, set loss_function_2 to ``mse`` and loss_weight_2 to + the value held in the old ``l2_reg_term`` item + + * masks - If `learn_mask` does not exist then it is set to ``True`` if `mask_type` is + not ``None`` otherwise it is set to ``False``. + + * masks type - Replace removed masks 'dfl_full' and 'facehull' with `components` mask + + * clipnorm - Only existed in 2 models (DFL-SAE + Unbalanced). Replaced with global + option autoclip + + Returns + ------- + bool + ``True`` if legacy items exist and state file has been updated, otherwise ``False`` + """ + logger.debug("Checking for legacy state file update") + priors = ["dssim_loss", "mask_type", "mask_type", "l2_reg_term", "clipnorm"] + new_items = ["loss_function", "learn_mask", "mask_type", "loss_function_2", + "autoclip"] + updated = False + for old, new in zip(priors, new_items): + if old not in self._config: + logger.debug("Legacy item '%s' not in config. Skipping update", old) + continue + + # dssim_loss > loss_function + if old == "dssim_loss": + self._config[new] = "ssim" if self._config[old] else "mae" + del self._config[old] + updated = True + logger.info("Updated config from legacy dssim format. New config loss " + "function: '%s'", self._config[new]) + continue + + # Add learn mask option and set to True if model has "penalized_mask_loss" specified + if old == "mask_type" and new == "learn_mask" and new not in self._config: + self._config[new] = self._config["mask_type"] is not None + updated = True + logger.info("Added new 'learn_mask' config item for this model. Value set to: %s", + self._config[new]) + continue + + # Replace removed masks with most similar equivalent + if old == "mask_type" and new == "mask_type" and self._config[old] in ("facehull", + "dfl_full"): + old_mask = self._config[old] + self._config[new] = "components" + updated = True + logger.info("Updated 'mask_type' from '%s' to '%s' for this model", + old_mask, self._config[new]) + + # Replace l2_reg_term with the correct loss_2_function and update the value of + # loss_2_weight + if old == "l2_reg_term": + self._config[new] = "mse" + self._config["loss_weight_2"] = self._config[old] + del self._config[old] + updated = True + logger.info("Updated config from legacy 'l2_reg_term' to 'loss_function_2'") + + # Replace clipnorm with correct gradient clipping type and value + if old == "clipnorm": + self._config[new] = self._config[old] + del self._config[old] + updated = True + logger.info("Updated config from legacy '%s' to '%s'", old, new) + + logger.debug("State file updated for legacy config: %s", updated) + return updated + + def _update_changed_config_items(self, config_changeable_items: dict) -> None: + """ Update any parameters which are not fixed and have been changed. + + Set the :attr:`model_needs_rebuild` to ``True`` if mixed precision state has changed + + Parameters + ---------- + config_changeable_items: dict + Configuration options that can be altered when resuming a model, and their current + values + """ + rebuild_tasks = ["mixed_precision"] + if not config_changeable_items: + logger.debug("No changeable parameters have been updated") + return + for key, val in config_changeable_items.items(): + old_val = self._config[key] + if old_val == val: + continue + self._config[key] = val + logger.info("Config item: '%s' has been updated from '%s' to '%s'", key, old_val, val) + self._rebuild_model = self._rebuild_model or key in rebuild_tasks + + +class _Inference(): # pylint:disable=too-few-public-methods + """ Calculates required layers and compiles a saved model for inference. + + Parameters + ---------- + saved_model: :class:`keras.models.Model` + The saved trained Faceswap model + switch_sides: bool + ``True`` if the swap should be performed "B" > "A" ``False`` if the swap should be + "A" > "B" + """ + def __init__(self, saved_model: tf.keras.models.Model, switch_sides: bool) -> None: + logger.debug("Initializing: %s (saved_model: %s, switch_sides: %s)", + self.__class__.__name__, saved_model, switch_sides) + self._config = saved_model.get_config() + + self._input_idx = 1 if switch_sides else 0 + self._output_idx = 0 if switch_sides else 1 + + self._input_names = [inp[0] for inp in self._config["input_layers"]] + self._model = self._make_inference_model(saved_model) + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def model(self) -> tf.keras.models.Model: + """ :class:`keras.models.Model`: The Faceswap model, compiled for inference. """ + return self._model + + def _get_nodes(self, nodes: np.ndarray) -> list[tuple[str, int]]: + """ Given in input list of nodes from a :attr:`keras.models.Model.get_config` dictionary, + filters the layer name(s) and output index of the node, splitting to the correct output + index in the event of multiple inputs. + + Parameters + ---------- + nodes: list + A node entry from the :attr:`keras.models.Model.get_config` dictionary + + Returns + ------- + list + The (node name, output index) for each node passed in + """ + anodes = np.array(nodes, dtype="object")[..., :3] + num_layers = anodes.shape[0] + anodes = anodes[self._output_idx] if num_layers == 2 else anodes[0] + + # Probably better checks for this, but this occurs when DNY preset is used and learn + # mask is enabled (i.e. the mask is created in fully connected layers) + anodes = anodes.squeeze() if anodes.ndim == 3 else anodes + + retval = [(node[0], node[2]) for node in anodes] + return retval + + def _make_inference_model(self, saved_model: tf.keras.models.Model) -> tf.keras.models.Model: + """ Extract the sub-models from the saved model that are required for inference. + + Parameters + ---------- + saved_model: :class:`keras.models.Model` + The saved trained Faceswap model + + Returns + ------- + :class:`keras.models.Model` + The model compiled for inference + """ + logger.debug("Compiling inference model. saved_model: %s", saved_model) + struct = self._get_filtered_structure() + model_inputs = self._get_inputs(saved_model.inputs) + compiled_layers: dict[str, tf.keras.layers.Layer] = {} + for layer in saved_model.layers: + if layer.name not in struct: + logger.debug("Skipping unused layer: '%s'", layer.name) + continue + inbound = struct[layer.name] + logger.debug("Processing layer '%s': (layer: %s, inbound_nodes: %s)", + layer.name, layer, inbound) + if not inbound: + model = model_inputs + logger.debug("Adding model inputs %s: %s", layer.name, model) + else: + layer_inputs = [] + for inp in inbound: + inbound_layer = compiled_layers[inp[0]] + if isinstance(inbound_layer, list) and len(inbound_layer) > 1: + # Multi output inputs + inbound_output_idx = inp[1] + next_input = inbound_layer[inbound_output_idx] + logger.debug("Selecting output index %s from multi output inbound layer: " + "%s (using: %s)", inbound_output_idx, inbound_layer, + next_input) + else: + next_input = inbound_layer + + layer_inputs.append(next_input) + + logger.debug("Compiling layer '%s': layer inputs: %s", layer.name, layer_inputs) + model = layer(layer_inputs) + compiled_layers[layer.name] = model + retval = keras.models.Model(model_inputs, model, name=f"{saved_model.name}_inference") + logger.debug("Compiled inference model '%s': %s", retval.name, retval) + return retval + + def _get_filtered_structure(self) -> OrderedDict: + """ Obtain the structure of the inference model. + + This parses the model config (in reverse) to obtain the required layers for an inference + model. + + Returns + ------- + :class:`collections.OrderedDict` + The layer name as key with the input name and output index as value. + """ + # Filter output layer + out = np.array(self._config["output_layers"], dtype="object") + if out.ndim == 2: + out = np.expand_dims(out, axis=1) # Needs to be expanded for _get_nodes + outputs = self._get_nodes(out) + + # Iterate backwards from the required output to get the reversed model structure + current_layers = [outputs[0]] + next_layers = [] + struct = OrderedDict() + drop_input = self._input_names[abs(self._input_idx - 1)] + switch_input = self._input_names[self._input_idx] + while True: + layer_info = current_layers.pop(0) + current_layer = next(lyr for lyr in self._config["layers"] + if lyr["name"] == layer_info[0]) + inbound = current_layer["inbound_nodes"] + + if not inbound: + break + + inbound_info = self._get_nodes(inbound) + + if any(inb[0] == drop_input for inb in inbound_info): # Switch inputs + inbound_info = [(switch_input if inb[0] == drop_input else inb[0], inb[1]) + for inb in inbound_info] + struct[layer_info[0]] = inbound_info + next_layers.extend(inbound_info) + + if not current_layers: + current_layers = next_layers + next_layers = [] + + struct[switch_input] = [] # Add the input layer + logger.debug("Model structure: %s", struct) + return struct + + def _get_inputs(self, inputs: list) -> list: + """ Obtain the inputs for the requested swap direction. + + Parameters + ---------- + inputs: list + The full list of input tensors to the saved faceswap training model + + Returns + ------- + list + List of input tensors to feed the model for the requested swap direction + """ + input_split = len(inputs) // 2 + start_idx = input_split * self._input_idx + retval = inputs[start_idx: start_idx + input_split] + logger.debug("model inputs: %s, input_split: %s, start_idx: %s, inference_inputs: %s", + inputs, input_split, start_idx, retval) + return retval diff --git a/plugins/train/model/_base/settings.py b/plugins/train/model/_base/settings.py new file mode 100644 index 0000000000..f2a9aba321 --- /dev/null +++ b/plugins/train/model/_base/settings.py @@ -0,0 +1,721 @@ +#!/usr/bin/env python3 +""" +Settings for the model base plugins. + +The objects in this module should not be called directly, but are called from +:class:`~plugins.train.model._base.ModelBase` + +Handles configuration of model plugins for: + - Loss configuration + - Optimizer settings + - General global model configuration settings +""" +from __future__ import annotations +from dataclasses import dataclass, field +import logging +import platform +import typing as T + +from contextlib import nullcontext + +import tensorflow as tf +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import losses as k_losses # pylint:disable=import-error +import tensorflow.keras.mixed_precision as mixedprecision # noqa pylint:disable=import-error + +from lib.model import losses, optimizers +from lib.model.autoclip import AutoClipper +from lib.utils import get_backend + +if T.TYPE_CHECKING: + from collections.abc import Callable + from contextlib import AbstractContextManager as ContextManager + from argparse import Namespace + from .model import State + +keras = tf.keras +K = keras.backend + +logger = logging.getLogger(__name__) + + +@dataclass +class LossClass: + """ Typing class for holding loss functions. + + Parameters + ---------- + function: Callable + The function that takes in the true/predicted images and returns the loss + init: bool, Optional + Whether the loss object ``True`` needs to be initialized (i.e. it's a class) or + ``False`` it does not require initialization (i.e. it's a function). + Default ``True`` + kwargs: dict + Any keyword arguments to supply to the loss function at initialization. + """ + function: Callable[[tf.Tensor, tf.Tensor], tf.Tensor] | T.Any = k_losses.mae + init: bool = True + kwargs: dict[str, T.Any] = field(default_factory=dict) + + +class Loss(): + """ Holds loss names and functions for an Autoencoder. + + Parameters + ---------- + config: dict + The configuration options for the current model plugin + color_order: str + Color order of the model. One of `"BGR"` or `"RGB"` + """ + def __init__(self, config: dict, color_order: T.Literal["bgr", "rgb"]) -> None: + logger.debug("Initializing %s: (color_order: %s)", self.__class__.__name__, color_order) + self._config = config + self._mask_channels = self._get_mask_channels() + self._inputs: list[tf.keras.layers.Layer] = [] + self._names: list[str] = [] + self._funcs: dict[str, Callable] = {} + + self._loss_dict = {"ffl": LossClass(function=losses.FocalFrequencyLoss), + "flip": LossClass(function=losses.LDRFLIPLoss, + kwargs={"color_order": color_order}), + "gmsd": LossClass(function=losses.GMSDLoss), + "l_inf_norm": LossClass(function=losses.LInfNorm), + "laploss": LossClass(function=losses.LaplacianPyramidLoss), + "logcosh": LossClass(function=k_losses.logcosh, init=False), + "lpips_alex": LossClass(function=losses.LPIPSLoss, + kwargs={"trunk_network": "alex"}), + "lpips_squeeze": LossClass(function=losses.LPIPSLoss, + kwargs={"trunk_network": "squeeze"}), + "lpips_vgg16": LossClass(function=losses.LPIPSLoss, + kwargs={"trunk_network": "vgg16"}), + "ms_ssim": LossClass(function=losses.MSSIMLoss), + "mae": LossClass(function=k_losses.mean_absolute_error, init=False), + "mse": LossClass(function=k_losses.mean_squared_error, init=False), + "pixel_gradient_diff": LossClass(function=losses.GradientLoss), + "ssim": LossClass(function=losses.DSSIMObjective), + "smooth_loss": LossClass(function=losses.GeneralizedLoss)} + + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def names(self) -> list[str]: + """ list: The list of loss names for the model. """ + return self._names + + @property + def functions(self) -> dict: + """ dict: The loss functions that apply to each model output. """ + return self._funcs + + @property + def _mask_inputs(self) -> list | None: + """ list: The list of input tensors to the model that contain the mask. Returns ``None`` + if there is no mask input to the model. """ + mask_inputs = [inp for inp in self._inputs if inp.name.startswith("mask")] + return None if not mask_inputs else mask_inputs + + @property + def _mask_shapes(self) -> list[tuple] | None: + """ list: The list of shape tuples for the mask input tensors for the model. Returns + ``None`` if there is no mask input. """ + if self._mask_inputs is None: + return None + return [K.int_shape(mask_input) for mask_input in self._mask_inputs] + + def configure(self, model: tf.keras.models.Model) -> None: + """ Configure the loss functions for the given inputs and outputs. + + Parameters + ---------- + model: :class:`keras.models.Model` + The model that is to be trained + """ + self._inputs = model.inputs + self._set_loss_names(model.outputs) + self._set_loss_functions(model.output_names) + self._names.insert(0, "total") + + def _set_loss_names(self, outputs: list[tf.Tensor]) -> None: + """ Name the losses based on model output. + + This is used for correct naming in the state file, for display purposes only. + + Adds the loss names to :attr:`names` + + Notes + ----- + TODO Currently there is an issue in Tensorflow that wraps all outputs in an Identity layer + when running in Eager Execution mode, which means we cannot use the name of the output + layers to name the losses (https://github.com/tensorflow/tensorflow/issues/32180). + With this in mind, losses are named based on their shapes + + Parameters + ---------- + outputs: list + A list of output tensors from the model plugin + """ + # TODO Use output names if/when these are fixed upstream + split_outputs = [outputs[:len(outputs) // 2], outputs[len(outputs) // 2:]] + for side, side_output in zip(("a", "b"), split_outputs): + output_names = [output.name for output in side_output] + output_shapes = [K.int_shape(output)[1:] for output in side_output] + output_types = ["mask" if shape[-1] == 1 else "face" for shape in output_shapes] + logger.debug("side: %s, output names: %s, output_shapes: %s, output_types: %s", + side, output_names, output_shapes, output_types) + for idx, name in enumerate(output_types): + suffix = "" if output_types.count(name) == 1 else f"_{idx}" + self._names.append(f"{name}_{side}{suffix}") + logger.debug(self._names) + + def _get_function(self, name: str) -> Callable[[tf.Tensor, tf.Tensor], tf.Tensor]: + """ Obtain the requested Loss function + + Parameters + ---------- + name: str + The name of the loss function from the training configuration file + + Returns + ------- + Keras Loss Function + The requested loss function + """ + func = self._loss_dict[name] + retval = func.function(**func.kwargs) if func.init else func.function # type:ignore + logger.debug("Obtained loss function `%s` (%s)", name, retval) + return retval + + def _set_loss_functions(self, output_names: list[str]): + """ Set the loss functions and their associated weights. + + Adds the loss functions to the :attr:`functions` dictionary. + + Parameters + ---------- + output_names: list + The output names from the model + """ + face_losses = [(lossname, self._config.get(f"loss_weight_{k[-1]}", 100)) + for k, lossname in sorted(self._config.items()) + if k.startswith("loss_function") + and self._config.get(f"loss_weight_{k[-1]}", 100) != 0 + and lossname is not None] + + for name, output_name in zip(self._names, output_names): + if name.startswith("mask"): + loss_func = self._get_function(self._config["mask_loss_function"]) + else: + loss_func = losses.LossWrapper() + for func, weight in face_losses: + self._add_face_loss_function(loss_func, func, weight / 100.) + + logger.debug("%s: (output_name: '%s', function: %s)", name, output_name, loss_func) + self._funcs[output_name] = loss_func + logger.debug("functions: %s", self._funcs) + + def _add_face_loss_function(self, + loss_wrapper: losses.LossWrapper, + loss_function: str, + weight: float) -> None: + """ Add the given face loss function at the given weight and apply any mouth and eye + multipliers + + Parameters + ---------- + loss_wrapper: :class:`lib.model.losses.LossWrapper` + The wrapper loss function that holds the face losses + loss_function: str + The loss function to add to the loss wrapper + weight: float + The amount of weight to apply to the given loss function + """ + logger.debug("Adding loss function: %s, weight: %s", loss_function, weight) + loss_wrapper.add_loss(self._get_function(loss_function), + weight=weight, + mask_channel=self._mask_channels[0]) + + channel_idx = 1 + for section in ("eye_multiplier", "mouth_multiplier"): + mask_channel = self._mask_channels[channel_idx] + multiplier = self._config[section] * 1. + if multiplier > 1.: + logger.debug("Adding section loss %s: %s", section, multiplier) + loss_wrapper.add_loss(self._get_function(loss_function), + weight=weight * multiplier, + mask_channel=mask_channel) + channel_idx += 1 + + def _get_mask_channels(self) -> list[int]: + """ Obtain the channels from the face targets that the masks reside in from the training + data generator. + + Returns + ------- + list: + A list of channel indices that contain the mask for the corresponding config item + """ + eye_multiplier = self._config["eye_multiplier"] + mouth_multiplier = self._config["mouth_multiplier"] + if not self._config["penalized_mask_loss"] and (eye_multiplier > 1 or + mouth_multiplier > 1): + logger.warning("You have selected eye/mouth loss multipliers greater than 1x, but " + "Penalized Mask Loss is disabled. Disabling all multipliers.") + eye_multiplier = 1 + mouth_multiplier = 1 + uses_masks = (self._config["penalized_mask_loss"], + eye_multiplier > 1, + mouth_multiplier > 1) + mask_channels = [-1 for _ in range(len(uses_masks))] + current_channel = 3 + for idx, mask_required in enumerate(uses_masks): + if mask_required: + mask_channels[idx] = current_channel + current_channel += 1 + logger.debug("uses_masks: %s, mask_channels: %s", uses_masks, mask_channels) + return mask_channels + + +class Optimizer(): # pylint:disable=too-few-public-methods + """ Obtain the selected optimizer with the appropriate keyword arguments. + + Parameters + ---------- + optimizer: str + The selected optimizer name for the plugin + learning_rate: float + The selected learning rate to use + autoclip: bool + ``True`` if AutoClip should be enabled otherwise ``False`` + epsilon: float + The value to use for the epsilon of the optimizer + """ + def __init__(self, + optimizer: str, + learning_rate: float, + autoclip: bool, + epsilon: float) -> None: + logger.debug("Initializing %s: (optimizer: %s, learning_rate: %s, autoclip: %s, " + ", epsilon: %s)", self.__class__.__name__, optimizer, learning_rate, + autoclip, epsilon) + valid_optimizers = {"adabelief": (optimizers.AdaBelief, + {"beta_1": 0.5, "beta_2": 0.99, "epsilon": epsilon}), + "adam": (optimizers.Adam, + {"beta_1": 0.5, "beta_2": 0.99, "epsilon": epsilon}), + "nadam": (optimizers.Nadam, + {"beta_1": 0.5, "beta_2": 0.99, "epsilon": epsilon}), + "rms-prop": (optimizers.RMSprop, {"epsilon": epsilon})} + optimizer_info = valid_optimizers[optimizer] + self._optimizer: Callable = optimizer_info[0] + self._kwargs: dict[str, T.Any] = optimizer_info[1] + + self._configure(learning_rate, autoclip) + logger.verbose("Using %s optimizer", optimizer.title()) # type:ignore[attr-defined] + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def optimizer(self) -> tf.keras.optimizers.Optimizer: + """ :class:`keras.optimizers.Optimizer`: The requested optimizer. """ + return self._optimizer(**self._kwargs) + + def _configure(self, + learning_rate: float, + autoclip: bool) -> None: + """ Configure the optimizer based on user settings. + + Parameters + ---------- + learning_rate: float + The selected learning rate to use + autoclip: bool + ``True`` if AutoClip should be enabled otherwise ``False`` + """ + self._kwargs["learning_rate"] = learning_rate + if not autoclip: + return + + logger.info("Enabling AutoClip") + self._kwargs["gradient_transformers"] = [AutoClipper(10, history_size=10000)] + logger.debug("optimizer kwargs: %s", self._kwargs) + + +class Settings(): + """ Tensorflow core training settings. + + Sets backend tensorflow settings prior to launching the model. + + Tensorflow 2 uses distribution strategies for multi-GPU/system training. These are context + managers. + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The arguments that were passed to the train or convert process as generated from + Faceswap's command line arguments + mixed_precision: bool + ``True`` if Mixed Precision training should be used otherwise ``False`` + allow_growth: bool + ``True`` if the Tensorflow allow_growth parameter should be set otherwise ``False`` + is_predict: bool, optional + ``True`` if the model is being loaded for inference, ``False`` if the model is being loaded + for training. Default: ``False`` + """ + def __init__(self, + arguments: Namespace, + mixed_precision: bool, + allow_growth: bool, + is_predict: bool) -> None: + logger.debug("Initializing %s: (arguments: %s, mixed_precision: %s, allow_growth: %s, " + "is_predict: %s)", self.__class__.__name__, arguments, mixed_precision, + allow_growth, is_predict) + self._set_tf_settings(allow_growth, arguments.exclude_gpus) + + use_mixed_precision = not is_predict and mixed_precision + self._use_mixed_precision = self._set_keras_mixed_precision(use_mixed_precision) + if self._use_mixed_precision: + logger.info("Enabling Mixed Precision Training.") + + if hasattr(arguments, "distribution_strategy"): + strategy = arguments.distribution_strategy + else: + strategy = "default" + self._strategy = self._get_strategy(strategy) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def use_mixed_precision(self) -> bool: + """ bool: ``True`` if mixed precision training has been enabled, otherwise ``False``. """ + return self._use_mixed_precision + + @classmethod + def loss_scale_optimizer( + cls, + optimizer: tf.keras.optimizers.Optimizer) -> mixedprecision.LossScaleOptimizer: + """ Optimize loss scaling for mixed precision training. + + Parameters + ---------- + optimizer: :class:`tf.keras.optimizers.Optimizer` + The optimizer instance to wrap + + Returns + -------- + :class:`tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer` + The original optimizer with loss scaling applied + """ + return mixedprecision.LossScaleOptimizer(optimizer) # pylint:disable=no-member + + @classmethod + def _set_tf_settings(cls, allow_growth: bool, exclude_devices: list[int]) -> None: + """ Specify Devices to place operations on and Allow TensorFlow to manage VRAM growth. + + Enables the Tensorflow allow_growth option if requested in the command line arguments + + Parameters + ---------- + allow_growth: bool + ``True`` if the Tensorflow allow_growth parameter should be set otherwise ``False`` + exclude_devices: list or ``None`` + List of GPU device indices that should not be made available to Tensorflow. Pass + ``None`` if all devices should be made available + """ + backend = get_backend() + if backend == "cpu": + logger.verbose("Hiding GPUs from Tensorflow") # type:ignore[attr-defined] + tf.config.set_visible_devices([], "GPU") + return + + if not exclude_devices and not allow_growth: + logger.debug("Not setting any specific Tensorflow settings") + return + + gpus = tf.config.list_physical_devices('GPU') + if exclude_devices: + gpus = [gpu for idx, gpu in enumerate(gpus) if idx not in exclude_devices] + logger.debug("Filtering devices to: %s", gpus) + tf.config.set_visible_devices(gpus, "GPU") + + if allow_growth and backend == "nvidia": + logger.debug("Setting Tensorflow 'allow_growth' option") + for gpu in gpus: + logger.info("Setting allow growth for GPU: %s", gpu) + tf.config.experimental.set_memory_growth(gpu, True) + logger.debug("Set Tensorflow 'allow_growth' option") + + @classmethod + def _set_keras_mixed_precision(cls, use_mixed_precision: bool) -> bool: + """ Enable the Keras experimental Mixed Precision API. + + Enables the Keras experimental Mixed Precision API if requested in the user configuration + file. + + Parameters + ---------- + use_mixed_precision: bool + ``True`` if experimental mixed precision support should be enabled for Nvidia GPUs + otherwise ``False``. + + Returns + ------- + bool + ``True`` if mixed precision has been enabled otherwise ``False`` + """ + logger.debug("use_mixed_precision: %s", use_mixed_precision) + if not use_mixed_precision: + policy = mixedprecision.Policy('float32') # pylint:disable=no-member + mixedprecision.set_global_policy(policy) # pylint:disable=no-member + logger.debug("Disabling mixed precision. (Compute dtype: %s, variable_dtype: %s)", + policy.compute_dtype, policy.variable_dtype) + return False + + policy = mixedprecision.Policy('mixed_float16') # pylint:disable=no-member + mixedprecision.set_global_policy(policy) # pylint:disable=no-member + logger.debug("Enabled mixed precision. (Compute dtype: %s, variable_dtype: %s)", + policy.compute_dtype, policy.variable_dtype) + return True + + def _get_strategy(self, + strategy: T.Literal["default", "central-storage", "mirrored"] + ) -> tf.distribute.Strategy | None: + """ If we are running on Nvidia backend and the strategy is not ``None`` then return + the correct tensorflow distribution strategy, otherwise return ``None``. + + Notes + ----- + By default Tensorflow defaults mirrored strategy to use the Nvidia NCCL method for + reductions, however this is only available in Linux, so the method used falls back to + `Hierarchical Copy All Reduce` if the OS is not Linux. + + Central Storage strategy is not compatible with Mixed Precision. However, in testing it + worked fine when using a single GPU, so we monkey-patch out the tests for Mixed-Precision + when using this strategy with a single GPU + + Parameters + ---------- + strategy: str + One of 'default', 'central-storage' or 'mirrored'. + + Returns + ------- + :class:`tensorflow.distribute.Strategy` or `None` + The request Tensorflow Strategy if the backend is Nvidia and the strategy is not + `"Default"` otherwise ``None`` + """ + if get_backend() not in ("nvidia", "directml", "rocm"): + retval = None + elif strategy == "mirrored": + retval = self._get_mirrored_strategy() + elif strategy == "central-storage": + retval = self._get_central_storage_strategy() + else: + retval = tf.distribute.get_strategy() + logger.debug("Using strategy: %s", retval) + return retval + + @classmethod + def _get_mirrored_strategy(cls) -> tf.distribute.MirroredStrategy: + """ Obtain an instance of a Tensorflow Mirrored Strategy, setting the cross device + operations appropriate for the OS in use. + + Returns + ------- + :class:`tensorflow.distribute.MirroredStrategy` + The Mirrored Distribution Strategy object with correct cross device operations set + """ + if platform.system().lower() == "linux": + cross_device_ops = tf.distribute.NcclAllReduce() + else: + cross_device_ops = tf.distribute.HierarchicalCopyAllReduce() + logger.debug("cross_device_ops: %s", cross_device_ops) + return tf.distribute.MirroredStrategy(cross_device_ops=cross_device_ops) + + @classmethod + def _get_central_storage_strategy(cls) -> tf.distribute.experimental.CentralStorageStrategy: + """ Obtain an instance of a Tensorflow Central Storage Strategy. If the strategy is being + run on a single GPU then monkey patch Tensorflows mixed-precision strategy checks to pass + successfully. + + Returns + ------- + :class:`tensorflow.distribute.experimental.CentralStorageStrategy` + The Central Storage Distribution Strategy object + """ + gpus = tf.config.get_visible_devices("GPU") + if len(gpus) == 1: + # TODO Remove these monkey patches when Strategy supports mixed-precision + from keras.mixed_precision import loss_scale_optimizer # noqa pylint:disable=import-outside-toplevel + + # Force a return of True on Loss Scale Optimizer Stategy check + loss_scale_optimizer.strategy_supports_loss_scaling = lambda: True + + # As LossScaleOptimizer aggregates gradients internally, it passes `False` as the value + # for `experimental_aggregate_gradients` in `OptimizerV2.apply_gradients`. This causes + # the optimizer to fail when checking against this strategy. We could monkey patch + # `Optimizer.apply_gradients`, but it is a lot more code to check, so we just switch + # the `experimental_aggregate_gradients` back to `True`. In brief testing this does not + # appear to have a negative impact. + func = lambda s, grads, wvars, name: s._optimizer.apply_gradients( # noqa pylint:disable=protected-access,unnecessary-lambda-assignment + list(zip(grads, wvars.value)), name, experimental_aggregate_gradients=True) + loss_scale_optimizer.LossScaleOptimizer._apply_gradients = func # noqa pylint:disable=protected-access + + return tf.distribute.experimental.CentralStorageStrategy(parameter_device="/cpu:0") + + def _get_mixed_precision_layers(self, layers: list[dict]) -> list[str]: + """ Obtain the names of the layers in a mixed precision model that have their dtype policy + explicitly set to mixed-float16. + + Parameters + ---------- + layers: List + The list of layers that appear in a keras's model configuration `dict` + + Returns + ------- + list + A list of layer names within the model that are assigned a float16 policy + """ + retval = [] + for layer in layers: + config = layer["config"] + + if layer["class_name"] in ("Functional", "Sequential"): # Recurse into sub-models + retval.extend(self._get_mixed_precision_layers(config["layers"])) + continue + + dtype = config["dtype"] + if isinstance(dtype, dict) and dtype["config"]["name"] == "mixed_float16": + logger.debug("Adding supported mixed precision layer: %s %s", layer["name"], dtype) + retval.append(layer["name"]) + else: + logger.debug("Skipping unsupported layer: %s %s", + layer.get("name", f"class_name: {layer['class_name']}"), dtype) + return retval + + def _switch_precision(self, layers: list[dict], compatible: list[str]) -> None: + """ Switch a model's datatype between mixed-float16 and float32. + + Parameters + ---------- + layers: List + The list of layers that appear in a keras's model configuration `dict` + compatible: List + A list of layer names that are compatible to have their datatype switched + """ + dtype = "mixed_float16" if self.use_mixed_precision else "float32" + policy = {"class_name": "Policy", "config": {"name": dtype}} + + for layer in layers: + config = layer["config"] + + if layer["class_name"] in ["Functional", "Sequential"]: # Recurse into sub-models + self._switch_precision(config["layers"], compatible) + continue + + if layer["name"] not in compatible: + logger.debug("Skipping incompatible layer: %s", layer["name"]) + continue + + logger.debug("Updating dtype for %s from: %s to: %s", + layer["name"], config["dtype"], policy) + config["dtype"] = policy + + def get_mixed_precision_layers(self, + build_func: Callable[[list[tf.keras.layers.Layer]], + tf.keras.models.Model], + inputs: list[tf.keras.layers.Layer] + ) -> tuple[tf.keras.models.Model, list[str]]: + """ Get and store the mixed precision layers from a full precision enabled model. + + Parameters + ---------- + build_func: Callable + The function to be called to compile the newly created model + inputs: + The inputs to the model to be compiled + + Returns + ------- + model: :class:`tensorflow.keras.model` + The built model in fp32 + list + The list of layer names within the full precision model that can be switched + to mixed precision + """ + logger.info("Storing Mixed Precision compatible layers. Please ignore any following " + "warnings about using mixed precision.") + self._set_keras_mixed_precision(True) + with tf.device("CPU"): + model = build_func(inputs) + layers = self._get_mixed_precision_layers(model.get_config()["layers"]) + + tf.keras.backend.clear_session() + self._set_keras_mixed_precision(False) + + config = model.get_config() + self._switch_precision(config["layers"], layers) + new_model = model.from_config(config) + del model + return new_model, layers + + def check_model_precision(self, + model: tf.keras.models.Model, + state: "State") -> tf.keras.models.Model: + """ Check the model's precision. + + If this is a new model, then + Rewrite an existing model's training precsion mode from mixed-float16 to float32 or + vice versa. + + This is not easy to do in keras, so we edit the model's config to change the dtype policy + for compatible layers. Create a new model from this config, then port the weights from the + old model to the new model. + + Parameters + ---------- + model: :class:`keras.models.Model` + The original saved keras model to rewrite the dtype + state: ~:class:`plugins.train.model._base.model.State` + The State information for the model + + Returns + ------- + :class:`keras.models.Model` + The original model with the datatype updated + """ + if self.use_mixed_precision and not state.mixed_precision_layers: + # Switching to mixed precision on a model which was started in FP32 prior to the + # ability to switch between precisions on a saved model is not supported as we + # do not have the compatible layer names + logger.warning("Switching from Full Precision to Mixed Precision is not supported on " + "older model files. Reverting to Full Precision.") + return model + + config = model.get_config() + + if not self.use_mixed_precision and not state.mixed_precision_layers: + # Switched to Full Precision, get compatible layers from model if not already stored + state.add_mixed_precision_layers(self._get_mixed_precision_layers(config["layers"])) + + self._switch_precision(config["layers"], state.mixed_precision_layers) + + new_model = keras.models.Model().from_config(config) + new_model.set_weights(model.get_weights()) + logger.info("Mixed precision has been updated from '%s' to '%s'", + not self.use_mixed_precision, self.use_mixed_precision) + del model + return new_model + + def strategy_scope(self) -> ContextManager: + """ Return the strategy scope if we have set a strategy, otherwise return a null + context. + + Returns + ------- + :func:`tensorflow.python.distribute.Strategy.scope` or :func:`contextlib.nullcontext` + The tensorflow strategy scope if a strategy is valid in the current scenario. A null + context manager if the strategy is not valid in the current scenario + """ + retval = nullcontext() if self._strategy is None else self._strategy.scope() + logger.debug("Using strategy scope: %s", retval) + return retval diff --git a/plugins/train/model/dfaker.py b/plugins/train/model/dfaker.py index 221ff179b5..0ad08357fd 100644 --- a/plugins/train/model/dfaker.py +++ b/plugins/train/model/dfaker.py @@ -4,20 +4,15 @@ import logging import sys -from lib.model.nn_blocks import Conv2DOutput, UpscaleBlock, ResidualBlock -from lib.utils import get_backend -from .original import Model as OriginalModel, KerasModel - -if get_backend() == "amd": - from keras.initializers import RandomNormal - from keras.layers import Input, LeakyReLU -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.initializers import RandomNormal # noqa pylint:disable=import-error,no-name-in-module - from tensorflow.keras.layers import Input, LeakyReLU # noqa pylint:disable=import-error,no-name-in-module +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.initializers import RandomNormal # pylint:disable=import-error +from tensorflow.keras.layers import Input, LeakyReLU # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error +from lib.model.nn_blocks import Conv2DOutput, UpscaleBlock, ResidualBlock +from .original import Model as OriginalModel -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Model(OriginalModel): @@ -64,4 +59,4 @@ def decoder(self, side): var_y = UpscaleBlock(64, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel([input_], outputs=outputs, name=f"decoder_{side}") + return KModel([input_], outputs=outputs, name=f"decoder_{side}") diff --git a/plugins/train/model/dfl_h128.py b/plugins/train/model/dfl_h128.py index 7d159c6e77..2bc1e61709 100644 --- a/plugins/train/model/dfl_h128.py +++ b/plugins/train/model/dfl_h128.py @@ -3,15 +3,12 @@ Based on https://github.com/iperov/DeepFaceLab """ -from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock -from lib.utils import get_backend -from .original import Model as OriginalModel, KerasModel +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Dense, Flatten, Input, Reshape # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error -if get_backend() == "amd": - from keras.layers import Dense, Flatten, Input, Reshape -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Dense, Flatten, Input, Reshape # noqa pylint:disable=import-error,no-name-in-module +from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock +from .original import Model as OriginalModel class Model(OriginalModel): @@ -32,7 +29,7 @@ def encoder(self): var_x = Dense(8 * 8 * self.encoder_dim)(var_x) var_x = Reshape((8, 8, self.encoder_dim))(var_x) var_x = UpscaleBlock(self.encoder_dim, activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder(self, side): """ DFL H128 Decoder """ @@ -51,4 +48,4 @@ def decoder(self, side): var_y = UpscaleBlock(self.encoder_dim // 4, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name=f"decoder_{side}") + return KModel(input_, outputs=outputs, name=f"decoder_{side}") diff --git a/plugins/train/model/dfl_sae.py b/plugins/train/model/dfl_sae.py index 6f00a96b8f..0c54e0031d 100644 --- a/plugins/train/model/dfl_sae.py +++ b/plugins/train/model/dfl_sae.py @@ -2,19 +2,19 @@ """ DeepFaceLab SAE Model Based on https://github.com/iperov/DeepFaceLab """ +import logging import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error + from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock -from lib.utils import get_backend -from ._base import ModelBase, KerasModel, logger +from ._base import ModelBase -if get_backend() == "amd": - from keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape # noqa pylint:disable=import-error,no-name-in-module +logger = logging.getLogger(__name__) class Model(ModelBase): @@ -73,9 +73,7 @@ def build_model(self, inputs): else: outputs = [self.decoder("a", enc_output_shape)(encoder_a), self.decoder("b", enc_output_shape)(encoder_b)] - autoencoder = KerasModel(inputs, - outputs, - name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder_df(self): @@ -91,7 +89,7 @@ def encoder_df(self): var_x = Dense(lowest_dense_res * lowest_dense_res * self.ae_dims)(var_x) var_x = Reshape((lowest_dense_res, lowest_dense_res, self.ae_dims))(var_x) var_x = UpscaleBlock(self.ae_dims, activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name="encoder_df") + return KModel(input_, var_x, name="encoder_df") def encoder_liae(self): """ DFL SAE LIAE Encoder Network """ @@ -102,7 +100,7 @@ def encoder_liae(self): var_x = Conv2DBlock(dims * 4, activation="leakyrelu")(var_x) var_x = Conv2DBlock(dims * 8, activation="leakyrelu")(var_x) var_x = Flatten()(var_x) - return KerasModel(input_, var_x, name="encoder_liae") + return KModel(input_, var_x, name="encoder_liae") def inter_liae(self, side, input_shape): """ DFL SAE LIAE Intermediate Network """ @@ -113,7 +111,7 @@ def inter_liae(self, side, input_shape): var_x = Dense(lowest_dense_res * lowest_dense_res * self.ae_dims * 2)(var_x) var_x = Reshape((lowest_dense_res, lowest_dense_res, self.ae_dims * 2))(var_x) var_x = UpscaleBlock(self.ae_dims * 2, activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name=f"intermediate_{side}") + return KModel(input_, var_x, name=f"intermediate_{side}") def decoder(self, side, input_shape): """ DFL SAE Decoder Network""" @@ -151,15 +149,15 @@ def decoder(self, side, input_shape): var_y = UpscaleBlock(self.decoder_dim * 2, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name=f"decoder_{side}") + return KModel(input_, outputs=outputs, name=f"decoder_{side}") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ - mappings = dict(df={f"{self.name}_encoder.h5": "encoder_df", - f"{self.name}_decoder_A.h5": "decoder_a", - f"{self.name}_decoder_B.h5": "decoder_b"}, - liae={f"{self.name}_encoder.h5": "encoder_liae", - f"{self.name}_intermediate_B.h5": "intermediate_both", - f"{self.name}_intermediate.h5": "intermediate_b", - f"{self.name}_decoder.h5": "decoder_both"}) + mappings = {"df": {f"{self.name}_encoder.h5": "encoder_df", + f"{self.name}_decoder_A.h5": "decoder_a", + f"{self.name}_decoder_B.h5": "decoder_b"}, + "liae": {f"{self.name}_encoder.h5": "encoder_liae", + f"{self.name}_intermediate_B.h5": "intermediate_both", + f"{self.name}_intermediate.h5": "intermediate_b", + f"{self.name}_decoder.h5": "decoder_both"}} return mappings[self.config["architecture"]] diff --git a/plugins/train/model/dfl_sae_defaults.py b/plugins/train/model/dfl_sae_defaults.py index 34fc916257..38c43d3b66 100644 --- a/plugins/train/model/dfl_sae_defaults.py +++ b/plugins/train/model/dfl_sae_defaults.py @@ -44,74 +44,59 @@ _HELPTEXT = "DFL SAE Model (Adapted from https://github.com/iperov/DeepFaceLab)" -_DEFAULTS = { - "input_size": { - "default": 128, - "info": "Resolution (in pixels) of the input image to train on.\n" - "BE AWARE Larger resolution will dramatically increase VRAM requirements.\n" - "\nMust be divisible by 16.", - "datatype": int, - "rounding": 16, - "min_max": (64, 256), - "group": "size", - "fixed": True, - }, - "clipnorm": { - "default": True, - "info": "Controls gradient clipping of the optimizer. Can prevent model corruption at " - "the expense of VRAM.", - "datatype": bool, - "fixed": False, - "group": "settings", - }, - "architecture": { - "default": "df", - "info": "Model architecture:" - "\n\t'df': Keeps the faces more natural." - "\n\t'liae': Can help fix overly different face shapes.", - "datatype": str, - "choices": ["df", "liae"], - "gui_radio": True, - "fixed": True, - "group": "network", - }, - "autoencoder_dims": { - "default": 0, - "info": "Face information is stored in AutoEncoder dimensions. If there are not enough " - "dimensions then certain facial features may not be recognized." - "\nHigher number of dimensions are better, but require more VRAM." - "\nSet to 0 to use the architecture defaults (256 for liae, 512 for df).", - "datatype": int, - "rounding": 32, - "min_max": (0, 1024), - "fixed": True, - "group": "network", - }, - "encoder_dims": { - "default": 42, - "info": "Encoder dimensions per channel. Higher number of encoder dimensions will help " - "the model to recognize more facial features, but will require more VRAM.", - "datatype": int, - "rounding": 1, - "min_max": (21, 85), - "fixed": True, - "group": "network", - }, - "decoder_dims": { - "default": 21, - "info": "Decoder dimensions per channel. Higher number of decoder dimensions will help " - "the model to improve details, but will require more VRAM.", - "datatype": int, - "rounding": 1, - "min_max": (10, 85), - "fixed": True, - "group": "network", - }, - "multiscale_decoder": { - "default": False, - "info": "Multiscale decoder can help to obtain better details.", - "datatype": bool, - "fixed": True, - "group": "network", - }, -} +_DEFAULTS = dict( + input_size=dict( + default=128, + info="Resolution (in pixels) of the input image to train on.\n" + "BE AWARE Larger resolution will dramatically increase VRAM requirements.\n" + "\nMust be divisible by 16.", + datatype=int, + rounding=16, + min_max=(64, 256), + group="size", + fixed=True), + architecture=dict( + default="df", + info="Model architecture:" + "\n\t'df': Keeps the faces more natural." + "\n\t'liae': Can help fix overly different face shapes.", + datatype=str, + choices=["df", "liae"], + gui_radio=True, + fixed=True, + group="network"), + autoencoder_dims=dict( + default=0, + info="Face information is stored in AutoEncoder dimensions. If there are not enough " + "dimensions then certain facial features may not be recognized." + "\nHigher number of dimensions are better, but require more VRAM." + "\nSet to 0 to use the architecture defaults (256 for liae, 512 for df).", + datatype=int, + rounding=32, + min_max=(0, 1024), + fixed=True, + group="network"), + encoder_dims=dict( + default=42, + info="Encoder dimensions per channel. Higher number of encoder dimensions will help " + "the model to recognize more facial features, but will require more VRAM.", + datatype=int, + rounding=1, + min_max=(21, 85), + fixed=True, + group="network"), + decoder_dims=dict( + default=21, + info="Decoder dimensions per channel. Higher number of decoder dimensions will help " + "the model to improve details, but will require more VRAM.", + datatype=int, + rounding=1, + min_max=(10, 85), + fixed=True, + group="network"), + multiscale_decoder=dict( + default=False, + info="Multiscale decoder can help to obtain better details.", + datatype=bool, + fixed=True, + group="network")) diff --git a/plugins/train/model/dlight.py b/plugins/train/model/dlight.py index 18808122ff..154b090466 100644 --- a/plugins/train/model/dlight.py +++ b/plugins/train/model/dlight.py @@ -7,22 +7,22 @@ kvrooman for numerous insights and invaluable aid DeepHomage for lots of testing """ +import logging + +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import ( # pylint:disable=import-error + AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout, Flatten, Input, Reshape, + LeakyReLU, UpSampling2D) +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error from lib.model.nn_blocks import (Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock, Upscale2xBlock) -from lib.utils import FaceswapError, get_backend +from lib.utils import FaceswapError + +from ._base import ModelBase -from ._base import ModelBase, KerasModel, logger -if get_backend() == "amd": - from keras.layers import ( - AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout, Flatten, Input, Reshape, - LeakyReLU, UpSampling2D) -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import ( # pylint:disable=import-error,no-name-in-module - AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout, Flatten, Input, Reshape, - LeakyReLU, UpSampling2D) +logger = logging.getLogger(__name__) class Model(ModelBase): @@ -32,21 +32,22 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.input_shape = (128, 128, 3) - self.features = dict(lowmem=0, fair=1, best=2)[self.config["features"]] + self.features = {"lowmem": 0, "fair": 1, "best": 2}[self.config["features"]] self.encoder_filters = 64 if self.features > 0 else 48 bonum_fortunam = 128 self.encoder_dim = {0: 512 + bonum_fortunam, 1: 1024 + bonum_fortunam, 2: 1536 + bonum_fortunam}[self.features] - self.details = dict(fast=0, good=1)[self.config["details"]] + self.details = {"fast": 0, "good": 1}[self.config["details"]] try: self.upscale_ratio = {128: 2, 256: 4, 384: 6}[self.config["output_size"]] - except KeyError: + except KeyError as err: logger.error("Config error: output_size must be one of: 128, 256, or 384.") - raise FaceswapError("Config error: output_size must be one of: 128, 256, or 384.") + raise FaceswapError("Config error: output_size must be one of: " + "128, 256, or 384.") from err logger.debug("output_size: %s, features: %s, encoder_filters: %s, encoder_dim: %s, " " details: %s, upscale_ratio: %s", self.config["output_size"], self.features, @@ -62,7 +63,7 @@ def build_model(self, inputs): outputs = [self.decoder_a()(encoder_a), decoder_b()(encoder_b)] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder(self): @@ -101,7 +102,7 @@ def encoder(self): var_x = Dropout(0.05)(var_x) var_x = Reshape((4, 4, 1024))(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder_a(self): """ DeLight Decoder A(old face) Network """ @@ -133,7 +134,7 @@ def decoder_a(self): outputs.append(var_y) - return KerasModel([input_], outputs=outputs, name="decoder_a") + return KModel([input_], outputs=outputs, name="decoder_a") def decoder_b_fast(self): """ DeLight Fast Decoder B(new face) Network """ @@ -168,7 +169,7 @@ def decoder_b_fast(self): outputs.append(var_y) - return KerasModel([input_], outputs=outputs, name="decoder_b_fast") + return KModel([input_], outputs=outputs, name="decoder_b_fast") def decoder_b(self): """ DeLight Decoder B(new face) Network """ @@ -220,7 +221,7 @@ def decoder_b(self): outputs.append(var_y) - return KerasModel([input_], outputs=outputs, name="decoder_b") + return KModel([input_], outputs=outputs, name="decoder_b") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ diff --git a/plugins/train/model/iae.py b/plugins/train/model/iae.py index dbbb982e30..d2690dd3be 100644 --- a/plugins/train/model/iae.py +++ b/plugins/train/model/iae.py @@ -1,17 +1,13 @@ #!/usr/bin/env python3 """ Improved autoencoder for faceswap """ -from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock -from lib.utils import get_backend - -from ._base import ModelBase, KerasModel +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, Reshape # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error -if get_backend() == "amd": - from keras.layers import Concatenate, Dense, Flatten, Input, Reshape +from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, Reshape # noqa pylint:disable=import-error,no-name-in-module +from ._base import ModelBase class Model(ModelBase): @@ -35,7 +31,7 @@ def build_model(self, inputs): outputs = [decoder(Concatenate()([inter_a(encoder_a), inter_both(encoder_a)])), decoder(Concatenate()([inter_b(encoder_b), inter_both(encoder_b)]))] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder(self): @@ -47,7 +43,7 @@ def encoder(self): var_x = Conv2DBlock(512, activation="leakyrelu")(var_x) var_x = Conv2DBlock(1024, activation="leakyrelu")(var_x) var_x = Flatten()(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def intermediate(self, side): """ Intermediate Network """ @@ -55,7 +51,7 @@ def intermediate(self, side): var_x = Dense(self.encoder_dim)(input_) var_x = Dense(4 * 4 * int(self.encoder_dim/2))(var_x) var_x = Reshape((4, 4, int(self.encoder_dim/2)))(var_x) - return KerasModel(input_, var_x, name=f"inter_{side}") + return KModel(input_, var_x, name=f"inter_{side}") def decoder(self): """ Decoder Network """ @@ -76,7 +72,7 @@ def decoder(self): var_y = UpscaleBlock(64, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name="mask_out")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name="decoder") + return KModel(input_, outputs=outputs, name="decoder") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ diff --git a/plugins/train/model/lightweight.py b/plugins/train/model/lightweight.py index 7dc0c69880..4feca05ab1 100644 --- a/plugins/train/model/lightweight.py +++ b/plugins/train/model/lightweight.py @@ -4,8 +4,10 @@ Based on the original https://www.reddit.com/r/deepfakes/ code sample + contributions """ +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error + from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock -from .original import Model as OriginalModel, KerasModel, Dense, Flatten, Input, Reshape +from .original import Model as OriginalModel, Dense, Flatten, Input, Reshape class Model(OriginalModel): @@ -25,7 +27,7 @@ def encoder(self): var_x = Dense(4 * 4 * 512)(var_x) var_x = Reshape((4, 4, 512))(var_x) var_x = UpscaleBlock(256, activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder(self, side): """ Decoder Network """ @@ -46,4 +48,4 @@ def decoder(self, side): activation="sigmoid", name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name=f"decoder_{side}") + return KModel(input_, outputs=outputs, name=f"decoder_{side}") diff --git a/plugins/train/model/original.py b/plugins/train/model/original.py index d23b59c532..0613a5d55e 100644 --- a/plugins/train/model/original.py +++ b/plugins/train/model/original.py @@ -6,15 +6,12 @@ from. """ -from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock -from lib.utils import get_backend -from ._base import KerasModel, ModelBase +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.layers import Dense, Flatten, Reshape, Input # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error -if get_backend() == "amd": - from keras.layers import Dense, Flatten, Reshape, Input -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.layers import Dense, Flatten, Reshape, Input # noqa pylint:disable=import-error,no-name-in-module +from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock +from ._base import ModelBase class Model(ModelBase): @@ -66,12 +63,6 @@ def build_model(self, inputs): 2 Decoders are then defined (one for each side) with the encoder instances passed in as input to the corresponding decoders. - It is important to note that any models and sub-models should not call - :class:`keras.models.Model` directly, but rather call - :class:`plugins.train.model._base.KerasModel`. This acts as a wrapper for Keras' Model - class, but handles some minor differences which need to be handled between Nvidia and AMD - backends. - The final output of the model should always call :class:`lib.model.nn_blocks.Conv2DOutput` so that the correct data type is set for the final activation, to support Mixed Precision Training. Failure to do so is likely to lead to issues when Mixed Precision is enabled. @@ -85,8 +76,7 @@ def build_model(self, inputs): Returns ------- :class:`keras.models.Model` - The output of this function must be a keras model generated from - :class:`plugins.train.model._base.KerasModel`. See Keras documentation for the correct + See Keras documentation for the correct structure, but note that parameter :attr:`name` is a required rather than an optional argument in Faceswap. You should assign this to the attribute ``self.name`` that is automatically generated from the plugin's filename. @@ -100,7 +90,7 @@ def build_model(self, inputs): outputs = [self.decoder("a")(encoder_a), self.decoder("b")(encoder_b)] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder(self): @@ -127,7 +117,7 @@ def encoder(self): var_x = Dense(4 * 4 * 1024)(var_x) var_x = Reshape((4, 4, 1024))(var_x) var_x = UpscaleBlock(512, activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder(self, side): """ The original Faceswap Decoder Network. @@ -160,7 +150,7 @@ def decoder(self, side): var_y = UpscaleBlock(64, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name=f"decoder_{side}") + return KModel(input_, outputs=outputs, name=f"decoder_{side}") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ diff --git a/plugins/train/model/phaze_a.py b/plugins/train/model/phaze_a.py index 62468c53e0..d6169fe252 100644 --- a/plugins/train/model/phaze_a.py +++ b/plugins/train/model/phaze_a.py @@ -1,50 +1,31 @@ #!/usr/bin/env python3 """ Phaze-A Model by TorzDF with thanks to BirbFakes and the myriad of testers. """ -# pylint: disable=too-many-lines -import sys +# pylint:disable=too-many-lines +from __future__ import annotations +import logging +import typing as T from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union - -if sys.version_info < (3, 8): - from typing_extensions import Literal -else: - from typing import Literal - import numpy as np +import tensorflow as tf from lib.model.nn_blocks import ( Conv2D, Conv2DBlock, Conv2DOutput, ResidualBlock, UpscaleBlock, Upscale2xBlock, UpscaleResizeImagesBlock, UpscaleDNYBlock) from lib.model.normalization import ( - AdaInstanceNormalization, GroupNormalization, InstanceNormalization, LayerNormalization, - RMSNormalization) -from lib.utils import get_backend, get_tf_version, FaceswapError - -from ._base import KerasModel, ModelBase, logger, _get_all_sub_models - -if get_backend() == "amd": - from keras import applications as kapp, backend as K - from keras.layers import ( - Add, BatchNormalization, Concatenate, Dense, Dropout, Flatten, GaussianNoise, MaxPool2D, - GlobalAveragePooling2D, GlobalMaxPooling2D, Input, LeakyReLU, Reshape, UpSampling2D, - Conv2D as KConv2D) - from keras.models import clone_model - # typing checks - import keras - from plaidml.tile import Value as Tensor # pylint:disable=import-error -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import applications as kapp, backend as K # pylint:disable=import-error - from tensorflow.keras.layers import ( # pylint:disable=import-error,no-name-in-module - Add, BatchNormalization, Concatenate, Dense, Dropout, Flatten, GaussianNoise, MaxPool2D, - GlobalAveragePooling2D, GlobalMaxPooling2D, Input, LeakyReLU, Reshape, UpSampling2D, - Conv2D as KConv2D) - from tensorflow.keras.models import clone_model # noqa pylint:disable=import-error,no-name-in-module - # typing checks - from tensorflow import keras - from tensorflow import Tensor + AdaInstanceNormalization, GroupNormalization, InstanceNormalization, RMSNormalization) +from lib.model.networks import ViT, TypeModelsViT +from lib.utils import get_tf_version, FaceswapError + +from ._base import ModelBase, get_all_sub_models + +logger = logging.getLogger(__name__) + +K = tf.keras.backend +kapp = tf.keras.applications +kl = tf.keras.layers +keras = tf.keras @dataclass @@ -58,9 +39,6 @@ class _EncoderInfo: exist in Keras Applications default_size: int The default input size of the encoder - no_amd: bool, optional - ``True`` if the encoder is not compatible with the PlaidML backend otherwise ``False``. - Default: ``False`` tf_min: float, optional The lowest version of Tensorflow that the encoder can be used for. Default: `2.0` scaling: tuple, optional @@ -75,87 +53,98 @@ class _EncoderInfo: """ keras_name: str default_size: int - no_amd: bool = False - tf_min: float = 2.0 - scaling: Tuple[int, int] = (0, 1) + tf_min: tuple[int, int] = (2, 0) + scaling: tuple[int, int] = (0, 1) min_size: int = 32 enforce_for_weights: bool = False - color_order: Literal["bgr", "rgb"] = "rgb" - - -_MODEL_MAPPING: Dict[str, _EncoderInfo] = dict( - densenet121=_EncoderInfo( + color_order: T.Literal["bgr", "rgb"] = "rgb" + + +_MODEL_MAPPING: dict[str, _EncoderInfo] = { + "clipv_farl-b-16-16": _EncoderInfo( + keras_name="FaRL-B-16-16", default_size=224), + "clipv_farl-b-16-64": _EncoderInfo( + keras_name="FaRL-B-16-64", default_size=224), + "clipv_vit-b-16": _EncoderInfo( + keras_name="ViT-B-16", default_size=224), + "clipv_vit-b-32": _EncoderInfo( + keras_name="ViT-B-32", default_size=224), + "clipv_vit-l-14": _EncoderInfo( + keras_name="ViT-L-14", default_size=224), + "clipv_vit-l-14-336px": _EncoderInfo( + keras_name="ViT-L-14-336px", default_size=336), + "densenet121": _EncoderInfo( keras_name="DenseNet121", default_size=224), - densenet169=_EncoderInfo( + "densenet169": _EncoderInfo( keras_name="DenseNet169", default_size=224), - densenet201=_EncoderInfo( + "densenet201": _EncoderInfo( keras_name="DenseNet201", default_size=224), - efficientnet_b0=_EncoderInfo( - keras_name="EfficientNetB0", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=224), - efficientnet_b1=_EncoderInfo( - keras_name="EfficientNetB1", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=240), - efficientnet_b2=_EncoderInfo( - keras_name="EfficientNetB2", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=260), - efficientnet_b3=_EncoderInfo( - keras_name="EfficientNetB3", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=300), - efficientnet_b4=_EncoderInfo( - keras_name="EfficientNetB4", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=380), - efficientnet_b5=_EncoderInfo( - keras_name="EfficientNetB5", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=456), - efficientnet_b6=_EncoderInfo( - keras_name="EfficientNetB6", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=528), - efficientnet_b7=_EncoderInfo( - keras_name="EfficientNetB7", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=600), - efficientnet_v2_b0=_EncoderInfo( - keras_name="EfficientNetV2B0", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=224), - efficientnet_v2_b1=_EncoderInfo( - keras_name="EfficientNetV2B1", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=240), - efficientnet_v2_b2=_EncoderInfo( - keras_name="EfficientNetV2B2", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=260), - efficientnet_v2_b3=_EncoderInfo( - keras_name="EfficientNetV2B3", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=300), - efficientnet_v2_s=_EncoderInfo( - keras_name="EfficientNetV2S", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=384), - efficientnet_v2_m=_EncoderInfo( - keras_name="EfficientNetV2M", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480), - efficientnet_v2_l=_EncoderInfo( - keras_name="EfficientNetV2L", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480), - inception_resnet_v2=_EncoderInfo( + "efficientnet_b0": _EncoderInfo( + keras_name="EfficientNetB0", tf_min=(2, 3), scaling=(0, 255), default_size=224), + "efficientnet_b1": _EncoderInfo( + keras_name="EfficientNetB1", tf_min=(2, 3), scaling=(0, 255), default_size=240), + "efficientnet_b2": _EncoderInfo( + keras_name="EfficientNetB2", tf_min=(2, 3), scaling=(0, 255), default_size=260), + "efficientnet_b3": _EncoderInfo( + keras_name="EfficientNetB3", tf_min=(2, 3), scaling=(0, 255), default_size=300), + "efficientnet_b4": _EncoderInfo( + keras_name="EfficientNetB4", tf_min=(2, 3), scaling=(0, 255), default_size=380), + "efficientnet_b5": _EncoderInfo( + keras_name="EfficientNetB5", tf_min=(2, 3), scaling=(0, 255), default_size=456), + "efficientnet_b6": _EncoderInfo( + keras_name="EfficientNetB6", tf_min=(2, 3), scaling=(0, 255), default_size=528), + "efficientnet_b7": _EncoderInfo( + keras_name="EfficientNetB7", tf_min=(2, 3), scaling=(0, 255), default_size=600), + "efficientnet_v2_b0": _EncoderInfo( + keras_name="EfficientNetV2B0", tf_min=(2, 8), scaling=(-1, 1), default_size=224), + "efficientnet_v2_b1": _EncoderInfo( + keras_name="EfficientNetV2B1", tf_min=(2, 8), scaling=(-1, 1), default_size=240), + "efficientnet_v2_b2": _EncoderInfo( + keras_name="EfficientNetV2B2", tf_min=(2, 8), scaling=(-1, 1), default_size=260), + "efficientnet_v2_b3": _EncoderInfo( + keras_name="EfficientNetV2B3", tf_min=(2, 8), scaling=(-1, 1), default_size=300), + "efficientnet_v2_s": _EncoderInfo( + keras_name="EfficientNetV2S", tf_min=(2, 8), scaling=(-1, 1), default_size=384), + "efficientnet_v2_m": _EncoderInfo( + keras_name="EfficientNetV2M", tf_min=(2, 8), scaling=(-1, 1), default_size=480), + "efficientnet_v2_l": _EncoderInfo( + keras_name="EfficientNetV2L", tf_min=(2, 8), scaling=(-1, 1), default_size=480), + "inception_resnet_v2": _EncoderInfo( keras_name="InceptionResNetV2", scaling=(-1, 1), min_size=75, default_size=299), - inception_v3=_EncoderInfo( + "inception_v3": _EncoderInfo( keras_name="InceptionV3", scaling=(-1, 1), min_size=75, default_size=299), - mobilenet=_EncoderInfo( + "mobilenet": _EncoderInfo( keras_name="MobileNet", scaling=(-1, 1), default_size=224), - mobilenet_v2=_EncoderInfo( + "mobilenet_v2": _EncoderInfo( keras_name="MobileNetV2", scaling=(-1, 1), default_size=224), - mobilenet_v3_large=_EncoderInfo( - keras_name="MobileNetV3Large", no_amd=True, tf_min=2.4, scaling=(-1, 1), default_size=224), - mobilenet_v3_small=_EncoderInfo( - keras_name="MobileNetV3Small", no_amd=True, tf_min=2.4, scaling=(-1, 1), default_size=224), - nasnet_large=_EncoderInfo( + "mobilenet_v3_large": _EncoderInfo( + keras_name="MobileNetV3Large", tf_min=(2, 4), scaling=(-1, 1), default_size=224), + "mobilenet_v3_small": _EncoderInfo( + keras_name="MobileNetV3Small", tf_min=(2, 4), scaling=(-1, 1), default_size=224), + "nasnet_large": _EncoderInfo( keras_name="NASNetLarge", scaling=(-1, 1), default_size=331, enforce_for_weights=True), - nasnet_mobile=_EncoderInfo( + "nasnet_mobile": _EncoderInfo( keras_name="NASNetMobile", scaling=(-1, 1), default_size=224, enforce_for_weights=True), - resnet50=_EncoderInfo( + "resnet50": _EncoderInfo( keras_name="ResNet50", scaling=(-1, 1), min_size=32, default_size=224), - resnet50_v2=_EncoderInfo( - keras_name="ResNet50V2", no_amd=True, scaling=(-1, 1), default_size=224), - resnet101=_EncoderInfo( - keras_name="ResNet101", no_amd=True, scaling=(-1, 1), default_size=224), - resnet101_v2=_EncoderInfo( - keras_name="ResNet101V2", no_amd=True, scaling=(-1, 1), default_size=224), - resnet152=_EncoderInfo( - keras_name="ResNet152", no_amd=True, scaling=(-1, 1), default_size=224), - resnet152_v2=_EncoderInfo( - keras_name="ResNet152V2", no_amd=True, scaling=(-1, 1), default_size=224), - vgg16=_EncoderInfo( + "resnet50_v2": _EncoderInfo( + keras_name="ResNet50V2", scaling=(-1, 1), default_size=224), + "resnet101": _EncoderInfo( + keras_name="ResNet101", scaling=(-1, 1), default_size=224), + "resnet101_v2": _EncoderInfo( + keras_name="ResNet101V2", scaling=(-1, 1), default_size=224), + "resnet152": _EncoderInfo( + keras_name="ResNet152", scaling=(-1, 1), default_size=224), + "resnet152_v2": _EncoderInfo( + keras_name="ResNet152V2", scaling=(-1, 1), default_size=224), + "vgg16": _EncoderInfo( keras_name="VGG16", color_order="bgr", scaling=(0, 255), default_size=224), - vgg19=_EncoderInfo( + "vgg19": _EncoderInfo( keras_name="VGG19", color_order="bgr", scaling=(0, 255), default_size=224), - xception=_EncoderInfo( + "xception": _EncoderInfo( keras_name="Xception", scaling=(-1, 1), min_size=71, default_size=299), - fs_original=_EncoderInfo( - keras_name="", color_order="bgr", min_size=32, default_size=1024)) + "fs_original": _EncoderInfo( + keras_name="", color_order="bgr", min_size=32, default_size=1024)} class Model(ModelBase): @@ -174,13 +163,13 @@ class Model(ModelBase): """ def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) - if self.config["output_size"] % 64 != 0: - raise FaceswapError("Phaze-A output shape must be a multiple of 64") + if self.config["output_size"] % 16 != 0: + raise FaceswapError("Phaze-A output shape must be a multiple of 16") self._validate_encoder_architecture() self.config["freeze_layers"] = self._select_freeze_layers() - self.input_shape = self._get_input_shape() + self.input_shape: tuple[int, int, int] = self._get_input_shape() self.color_order = _MODEL_MAPPING[self.config["enc_architecture"]].color_order def build(self) -> None: @@ -197,13 +186,13 @@ def build(self) -> None: super().build() return with self._settings.strategy_scope(): - model = self._io._load() # pylint:disable=protected-access + model = self.io.load() model = self._update_dropouts(model) self._model = model self._compile_model() self._output_summary() - def _update_dropouts(self, model: keras.models.Model) -> keras.models.Model: + def _update_dropouts(self, model: tf.keras.models.Model) -> tf.keras.models.Model: """ Update the saved model with new dropout rates. Keras, annoyingly, does not actually change the dropout of the underlying layer, so we need @@ -219,18 +208,18 @@ def _update_dropouts(self, model: keras.models.Model) -> keras.models.Model: :class:`keras.models.Model` The loaded Keras Model with the dropout rates updated """ - dropouts = dict(fc=self.config["fc_dropout"], - gblock=self.config["fc_gblock_dropout"]) + dropouts = {"fc": self.config["fc_dropout"], + "gblock": self.config["fc_gblock_dropout"]} logger.debug("Config dropouts: %s", dropouts) updated = False - for mod in _get_all_sub_models(model): + for mod in get_all_sub_models(model): if not mod.name.startswith("fc_"): continue key = "gblock" if "gblock" in mod.name else mod.name.split("_")[0] rate = dropouts[key] log_once = False for layer in mod.layers: - if not isinstance(layer, Dropout): + if not isinstance(layer, kl.Dropout): continue if layer.rate != rate: logger.debug("Updating dropout rate for %s from %s to %s", @@ -243,13 +232,13 @@ def _update_dropouts(self, model: keras.models.Model) -> keras.models.Model: updated = True if updated: logger.debug("Dropout rate updated. Cloning model") - new_model = clone_model(model) + new_model = keras.models.clone_model(model) new_model.set_weights(model.get_weights()) del model model = new_model return model - def _select_freeze_layers(self) -> List[str]: + def _select_freeze_layers(self) -> list[str]: """ Process the selected frozen layers and replace the `keras_encoder` option with the actual keras model name @@ -262,6 +251,8 @@ def _select_freeze_layers(self) -> List[str]: layers = self.config["freeze_layers"] # EfficientNetV2 is inconsistent with other model's naming conventions keras_name = _MODEL_MAPPING[arch].keras_name.replace("EfficientNetV2", "EfficientNetV2-") + # CLIPv model is always called 'visual' regardless of weights/format loaded + keras_name = "visual" if arch.startswith("clipv_") else keras_name if "keras_encoder" not in self.config["freeze_layers"]: retval = layers @@ -271,9 +262,10 @@ def _select_freeze_layers(self) -> List[str]: else: retval = [layer for layer in layers if layer != "keras_encoder"] logger.debug("Removing 'keras_encoder' for '%s'", arch) + return retval - def _get_input_shape(self) -> Tuple[int, int, int]: + def _get_input_shape(self) -> tuple[int, int, int]: """ Obtain the input shape for the model. Input shape is calculated from the selected Encoder's input size, scaled to the user @@ -295,7 +287,7 @@ def _get_input_shape(self) -> Tuple[int, int, int]: scaling = self.config["enc_scaling"] / 100 min_size = _MODEL_MAPPING[arch].min_size - size = int(max(min_size, min(default_size, ((default_size * scaling) // 16) * 16))) + size = int(max(min_size, ((default_size * scaling) // 16) * 16)) if self.config["enc_load_weights"] and enforce_size and scaling != 1.0: logger.warning("%s requires input size to be %spx when loading imagenet weights. " @@ -320,19 +312,14 @@ def _validate_encoder_architecture(self) -> None: raise FaceswapError(f"'{arch}' is not a valid choice for encoder architecture. Choose " f"one of {list(_MODEL_MAPPING.keys())}.") - if get_backend() == "amd" and model.no_amd: - valid = [k for k, v in _MODEL_MAPPING.items() if not v.no_amd] - raise FaceswapError(f"'{arch}' is not compatible with the AMD backend. Choose one of " - f"{valid}.") - tf_ver = get_tf_version() tf_min = model.tf_min - if get_backend() != "amd" and tf_ver < tf_min: + if tf_ver < tf_min: raise FaceswapError(f"{arch}' is not compatible with your version of Tensorflow. The " f"minimum version required is {tf_min} whilst you have version " f"{tf_ver} installed.") - def build_model(self, inputs: List[Tensor]) -> keras.models.Model: + def build_model(self, inputs: list[tf.Tensor]) -> tf.keras.models.Model: """ Create the model's structure. Parameters @@ -344,11 +331,7 @@ def build_model(self, inputs: List[Tensor]) -> keras.models.Model: Returns ------- :class:`keras.models.Model` - The output of this function must be a keras model generated from - :class:`plugins.train.model._base.KerasModel`. See Keras documentation for the correct - structure, but note that parameter :attr:`name` is a required rather than an optional - argument in Faceswap. You should assign this to the attribute ``self.name`` that is - automatically generated from the plugin's filename. + The generated model """ # Create sub-Models encoders = self._build_encoders(inputs) @@ -358,10 +341,10 @@ def build_model(self, inputs: List[Tensor]) -> keras.models.Model: # Create Autoencoder outputs = [decoders["a"], decoders["b"]] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = keras.models.Model(inputs, outputs, name=self.model_name) return autoencoder - def _build_encoders(self, inputs: List[Tensor]) -> Dict[str, keras.models.Model]: + def _build_encoders(self, inputs: list[tf.Tensor]) -> dict[str, tf.keras.models.Model]: """ Build the encoders for Phaze-A Parameters @@ -376,13 +359,13 @@ def _build_encoders(self, inputs: List[Tensor]) -> Dict[str, keras.models.Model] side as key ('a' or 'b'), encoder for side as value """ encoder = Encoder(self.input_shape, self.config)() - retval = dict(a=encoder(inputs[0]), b=encoder(inputs[1])) + retval = {"a": encoder(inputs[0]), "b": encoder(inputs[1])} logger.debug("Encoders: %s", retval) return retval def _build_fully_connected( self, - inputs: Dict[str, keras.models.Model]) -> Dict[str, List[keras.models.Model]]: + inputs: dict[str, tf.keras.models.Model]) -> dict[str, list[tf.keras.models.Model]]: """ Build the fully connected layers for Phaze-A Parameters @@ -413,22 +396,22 @@ def _build_fully_connected( fc_shared = fc_a else: fc_shared = fc_both - inter_a = [Concatenate(name="inter_a")([inter_a[0], fc_shared(inputs["a"])])] - inter_b = [Concatenate(name="inter_b")([inter_b[0], fc_shared(inputs["b"])])] + inter_a = [kl.Concatenate(name="inter_a")([inter_a[0], fc_shared(inputs["a"])])] + inter_b = [kl.Concatenate(name="inter_b")([inter_b[0], fc_shared(inputs["b"])])] if self.config["enable_gblock"]: fc_gblock = FullyConnected("gblock", input_shapes, self.config)() inter_a.append(fc_gblock(inputs["a"])) inter_b.append(fc_gblock(inputs["b"])) - retval = dict(a=inter_a, b=inter_b) + retval = {"a": inter_a, "b": inter_b} logger.debug("Fully Connected: %s", retval) return retval def _build_g_blocks( self, - inputs: Dict[str, List[keras.models.Model]] - ) -> Dict[str, Union[List[keras.models.Model], keras.models.Model]]: + inputs: dict[str, list[tf.keras.models.Model]] + ) -> dict[str, list[tf.keras.models.Model] | tf.keras.models.Model]: """ Build the g-block layers for Phaze-A. If a g-block has not been selected for this model, then the original `inters` models are @@ -451,19 +434,18 @@ def _build_g_blocks( input_shapes = [K.int_shape(inter)[1:] for inter in inputs["a"]] if self.config["split_gblock"]: - retval = dict(a=GBlock("a", input_shapes, self.config)()(inputs["a"]), - b=GBlock("b", input_shapes, self.config)()(inputs["b"])) + retval = {"a": GBlock("a", input_shapes, self.config)()(inputs["a"]), + "b": GBlock("b", input_shapes, self.config)()(inputs["b"])} else: g_block = GBlock("both", input_shapes, self.config)() - retval = dict(a=g_block((inputs["a"])), b=g_block((inputs["b"]))) + retval = {"a": g_block((inputs["a"])), "b": g_block((inputs["b"]))} logger.debug("G-Blocks: %s", retval) return retval - def _build_decoders( - self, - inputs: Dict[str, Union[List[keras.models.Model], keras.models.Model]] - ) -> Dict[str, keras.models.Model]: + def _build_decoders(self, + inputs: dict[str, list[tf.keras.models.Model] | tf.keras.models.Model] + ) -> dict[str, tf.keras.models.Model]: """ Build the encoders for Phaze-A Parameters @@ -482,28 +464,34 @@ def _build_decoders( # There will only ever be 1 input. For inters: either inter out, or concatenate of inters # For g-block, this only ever has one output input_ = input_[0] if isinstance(input_, list) else input_ + + # If learning a mask and upscales have been placed into FC layer, then the mask will also + # come as an input + if self.config["learn_mask"] and self.config["dec_upscales_in_fc"]: + input_ = input_[0] + input_shape = K.int_shape(input_)[1:] if self.config["split_decoders"]: - retval = dict(a=Decoder("a", input_shape, self.config)()(inputs["a"]), - b=Decoder("b", input_shape, self.config)()(inputs["b"])) + retval = {"a": Decoder("a", input_shape, self.config)()(inputs["a"]), + "b": Decoder("b", input_shape, self.config)()(inputs["b"])} else: decoder = Decoder("both", input_shape, self.config)() - retval = dict(a=decoder(inputs["a"]), b=decoder(inputs["b"])) + retval = {"a": decoder(inputs["a"]), "b": decoder(inputs["b"])} logger.debug("Decoders: %s", retval) return retval -def _bottleneck(inputs: Tensor, bottleneck: str, size: int, normalization: str) -> Tensor: +def _bottleneck(inputs: tf.Tensor, bottleneck: str, size: int, normalization: str) -> tf.Tensor: """ The bottleneck fully connected layer. Can be called from Encoder or FullyConnected layers. Parameters ---------- inputs: tensor The input to the bottleneck layer - bottleneck: str - The type of layer to use for the bottleneck + bottleneck: str or ``None`` + The type of layer to use for the bottleneck. ``None`` to not use a bottleneck size: int The number of nodes for the dense layer (if selected) normalization: str @@ -514,31 +502,31 @@ def _bottleneck(inputs: Tensor, bottleneck: str, size: int, normalization: str) tensor The output from the bottleneck """ - norms = dict(layer=LayerNormalization, - rms=RMSNormalization, - instance=InstanceNormalization) - bottlenecks = dict(average_pooling=GlobalAveragePooling2D(), - dense=Dense(size), - max_pooling=GlobalMaxPooling2D()) + norms = {"layer": kl.LayerNormalization, + "rms": RMSNormalization, + "instance": InstanceNormalization} + bottlenecks = {"average_pooling": kl.GlobalAveragePooling2D(), + "dense": kl.Dense(size), + "max_pooling": kl.GlobalMaxPooling2D()} var_x = inputs if normalization: var_x = norms[normalization]()(var_x) - if bottleneck == "dense" and len(K.int_shape(var_x)[1:]) > 1: - # Flatten non-1D inputs for dense bottleneck - var_x = Flatten()(var_x) - var_x = bottlenecks[bottleneck](var_x) - if len(K.int_shape(var_x)[1:]) > 1: + if bottleneck == "dense" and K.ndim(var_x) > 2: # Flatten non-1D inputs for dense + var_x = kl.Flatten()(var_x) + if bottleneck != "flatten": + var_x = bottlenecks[bottleneck](var_x) + if K.ndim(var_x) > 2: # Flatten prior to fc layers - var_x = Flatten()(var_x) + var_x = kl.Flatten()(var_x) return var_x -def _get_upscale_layer(method: Literal["resize_images", "subpixel", "upscale_dny", "upscale_fast", - "upscale_hybrid", "upsample2d"], +def _get_upscale_layer(method: T.Literal["resize_images", "subpixel", "upscale_dny", + "upscale_fast", "upscale_hybrid", "upsample2d"], filters: int, - activation: Optional[str] = None, - upsamples: Optional[int] = None, - interpolation: Optional[str] = None) -> keras.layers.Layer: + activation: str | None = None, + upsamples: int | None = None, + interpolation: str | None = None) -> tf.keras.layers.Layer: """ Obtain an instance of the requested upscale method. Parameters @@ -564,12 +552,12 @@ def _get_upscale_layer(method: Literal["resize_images", "subpixel", "upscale_dny The selected configured upscale layer """ if method == "upsample2d": - kwargs: Dict[str, Union[str, int]] = {} + kwargs: dict[str, str | int] = {} if upsamples: kwargs["size"] = upsamples if interpolation: kwargs["interpolation"] = interpolation - return UpSampling2D(**kwargs) + return kl.UpSampling2D(**kwargs) if method == "subpixel": return UpscaleBlock(filters, activation=activation) if method == "upscale_fast": @@ -585,7 +573,7 @@ def _get_curve(start_y: int, end_y: int, num_points: int, scale: float, - mode: Literal["full", "cap_max", "cap_min"] = "full") -> List[int]: + mode: T.Literal["full", "cap_max", "cap_min"] = "full") -> list[int]: """ Obtain a curve. For the given start and end y values, return the y co-ordinates of a curve for the given @@ -625,7 +613,7 @@ def _get_curve(start_y: int, else: y_axis = [start_y] scale = 1. - abs(scale) - for _ in range(num_points): + for _ in range(num_points - 1): current_value = max(end_y, int(((y_axis[-1] * scale) // 8) * 8)) y_axis.append(current_value) if current_value == end_y: @@ -674,24 +662,24 @@ class Encoder(): # pylint:disable=too-few-public-methods config: dict The model configuration options """ - def __init__(self, input_shape: Tuple[int, int, int], config: dict) -> None: + def __init__(self, input_shape: tuple[int, int, int], config: dict) -> None: self.input_shape = input_shape self._config = config self._input_shape = input_shape @property - def _model_kwargs(self) -> Dict[str, Dict[str, Union[str, bool]]]: + def _model_kwargs(self) -> dict[str, dict[str, str | bool]]: """ dict: Configuration option for architecture mapped to optional kwargs. """ - return dict(mobilenet=dict(alpha=self._config["mobilenet_width"], - depth_multiplier=self._config["mobilenet_depth"], - dropout=self._config["mobilenet_dropout"]), - mobilenet_v2=dict(alpha=self._config["mobilenet_width"]), - mobilenet_v3=dict(alpha=self._config["mobilenet_width"], - minimalist=self._config["mobilenet_minimalistic"], - include_preprocessing=False)) + return {"mobilenet": {"alpha": self._config["mobilenet_width"], + "depth_multiplier": self._config["mobilenet_depth"], + "dropout": self._config["mobilenet_dropout"]}, + "mobilenet_v2": {"alpha": self._config["mobilenet_width"]}, + "mobilenet_v3": {"alpha": self._config["mobilenet_width"], + "minimalist": self._config["mobilenet_minimalistic"], + "include_preprocessing": False}} @property - def _selected_model(self) -> Tuple[_EncoderInfo, dict]: + def _selected_model(self) -> tuple[_EncoderInfo, dict]: """ tuple(dict, :class:`_EncoderInfo`): The selected encoder model and it's associated keyword arguments """ arch = self._config["enc_architecture"] @@ -701,7 +689,7 @@ def _selected_model(self) -> Tuple[_EncoderInfo, dict]: kwargs["include_preprocessing"] = False return model, kwargs - def __call__(self) -> keras.models.Model: + def __call__(self) -> tf.keras.models.Model: """ Create the Phaze-A Encoder Model. Returns @@ -709,10 +697,11 @@ def __call__(self) -> keras.models.Model: :class:`keras.models.Model` The selected Encoder Model """ - input_ = Input(shape=self._input_shape) + input_ = kl.Input(shape=self._input_shape) var_x = input_ scaling = self._selected_model[0].scaling + if scaling: # Some models expect different scaling. logger.debug("Scaling to %s for '%s'", scaling, self._config["enc_architecture"]) @@ -724,6 +713,19 @@ def __call__(self) -> keras.models.Model: var_x = var_x * 2. var_x = var_x - 1.0 + if (self._config["enc_architecture"].startswith("efficientnet_b") + and self._config["mixed_precision"]): + # There is a bug in EfficientNet pre-processing where the normalized mean for the + # imagenet rgb values are not cast to float16 when mixed precision is enabled. + # We monkeypatch in a cast constant until the issue is resolved + # TODO revert if/when applying Imagenet Normalization works with mixed precision + # confirmed bugged: TF2.10 + logger.debug("Patching efficientnet.IMAGENET_STDDEV_RGB to float16 constant") + from keras.applications import efficientnet # pylint:disable=import-outside-toplevel + setattr(efficientnet, + "IMAGENET_STDDEV_RGB", + K.constant(efficientnet.IMAGENET_STDDEV_RGB, dtype="float16")) + var_x = self._get_encoder_model()(var_x) if self._config["bottleneck_in_encoder"]: @@ -732,9 +734,9 @@ def __call__(self) -> keras.models.Model: self._config["bottleneck_size"], self._config["bottleneck_norm"]) - return KerasModel(input_, var_x, name="encoder") + return keras.models.Model(input_, var_x, name="encoder") - def _get_encoder_model(self) -> keras.models.Model: + def _get_encoder_model(self) -> tf.keras.models.Model: """ Return the model defined by the selected architecture. Returns @@ -743,7 +745,14 @@ def _get_encoder_model(self) -> keras.models.Model: The selected keras model for the chosen encoder architecture """ model, kwargs = self._selected_model - if model.keras_name: + if model.keras_name and self._config["enc_architecture"].startswith("clipv_"): + assert model.keras_name in T.get_args(TypeModelsViT) + kwargs["input_shape"] = self._input_shape + kwargs["load_weights"] = self._config["enc_load_weights"] + retval = ViT(T.cast(TypeModelsViT, model.keras_name), + input_size=self._input_shape[0], + load_weights=self._config["enc_load_weights"])() + elif model.keras_name: kwargs["input_shape"] = self._input_shape kwargs["include_top"] = False kwargs["weights"] = "imagenet" if self._config["enc_load_weights"] else None @@ -772,7 +781,7 @@ def __init__(self, config: dict) -> None: self._kernel_size = 3 if self._is_alt else 5 self._strides = 1 if self._is_alt else 2 - def __call__(self, inputs: Tensor) -> Tensor: + def __call__(self, inputs: tf.Tensor) -> tf.Tensor: """ Call the original Faceswap Encoder Parameters @@ -815,7 +824,7 @@ def __call__(self, inputs: Tensor) -> Tensor: strides=self._strides, relu_alpha=self._relu_alpha, name=f"{name}_convblk_{i}_1")(var_x) - var_x = MaxPool2D(2, name=f"{name}_pool_{i}")(var_x) + var_x = kl.MaxPool2D(2, name=f"{name}_pool_{i}")(var_x) return var_x @@ -832,7 +841,7 @@ class FullyConnected(): # pylint:disable=too-few-public-methods The user configuration dictionary """ def __init__(self, - side: Literal["a", "b", "both", "gblock", "shared"], + side: T.Literal["a", "b", "both", "gblock", "shared"], input_shape: tuple, config: dict) -> None: logger.debug("Initializing: %s (side: %s, input_shape: %s)", @@ -896,7 +905,7 @@ def _scale_filters(self, original_filters: int) -> int: logger.debug("original_filters: %s, scaled_filters: %s", original_filters, retval) return retval - def _do_upsampling(self, inputs: Tensor) -> Tensor: + def _do_upsampling(self, inputs: tf.Tensor) -> tf.Tensor: """ Perform the upsampling at the end of the fully connected layers. Parameters @@ -926,10 +935,10 @@ def _do_upsampling(self, inputs: Tensor) -> Tensor: activation="leakyrelu") var_x = upscaler(var_x) if upsampler == "upsample2d": - var_x = LeakyReLU(alpha=0.1)(var_x) + var_x = kl.LeakyReLU(alpha=0.1)(var_x) return var_x - def __call__(self) -> keras.models.Model: + def __call__(self) -> tf.keras.models.Model: """ Call the intermediate layer. Returns @@ -937,7 +946,7 @@ def __call__(self) -> keras.models.Model: :class:`keras.models.Model` The Fully connected model """ - input_ = Input(shape=self._input_shape) + input_ = kl.Input(shape=self._input_shape) var_x = input_ node_curve = _get_curve(self._min_nodes, @@ -953,25 +962,24 @@ def __call__(self) -> keras.models.Model: dropout = f"{self._prefix}_dropout" for idx, nodes in enumerate(node_curve): - var_x = Dropout(self._config[dropout], name=f"{dropout}_{idx + 1}")(var_x) - var_x = Dense(nodes)(var_x) + var_x = kl.Dropout(self._config[dropout], name=f"{dropout}_{idx + 1}")(var_x) + var_x = kl.Dense(nodes)(var_x) if self._side != "gblock": dim = self._config["fc_dimensions"] - var_x = Reshape((dim, dim, int(self._max_nodes / (dim ** 2))))(var_x) + var_x = kl.Reshape((dim, dim, int(self._max_nodes / (dim ** 2))))(var_x) var_x = self._do_upsampling(var_x) num_upscales = self._config["dec_upscales_in_fc"] if num_upscales: var_x = UpscaleBlocks(self._side, - K.int_shape(var_x)[1:], self._config, layer_indicies=(0, num_upscales))(var_x) - return KerasModel(input_, var_x, name=f"fc_{self._side}") + return keras.models.Model(input_, var_x, name=f"fc_{self._side}") -class UpscaleBlocks(): # pylint: disable=too-few-public-methods +class UpscaleBlocks(): # pylint:disable=too-few-public-methods """ Obtain a block of upscalers. This class exists outside of the :class:`Decoder` model, as it is possible to place some of @@ -985,8 +993,6 @@ class UpscaleBlocks(): # pylint: disable=too-few-public-methods ---------- side: ["a", "b", "both", "shared"] The side of the model that the Decoder belongs to. Used for naming - input_shape: tuple - The shape tuple for the input to the decoder. config: dict The user configuration dictionary layer_indices: tuple, optional @@ -995,23 +1001,21 @@ class UpscaleBlocks(): # pylint: disable=too-few-public-methods and the Decoder. ``None`` will generate the full Upscale chain. An end index of -1 will generate the layers from the starting index to the final upscale. Default: ``None`` """ - _filters: List[int] = [] + _filters: list[int] = [] def __init__(self, - side: Literal["a", "b", "both", "shared"], - input_shape: Tuple[int, int, int], + side: T.Literal["a", "b", "both", "shared"], config: dict, - layer_indicies: Optional[Tuple[int, int]] = None) -> None: - logger.debug("Initializing: %s (side: %s, input_shape: %s, layer_indicies: %s)", - self.__class__.__name__, side, input_shape, layer_indicies) + layer_indicies: tuple[int, int] | None = None) -> None: + logger.debug("Initializing: %s (side: %s, layer_indicies: %s)", + self.__class__.__name__, side, layer_indicies) self._side = side - self._input_shape = input_shape self._config = config self._is_dny = self._config["dec_upscale_method"].lower() == "upscale_dny" self._layer_indicies = layer_indicies logger.debug("Initialized: %s", self.__class__.__name__,) - def _reshape_for_output(self, inputs: Tensor) -> Tensor: + def _reshape_for_output(self, inputs: tf.Tensor) -> tf.Tensor: """ Reshape the input for arbitrary output sizes. The number of filters in the input will have been scaled to the model output size allowing @@ -1035,14 +1039,14 @@ def _reshape_for_output(self, inputs: Tensor) -> Tensor: new_shape = (new_dim, new_dim, np.prod(old_shape) // new_dim ** 2) logger.debug("Reshaping tensor from %s to %s for output size %s", K.int_shape(inputs)[1:], new_shape, self._config["output_size"]) - var_x = Reshape(new_shape)(var_x) + var_x = kl.Reshape(new_shape)(var_x) return var_x def _upscale_block(self, - inputs: Tensor, + inputs: tf.Tensor, filters: int, skip_residual: bool = False, - is_mask: bool = False) -> Tensor: + is_mask: bool = False) -> tf.Tensor: """ Upscale block for Phaze-A Decoder. Uses requested upscale method, adds requested regularization and activation function. @@ -1072,19 +1076,19 @@ def _upscale_block(self, var_x = upscaler(inputs) if not is_mask and self._config["dec_gaussian"]: - var_x = GaussianNoise(1.0)(var_x) + var_x = kl.GaussianNoise(1.0)(var_x) if not is_mask and self._config["dec_res_blocks"] and not skip_residual: var_x = self._normalization(var_x) - var_x = LeakyReLU(alpha=0.2)(var_x) + var_x = kl.LeakyReLU(alpha=0.2)(var_x) for _ in range(self._config["dec_res_blocks"]): var_x = ResidualBlock(filters)(var_x) else: var_x = self._normalization(var_x) if not self._is_dny: - var_x = LeakyReLU(alpha=0.1)(var_x) + var_x = kl.LeakyReLU(alpha=0.1)(var_x) return var_x - def _normalization(self, inputs: Tensor) -> Tensor: + def _normalization(self, inputs: tf.Tensor) -> tf.Tensor: """ Add a normalization layer if requested. Parameters @@ -1099,14 +1103,14 @@ def _normalization(self, inputs: Tensor) -> Tensor: """ if not self._config["dec_norm"]: return inputs - norms = dict(batch=BatchNormalization, - group=GroupNormalization, - instance=InstanceNormalization, - layer=LayerNormalization, - rms=RMSNormalization) + norms = {"batch": kl.BatchNormalization, + "group": GroupNormalization, + "instance": InstanceNormalization, + "layer": kl.LayerNormalization, + "rms": RMSNormalization} return norms[self._config["dec_norm"]]()(inputs) - def _dny_entry(self, inputs: Tensor) -> Tensor: + def _dny_entry(self, inputs: tf.Tensor) -> tf.Tensor: """ Entry convolutions for using the upscale_dny method. Parameters @@ -1131,29 +1135,40 @@ def _dny_entry(self, inputs: Tensor) -> Tensor: relu_alpha=0.2)(var_x) return var_x - def __call__(self, inputs: Optional[Tensor] = None) -> Tensor: - """ Decoder Network. + def __call__(self, inputs: tf.Tensor | list[tf.Tensor]) -> tf.Tensor | list[tf.Tensor]: + """ Upscale Network. Parameters - inputs: Tensor, optional - If the input is an output from another model (such as the Fully Connected Model) this - should be ``None`` otherwise it should be a Tensor + inputs: Tensor or list of tensors + Input tensor(s) to upscale block. This will be a single tensor if learn mask is not + selected or if this is the first call to the upscale blocks. If learn mask is selected + and this is not the first call to upscale blocks, then this will be a list of the face + and mask tensors. Returns ------- - :class:`keras.models.Model` - The Decoder model + Tensor or list of tensors + The output of encoder blocks. Either a single tensor (if learn mask is not enabled) or + list of tensors (if learn mask is enabled) """ - inputs = Input(shape=self._input_shape) if inputs is None else inputs - var_x = inputs start_idx, end_idx = (0, None) if self._layer_indicies is None else self._layer_indicies end_idx = None if end_idx == -1 else end_idx + if self._config["learn_mask"] and start_idx == 0: + # Mask needs to be created + var_x = inputs + var_y = inputs + elif self._config["learn_mask"]: + # Mask has already been created and is an input to upscale blocks + var_x, var_y = inputs + else: + # No mask required + var_x = inputs + if start_idx == 0: var_x = self._reshape_for_output(var_x) if self._config["learn_mask"]: - var_y = inputs var_y = self._reshape_for_output(var_y) if self._is_dny: @@ -1197,20 +1212,24 @@ class GBlock(): # pylint:disable=too-few-public-methods The user configuration dictionary """ def __init__(self, - side: Literal["a", "b", "both"], - input_shapes: Union[list, tuple], + side: T.Literal["a", "b", "both"], + input_shapes: list | tuple, config: dict) -> None: logger.debug("Initializing: %s (side: %s, input_shapes: %s)", self.__class__.__name__, side, input_shapes) self._side = side self._config = config - self._inputs = [Input(shape=shape) for shape in input_shapes] + self._inputs = [kl.Input(shape=shape) for shape in input_shapes] self._dense_nodes = 512 self._dense_recursions = 3 logger.debug("Initialized: %s", self.__class__.__name__) @classmethod - def _g_block(cls, inputs: Tensor, style: Tensor, filters: int, recursions: int = 2) -> Tensor: + def _g_block(cls, + inputs: tf.Tensor, + style: tf.Tensor, + filters: int, + recursions: int = 2) -> tf.Tensor: """ G_block adapted from ADAIN StyleGAN. Parameters @@ -1231,19 +1250,19 @@ def _g_block(cls, inputs: Tensor, style: Tensor, filters: int, recursions: int = """ var_x = inputs for i in range(recursions): - styles = [Reshape([1, 1, filters])(Dense(filters)(style)) for _ in range(2)] - noise = KConv2D(filters, 1, padding="same")(GaussianNoise(1.0)(var_x)) + styles = [kl.Reshape([1, 1, filters])(kl.Dense(filters)(style)) for _ in range(2)] + noise = kl.Conv2D(filters, 1, padding="same")(kl.GaussianNoise(1.0)(var_x)) if i == recursions - 1: - var_x = KConv2D(filters, 3, padding="same")(var_x) + var_x = kl.Conv2D(filters, 3, padding="same")(var_x) var_x = AdaInstanceNormalization(dtype="float32")([var_x, *styles]) - var_x = Add()([var_x, noise]) - var_x = LeakyReLU(0.2)(var_x) + var_x = kl.Add()([var_x, noise]) + var_x = kl.LeakyReLU(0.2)(var_x) return var_x - def __call__(self) -> keras.models.Model: + def __call__(self) -> tf.keras.models.Model: """ G-Block Network. Returns @@ -1253,16 +1272,16 @@ def __call__(self) -> keras.models.Model: """ var_x, style = self._inputs for i in range(self._dense_recursions): - style = Dense(self._dense_nodes, kernel_initializer="he_normal")(style) + style = kl.Dense(self._dense_nodes, kernel_initializer="he_normal")(style) if i != self._dense_recursions - 1: # Don't add leakyReLu to final output - style = LeakyReLU(0.1)(style) + style = kl.LeakyReLU(0.1)(style) # Scale g_block filters to side dense g_filts = K.int_shape(var_x)[-1] var_x = Conv2D(g_filts, 3, strides=1, padding="same")(var_x) - var_x = GaussianNoise(1.0)(var_x) + var_x = kl.GaussianNoise(1.0)(var_x) var_x = self._g_block(var_x, style, g_filts) - return KerasModel(self._inputs, var_x, name=f"g_block_{self._side}") + return keras.models.Model(self._inputs, var_x, name=f"g_block_{self._side}") class Decoder(): # pylint:disable=too-few-public-methods @@ -1278,8 +1297,8 @@ class Decoder(): # pylint:disable=too-few-public-methods The user configuration dictionary """ def __init__(self, - side: Literal["a", "b", "both"], - input_shape: Tuple[int, int, int], + side: T.Literal["a", "b", "both"], + input_shape: tuple[int, int, int], config: dict) -> None: logger.debug("Initializing: %s (side: %s, input_shape: %s)", self.__class__.__name__, side, input_shape) @@ -1288,7 +1307,7 @@ def __init__(self, self._config = config logger.debug("Initialized: %s", self.__class__.__name__,) - def __call__(self) -> keras.models.Model: + def __call__(self) -> tf.keras.models.Model: """ Decoder Network. Returns @@ -1296,15 +1315,18 @@ def __call__(self) -> keras.models.Model: :class:`keras.models.Model` The Decoder model """ - inputs = Input(shape=self._input_shape) - var_x = inputs + inputs = kl.Input(shape=self._input_shape) + num_ups_in_fc = self._config["dec_upscales_in_fc"] - indicies = None if not num_ups_in_fc else (num_ups_in_fc, -1) + if self._config["learn_mask"] and num_ups_in_fc: + # Mask has already been created in FC and is an output of that model + inputs = [inputs, kl.Input(shape=self._input_shape)] + + indicies = None if not num_ups_in_fc else (num_ups_in_fc, -1) upscales = UpscaleBlocks(self._side, - self._input_shape, self._config, - layer_indicies=indicies)(var_x) + layer_indicies=indicies)(inputs) if self._config["learn_mask"]: var_x, var_y = upscales @@ -1317,4 +1339,4 @@ def __call__(self) -> keras.models.Model: self._config["dec_output_kernel"], name="mask_out")(var_y)) - return KerasModel(inputs, outputs=outputs, name=f"decoder_{self._side}") + return keras.models.Model(inputs, outputs=outputs, name=f"decoder_{self._side}") diff --git a/plugins/train/model/phaze_a_defaults.py b/plugins/train/model/phaze_a_defaults.py index 03b1ae7c40..8cff3d6458 100644 --- a/plugins/train/model/phaze_a_defaults.py +++ b/plugins/train/model/phaze_a_defaults.py @@ -7,41 +7,38 @@ within the faceswap/config folder. The following variables should be defined: - _HELPTEXT: A string describing what this plugin does - _DEFAULTS: A dictionary containing the options, defaults and meta information. The - dictionary should be defined as: - {: {}} + "_HELPTEXT: A string describing what this plugin does + "_DEFAULTS: A dictionary containing the options, defaults and meta information. The + " dictionary should be defined as: + " {: {}} - should always be lower text. - dictionary requirements are listed below. + " should always be lower text. + " dictionary requirements are listed below. The following keys are expected for the _DEFAULTS dict: - datatype: [required] A python type class. This limits the type of data that can be - provided in the .ini file and ensures that the value is returned in the - correct type to faceswap. Valid data types are: , , - , . - default: [required] The default value for this option. - info: [required] A string describing what this option does. - choices: [optional] If this option's datatype is of then valid - selections can be defined here. This validates the option and also enables - a combobox / radio option in the GUI. - gui_radio: [optional] If are defined, this indicates that the GUI should use - radio buttons rather than a combobox to display this option. - min_max: [partial] For and data types this is required - otherwise it is ignored. Should be a tuple of min and max accepted values. - This is used for controlling the GUI slider range. Values are not enforced. - rounding: [partial] For and data types this is - required otherwise it is ignored. Used for the GUI slider. For floats, this - is the number of decimal places to display. For ints this is the step size. - fixed: [optional] [train only]. Training configurations are fixed when the model is - created, and then reloaded from the state file. Marking an item as fixed=False - indicates that this value can be changed for existing models, and will override - the value saved in the state file with the updated value in config. If not - provided this will default to True. + "datatype: [required] A python type class. This limits the type of data that can be + " provided in the .ini file and ensures that the value is returned in the + " correct type to faceswap. Valid data types are: , , + " , . + "default: [required] The default value for this option. + "info: [required] A string describing what this option does. + "choices: [optional] If this option's datatype is of then valid + " selections can be defined here. This validates the option and also enables + " a combobox / radio option in the GUI. + "gui_radio: [optional] If are defined, this indicates that the GUI should use + " radio buttons rather than a combobox to display this option. + "min_max: [partial] For and data types this is required + " otherwise it is ignored. Should be a tuple of min and max accepted values. + " This is used for controlling the GUI slider range. Values are not enforced. + "rounding: [partial] For and data types this is + " required otherwise it is ignored. Used for the GUI slider. For floats, this + " is the number of decimal places to display. For ints this is the step size. + "fixed: [optional] [train only]. Training configurations are fixed when the model is + " created, and then reloaded from the state file. Marking an item as fixed=False + " indicates that this value can be changed for existing models, and will override + " the value saved in the state file with the updated value in config. If not + " provided this will default to True. """ -from typing import List - -from lib.utils import get_backend _HELPTEXT: str = ( "Phaze-A Model by TorzDF, with thanks to BirbFakes.\n" @@ -49,634 +46,683 @@ "inspiration from Nvidia's StyleGAN for the Decoder. It is highly recommended to research to " "understand the parameters better.") -_ENCODERS: List[str] = [ - "densenet121", "densenet169", "densenet201", "inception_resnet_v2", "inception_v3", - "mobilenet", "mobilenet_v2", "nasnet_large", "nasnet_mobile", "resnet50", "vgg16", "vgg19", - "xception", "fs_original"] - -if get_backend() != "amd": - _ENCODERS.extend(["efficientnet_b0", "efficientnet_b1", "efficientnet_b2", "efficientnet_b3", - "efficientnet_b4", "efficientnet_b5", "efficientnet_b6", "efficientnet_b7", - "efficientnet_v2_b0", "efficientnet_v2_b1", "efficientnet_v2_b2", - "efficientnet_v2_b3", "efficientnet_v2_l", "efficientnet_v2_m", - "efficientnet_v2_s", "mobilenet_v3_large", "mobilenet_v3_small", - "resnet50_v2", "resnet101", "resnet101_v2", "resnet152", "resnet152_v2"]) -_ENCODERS = sorted(_ENCODERS) - +_ENCODERS: list[str] = sorted([ + "clipv_vit-b-16", "clipv_vit-b-32", "clipv_vit-l-14", "clipv_vit-l-14-336px", + "clipv_farl-b-16-16", "clipv_farl-b-16-64", + "densenet121", "densenet169", "densenet201", "efficientnet_b0", "efficientnet_b1", + "efficientnet_b2", "efficientnet_b3", "efficientnet_b4", "efficientnet_b5", "efficientnet_b6", + "efficientnet_b7", "efficientnet_v2_b0", "efficientnet_v2_b1", "efficientnet_v2_b2", + "efficientnet_v2_b3", "efficientnet_v2_l", "efficientnet_v2_m", "efficientnet_v2_s", + "inception_resnet_v2", "inception_v3", "mobilenet", "mobilenet_v2", "mobilenet_v3_large", + "mobilenet_v3_small", "nasnet_large", "nasnet_mobile", "resnet50", "resnet50_v2", "resnet101", + "resnet101_v2", "resnet152", "resnet152_v2", "vgg16", "vgg19", "xception", "fs_original"]) -_DEFAULTS = dict( +_DEFAULTS = { # General - output_size=dict( - default=128, - info="Resolution (in pixels) of the output image to generate.\n" - "BE AWARE Larger resolution will dramatically increase VRAM requirements.", - datatype=int, - rounding=64, - min_max=(64, 2048), - group="general", - fixed=True), - shared_fc=dict( - default="none", - info="Whether to create a shared fully connected layer. This layer will have the same " - "structure as the fully connected layers used for each side of the model. A shared " - "fully connected layer looks for patterns that are common to both sides. NB: " - "Enabling this option only makes sense if 'split fc' is selected." - "\n\tnone - Do not create a Fully Connected layer for shared data. (Original method)" - "\n\tfull - Create an exclusive Fully Connected layer for shared data. (IAE method)" - "\n\thalf - Use the 'fc_a' layer for shared data. This saves VRAM by re-using the " - "'A' side's fully connected model for the shared data. However, this will lead to " - "an 'unbalanced' model and can lead to more identity bleed (DFL method)", - datatype=str, - choices=["none", "full", "half"], - gui_radio=True, - group="general", - fixed=True), - enable_gblock=dict( - default=True, - info="Whether to enable the G-Block. If enabled, this will create a shared fully " - "connected layer (configurable in the 'G-Block hidden layers' section) to look for " - "patterns in the combined data, before feeding a block prior to the decoder for " - "merging this shared and combined data." - "\n\tTrue - Use the G-Block in the Decoder. A combined fully connected layer will be " - "created to feed this block which can be configured below." - "\n\tFalse - Don't use the G-Block in the decoder. No combined fully connected layer " - "will be created.", - datatype=bool, - group="general", - fixed=True), - split_fc=dict( - default=True, - info="Whether to use a single shared Fully Connected layer or separate Fully Connected " - "layers for each side." - "\n\tTrue - Use separate Fully Connected layers for Face A and Face B. This is more " - "similar to the 'IAE' style of model." - "\n\tFalse - Use combined Fully Connected layers for both sides. This is more " - "similar to the original Faceswap architecture.", - datatype=bool, - group="general", - fixed=True), - split_gblock=dict( - default=False, - info="If the G-Block is enabled, Whether to use a single G-Block shared between both " - "sides, or whether to have a separate G-Block (one for each side). NB: The Fully " - "Connected layer that feeds the G-Block will always be shared." - "\n\tTrue - Use separate G-Blocks for Face A and Face B." - "\n\tFalse - Use a combined G-Block layers for both sides.", - datatype=bool, - group="general", - fixed=True), - split_decoders=dict( - default=False, - info="Whether to use a single decoder or split decoders." - "\n\tTrue - Use a separate decoder for Face A and Face B. This is more similar to " - "the original Faceswap architecture." - "\n\tFalse - Use a combined Decoder. This is more similar to 'IAE' style " - "architecture.", - datatype=bool, - group="general", - fixed=True), + "output_size": { + "default": 128, + "info": ( + "Resolution (in pixels) of the output image to generate.\n" + "BE AWARE Larger resolution will dramatically increase VRAM requirements."), + "datatype": int, + "rounding": 16, + "min_max": (64, 2048), + "group": "general", + "fixed": True}, + "shared_fc": { + "default": "none", + "info": ( + "Whether to create a shared fully connected layer. This layer will have the same " + "structure as the fully connected layers used for each side of the model. A shared " + "fully connected layer looks for patterns that are common to both sides. NB: " + "Enabling this option only makes sense if 'split fc' is selected." + "\n\tnone - Do not create a Fully Connected layer for shared data. (Original method)" + "\n\tfull - Create an exclusive Fully Connected layer for shared data. (IAE method)" + "\n\thalf - Use the 'fc_a' layer for shared data. This saves VRAM by re-using the " + "'A' side's fully connected model for the shared data. However, this will lead to " + "an 'unbalanced' model and can lead to more identity bleed (DFL method)"), + "datatype": str, + "choices": ["none", "full", "half"], + "gui_radio": True, + "group": "general", + "fixed": True}, + "enable_gblock": { + "default": True, + "info": ( + "Whether to enable the G-Block. If enabled, this will create a shared fully " + "connected layer (configurable in the 'G-Block hidden layers' section) to look for " + "patterns in the combined data, before feeding a block prior to the decoder for " + "merging this shared and combined data." + "\n\tTrue - Use the G-Block in the Decoder. A combined fully connected layer will be " + "created to feed this block which can be configured below." + "\n\tFalse - Don't use the G-Block in the decoder. No combined fully connected layer " + "will be created."), + "datatype": bool, + "group": "general", + "fixed": True}, + "split_fc": { + "default": True, + "info": ( + "Whether to use a single shared Fully Connected layer or separate Fully Connected " + "layers for each side." + "\n\tTrue - Use separate Fully Connected layers for Face A and Face B. This is more " + "similar to the 'IAE' style of model." + "\n\tFalse - Use combined Fully Connected layers for both sides. This is more " + "similar to the original Faceswap architecture."), + "datatype": bool, + "group": "general", + "fixed": True}, + "split_gblock": { + "default": False, + "info": ( + "If the G-Block is enabled, Whether to use a single G-Block shared between both " + "sides, or whether to have a separate G-Block (one for each side). NB: The Fully " + "Connected layer that feeds the G-Block will always be shared." + "\n\tTrue - Use separate G-Blocks for Face A and Face B." + "\n\tFalse - Use a combined G-Block layers for both sides."), + "datatype": bool, + "group": "general", + "fixed": True}, + "split_decoders": { + "default": False, + "info": ( + "Whether to use a single decoder or split decoders." + "\n\tTrue - Use a separate decoder for Face A and Face B. This is more similar to " + "the original Faceswap architecture." + "\n\tFalse - Use a combined Decoder. This is more similar to 'IAE' style " + "architecture."), + "datatype": bool, + "group": "general", + "fixed": True}, # Encoder - enc_architecture=dict( - default="fs_original", - info="The encoder architecture to use. See the relevant config sections for specific " - "architecture tweaking.\nNB: For keras based pre-built models, the global " - "initializers and padding options will be ignored for the selected encoder." - "\n\tdensenet: (32px -224px). Ref: Densely Connected Convolutional Networks (2016): " - "https://arxiv.org/abs/1608.06993?source=post_page" - "\n\tefficientnet: [Tensorflow 2.3+ only] EfficientNet has numerous variants (B0 - " - "B8) that increases the model width, depth and dimensional space at each step. The " - "minimum input resolution is 32px for all variants. The maximum input resolution for " - "each variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, b4: 380px, b5: 456px, " - "b6: 528px, b7 600px. Ref: Rethinking Model Scaling for Convolutional Neural " - "Networks (2020): https://arxiv.org/abs/1905.11946" - "\n\tefficientnet_v2: [Tensorflow 2.8+ only] EfficientNetV2 is the follow up to " - "efficientnet. It has numerous variants (B0 - B3 and Small, Medium and Large) that " - "increases the model width, depth and dimensional space at each step. The minimum " - "input resolution is 32px for all variants. The maximum input resolution for each " - "variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, s: 384px, m: 480px, l: " - "480px. Ref: EfficientNetV2: Smaller Models and Faster Training (2021): " - "https://arxiv.org/abs/2104.00298" - "\n\tfs_original: (32px - 1024px). A configurable variant of the original facewap " - "encoder. ImageNet weights cannot be loaded for this model. Additional parameters " - "can be configured with the 'fs_enc' options. A version of this encoder is used in " - "the following models: Original, Original (lowmem), Dfaker, DFL-H128, DFL-SAE, IAE, " - "Lightweight." - "\n\tinception_resnet_v2: (75px - 299px). Ref: Inception-ResNet and the Impact of " - "Residual Connections on Learning (2016): https://arxiv.org/abs/1602.07261" - "\n\tinceptionV3: (75px - 299px). Ref: Rethinking the Inception Architecture for " - "Computer Vision (2015): https://arxiv.org/abs/1512.00567" - "\n\tmobilenet: (32px - 224px). Additional MobileNet parameters can be set with the " - "'mobilenet' options. Ref: MobileNets: Efficient Convolutional Neural Networks for " - "Mobile Vision Applications (2017): https://arxiv.org/abs/1704.04861" - "\n\tmobilenet_v2: (32px - 224px). Additional MobileNet parameters can be set with " - "the 'mobilenet' options. Ref: MobileNetV2: Inverted Residuals and Linear " - "Bottlenecks (2018): https://arxiv.org/abs/1801.04381" - "\n\tmobilenet_v3: (32px - 224px). Additional MobileNet parameters can be set with " - "the 'mobilenet' options. Ref: Searching for MobileNetV3 (2019): " - "https://arxiv.org/pdf/1905.02244.pdf" - "\n\tnasnet: (32px - 331px (large) or 224px (mobile)). Ref: Learning Transferable " - "Architectures for Scalable Image Recognition (2017): " - "https://arxiv.org/abs/1707.07012" - "\n\tresnet: (32px - 224px). Deep Residual Learning for Image Recognition (2015): " - "https://arxiv.org/abs/1512.03385" - "\n\tvgg: (32px - 224px). Very Deep Convolutional Networks for Large-Scale Image " - "Recognition (2014): https://arxiv.org/abs/1409.1556" - "\n\txception: (71px - 229px). Ref: Deep Learning with Depthwise Separable " - "Convolutions (2017): https://arxiv.org/abs/1409.1556.\n", - datatype=str, - choices=_ENCODERS, - gui_radio=False, - group="encoder", - fixed=True), - enc_scaling=dict( - default=7, - info="Input scaling for the encoder. Some of the encoders have large input sizes, which " - "often are not helpful for Faceswap. This setting scales the dimensional space that " - "the encoder works in. For example an encoder with a maximum input size of 224px " - "will be input an image of 112px at 50%% scaling. See the Architecture tooltip for " - "the minimum and maximum sizes for each encoder. NB: The input size will be rounded " - "down to the nearest 16 pixels.", - datatype=int, - min_max=(0, 100), - rounding=1, - group="encoder", - fixed=True), - enc_load_weights=dict( - default=True, - info="Load pre-trained weights trained on ImageNet data. Only available for non-Faceswap " - "encoders (i.e. those not beginning with 'fs'). NB: If you use the global 'load " - "weights' option and have selected to load weights from a previous model's 'encoder' " - "or 'keras_encoder' then the weights loaded here will be replaced by the weights " - "loaded from your saved model.", - datatype=bool, - group="encoder", - fixed=True), + "enc_architecture": { + "default": "fs_original", + "info": ( + "The encoder architecture to use. See the relevant config sections for specific " + "architecture tweaking.\nNB: For keras based pre-built models, the global " + "initializers and padding options will be ignored for the selected encoder." + "\n\n\tCLIPv: This is an implementation of the Visual encoder from the CLIP " + "transformer. The ViT weights are trained on imagenet whilst the FaRL weights are " + "trained on face related tasks. All have a default input size of 224px except for " + "ViT-L-14-336px that has an input size of 336px. Ref: Learning Transferable Visual " + "Models From Natural Language Supervision (2021): https://arxiv.org/abs/2103.00020" + "\n\n\tdensenet: (32px -224px). Ref: Densely Connected Convolutional Networks " + "(2016): https://arxiv.org/abs/1608.06993?source=post_page" + "\n\n\tefficientnet: [Tensorflow 2.3+ only] EfficientNet has numerous variants (B0 - " + "B8) that increases the model width, depth and dimensional space at each step. The " + "minimum input resolution is 32px for all variants. The maximum input resolution for " + "each variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, b4: 380px, b5: 456px, " + "b6: 528px, b7 600px. Ref: Rethinking Model Scaling for Convolutional Neural " + "Networks (2020): https://arxiv.org/abs/1905.11946" + "\n\n\tefficientnet_v2: [Tensorflow 2.8+ only] EfficientNetV2 is the follow up to " + "efficientnet. It has numerous variants (B0 - B3 and Small, Medium and Large) that " + "increases the model width, depth and dimensional space at each step. The minimum " + "input resolution is 32px for all variants. The maximum input resolution for each " + "variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, s: 384px, m: 480px, l: " + "480px. Ref: EfficientNetV2: Smaller Models and Faster Training (2021): " + "https://arxiv.org/abs/2104.00298" + "\n\n\tfs_original: (32px - 1024px). A configurable variant of the original facewap " + "encoder. ImageNet weights cannot be loaded for this model. Additional parameters " + "can be configured with the 'fs_enc' options. A version of this encoder is used in " + "the following models: Original, Original (lowmem), Dfaker, DFL-H128, DFL-SAE, IAE, " + "Lightweight." + "\n\n\tinception_resnet_v2: (75px - 299px). Ref: Inception-ResNet and the Impact of " + "Residual Connections on Learning (2016): https://arxiv.org/abs/1602.07261" + "\n\n\tinceptionV3: (75px - 299px). Ref: Rethinking the Inception Architecture for " + "Computer Vision (2015): https://arxiv.org/abs/1512.00567" + "\n\n\tmobilenet: (32px - 224px). Additional MobileNet parameters can be set with " + "the 'mobilenet' options. Ref: MobileNets: Efficient Convolutional Neural Networks " + "for Mobile Vision Applications (2017): https://arxiv.org/abs/1704.04861" + "\n\n\tmobilenet_v2: (32px - 224px). Additional MobileNet parameters can be set with " + "the 'mobilenet' options. Ref: MobileNetV2: Inverted Residuals and Linear " + "Bottlenecks (2018): https://arxiv.org/abs/1801.04381" + "\n\n\tmobilenet_v3: (32px - 224px). Additional MobileNet parameters can be set with " + "the 'mobilenet' options. Ref: Searching for MobileNetV3 (2019): " + "https://arxiv.org/pdf/1905.02244.pdf" + "\n\n\tnasnet: (32px - 331px (large) or 224px (mobile)). Ref: Learning Transferable " + "Architectures for Scalable Image Recognition (2017): " + "https://arxiv.org/abs/1707.07012" + "\n\n\tresnet: (32px - 224px). Deep Residual Learning for Image Recognition (2015): " + "https://arxiv.org/abs/1512.03385" + "\n\n\tvgg: (32px - 224px). Very Deep Convolutional Networks for Large-Scale Image " + "Recognition (2014): https://arxiv.org/abs/1409.1556" + "\n\n\txception: (71px - 229px). Ref: Deep Learning with Depthwise Separable " + "Convolutions (2017): https://arxiv.org/abs/1409.1556.\n"), + "datatype": str, + "choices": _ENCODERS, + "gui_radio": False, + "group": "encoder", + "fixed": True}, + "enc_scaling": { + "default": 7, + "info": ( + "Input scaling for the encoder. Some of the encoders have large input sizes, which " + "often are not helpful for Faceswap. This setting scales the dimensional space that " + "the encoder works in. For example an encoder with a maximum input size of 224px " + "will be input an image of 112px at 50%% scaling. See the Architecture tooltip for " + "the minimum and maximum sizes for each encoder. NB: The input size will be rounded " + "down to the nearest 16 pixels."), + "datatype": int, + "min_max": (0, 200), + "rounding": 1, + "group": "encoder", + "fixed": True}, + "enc_load_weights": { + "default": True, + "info": ( + "Load pre-trained weights trained on ImageNet data. Only available for non-" + "Faceswap encoders (i.e. those not beginning with 'fs'). NB: If you use the global " + "'load weights' option and have selected to load weights from a previous model's " + "'encoder' or 'keras_encoder' then the weights loaded here will be replaced by the " + "weights loaded from your saved model."), + "datatype": bool, + "group": "encoder", + "fixed": True}, # Bottleneck - bottleneck_type=dict( - default="dense", - info="The type of layer to use for the bottleneck." - "\n\taverage_pooling: Use a Global Average Pooling 2D layer for the bottleneck." - "\n\tdense: Use a Dense layer for the bottleneck (the traditional Faceswap method). " - "You can set the size of the Dense layer with the 'bottleneck_size' parameter." - "\n\tmax_pooling: Use a Global Max Pooling 2D layer for the bottleneck.", - datatype=str, - group="bottleneck", - gui_radio=True, - choices=["average_pooling", "dense", "max_pooling"], - fixed=True), - bottleneck_norm=dict( - default="none", - info="Apply a normalization layer after encoder output and prior to the bottleneck." - "\n\tnone - Do not apply a normalization layer" - "\n\tinstance - Apply Instance Normalization" - "\n\tlayer - Apply Layer Normalization (Ba et al., 2016)" - "\n\trms - Apply Root Mean Squared Layer Normalization (Zhang et al., 2019). A " - "simplified version of Layer Normalization with reduced overhead.", - datatype=str, - gui_radio=True, - choices=["none", "instance", "layer", "rms"], - group="bottleneck", - fixed=True), - bottleneck_size=dict( - default=1024, - info="If using a Dense layer for the bottleneck, then this is the number of nodes to use.", - datatype=int, - rounding=128, - min_max=(128, 4096), - group="bottleneck", - fixed=True), - bottleneck_in_encoder=dict( - default=True, - info="Whether to place the bottleneck in the Encoder or to place it with the other hidden " - "layers. Placing the bottleneck in the encoder means that both sides will share the " - "same bottleneck. Placing it with the other fully connected layers means that each " - "fully connected layer will each get their own bottleneck. This may be combined or " - "split depending on your overall architecture configuration settings.", - datatype=bool, - group="bottleneck", - fixed=True), + "bottleneck_type": { + "default": "dense", + "info": ( + "The type of layer to use for the bottleneck." + "\n\taverage_pooling: Use a Global Average Pooling 2D layer for the bottleneck." + "\n\tdense: Use a Dense layer for the bottleneck (the traditional Faceswap method). " + "You can set the size of the Dense layer with the 'bottleneck_size' parameter." + "\n\tmax_pooling: Use a Global Max Pooling 2D layer for the bottleneck." + "\n\flatten: Don't use a bottleneck at all. Some encoders output in a size that make " + "a bottleneck unnecessary. This option flattens the output from the encoder, with no " + "further operations"), + "datatype": str, + "group": "bottleneck", + "gui_radio": True, + "choices": ["average_pooling", "dense", "max_pooling", "flatten"], + "fixed": True}, + "bottleneck_norm": { + "default": "none", + "info": ( + "Apply a normalization layer after encoder output and prior to the bottleneck." + "\n\tnone - Do not apply a normalization layer" + "\n\tinstance - Apply Instance Normalization" + "\n\tlayer - Apply Layer Normalization (Ba et al., 2016)" + "\n\trms - Apply Root Mean Squared Layer Normalization (Zhang et al., 2019). A " + "simplified version of Layer Normalization with reduced overhead."), + "datatype": str, + "gui_radio": True, + "choices": ["none", "instance", "layer", "rms"], + "group": "bottleneck", + "fixed": True}, + "bottleneck_size": { + "default": 1024, + "info": ( + "If using a Dense layer for the bottleneck, then this is the number of nodes to " + "use."), + "datatype": int, + "rounding": 128, + "min_max": (128, 4096), + "group": "bottleneck", + "fixed": True}, + "bottleneck_in_encoder": { + "default": True, + "info": ( + "Whether to place the bottleneck in the Encoder or to place it with the other " + "hidden layers. Placing the bottleneck in the encoder means that both sides will " + "share the same bottleneck. Placing it with the other fully connected layers means " + "that each fully connected layer will each get their own bottleneck. This may be " + "combined or split depending on your overall architecture configuration settings."), + "datatype": bool, + "group": "bottleneck", + "fixed": True}, # Intermediate Layers - fc_depth=dict( - default=1, - info="The number of consecutive Dense (fully connected) layers to include in each side's " - "intermediate layer.", - datatype=int, - rounding=1, - min_max=(0, 16), - group="hidden layers", - fixed=True), - fc_min_filters=dict( - default=1024, - info="The number of filters to use for the initial fully connected layer. The number of " - "nodes actually used is: fc_min_filters x fc_dimensions x fc_dimensions.\nNB: This " - "value may be scaled down, depending on output resolution.", - datatype=int, - rounding=16, - min_max=(16, 5120), - group="hidden layers", - fixed=True), - fc_max_filters=dict( - default=1024, - info="This is the number of filters to be used in the final reshape layer at the end of " - "the fully connected layers. The actual number of nodes used for the final fully " - "connected layer is: fc_min_filters x fc_dimensions x fc_dimensions.\nNB: This value " - "may be scaled down, depending on output resolution.", - datatype=int, - rounding=64, - min_max=(128, 5120), - group="hidden layers", - fixed=True), - fc_dimensions=dict( - default=4, - info="The height and width dimension for the final reshape layer at the end of the fully " - "connected layers.\nNB: The total number of nodes within the final fully connected " - "layer will be: fc_dimensions x fc_dimensions x fc_max_filters.", - datatype=int, - rounding=1, - min_max=(1, 16), - group="hidden layers", - fixed=True), - fc_filter_slope=dict( - default=-0.5, - info="The rate that the filters move from the minimum number of filters to the maximum " - "number of filters. EG:\n" - "Negative numbers will change the number of filters quicker at first and slow down " - "each layer.\n" - "Positive numbers will change the number of filters slower at first but then speed " - "up each layer.\n" - "0.0 - This will change at a linear rate (i.e. the same number of filters will be " - "changed at each layer).", - datatype=float, - min_max=(-.99, .99), - rounding=2, - group="hidden layers", - fixed=True), - fc_dropout=dict( - default=0.0, - info="Dropout is a form of regularization that can prevent a model from over-fitting and " - "help to keep neurons 'alive'. 0.5 will dropout half the connections between each " - "fully connected layer, 0.25 will dropout a quarter of the connections etc. Set to " - "0.0 to disable.", - datatype=float, - rounding=2, - min_max=(0.0, 0.99), - group="hidden layers", - fixed=False), - fc_upsampler=dict( - default="upsample2d", - info="The type of dimensional upsampling to perform at the end of the fully connected " - "layers, if upsamples > 0. The number of filters used for the upscale layers will be " - "the value given in 'fc_upsample_filters'." - "\n\tupsample2d - A lightweight and VRAM friendly method. 'quick and dirty' but does " - "not learn any parameters" - "\n\tsubpixel - Sub-pixel upscaler using depth-to-space which may require more " - "VRAM." - "\n\tresize_images - Uses the Keras resize_image function to save about half as much " - "vram as the heaviest methods." - "\n\tupscale_fast - Developed by Andenixa. Focusses on speed to upscale, but " - "requires more VRAM." - "\n\tupscale_hybrid - Developed by Andenixa. Uses a combination of PixelShuffler and " - "Upsampling2D to upscale, saving about 1/3rd of VRAM of the heaviest methods.", - datatype=str, - choices=["resize_images", "subpixel", "upscale_fast", "upscale_hybrid", "upsample2d"], - group="hidden layers", - gui_radio=False, - fixed=True), - fc_upsamples=dict( - default=1, - info="Some upsampling can occur within the Fully Connected layers rather than in the " - "Decoder to increase the dimensional space. Set how many upscale layers should occur " - "within the Fully Connected layers.", - datatype=int, - min_max=(0, 4), - rounding=1, - group="hidden layers", - fixed=True), - fc_upsample_filters=dict( - default=512, - info="If you have selected an upsampler which requires filters (i.e. any upsampler with " - "the exception of Upsampling2D), then this is the number of filters to be used for " - "the upsamplers within the fully connected layers, NB: This value may be scaled " - "down, depending on output resolution. Also note, that this figure will dictate the " - "number of filters used for the G-Block, if selected.", - datatype=int, - rounding=64, - min_max=(128, 5120), - group="hidden layers", - fixed=True), + "fc_depth": { + "default": 1, + "info": ( + "The number of consecutive Dense (fully connected) layers to include in each " + "side's intermediate layer."), + "datatype": int, + "rounding": 1, + "min_max": (0, 16), + "group": "hidden layers", + "fixed": True}, + "fc_min_filters": { + "default": 1024, + "info": ( + "The number of filters to use for the initial fully connected layer. The number of " + "nodes actually used is: fc_min_filters x fc_dimensions x fc_dimensions.\nNB: This " + "value may be scaled down, depending on output resolution."), + "datatype": int, + "rounding": 16, + "min_max": (16, 5120), + "group": "hidden layers", + "fixed": True}, + "fc_max_filters": { + "default": 1024, + "info": ( + "This is the number of filters to be used in the final reshape layer at the end of " + "the fully connected layers. The actual number of nodes used for the final fully " + "connected layer is: fc_min_filters x fc_dimensions x fc_dimensions.\nNB: This value " + "may be scaled down, depending on output resolution."), + "datatype": int, + "rounding": 64, + "min_max": (128, 5120), + "group": "hidden layers", + "fixed": True}, + "fc_dimensions": { + "default": 4, + "info": ( + "The height and width dimension for the final reshape layer at the end of the " + "fully connected layers.\nNB: The total number of nodes within the final fully " + "connected layer will be: fc_dimensions x fc_dimensions x fc_max_filters."), + "datatype": int, + "rounding": 1, + "min_max": (1, 16), + "group": "hidden layers", + "fixed": True}, + "fc_filter_slope": { + "default": -0.5, + "info": ( + "The rate that the filters move from the minimum number of filters to the maximum " + "number of filters. EG:\n" + "Negative numbers will change the number of filters quicker at first and slow down " + "each layer.\n" + "Positive numbers will change the number of filters slower at first but then speed " + "up each layer.\n" + "0.0 - This will change at a linear rate (i.e. the same number of filters will be " + "changed at each layer)."), + "datatype": float, + "min_max": (-.99, .99), + "rounding": 2, + "group": "hidden layers", + "fixed": True}, + "fc_dropout": { + "default": 0.0, + "info": ( + "Dropout is a form of regularization that can prevent a model from over-fitting " + "and help to keep neurons 'alive'. 0.5 will dropout half the connections between each " + "fully connected layer, 0.25 will dropout a quarter of the connections etc. Set to " + "0.0 to disable."), + "datatype": float, + "rounding": 2, + "min_max": (0.0, 0.99), + "group": "hidden layers", + "fixed": False}, + "fc_upsampler": { + "default": "upsample2d", + "info": ( + "The type of dimensional upsampling to perform at the end of the fully connected " + "layers, if upsamples > 0. The number of filters used for the upscale layers will be " + "the value given in 'fc_upsample_filters'." + "\n\tupsample2d - A lightweight and VRAM friendly method. 'quick and dirty' but does " + "not learn any parameters" + "\n\tsubpixel - Sub-pixel upscaler using depth-to-space which may require more " + "VRAM." + "\n\tresize_images - Uses the Keras resize_image function to save about half as much " + "vram as the heaviest methods." + "\n\tupscale_fast - Developed by Andenixa. Focusses on speed to upscale, but " + "requires more VRAM." + "\n\tupscale_hybrid - Developed by Andenixa. Uses a combination of PixelShuffler and " + "Upsampling2D to upscale, saving about 1/3rd of VRAM of the heaviest methods."), + "datatype": str, + "choices": ["resize_images", "subpixel", "upscale_fast", "upscale_hybrid", "upsample2d"], + "group": "hidden layers", + "gui_radio": False, + "fixed": True}, + "fc_upsamples": { + "default": 1, + "info": ( + "Some upsampling can occur within the Fully Connected layers rather than in the " + "Decoder to increase the dimensional space. Set how many upscale layers should occur " + "within the Fully Connected layers."), + "datatype": int, + "min_max": (0, 4), + "rounding": 1, + "group": "hidden layers", + "fixed": True}, + "fc_upsample_filters": { + "default": 512, + "info": ( + "If you have selected an upsampler which requires filters (i.e. any upsampler with " + "the exception of Upsampling2D), then this is the number of filters to be used for " + "the upsamplers within the fully connected layers, NB: This value may be scaled " + "down, depending on output resolution. Also note, that this figure will dictate the " + "number of filters used for the G-Block, if selected."), + "datatype": int, + "rounding": 64, + "min_max": (128, 5120), + "group": "hidden layers", + "fixed": True}, # G-Block - fc_gblock_depth=dict( - default=3, - info="The number of consecutive Dense (fully connected) layers to include in the G-Block " - "shared layer.", - datatype=int, - rounding=1, - min_max=(1, 16), - group="g-block hidden layers", - fixed=True), - fc_gblock_min_nodes=dict( - default=512, - info="The number of nodes to use for the initial G-Block shared fully connected layer.", - datatype=int, - rounding=64, - min_max=(128, 5120), - group="g-block hidden layers", - fixed=True), - fc_gblock_max_nodes=dict( - default=512, - info="The number of nodes to use for the final G-Block shared fully connected layer.", - datatype=int, - rounding=64, - min_max=(128, 5120), - group="g-block hidden layers", - fixed=True), - fc_gblock_filter_slope=dict( - default=-0.5, - info="The rate that the filters move from the minimum number of filters to the maximum " - "number of filters for the G-Block shared layers. EG:\n" - "Negative numbers will change the number of filters quicker at first and slow down " - "each layer.\n" - "Positive numbers will change the number of filters slower at first but then speed " - "up each layer.\n" - "0.0 - This will change at a linear rate (i.e. the same number of filters will be " - "changed at each layer).", - datatype=float, - min_max=(-.99, .99), - rounding=2, - group="g-block hidden layers", - fixed=True), - fc_gblock_dropout=dict( - default=0.0, - info="Dropout is a regularization technique that can prevent a model from over-fitting " - "and help to keep neurons 'alive'. 0.5 will dropout half the connections between " - "each fully connected layer, 0.25 will dropout a quarter of the connections etc. Set " - "to 0.0 to disable.", - datatype=float, - rounding=2, - min_max=(0.0, 0.99), - group="g-block hidden layers", - fixed=False), + "fc_gblock_depth": { + "default": 3, + "info": ( + "The number of consecutive Dense (fully connected) layers to include in the " + "G-Block shared layer."), + "datatype": int, + "rounding": 1, + "min_max": (1, 16), + "group": "g-block hidden layers", + "fixed": True}, + "fc_gblock_min_nodes": { + "default": 512, + "info": "The number of nodes to use for the initial G-Block shared fully connected layer.", + "datatype": int, + "rounding": 64, + "min_max": (128, 5120), + "group": "g-block hidden layers", + "fixed": True}, + "fc_gblock_max_nodes": { + "default": 512, + "info": "The number of nodes to use for the final G-Block shared fully connected layer.", + "datatype": int, + "rounding": 64, + "min_max": (128, 5120), + "group": "g-block hidden layers", + "fixed": True}, + "fc_gblock_filter_slope": { + "default": -0.5, + "info": ( + "The rate that the filters move from the minimum number of filters to the maximum " + "number of filters for the G-Block shared layers. EG:\n" + "Negative numbers will change the number of filters quicker at first and slow down " + "each layer.\n" + "Positive numbers will change the number of filters slower at first but then speed " + "up each layer.\n" + "0.0 - This will change at a linear rate (i.e. the same number of filters will be " + "changed at each layer)."), + "datatype": float, + "min_max": (-.99, .99), + "rounding": 2, + "group": "g-block hidden layers", + "fixed": True}, + "fc_gblock_dropout": { + "default": 0.0, + "info": ( + "Dropout is a regularization technique that can prevent a model from over-fitting " + "and help to keep neurons 'alive'. 0.5 will dropout half the connections between " + "each fully connected layer, 0.25 will dropout a quarter of the connections etc. Set " + "to 0.0 to disable."), + "datatype": float, + "rounding": 2, + "min_max": (0.0, 0.99), + "group": "g-block hidden layers", + "fixed": False}, # Decoder - dec_upscale_method=dict( - default="subpixel", - info="The method to use for the upscales within the decoder. Images are upscaled multiple " - "times within the decoder as the network learns to reconstruct the face." - "\n\tsubpixel - Sub-pixel upscaler using depth-to-space which requires more " - "VRAM." - "\n\tresize_images - Uses the Keras resize_image function to save about half as much " - "vram as the heaviest methods." - "\n\tupscale_fast - Developed by Andenixa. Focusses on speed to upscale, but " - "requires more VRAM." - "\n\tupscale_hybrid - Developed by Andenixa. Uses a combination of PixelShuffler and " - "Upsampling2D to upscale, saving about 1/3rd of VRAM of the heaviest methods." - "\n\tupscale_dny - An alternative upscale implementation using Upsampling2D to " - "upsale.", - datatype=str, - choices=["subpixel", "resize_images", "upscale_fast", "upscale_hybrid", "upscale_dny"], - gui_radio=True, - group="decoder", - fixed=True), - dec_upscales_in_fc=dict( - default=0, - min_max=(0, 6), - rounding=1, - info="It is possible to place some of the upscales at the end of the fully connected " - "model. For models with split decoders, but a shared fully connected layer, this would " - "have the effect of saving some VRAM but possibly at the cost of introducing artefacts. " - "For models with a shared decoder but split fully connected layers, this would have the " - "effect of increasing VRAM usage by processing some of the upscales for each side rather " - "than together.", - datatype=int, - group="decoder", - fixed=True), - dec_norm=dict( - default="none", - info="Normalization to apply to apply after each upscale." - "\n\tnone - Do not apply a normalization layer" - "\n\tbatch - Apply Batch Normalization" - "\n\tgroup - Apply Group Normalization" - "\n\tinstance - Apply Instance Normalization" - "\n\tlayer - Apply Layer Normalization (Ba et al., 2016)" - "\n\trms - Apply Root Mean Squared Layer Normalization (Zhang et al., 2019). A " - "simplified version of Layer Normalization with reduced overhead.", - datatype=str, - gui_radio=True, - choices=["none", "batch", "group", "instance", "layer", "rms"], - group="decoder", - fixed=True), - dec_min_filters=dict( - default=64, - info="The minimum number of filters to use in decoder upscalers (i.e. the number of " - "filters to use for the final upscale layer).", - datatype=int, - min_max=(16, 512), - rounding=16, - group="decoder", - fixed=True), - dec_max_filters=dict( - default=512, - info="The maximum number of filters to use in decoder upscalers (i.e. the number of " - "filters to use for the first upscale layer).", - datatype=int, - min_max=(256, 5120), - rounding=64, - group="decoder", - fixed=True), - dec_slope_mode=dict( - default="full", - info="Alters the action of the filter slope.\n" - "\n\tfull: The number of filters at each upscale layer will reduce from the chosen " - "max_filters at the first layer to the chosen min_filters at the last layer as " - "dictated by the dec_filter_slope." - "\n\tcap_max: The filters will decline at a fixed rate from each upscale to the next " - "based on the filter_slope setting. If there are more upscales than filters, " - "then the earliest upscales will be capped at the max_filter value until the filters " - "can reduce to the min_filters value at the final upscale. (EG: 512 -> 512 -> 512 -> " - "256 -> 128 -> 64)." - "\n\tcap_min: The filters will decline at a fixed rate from each upscale to the next " - "based on the filter_slope setting. If there are more upscales than filters, then " - "the earliest upscales will drop their filters until the min_filter value is met and " - "repeat the min_filter value for the remaining upscales. (EG: 512 -> 256 -> 128 -> " - "64 -> 64 -> 64).", - choices=["full", "cap_max", "cap_min"], - group="decoder", - fixed=True, - gui_radio=True), - dec_filter_slope=dict( - default=-0.45, - info="The rate that the filters reduce at each upscale layer.\n" - "\n\tFull Slope Mode: Negative numbers will drop the number of filters quicker at " - "first and slow down each upscale. Positive numbers will drop the number of filters " - "slower at first but then speed up each upscale. A value of 0.0 will reduce at a " - "linear rate (i.e. the same number of filters will be reduced at each upscale).\n" - "\n\tCap Min/Max Slope Mode: Only positive values will work here. Negative values " - "will automatically be converted to their positive counterpart. A value of 0.5 will " - "halve the number of filters at each upscale until the minimum value is reached. A " - "value of 0.33 will be reduce the number of filters by a third until the minimum " - "value is reached etc.", - datatype=float, - min_max=(-.99, .99), - rounding=2, - group="decoder", - fixed=True), - dec_res_blocks=dict( - default=1, - info="The number of Residual Blocks to apply to each upscale layer. Set to 0 to disable " - "residual blocks entirely.", - datatype=int, - rounding=1, - min_max=(0, 8), - group="decoder", - fixed=True), - dec_output_kernel=dict( - default=5, - info="The kernel size to apply to the final Convolution layer.", - datatype=int, - rounding=2, - min_max=(1, 9), - group="decoder", - fixed=True), - dec_gaussian=dict( - default=True, - info="Gaussian Noise acts as a regularization technique for preventing overfitting of " - "data." - "\n\tTrue - Apply a Gaussian Noise layer to each upscale." - "\n\tFalse - Don't apply a Gaussian Noise layer to each upscale.", - datatype=bool, - group="decoder", - fixed=True), - dec_skip_last_residual=dict( - default=True, - info="If Residual blocks have been enabled, enabling this option will not apply a " - "Residual block to the final upscaler." - "\n\tTrue - Don't apply a Residual block to the final upscale." - "\n\tFalse - Apply a Residual block to all upscale layers.", - datatype=bool, - group="decoder", - fixed=True), + "dec_upscale_method": { + "default": "subpixel", + "info": ( + "The method to use for the upscales within the decoder. Images are upscaled " + "multiple times within the decoder as the network learns to reconstruct the face." + "\n\tsubpixel - Sub-pixel upscaler using depth-to-space which requires more " + "VRAM." + "\n\tresize_images - Uses the Keras resize_image function to save about half as much " + "vram as the heaviest methods." + "\n\tupscale_fast - Developed by Andenixa. Focusses on speed to upscale, but " + "requires more VRAM." + "\n\tupscale_hybrid - Developed by Andenixa. Uses a combination of PixelShuffler and " + "Upsampling2D to upscale, saving about 1/3rd of VRAM of the heaviest methods." + "\n\tupscale_dny - An alternative upscale implementation using Upsampling2D to " + "upsale."), + "datatype": str, + "choices": ["subpixel", "resize_images", "upscale_fast", "upscale_hybrid", "upscale_dny"], + "gui_radio": True, + "group": "decoder", + "fixed": True}, + "dec_upscales_in_fc": { + "default": 0, + "min_max": (0, 6), + "rounding": 1, + "info": ( + "It is possible to place some of the upscales at the end of the fully connected " + "model. For models with split decoders, but a shared fully connected layer, this " + "would have the effect of saving some VRAM but possibly at the cost of introducing " + "artefacts. For models with a shared decoder but split fully connected layers, this " + "would have the effect of increasing VRAM usage by processing some of the upscales " + "for each side rather than together."), + "datatype": int, + "group": "decoder", + "fixed": True}, + "dec_norm": { + "default": "none", + "info": ( + "Normalization to apply to apply after each upscale." + "\n\tnone - Do not apply a normalization layer" + "\n\tbatch - Apply Batch Normalization" + "\n\tgroup - Apply Group Normalization" + "\n\tinstance - Apply Instance Normalization" + "\n\tlayer - Apply Layer Normalization (Ba et al., 2016)" + "\n\trms - Apply Root Mean Squared Layer Normalization (Zhang et al., 2019). A " + "simplified version of Layer Normalization with reduced overhead."), + "datatype": str, + "gui_radio": True, + "choices": ["none", "batch", "group", "instance", "layer", "rms"], + "group": "decoder", + "fixed": True}, + "dec_min_filters": { + "default": 64, + "info": ( + "The minimum number of filters to use in decoder upscalers (i.e. the number of " + "filters to use for the final upscale layer)."), + "datatype": int, + "min_max": (16, 512), + "rounding": 16, + "group": "decoder", + "fixed": True}, + "dec_max_filters": { + "default": 512, + "info": ( + "The maximum number of filters to use in decoder upscalers (i.e. the number of " + "filters to use for the first upscale layer)."), + "datatype": int, + "min_max": (256, 5120), + "rounding": 64, + "group": "decoder", + "fixed": True}, + "dec_slope_mode": { + "default": "full", + "info": ( + "Alters the action of the filter slope.\n" + "\n\tfull: The number of filters at each upscale layer will reduce from the chosen " + "max_filters at the first layer to the chosen min_filters at the last layer as " + "dictated by the dec_filter_slope." + "\n\tcap_max: The filters will decline at a fixed rate from each upscale to the next " + "based on the filter_slope setting. If there are more upscales than filters, " + "then the earliest upscales will be capped at the max_filter value until the filters " + "can reduce to the min_filters value at the final upscale. (EG: 512 -> 512 -> 512 -> " + "256 -> 128 -> 64)." + "\n\tcap_min: The filters will decline at a fixed rate from each upscale to the next " + "based on the filter_slope setting. If there are more upscales than filters, then " + "the earliest upscales will drop their filters until the min_filter value is met and " + "repeat the min_filter value for the remaining upscales. (EG: 512 -> 256 -> 128 -> " + "64 -> 64 -> 64)."), + "choices": ["full", "cap_max", "cap_min"], + "group": "decoder", + "fixed": True, + "gui_radio": True}, + "dec_filter_slope": { + "default": -0.45, + "info": ( + "The rate that the filters reduce at each upscale layer.\n" + "\n\tFull Slope Mode: Negative numbers will drop the number of filters quicker at " + "first and slow down each upscale. Positive numbers will drop the number of filters " + "slower at first but then speed up each upscale. A value of 0.0 will reduce at a " + "linear rate (i.e. the same number of filters will be reduced at each upscale).\n" + "\n\tCap Min/Max Slope Mode: Only positive values will work here. Negative values " + "will automatically be converted to their positive counterpart. A value of 0.5 will " + "halve the number of filters at each upscale until the minimum value is reached. A " + "value of 0.33 will be reduce the number of filters by a third until the minimum " + "value is reached etc."), + "datatype": float, + "min_max": (-.99, .99), + "rounding": 2, + "group": "decoder", + "fixed": True}, + "dec_res_blocks": { + "default": 1, + "info": ( + "The number of Residual Blocks to apply to each upscale layer. Set to 0 to disable " + "residual blocks entirely."), + "datatype": int, + "rounding": 1, + "min_max": (0, 8), + "group": "decoder", + "fixed": True}, + "dec_output_kernel": { + "default": 5, + "info": "The kernel size to apply to the final Convolution layer.", + "datatype": int, + "rounding": 2, + "min_max": (1, 9), + "group": "decoder", + "fixed": True}, + "dec_gaussian": { + "default": True, + "info": ( + "Gaussian Noise acts as a regularization technique for preventing overfitting of " + "data." + "\n\tTrue - Apply a Gaussian Noise layer to each upscale." + "\n\tFalse - Don't apply a Gaussian Noise layer to each upscale."), + "datatype": bool, + "group": "decoder", + "fixed": True}, + "dec_skip_last_residual": { + "default": True, + "info": ( + "If Residual blocks have been enabled, enabling this option will not apply a " + "Residual block to the final upscaler." + "\n\tTrue - Don't apply a Residual block to the final upscale." + "\n\tFalse - Apply a Residual block to all upscale layers."), + "datatype": bool, + "group": "decoder", + "fixed": True}, # Weight management - freeze_layers=dict( - default="keras_encoder", - info="If the command line option 'freeze-weights' is enabled, then the layers indicated " - "here will be frozen the next time the model starts up. NB: Not all architectures " - "contain all of the layers listed here, so any layers marked for freezing that are " - "not within your chosen architecture will be ignored. EG:\n If 'split fc' has " - "been selected, then 'fc_a' and 'fc_b' are available for freezing. If it has " - "not been selected then 'fc_both' is available for freezing.", - datatype=list, - choices=["encoder", "keras_encoder", "fc_a", "fc_b", "fc_both", "fc_shared", "fc_gblock", - "g_block_a", "g_block_b", "g_block_both", "decoder_a", "decoder_b", - "decoder_both"], - group="weights", - fixed=False), - load_layers=dict( - default="encoder", - info="If the command line option 'load-weights' is populated, then the layers indicated " - "here will be loaded from the given weights file if starting a new model. NB Not all " - "architectures contain all of the layers listed here, so any layers marked for " - "loading that are not within your chosen architecture will be ignored. EG:\n If " - "'split fc' has been selected, then 'fc_a' and 'fc_b' are available for loading. If " - "it has not been selected then 'fc_both' is available for loading.", - datatype=list, - choices=["encoder", "fc_a", "fc_b", "fc_both", "fc_shared", "fc_gblock", "g_block_a", - "g_block_b", "g_block_both", "decoder_a", "decoder_b", "decoder_both"], - group="weights", - fixed=True), + "freeze_layers": { + "default": "keras_encoder", + "info": ( + "If the command line option 'freeze-weights' is enabled, then the layers indicated " + "here will be frozen the next time the model starts up. NB: Not all architectures " + "contain all of the layers listed here, so any layers marked for freezing that are " + "not within your chosen architecture will be ignored. EG:\n If 'split fc' has " + "been selected, then 'fc_a' and 'fc_b' are available for freezing. If it has " + "not been selected then 'fc_both' is available for freezing."), + "datatype": list, + "choices": ["encoder", "keras_encoder", "fc_a", "fc_b", "fc_both", "fc_shared", + "fc_gblock", "g_block_a", "g_block_b", "g_block_both", "decoder_a", + "decoder_b", "decoder_both"], + "group": "weights", + "fixed": False}, + "load_layers": { + "default": "encoder", + "info": ( + "If the command line option 'load-weights' is populated, then the layers indicated " + "here will be loaded from the given weights file if starting a new model. NB Not all " + "architectures contain all of the layers listed here, so any layers marked for " + "loading that are not within your chosen architecture will be ignored. EG:\n If " + "'split fc' has been selected, then 'fc_a' and 'fc_b' are available for loading. If " + "it has not been selected then 'fc_both' is available for loading."), + "datatype": list, + "choices": ["encoder", "fc_a", "fc_b", "fc_both", "fc_shared", "fc_gblock", "g_block_a", + "g_block_b", "g_block_both", "decoder_a", "decoder_b", "decoder_both"], + "group": "weights", + "fixed": True}, # # SPECIFIC ENCODER SETTINGS # # # Faceswap Original - fs_original_depth=dict( - default=4, - info="Faceswap Encoder only: The number of convolutions to perform within the encoder.", - datatype=int, - min_max=(2, 10), - rounding=1, - group="faceswap encoder configuration", - fixed=True), - fs_original_min_filters=dict( - default=128, - info="Faceswap Encoder only: The minumum number of filters to use for encoder " - "convolutions. (i.e. the number of filters to use for the first encoder layer).", - datatype=int, - min_max=(16, 2048), - rounding=64, - group="faceswap encoder configuration", - fixed=True), - fs_original_max_filters=dict( - default=1024, - info="Faceswap Encoder only: The maximum number of filters to use for encoder " - "convolutions. (i.e. the number of filters to use for the final encoder layer).", - datatype=int, - min_max=(256, 8192), - rounding=128, - group="faceswap encoder configuration", - fixed=True), - fs_original_use_alt=dict( - default=False, - info="Use a slightly alternate version of the Faceswap Encoder." - "\n\tTrue - Use the alternate variation of the Faceswap Encoder." - "\n\tFalse - Use the original Faceswap Encoder.", - datatype=bool, - group="faceswap encoder configuration", - fixed=True), + "fs_original_depth": { + "default": 4, + "info": "Faceswap Encoder only: The number of convolutions to perform within the encoder.", + "datatype": int, + "min_max": (2, 10), + "rounding": 1, + "group": "faceswap encoder configuration", + "fixed": True}, + "fs_original_min_filters": { + "default": 128, + "info": ( + "Faceswap Encoder only: The minumum number of filters to use for encoder " + "convolutions. (i.e. the number of filters to use for the first encoder layer)."), + "datatype": int, + "min_max": (16, 2048), + "rounding": 64, + "group": "faceswap encoder configuration", + "fixed": True}, + "fs_original_max_filters": { + "default": 1024, + "info": ( + "Faceswap Encoder only: The maximum number of filters to use for encoder " + "convolutions. (i.e. the number of filters to use for the final encoder layer)."), + "datatype": int, + "min_max": (256, 8192), + "rounding": 128, + "group": "faceswap encoder configuration", + "fixed": True}, + "fs_original_use_alt": { + "default": False, + "info": ( + "Use a slightly alternate version of the Faceswap Encoder." + "\n\tTrue - Use the alternate variation of the Faceswap Encoder." + "\n\tFalse - Use the original Faceswap Encoder."), + "datatype": bool, + "group": "faceswap encoder configuration", + "fixed": True}, # MobileNet - mobilenet_width=dict( - default=1.0, - info="The width multiplier for mobilenet encoders. Controls the width of the " - "network. Values less than 1.0 proportionally decrease the number of filters within " - "each layer. Values greater than 1.0 proportionally increase the number of filters " - "within each layer. 1.0 is the default number of layers used within the paper.\n" - "NB: This option is ignored for any non-mobilenet encoders.\n" - "NB: If loading ImageNet weights, then for MobilenetV1 only values of '0.25', " - "'0.5', '0.75' or '1.0 can be selected. For MobilenetV2 only values of '0.35', " - "'0.50', '0.75', '1.0', '1.3' or '1.4' can be selected. For mobilenet_v3 only values " - "of '0.75' or '1.0' can be selected", - datatype=float, - min_max=(0.1, 2.0), - rounding=2, - group="mobilenet encoder configuration", - fixed=True), - mobilenet_depth=dict( - default=1, - info="The depth multiplier for MobilenetV1 encoder. This is the depth multiplier " - "for depthwise convolution (known as the resolution multiplier within the original " - "paper).\n" - "NB: This option is only used for MobilenetV1 and is ignored for all other " - "encoders.\n" - "NB: If loading ImageNet weights, this must be set to 1.", - datatype=int, - min_max=(1, 10), - rounding=1, - group="mobilenet encoder configuration", - fixed=True), - mobilenet_dropout=dict( - default=0.001, - info="The dropout rate for MobilenetV1 encoder.\n" - "NB: This option is only used for MobilenetV1 and is ignored for all other " - "encoders.", - datatype=float, - min_max=(0.001, 2.0), - rounding=3, - group="mobilenet encoder configuration", - fixed=True), - mobilenet_minimalistic=dict( - default=False, - info="Use a minimilist version of MobilenetV3.\n" - "In addition to large and small models MobilenetV3 also contains so-called " - "minimalistic models, these models have the same per-layer dimensions characteristic " - "as MobilenetV3 however, they don't utilize any of the advanced blocks " - "(squeeze-and-excite units, hard-swish, and 5x5 convolutions). While these models " - "are less efficient on CPU, they are much more performant on GPU/DSP.\n" - "NB: This option is only used for MobilenetV3 and is ignored for all other " - "encoders.\n", - datatype=bool, - group="mobilenet encoder configuration", - fixed=True), - ) + "mobilenet_width": { + "default": 1.0, + "info": ( + "The width multiplier for mobilenet encoders. Controls the width of the " + "network. Values less than 1.0 proportionally decrease the number of filters within " + "each layer. Values greater than 1.0 proportionally increase the number of filters " + "within each layer. 1.0 is the default number of layers used within the paper.\n" + "NB: This option is ignored for any non-mobilenet encoders.\n" + "NB: If loading ImageNet weights, then for MobilenetV1 only values of '0.25', " + "'0.5', '0.75' or '1.0 can be selected. For MobilenetV2 only values of '0.35', " + "'0.50', '0.75', '1.0', '1.3' or '1.4' can be selected. For mobilenet_v3 only values " + "of '0.75' or '1.0' can be selected"), + "datatype": float, + "min_max": (0.1, 2.0), + "rounding": 2, + "group": "mobilenet encoder configuration", + "fixed": True}, + "mobilenet_depth": { + "default": 1, + "info": ( + "The depth multiplier for MobilenetV1 encoder. This is the depth multiplier " + "for depthwise convolution (known as the resolution multiplier within the original " + "paper).\n" + "NB: This option is only used for MobilenetV1 and is ignored for all other " + "encoders.\n" + "NB: If loading ImageNet weights, this must be set to 1."), + "datatype": int, + "min_max": (1, 10), + "rounding": 1, + "group": "mobilenet encoder configuration", + "fixed": True}, + "mobilenet_dropout": { + "default": 0.001, + "info": ( + "The dropout rate for MobilenetV1 encoder.\n" + "NB: This option is only used for MobilenetV1 and is ignored for all other " + "encoders."), + "datatype": float, + "min_max": (0.001, 2.0), + "rounding": 3, + "group": "mobilenet encoder configuration", + "fixed": True}, + "mobilenet_minimalistic": { + "default": False, + "info": ( + "Use a minimilist version of MobilenetV3.\n" + "In addition to large and small models MobilenetV3 also contains so-called " + "minimalistic models, these models have the same per-layer dimensions characteristic " + "as MobilenetV3 however, they don't utilize any of the advanced blocks " + "(squeeze-and-excite units, hard-swish, and 5x5 convolutions). While these models " + "are less efficient on CPU, they are much more performant on GPU/DSP.\n" + "NB: This option is only used for MobilenetV3 and is ignored for all other " + "encoders.\n"), + "datatype": bool, + "group": "mobilenet encoder configuration", + "fixed": True}, + } diff --git a/plugins/train/model/realface.py b/plugins/train/model/realface.py index fd1aa73569..30d0d7f8f1 100644 --- a/plugins/train/model/realface.py +++ b/plugins/train/model/realface.py @@ -7,19 +7,18 @@ Additional thanks: Birb - source of inspiration, great Encoder ideas Kvrooman - additional counseling on auto-encoders and practical advice """ +import logging import sys +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.initializers import RandomNormal # pylint:disable=import-error +from tensorflow.keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error + from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock -from lib.utils import get_backend -from ._base import ModelBase, KerasModel, logger +from ._base import ModelBase -if get_backend() == "amd": - from keras.initializers import RandomNormal - from keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.initializers import RandomNormal # noqa pylint:disable=import-error,no-name-in-module - from tensorflow.keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape # noqa pylint:disable=import-error,no-name-in-module +logger = logging.getLogger(__name__) class Model(ModelBase): @@ -77,7 +76,7 @@ def build_model(self, inputs): outputs = [self.decoder_a()(encoder_a), self.decoder_b()(encoder_b)] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder(self): @@ -95,7 +94,7 @@ def encoder(self): var_x = Conv2DBlock(encoder_complexity * 2**(idx + 1), activation="leakyrelu")(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder_b(self): """ RealFace Decoder Network """ @@ -139,7 +138,7 @@ def decoder_b(self): outputs += [var_y] - return KerasModel(input_, outputs=outputs, name="decoder_b") + return KModel(input_, outputs=outputs, name="decoder_b") def decoder_a(self): """ RealFace Decoder (A) Network """ @@ -184,7 +183,7 @@ def decoder_a(self): outputs += [var_y] - return KerasModel(input_, outputs=outputs, name="decoder_a") + return KModel(input_, outputs=outputs, name="decoder_a") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ diff --git a/plugins/train/model/unbalanced.py b/plugins/train/model/unbalanced.py index 9330f0f4c0..6f83166305 100644 --- a/plugins/train/model/unbalanced.py +++ b/plugins/train/model/unbalanced.py @@ -3,17 +3,14 @@ Based on the original https://www.reddit.com/r/deepfakes/ code sample + contributions """ -from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock -from lib.utils import get_backend -from ._base import ModelBase, KerasModel +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.initializers import RandomNormal # pylint:disable=import-error +from tensorflow.keras.layers import ( # pylint:disable=import-error + Dense, Flatten, Input, LeakyReLU, Reshape, SpatialDropout2D) +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error -if get_backend() == "amd": - from keras.initializers import RandomNormal - from keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape, SpatialDropout2D -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.initializers import RandomNormal # noqa pylint:disable=import-error,no-name-in-module - from tensorflow.keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape, SpatialDropout2D # noqa pylint:disable=import-error,no-name-in-module +from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock +from ._base import ModelBase class Model(ModelBase): @@ -33,12 +30,12 @@ def build_model(self, inputs): outputs = [self.decoder_a()(encoder_a), self.decoder_b()(encoder_b)] - autoencoder = KerasModel(inputs, outputs, name=self.model_name) + autoencoder = KModel(inputs, outputs, name=self.model_name) return autoencoder def encoder(self): """ Unbalanced Encoder """ - kwargs = dict(kernel_initializer=self.kernel_initializer) + kwargs = {"kernel_initializer": self.kernel_initializer} encoder_complexity = 128 if self.low_mem else self.config["complexity_encoder"] dense_dim = 384 if self.low_mem else 512 dense_shape = self.input_shape[0] // 16 @@ -61,11 +58,11 @@ def encoder(self): var_x = Dense(dense_shape * dense_shape * dense_dim, kernel_initializer=self.kernel_initializer)(var_x) var_x = Reshape((dense_shape, dense_shape, dense_dim))(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder_a(self): """ Decoder for side A """ - kwargs = dict(kernel_size=5, kernel_initializer=self.kernel_initializer) + kwargs = {"kernel_size": 5, "kernel_initializer": self.kernel_initializer} decoder_complexity = 320 if self.low_mem else self.config["complexity_decoder_a"] dense_dim = 384 if self.low_mem else 512 decoder_shape = self.input_shape[0] // 16 @@ -93,13 +90,13 @@ def decoder_a(self): var_y = UpscaleBlock(decoder_complexity // 4, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name="mask_out_a")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name="decoder_a") + return KModel(input_, outputs=outputs, name="decoder_a") def decoder_b(self): """ Decoder for side B """ - kwargs = dict(kernel_size=5, kernel_initializer=self.kernel_initializer) - dense_dim = 384 if self.low_mem else self.config["complexity_decoder_b"] - decoder_complexity = 384 if self.low_mem else 512 + kwargs = {"kernel_size": 5, "kernel_initializer": self.kernel_initializer} + decoder_complexity = 384 if self.low_mem else self.config["complexity_decoder_b"] + dense_dim = 384 if self.low_mem else 512 decoder_shape = self.input_shape[0] // 16 input_ = Input(shape=(decoder_shape, decoder_shape, dense_dim)) @@ -137,7 +134,7 @@ def decoder_b(self): var_y = UpscaleBlock(decoder_complexity // 8, activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name="mask_out_b")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name="decoder_b") + return KModel(input_, outputs=outputs, name="decoder_b") def _legacy_mapping(self): """ The mapping of legacy separate model names to single model names """ diff --git a/plugins/train/model/unbalanced_defaults.py b/plugins/train/model/unbalanced_defaults.py index 317aec23ff..28bbcfd5da 100755 --- a/plugins/train/model/unbalanced_defaults.py +++ b/plugins/train/model/unbalanced_defaults.py @@ -47,90 +47,70 @@ ) -_DEFAULTS = { - "input_size": { - "default": 128, - "info": "Resolution (in pixels) of the image to train on.\n" - "BE AWARE Larger resolution will dramatically increaseVRAM requirements.\n" - "Make sure your resolution is divisible by 64 (e.g. 64, 128, 256 etc.).\n" - "NB: Your faceset must be at least 1.6x larger than your required input " - "size.\n(e.g. 160 is the maximum input size for a 256x256 faceset).", - "datatype": int, - "rounding": 64, - "min_max": (64, 512), - "choices": [], - "gui_radio": False, - "group": "size", - "fixed": True, - }, - "lowmem": { - "default": False, - "info": "Lower memory mode. Set to 'True' if having issues with VRAM useage.\n" - "NB: Models with a changed lowmem mode are not compatible with each other.\n" - "NB: lowmem will override cutom nodes and complexity settings.", - "datatype": bool, - "rounding": None, - "min_max": None, - "choices": [], - "gui_radio": False, - "group": "settings", - "fixed": True, - }, - "clipnorm": { - "default": True, - "info": "Controls gradient clipping of the optimizer. Can prevent model corruption at " - "the expense of VRAM.", - "datatype": bool, - "rounding": None, - "min_max": None, - "choices": [], - "gui_radio": False, - "group": "settings", - "fixed": True, - }, - "nodes": { - "default": 1024, - "info": "Number of nodes for decoder. Don't change this unless you know what you are " - "doing!", - "datatype": int, - "rounding": 64, - "min_max": (512, 4096), - "choices": [], - "gui_radio": False, - "fixed": True, - "group": "network", - }, - "complexity_encoder": { - "default": 128, - "info": "Encoder Convolution Layer Complexity. sensible ranges: 128 to 160.", - "datatype": int, - "rounding": 16, - "min_max": (64, 1024), - "choices": [], - "gui_radio": False, - "fixed": True, - "group": "network", - }, - "complexity_decoder_a": { - "default": 384, - "info": "Decoder A Complexity.", - "datatype": int, - "rounding": 16, - "min_max": (64, 1024), - "choices": [], - "gui_radio": False, - "fixed": True, - "group": "network", - }, - "complexity_decoder_b": { - "default": 512, - "info": "Decoder B Complexity.", - "datatype": int, - "rounding": 16, - "min_max": (64, 1024), - "choices": [], - "gui_radio": False, - "fixed": True, - "group": "network", - }, -} +_DEFAULTS = dict( + input_size=dict( + default=128, + info="Resolution (in pixels) of the image to train on.\n" + "BE AWARE Larger resolution will dramatically increaseVRAM requirements.\n" + "Make sure your resolution is divisible by 64 (e.g. 64, 128, 256 etc.).\n" + "NB: Your faceset must be at least 1.6x larger than your required input " + "size.\n(e.g. 160 is the maximum input size for a 256x256 faceset).", + datatype=int, + rounding=64, + min_max=(64, 512), + choices=[], + gui_radio=False, + group="size", + fixed=True), + lowmem=dict( + default=False, + info="Lower memory mode. Set to 'True' if having issues with VRAM useage.\n" + "NB: Models with a changed lowmem mode are not compatible with each other.\n" + "NB: lowmem will override cutom nodes and complexity settings.", + datatype=bool, + rounding=None, + min_max=None, + choices=[], + gui_radio=False, + group="settings", + fixed=True), + nodes=dict( + default=1024, + info="Number of nodes for decoder. Don't change this unless you know what you are doing!", + datatype=int, + rounding=64, + min_max=(512, 4096), + choices=[], + gui_radio=False, + fixed=True, + group="network"), + complexity_encoder=dict( + default=128, + info="Encoder Convolution Layer Complexity. sensible ranges: 128 to 160.", + datatype=int, + rounding=16, + min_max=(64, 1024), + choices=[], + gui_radio=False, + fixed=True, + group="network"), + complexity_decoder_a=dict( + default=384, + info="Decoder A Complexity.", + datatype=int, + rounding=16, + min_max=(64, 1024), + choices=[], + gui_radio=False, + fixed=True, + group="network"), + complexity_decoder_b=dict( + default=512, + info="Decoder B Complexity.", + datatype=int, + rounding=16, + min_max=(64, 1024), + choices=[], + gui_radio=False, + fixed=True, + group="network")) diff --git a/plugins/train/model/villain.py b/plugins/train/model/villain.py index 16efdc6c4d..1d6bfc7f10 100644 --- a/plugins/train/model/villain.py +++ b/plugins/train/model/villain.py @@ -3,20 +3,16 @@ Based on the original https://www.reddit.com/r/deepfakes/ code sample + contributions Adapted from a model by VillainGuy (https://github.com/VillainGuy) """ +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras.initializers import RandomNormal # pylint:disable=import-error +from tensorflow.keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape # noqa:E501 # pylint:disable=import-error +from tensorflow.keras.models import Model as KModel # pylint:disable=import-error + from lib.model.layers import PixelShuffler from lib.model.nn_blocks import (Conv2DOutput, Conv2DBlock, ResidualBlock, SeparableConv2DBlock, UpscaleBlock) -from lib.utils import get_backend - -from .original import Model as OriginalModel, KerasModel -if get_backend() == "amd": - from keras.initializers import RandomNormal - from keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras.initializers import RandomNormal # noqa pylint:disable=import-error,no-name-in-module - from tensorflow.keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape # noqa pylint:disable=import-error,no-name-in-module +from .original import Model as OriginalModel class Model(OriginalModel): @@ -29,7 +25,7 @@ def __init__(self, *args, **kwargs): def encoder(self): """ Encoder Network """ - kwargs = dict(kernel_initializer=self.kernel_initializer) + kwargs = {"kernel_initializer": self.kernel_initializer} input_ = Input(shape=self.input_shape) in_conv_filters = self.input_shape[0] if self.input_shape[0] > 128: @@ -61,11 +57,11 @@ def encoder(self): var_x = Dense(dense_shape * dense_shape * 1024, **kwargs)(var_x) var_x = Reshape((dense_shape, dense_shape, 1024))(var_x) var_x = UpscaleBlock(512, activation="leakyrelu", **kwargs)(var_x) - return KerasModel(input_, var_x, name="encoder") + return KModel(input_, var_x, name="encoder") def decoder(self, side): """ Decoder Network """ - kwargs = dict(kernel_initializer=self.kernel_initializer) + kwargs = {"kernel_initializer": self.kernel_initializer} decoder_shape = self.input_shape[0] // 8 input_ = Input(shape=(decoder_shape, decoder_shape, 512)) @@ -89,4 +85,4 @@ def decoder(self, side): var_y = UpscaleBlock(self.input_shape[0], activation="leakyrelu")(var_y) var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y) outputs.append(var_y) - return KerasModel(input_, outputs=outputs, name=f"decoder_{side}") + return KModel(input_, outputs=outputs, name=f"decoder_{side}") diff --git a/plugins/train/trainer/_base.py b/plugins/train/trainer/_base.py index cd0640867d..dbbc74c9fd 100644 --- a/plugins/train/trainer/_base.py +++ b/plugins/train/trainer/_base.py @@ -6,11 +6,11 @@ inherits from this class. If further plugins are developed, then common code should be kept here, with "original" unique code split out to the original plugin. """ - -# pylint:disable=too-many-lines +from __future__ import annotations import logging import os import time +import typing as T import cv2 import numpy as np @@ -19,14 +19,21 @@ from tensorflow.python.framework import ( # pylint:disable=no-name-in-module errors_impl as tf_errors) -from lib.training import TrainingDataGenerator -from lib.utils import FaceswapError, get_backend, get_folder, get_image_paths, get_tf_version +from lib.image import hex_to_rgb +from lib.training import Feeder, LearningRateFinder +from lib.utils import FaceswapError, get_folder, get_image_paths from plugins.train._config import Config -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from collections.abc import Callable + from plugins.train.model._base import ModelBase + from lib.config import ConfigValueType + +logger = logging.getLogger(__name__) -def _get_config(plugin_name, configfile=None): +def _get_config(plugin_name: str, + configfile: str | None = None) -> dict[str, ConfigValueType]: """ Return the configuration for the requested trainer. Parameters @@ -39,8 +46,8 @@ def _get_config(plugin_name, configfile=None): Returns ------- - :class:`lib.config.FaceswapConfig` - The configuration file for the requested plugin + dict + The configuration dictionary for the requested plugin """ return Config(plugin_name, configfile=configfile).config_dict @@ -65,30 +72,49 @@ class TrainerBase(): from the default :file:`.config.train.ini` file. """ - def __init__(self, model, images, batch_size, configfile): + def __init__(self, + model: ModelBase, + images: dict[T.Literal["a", "b"], list[str]], + batch_size: int, + configfile: str | None) -> None: logger.debug("Initializing %s: (model: '%s', batch_size: %s)", self.__class__.__name__, model, batch_size) self._model = model self._config = self._get_config(configfile) + self._feeder = Feeder(images, model, batch_size, self._config) + + self._exit_early = self._handle_lr_finder() + if self._exit_early: + return + self._model.state.add_session_batchsize(batch_size) self._images = images self._sides = sorted(key for key in self._images.keys()) - self._feeder = _Feeder(images, self._model, batch_size, self._config) - self._tensorboard = self._set_tensorboard() self._samples = _Samples(self._model, self._model.coverage_ratio, - self._model.command_line_arguments.preview_scale / 100) + T.cast(int, self._config["mask_opacity"]), + T.cast(str, self._config["mask_color"])) + + num_images = self._config.get("preview_images", 14) + assert isinstance(num_images, int) self._timelapse = _Timelapse(self._model, self._model.coverage_ratio, - self._config.get("preview_images", 14), + num_images, + T.cast(int, self._config["mask_opacity"]), + T.cast(str, self._config["mask_color"]), self._feeder, self._images) logger.debug("Initialized %s", self.__class__.__name__) - def _get_config(self, configfile): + @property + def exit_early(self) -> bool: + """ True if the trainer should exit early, without perfoming any training steps """ + return self._exit_early + + def _get_config(self, configfile: str | None) -> dict[str, ConfigValueType]: """ Get the saved training config options. Override any global settings with the setting provided from the model's saved config. @@ -113,7 +139,35 @@ def _get_config(self, configfile): config[key] = new_val return config - def _set_tensorboard(self): + def _handle_lr_finder(self) -> bool: + """ Handle the learning rate finder. + + If this is a new model, then find the optimal learning rate and return ``True`` if user has + just requested the graph, otherwise return ``False`` to continue training + + If it as existing model, set the learning rate to the value found by the learing rate + finder and return ``False`` to continue training + + Returns + ------- + bool + ``True`` if the learning rate finder options dictate that training should not continue + after finding the optimal leaning rate + """ + if not self._model.command_line_arguments.use_lr_finder: + return False + + if self._model.state.iterations == 0 and self._model.state.session_id == 1: + lrf = LearningRateFinder(self._model, self._config, self._feeder) + success = lrf.find() + return self._config["lr_finder_mode"] == "graph_and_exit" or not success + + learning_rate = self._model.state.sessions[1]["config"]["learning_rate"] + logger.info("Setting learning rate from Learning Rate Finder to %s", + f"{learning_rate:.1e}") + return False + + def _set_tensorboard(self) -> tf.keras.callbacks.TensorBoard: """ Set up Tensorboard callback for logging loss. Bypassed if command line option "no-logs" has been selected. @@ -124,17 +178,17 @@ def _set_tensorboard(self): Tensorboard object for the the current training session. """ if self._model.state.current_session["no_logs"]: - logger.verbose("TensorBoard logging disabled") + logger.verbose("TensorBoard logging disabled") # type: ignore return None logger.debug("Enabling TensorBoard Logging") logger.debug("Setting up TensorBoard Logging") - log_dir = os.path.join(str(self._model.model_dir), + log_dir = os.path.join(str(self._model.io.model_dir), f"{self._model.name}_logs", f"session_{self._model.state.session_id}") tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=0, # Must be 0 or hangs - write_graph=get_backend() != "amd", + write_graph=True, write_images=False, update_freq="batch", profile_batch=0, @@ -142,14 +196,17 @@ def _set_tensorboard(self): embeddings_metadata=None) tensorboard.set_model(self._model.model) tensorboard.on_train_begin(0) - logger.verbose("Enabled TensorBoard Logging") + logger.verbose("Enabled TensorBoard Logging") # type: ignore return tensorboard - def toggle_mask(self): + def toggle_mask(self) -> None: """ Toggle the mask overlay on or off based on user input. """ self._samples.toggle_mask_display() - def train_one_step(self, viewer, timelapse_kwargs): + def train_one_step(self, + viewer: Callable[[np.ndarray, str], None] | None, + timelapse_kwargs: dict[T.Literal["input_a", "input_b", "output"], + str] | None) -> None: """ Running training on a batch of images for each side. Triggered from the training cycle in :class:`scripts.train.Train`. @@ -173,7 +230,7 @@ def train_one_step(self, viewer, timelapse_kwargs): Parameters ---------- - viewer: :func:`scripts.train.Train._show` + viewer: :func:`scripts.train.Train._show` or ``None`` The function that will display the preview image timelapse_kwargs: dict The keyword arguments for generating time-lapse previews. If a time-lapse preview is @@ -181,17 +238,16 @@ def train_one_step(self, viewer, timelapse_kwargs): the keys being `input_a`, `input_b`, `output`. """ self._model.state.increment_iterations() - logger.trace("Training one step: (iteration: %s)", self._model.iterations) - do_preview = viewer is not None - do_timelapse = timelapse_kwargs is not None + logger.trace("Training one step: (iteration: %s)", self._model.iterations) # type: ignore snapshot_interval = self._model.command_line_arguments.snapshot_interval do_snapshot = (snapshot_interval != 0 and self._model.iterations - 1 >= snapshot_interval and (self._model.iterations - 1) % snapshot_interval == 0) model_inputs, model_targets = self._feeder.get_batch() + try: - loss = self._model.model.train_on_batch(model_inputs, y=model_targets) + loss: list[float] = self._model.model.train_on_batch(model_inputs, y=model_targets) except tf_errors.ResourceExhaustedError as err: msg = ("You do not have enough GPU memory available to train the selected model at " "the selected settings. You can try a number of things:" @@ -203,43 +259,14 @@ def train_one_step(self, viewer, timelapse_kwargs): "\n4) Use a more lightweight model, or select the model's 'LowMem' option " "(in config) if it has one.") raise FaceswapError(msg) from err - except Exception as err: - if get_backend() == "amd": - # pylint:disable=import-outside-toplevel - from lib.plaidml_utils import is_plaidml_error - if (is_plaidml_error(err) and ( - "CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or - "enough memory for the current schedule" in str(err).lower())): - msg = ("You do not have enough GPU memory available to train the selected " - "model at the selected settings. You can try a number of things:" - "\n1) Close any other application that is using your GPU (web browsers " - "are particularly bad for this)." - "\n2) Lower the batchsize (the amount of images fed into the model " - "each iteration)." - "\n3) Use a more lightweight model, or select the model's 'LowMem' " - "option (in config) if it has one.") - raise FaceswapError(msg) from err - raise self._log_tensorboard(loss) loss = self._collate_and_store_loss(loss[1:]) self._print_loss(loss) - if do_snapshot: - self._model.snapshot() - - if do_preview: - self._feeder.generate_preview(do_preview) - self._samples.images = self._feeder.compile_sample(None) - samples = self._samples.show_sample() - if samples is not None: - viewer(samples, - "Training - 'S': Save Now. 'R': Refresh Preview. 'M': Toggle Mask. " - "'ENTER': Save and Quit") + self._model.io.snapshot() + self._update_viewers(viewer, timelapse_kwargs) - if do_timelapse: - self._timelapse.output_timelapse(timelapse_kwargs) - - def _log_tensorboard(self, loss): + def _log_tensorboard(self, loss: list[float]) -> None: """ Log current loss to Tensorboard log files Parameters @@ -249,21 +276,22 @@ def _log_tensorboard(self, loss): """ if not self._tensorboard: return - logger.trace("Updating TensorBoard log") + logger.trace("Updating TensorBoard log") # type: ignore logs = {log[0]: log[1] for log in zip(self._model.state.loss_names, loss)} - self._tensorboard.on_train_batch_end(self._model.iterations, logs=logs) - if get_tf_version() == 2.8: - # Bug in TF 2.8 where batch recording got deleted. - # ref: https://github.com/keras-team/keras/issues/16173 + # Bug in TF 2.8/2.9/2.10 where batch recording got deleted. + # ref: https://github.com/keras-team/keras/issues/16173 + with tf.summary.record_if(True), self._tensorboard._train_writer.as_default(): # noqa:E501 pylint:disable=protected-access,not-context-manager for name, value in logs.items(): tf.summary.scalar( "batch_" + name, value, - step=self._model._model._train_counter) # pylint:disable=protected-access + step=self._tensorboard._train_step) # pylint:disable=protected-access + # TODO revert this code if fixed in tensorflow + # self._tensorboard.on_train_batch_end(self._model.iterations, logs=logs) - def _collate_and_store_loss(self, loss): + def _collate_and_store_loss(self, loss: list[float]) -> list[float]: """ Collate the loss into totals for each side. The losses are summed into a total for each side. Loss totals are added to @@ -274,12 +302,13 @@ def _collate_and_store_loss(self, loss): Parameters ---------- loss: list - The list of loss ``floats`` for this iteration. + The list of loss ``floats`` for each side this iteration (excluding total combined + loss) Returns ------- list - List of 2 ``floats`` which is the total loss for each side + List of 2 ``floats`` which is the total loss for each side (eg sum of face + mask loss) Raises ------ @@ -295,10 +324,10 @@ def _collate_and_store_loss(self, loss): split = len(loss) // 2 combined_loss = [sum(loss[:split]), sum(loss[split:])] self._model.add_history(combined_loss) - logger.trace("original loss: %s, comibed_loss: %s", loss, combined_loss) + logger.trace("original loss: %s, combined_loss: %s", loss, combined_loss) # type: ignore return combined_loss - def _print_loss(self, loss): + def _print_loss(self, loss: list[float]) -> None: """ Outputs the loss for the current iteration to the console. Parameters @@ -317,7 +346,33 @@ def _print_loss(self, loss): logger.warning("Swallowed OS Error caused by Tensorflow distributed training. output " "line: %s, error: %s", output, str(err)) - def clear_tensorboard(self): + def _update_viewers(self, + viewer: Callable[[np.ndarray, str], None] | None, + timelapse_kwargs: dict[T.Literal["input_a", "input_b", "output"], + str] | None) -> None: + """ Update the preview viewer and timelapse output + + Parameters + ---------- + viewer: :func:`scripts.train.Train._show` or ``None`` + The function that will display the preview image + timelapse_kwargs: dict + The keyword arguments for generating time-lapse previews. If a time-lapse preview is + not required then this should be ``None``. Otherwise all values should be full paths + the keys being `input_a`, `input_b`, `output`. + """ + if viewer is not None: + self._samples.images = self._feeder.generate_preview() + samples = self._samples.show_sample() + if samples is not None: + viewer(samples, + "Training - 'S': Save Now. 'R': Refresh Preview. 'M': Toggle Mask. 'F': " + "Toggle Screen Fit-Actual Size. 'ENTER': Save and Quit") + + if timelapse_kwargs: + self._timelapse.output_timelapse(timelapse_kwargs) + + def clear_tensorboard(self) -> None: """ Stop Tensorboard logging. Tensorboard logging needs to be explicitly shutdown on training termination. Called from @@ -329,274 +384,6 @@ def clear_tensorboard(self): self._tensorboard.on_train_end(None) -class _Feeder(): - """ Handles the processing of a Batch for training the model and generating samples. - - Parameters - ---------- - images: dict - The list of full paths to the training images for this :class:`_Feeder` for each side - model: plugin from :mod:`plugins.train.model` - The selected model that will be running this trainer - batch_size: int - The size of the batch to be processed for each side at each iteration - config: :class:`lib.config.FaceswapConfig` - The configuration for this trainer - """ - def __init__(self, images, model, batch_size, config): - logger.debug("Initializing %s: num_images: %s, batch_size: %s, config: %s)", - self.__class__.__name__, len(images), batch_size, config) - self._model = model - self._images = images - self._config = config - self._target = {} - self._samples = {} - self._masks = {} - - self._feeds = {side: self._load_generator(idx).minibatch_ab(images[side], batch_size, side) - for idx, side in enumerate(("a", "b"))} - - self._display_feeds = dict(preview=self._set_preview_feed(), timelapse={}) - logger.debug("Initialized %s:", self.__class__.__name__) - - def _load_generator(self, output_index): - """ Load the :class:`~lib.training_data.TrainingDataGenerator` for this feeder. - - Parameters - ---------- - output_index: int - The output index from the model to get output shapes for - - Returns - ------- - :class:`~lib.training_data.TrainingDataGenerator` - The training data generator - """ - logger.debug("Loading generator") - input_size = self._model.model.input_shape[output_index][1] - output_shapes = self._model.output_shapes[output_index] - logger.debug("input_size: %s, output_shapes: %s", input_size, output_shapes) - generator = TrainingDataGenerator(input_size, - output_shapes, - self._model.coverage_ratio, - self._model.color_order, - not self._model.command_line_arguments.no_augment_color, - self._model.command_line_arguments.no_flip, - self._model.command_line_arguments.no_warp, - self._model.command_line_arguments.warp_to_landmarks, - self._config) - return generator - - def _set_preview_feed(self): - """ Set the preview feed for this feeder. - - Creates a generator from :class:`lib.training_data.TrainingDataGenerator` specifically - for previews for the feeder. - - Returns - ------- - dict - The side ("a" or "b") as key, :class:`~lib.training_data.TrainingDataGenerator` as - value. - """ - retval = {} - for idx, side in enumerate(("a", "b")): - logger.debug("Setting preview feed: (side: '%s')", side) - preview_images = self._config.get("preview_images", 14) - preview_images = min(max(preview_images, 2), 16) - batchsize = min(len(self._images[side]), preview_images) - retval[side] = self._load_generator(idx).minibatch_ab(self._images[side], - batchsize, - side, - do_shuffle=True, - is_preview=True) - logger.debug("Set preview feed. Batchsize: %s", batchsize) - return retval - - def get_batch(self): - """ Get the feed data and the targets for each training side for feeding into the model's - train function. - - Returns - ------- - model_inputs: list - The inputs to the model for each side A and B - model_targets: list - The targets for the model for each side A and B - """ - model_inputs = [] - model_targets = [] - for side in ("a", "b"): - batch = next(self._feeds[side]) - side_inputs = batch["feed"] - side_targets = self._compile_mask_targets(batch["targets"], - batch["masks"], - batch.get("additional_masks", None)) - if self._model.config["learn_mask"]: - side_targets = side_targets + [batch["masks"]] - logger.trace("side: %s, input_shapes: %s, target_shapes: %s", - side, [i.shape for i in side_inputs], [i.shape for i in side_targets]) - if get_backend() == "amd": - model_inputs.extend(side_inputs) - model_targets.extend(side_targets) - else: - model_inputs.append(side_inputs) - model_targets.append(side_targets) - return model_inputs, model_targets - - def _compile_mask_targets(self, targets, masks, additional_masks): - """ Compile the masks into the targets for penalized loss and for targeted learning. - - Penalized loss expects the target mask to be included for all outputs in the 4th channel - of the targets. Any additional masks are placed into subsequent channels for extraction - by the relevant loss functions. - - Parameters - ---------- - targets: list - The targets for the model, with the mask as the final entry in the list - masks: list - The masks for the model - additional_masks: list or ``None`` - Any additional masks for the model, or ``None`` if no additional masks are required - - Returns - ------- - list - The targets for the model with the mask compiled into the 4th channel. The original - mask is still output as the final item in the list - """ - if not self._model.config["penalized_mask_loss"] and additional_masks is None: - logger.trace("No masks to compile. Returning targets") - return targets - - if not self._model.config["penalized_mask_loss"] and additional_masks is not None: - masks = additional_masks - elif additional_masks is not None: - masks = np.concatenate((masks, additional_masks), axis=-1) - - for idx, tgt in enumerate(targets): - tgt_dim = tgt.shape[1] - if tgt_dim == masks.shape[1]: - add_masks = masks - else: - add_masks = np.array([cv2.resize(mask, (tgt_dim, tgt_dim)) - for mask in masks]) - if add_masks.ndim == 3: - add_masks = add_masks[..., None] - targets[idx] = np.concatenate((tgt, add_masks), axis=-1) - logger.trace("masks added to targets: %s", [tgt.shape for tgt in targets]) - return targets - - def generate_preview(self, do_preview): - """ Generate the preview images. - - Parameters - ---------- - do_preview: bool - Whether the previews should be generated. ``True`` if they should ``False`` if they - should not be generated, in which case currently stored previews should be deleted. - """ - if not do_preview: - self._samples = {} - self._target = {} - self._masks = {} - return - logger.debug("Generating preview") - for side in ("a", "b"): - batch = next(self._display_feeds["preview"][side]) - self._samples[side] = batch["samples"] - self._target[side] = batch["targets"][-1] - self._masks[side] = batch["masks"] - - def compile_sample(self, batch_size, samples=None, images=None, masks=None): - """ Compile the preview samples for display. - - Parameters - ---------- - batch_size: int - The requested batch size for each training iterations - samples: dict, optional - Dictionary for side "a", "b" of :class:`numpy.ndarray`. The sample images that should - be used for creating the preview. If ``None`` then the samples will be generated from - the internal random image generator. Default: ``None`` - images: dict, optional - Dictionary for side "a", "b" of :class:`numpy.ndarray`. The target images that should - be used for creating the preview. If ``None`` then the targets will be generated from - the internal random image generator. Default: ``None`` - masks: dict, optional - Dictionary for side "a", "b" of :class:`numpy.ndarray`. The masks that should be used - for creating the preview. If ``None`` then the masks will be generated from the - internal random image generator. Default: ``None`` - - Returns - ------- - list - The list of samples, targets and masks as :class:`numpy.ndarrays` for creating a - preview image - """ - num_images = self._config.get("preview_images", 14) - num_images = min(batch_size, num_images) if batch_size is not None else num_images - retval = {} - for side in ("a", "b"): - logger.debug("Compiling samples: (side: '%s', samples: %s)", side, num_images) - side_images = images[side] if images is not None else self._target[side] - side_masks = masks[side] if masks is not None else self._masks[side] - side_samples = samples[side] if samples is not None else self._samples[side] - retval[side] = [side_samples[0:num_images], - side_images[0:num_images], - side_masks[0:num_images]] - return retval - - def compile_timelapse_sample(self): - """ Compile the sample images for creating a time-lapse frame. - - Returns - ------- - dict - For sides "a" and "b"; The list of samples, targets and masks as - :class:`numpy.ndarrays` for creating a time-lapse frame - """ - batchsizes = [] - samples = {} - images = {} - masks = {} - for side in ("a", "b"): - batch = next(self._display_feeds["timelapse"][side]) - batchsizes.append(len(batch["samples"])) - samples[side] = batch["samples"] - images[side] = batch["targets"][-1] - masks[side] = batch["masks"] - batchsize = min(batchsizes) - sample = self.compile_sample(batchsize, samples=samples, images=images, masks=masks) - return sample - - def set_timelapse_feed(self, images, batch_size): - """ Set the time-lapse feed for this feeder. - - Creates a generator from :class:`lib.training_data.TrainingDataGenerator` specifically - for generating time-lapse previews for the feeder. - - Parameters - ---------- - images: list - The list of full paths to the images for creating the time-lapse for this - :class:`_Feeder` - batch_size: int - The number of images to be used to create the time-lapse preview. - """ - logger.debug("Setting time-lapse feed: (input_images: '%s', batch_size: %s)", - images, batch_size) - for idx, side in enumerate(("a", "b")): - self._display_feeds["timelapse"][side] = self._load_generator(idx).minibatch_ab( - images[side][:batch_size], - batch_size, - side, - do_shuffle=False, - is_timelapse=True) - logger.debug("Set time-lapse feed: %s", self._display_feeds["timelapse"]) - - class _Samples(): # pylint:disable=too-few-public-methods """ Compile samples for display for preview and time-lapse @@ -606,8 +393,10 @@ class _Samples(): # pylint:disable=too-few-public-methods The selected model that will be running this trainer coverage_ratio: float Ratio of face to be cropped out of the training image. - scaling: float, optional - The amount to scale the final preview image by. Default: `1.0` + mask_opacity: int + The opacity (as a percentage) to use for the mask overlay + mask_color: str + The hex RGB value to use the mask overlay Attributes ---------- @@ -616,26 +405,32 @@ class _Samples(): # pylint:disable=too-few-public-methods dictionary should contain 2 keys ("a" and "b") with the values being the training images for generating samples corresponding to each side. """ - def __init__(self, model, coverage_ratio, scaling=1.0): - logger.debug("Initializing %s: model: '%s', coverage_ratio: %s)", - self.__class__.__name__, model, coverage_ratio) + def __init__(self, + model: ModelBase, + coverage_ratio: float, + mask_opacity: int, + mask_color: str) -> None: + logger.debug("Initializing %s: model: '%s', coverage_ratio: %s, mask_opacity: %s, " + "mask_color: %s)", + self.__class__.__name__, model, coverage_ratio, mask_opacity, mask_color) self._model = model self._display_mask = model.config["learn_mask"] or model.config["penalized_mask_loss"] - self.images = {} + self.images: dict[T.Literal["a", "b"], list[np.ndarray]] = {} self._coverage_ratio = coverage_ratio - self._scaling = scaling + self._mask_opacity = mask_opacity / 100.0 + self._mask_color = np.array(hex_to_rgb(mask_color))[..., 2::-1] / 255. logger.debug("Initialized %s", self.__class__.__name__) - def toggle_mask_display(self): + def toggle_mask_display(self) -> None: """ Toggle the mask overlay on or off depending on user input. """ if not (self._model.config["learn_mask"] or self._model.config["penalized_mask_loss"]): return display_mask = not self._display_mask - print("\n") # Break to not garble loss output + print("") # Break to not garble loss output logger.info("Toggling mask display %s...", "on" if display_mask else "off") self._display_mask = display_mask - def show_sample(self): + def show_sample(self) -> np.ndarray: """ Compile a preview image. Returns @@ -644,49 +439,23 @@ def show_sample(self): A compiled preview image ready for display or saving """ logger.debug("Showing sample") - feeds = {} - figures = {} - headers = {} - for idx, side in enumerate(("a", "b")): - samples = self.images[side] - faces = samples[1] + feeds: dict[T.Literal["a", "b"], np.ndarray] = {} + for idx, side in enumerate(T.get_args(T.Literal["a", "b"])): + feed = self.images[side][0] input_shape = self._model.model.input_shape[idx][1:] - if input_shape[0] / faces.shape[1] != 1.0: - feeds[side] = self._resize_sample(side, faces, input_shape[0]) - feeds[side] = feeds[side].reshape((-1, ) + input_shape) + if input_shape[0] / feed.shape[1] != 1.0: + feeds[side] = self._resize_sample(side, feed, input_shape[0]) else: - feeds[side] = faces + feeds[side] = feed preds = self._get_predictions(feeds["a"], feeds["b"]) - - for side, samples in self.images.items(): - other_side = "a" if side == "b" else "b" - predictions = [preds[f"{side}_{side}"], - preds[f"{other_side}_{side}"]] - display = self._to_full_frame(side, samples, predictions) - headers[side] = self._get_headers(side, display[0].shape[1]) - figures[side] = np.stack([display[0], display[1], display[2], ], axis=1) - if self.images[side][0].shape[0] % 2 == 1: - figures[side] = np.concatenate([figures[side], - np.expand_dims(figures[side][0], 0)]) - - width = 4 - side_cols = width // 2 - if side_cols != 1: - headers = self._duplicate_headers(headers, side_cols) - - header = np.concatenate([headers["a"], headers["b"]], axis=1) - figure = np.concatenate([figures["a"], figures["b"]], axis=0) - height = int(figure.shape[0] / width) - figure = figure.reshape((width, height) + figure.shape[1:]) - figure = _stack_images(figure) - figure = np.concatenate((header, figure), axis=0) - - logger.debug("Compiled sample") - return np.clip(figure * 255, 0, 255).astype('uint8') + return self._compile_preview(preds) @classmethod - def _resize_sample(cls, side, sample, target_size): + def _resize_sample(cls, + side: T.Literal["a", "b"], + sample: np.ndarray, + target_size: int) -> np.ndarray: """ Resize a given image to the target size. Parameters @@ -705,24 +474,25 @@ def _resize_sample(cls, side, sample, target_size): """ scale = target_size / sample.shape[1] if scale == 1.0: - return sample + # cv2 complains if we don't do this :/ + return np.ascontiguousarray(sample) logger.debug("Resizing sample: (side: '%s', sample.shape: %s, target_size: %s, scale: %s)", side, sample.shape, target_size, scale) interpn = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA - retval = np.array([cv2.resize(img, (target_size, target_size), interpn) + retval = np.array([cv2.resize(img, (target_size, target_size), interpolation=interpn) for img in sample]) logger.debug("Resized sample: (side: '%s' shape: %s)", side, retval.shape) return retval - def _get_predictions(self, feed_a, feed_b): + def _get_predictions(self, feed_a: np.ndarray, feed_b: np.ndarray) -> dict[str, np.ndarray]: """ Feed the samples to the model and return predictions Parameters ---------- - feed_a: list - List of :class:`numpy.ndarray` of feed images for the "a" side - feed_a: list - List of :class:`numpy.ndarray` of feed images for the "b" side + feed_a: :class:`numpy.ndarray` + Feed images for the "a" side + feed_a: :class:`numpy.ndarray` + Feed images for the "b" side Returns ------- @@ -730,22 +500,23 @@ def _get_predictions(self, feed_a, feed_b): List of :class:`numpy.ndarray` of predictions received from the model """ logger.debug("Getting Predictions") - preds = {} - standard = self._model.model.predict([feed_a, feed_b]) - swapped = self._model.model.predict([feed_b, feed_a]) + preds: dict[str, np.ndarray] = {} - if self._model.config["learn_mask"] and get_backend() == "amd": - # Ravel results for plaidml - split = len(standard) // 2 - standard = [standard[:split], standard[split:]] - swapped = [swapped[:split], swapped[split:]] + # Calling model.predict() can lead to both VRAM and system memory leaks, so call model + # directly + standard = self._model.model([feed_a, feed_b]) + swapped = self._model.model([feed_b, feed_a]) if self._model.config["learn_mask"]: # Add mask to 4th channel of final output - standard = [np.concatenate(side[-2:], axis=-1) for side in standard] - swapped = [np.concatenate(side[-2:], axis=-1) for side in swapped] + standard = [np.concatenate(side[-2:], axis=-1) + for side in [[s.numpy() for s in t] for t in standard]] + swapped = [np.concatenate(side[-2:], axis=-1) + for side in [[s.numpy() for s in t] for t in swapped]] else: # Retrieve final output - standard = [side[-1] if isinstance(side, list) else side for side in standard] - swapped = [side[-1] if isinstance(side, list) else side for side in swapped] + standard = [side[-1] if isinstance(side, list) else side + for side in [t.numpy() for t in standard]] + swapped = [side[-1] if isinstance(side, list) else side + for side in [t.numpy() for t in swapped]] preds["a_a"] = standard[0] preds["b_b"] = standard[1] @@ -755,7 +526,51 @@ def _get_predictions(self, feed_a, feed_b): logger.debug("Returning predictions: %s", {key: val.shape for key, val in preds.items()}) return preds - def _to_full_frame(self, side, samples, predictions): + def _compile_preview(self, predictions: dict[str, np.ndarray]) -> np.ndarray: + """ Compile predictions and images into the final preview image. + + Parameters + ---------- + predictions: dict + The predictions from the model + + Returns + ------- + :class:`numpy.ndarry` + A compiled preview image ready for display or saving + """ + figures: dict[T.Literal["a", "b"], np.ndarray] = {} + headers: dict[T.Literal["a", "b"], np.ndarray] = {} + + for side, samples in self.images.items(): + other_side = "a" if side == "b" else "b" + preds = [predictions[f"{side}_{side}"], + predictions[f"{other_side}_{side}"]] + display = self._to_full_frame(side, samples, preds) + headers[side] = self._get_headers(side, display[0].shape[1]) + figures[side] = np.stack([display[0], display[1], display[2], ], axis=1) + if self.images[side][1].shape[0] % 2 == 1: + figures[side] = np.concatenate([figures[side], + np.expand_dims(figures[side][0], 0)]) + + width = 4 + if width // 2 != 1: + headers = self._duplicate_headers(headers, width // 2) + + header = np.concatenate([headers["a"], headers["b"]], axis=1) + figure = np.concatenate([figures["a"], figures["b"]], axis=0) + height = int(figure.shape[0] / width) + figure = figure.reshape((width, height) + figure.shape[1:]) + figure = _stack_images(figure) + figure = np.concatenate((header, figure), axis=0) + + logger.debug("Compiled sample") + return np.clip(figure * 255, 0, 255).astype('uint8') + + def _to_full_frame(self, + side: T.Literal["a", "b"], + samples: list[np.ndarray], + predictions: list[np.ndarray]) -> list[np.ndarray]: """ Patch targets and prediction images into images of model output size. Parameters @@ -763,7 +578,7 @@ def _to_full_frame(self, side, samples, predictions): side: {"a" or "b"} The side that these samples are for samples: list - List of :class:`numpy.ndarray` of feed images and target images + List of :class:`numpy.ndarray` of feed images and sample images predictions: list List of :class: `numpy.ndarray` of predictions from the model @@ -774,26 +589,31 @@ def _to_full_frame(self, side, samples, predictions): """ logger.debug("side: '%s', number of sample arrays: %s, prediction.shapes: %s)", side, len(samples), [pred.shape for pred in predictions]) - full, faces = samples[:2] + faces, full = samples[:2] if self._model.color_order.lower() == "rgb": # Switch color order for RGB model display full = full[..., ::-1] faces = faces[..., ::-1] - predictions = [pred[..., ::-1] if pred.shape[-1] == 3 else pred - for pred in predictions] + predictions = [pred[..., 2::-1] for pred in predictions] - full = self._process_full(side, full, predictions[0].shape[1], (0, 0, 255)) + full = self._process_full(side, full, predictions[0].shape[1], (0., 0., 1.0)) images = [faces] + predictions + if self._display_mask: images = self._compile_masked(images, samples[-1]) + elif self._model.config["learn_mask"]: + # Remove masks when learn mask is selected but mask toggle is off + images = [batch[..., :3] for batch in images] + images = [self._overlay_foreground(full.copy(), image) for image in images] - if self._scaling != 1.0: - new_size = int(images[0].shape[1] * self._scaling) - images = [self._resize_sample(side, image, new_size) for image in images] return images - def _process_full(self, side, images, prediction_size, color): + def _process_full(self, + side: T.Literal["a", "b"], + images: np.ndarray, + prediction_size: int, + color: tuple[float, float, float]) -> np.ndarray: """ Add a frame overlay to preview images indicating the region of interest. This applies the red border that appears in the preview images. @@ -834,8 +654,7 @@ def _process_full(self, side, images, prediction_size, color): logger.debug("Overlayed background. Shape: %s", images.shape) return images - @classmethod - def _compile_masked(cls, faces, masks): + def _compile_masked(self, faces: list[np.ndarray], masks: np.ndarray) -> list[np.ndarray]: """ Add the mask to the faces for masked preview. Places an opaque red layer over areas of the face that are masked out. @@ -853,34 +672,36 @@ def _compile_masked(cls, faces, masks): list List of :class:`numpy.ndarray` faces with the opaque mask layer applied """ - orig_masks = np.tile(1 - np.rint(masks), 3) - orig_masks[np.where((orig_masks == [1., 1., 1.]).all(axis=3))] = [0., 0., 1.] + orig_masks = 1. - masks + masks3: list[np.ndarray] | np.ndarray = [] if faces[-1].shape[-1] == 4: # Mask contained in alpha channel of predictions - pred_masks = [np.tile(1 - np.rint(face[..., -1])[..., None], 3) for face in faces[-2:]] - for swap_masks in pred_masks: - swap_masks[np.where((swap_masks == [1., 1., 1.]).all(axis=3))] = [0., 0., 1.] + pred_masks = [1. - face[..., -1][..., None] for face in faces[-2:]] faces[-2:] = [face[..., :-1] for face in faces[-2:]] masks3 = [orig_masks, *pred_masks] else: masks3 = np.repeat(np.expand_dims(orig_masks, axis=0), 3, axis=0) - retval = [np.array([cv2.addWeighted(img, 1.0, mask, 0.3, 0) - for img, mask in zip(previews, compiled_masks)]) - for previews, compiled_masks in zip(faces, masks3)] + retval: list[np.ndarray] = [] + overlays3 = np.ones_like(faces) * self._mask_color + for previews, overlays, compiled_masks in zip(faces, overlays3, masks3): + compiled_masks *= self._mask_opacity + overlays *= compiled_masks + previews *= (1. - compiled_masks) + retval.append(previews + overlays) logger.debug("masked shapes: %s", [faces.shape for faces in retval]) return retval @classmethod - def _overlay_foreground(cls, backgrounds, foregrounds): + def _overlay_foreground(cls, backgrounds: np.ndarray, foregrounds: np.ndarray) -> np.ndarray: """ Overlay the preview images into the center of the background images Parameters ---------- - backgrounds: list - List of :class:`numpy.ndarray` background images for placing the preview images onto - backgrounds: list - List of :class:`numpy.ndarray` preview images for placing onto the background images + backgrounds: :class:`numpy.ndarray` + Background images for placing the preview images onto + backgrounds: :class:`numpy.ndarray` + Preview images for placing onto the background images Returns ------- @@ -895,7 +716,7 @@ def _overlay_foreground(cls, backgrounds, foregrounds): return backgrounds @classmethod - def _get_headers(cls, side, width): + def _get_headers(cls, side: T.Literal["a", "b"], width: int) -> np.ndarray: """ Set header row for the final preview frame Parameters @@ -913,12 +734,11 @@ def _get_headers(cls, side, width): logger.debug("side: '%s', width: %s", side, width) titles = ("Original", "Swap") if side == "a" else ("Swap", "Original") - side = side.upper() height = int(width / 4.5) total_width = width * 3 logger.debug("height: %s, total_width: %s", height, total_width) font = cv2.FONT_HERSHEY_SIMPLEX - texts = [f"{titles[0]} ({side})", + texts = [f"{titles[0]} ({side.upper()})", f"{titles[0]} > {titles[0]}", f"{titles[0]} > {titles[1]}"] scaling = (width / 144) * 0.45 @@ -943,20 +763,22 @@ def _get_headers(cls, side, width): return header_box @classmethod - def _duplicate_headers(cls, headers, columns): + def _duplicate_headers(cls, + headers: dict[T.Literal["a", "b"], np.ndarray], + columns: int) -> dict[T.Literal["a", "b"], np.ndarray]: """ Duplicate headers for the number of columns displayed for each side. Parameters ---------- - headers: :class:`numpy.ndarray` - The header to be duplicated + headers: dict + The headers to be duplicated for each side columns: int The number of columns that the header needs to be duplicated for Returns ------- - :class:`numpy.ndarray` - The original headers duplicated by the number of columns + :class:dict + The original headers duplicated by the number of columns for each side """ for side, header in headers.items(): duped = tuple(header for _ in range(columns)) @@ -974,28 +796,38 @@ class _Timelapse(): # pylint:disable=too-few-public-methods The selected model that will be running this trainer coverage_ratio: float Ratio of face to be cropped out of the training image. - scaling: float, optional - The amount to scale the final preview image by. Default: `1.0` image_count: int The number of preview images to be displayed in the time-lapse - feeder: dict - The :class:`_Feeder` for generating the time-lapse images. + mask_opacity: int + The opacity (as a percentage) to use for the mask overlay + mask_color: str + The hex RGB value to use the mask overlay + feeder: :class:`~lib.training.generator.Feeder` + The feeder for generating the time-lapse images. image_paths: dict The full paths to the training images for each side of the model """ - def __init__(self, model, coverage_ratio, image_count, feeder, image_paths): + def __init__(self, + model: ModelBase, + coverage_ratio: float, + image_count: int, + mask_opacity: int, + mask_color: str, + feeder: Feeder, + image_paths: dict[T.Literal["a", "b"], list[str]]) -> None: logger.debug("Initializing %s: model: %s, coverage_ratio: %s, image_count: %s, " - "feeder: '%s', image_paths: %s)", self.__class__.__name__, model, - coverage_ratio, image_count, feeder, len(image_paths)) + "mask_opacity: %s, mask_color: %s, feeder: %s, image_paths: %s)", + self.__class__.__name__, model, coverage_ratio, image_count, mask_opacity, + mask_color, feeder, len(image_paths)) self._num_images = image_count - self._samples = _Samples(model, coverage_ratio) + self._samples = _Samples(model, coverage_ratio, mask_opacity, mask_color) self._model = model self._feeder = feeder self._image_paths = image_paths - self._output_file = None + self._output_file = "" logger.debug("Initialized %s", self.__class__.__name__) - def _setup(self, input_a=None, input_b=None, output=None): + def _setup(self, input_a: str, input_b: str, output: str) -> None: """ Setup the time-lapse folder locations and the time-lapse feed. Parameters @@ -1009,15 +841,15 @@ def _setup(self, input_a=None, input_b=None, output=None): default to the model folder """ logger.debug("Setting up time-lapse") - if output is None: - output = get_folder(os.path.join(str(self._model.model_dir), + if not output: + output = get_folder(os.path.join(str(self._model.io.model_dir), f"{self._model.name}_timelapse")) - self._output_file = str(output) + self._output_file = output logger.debug("Time-lapse output set to '%s'", self._output_file) # Rewrite paths to pull from the training images so mask and face data can be accessed - images = {} - for side, input_ in zip(("a", "b"), (input_a, input_b)): + images: dict[T.Literal["a", "b"], list[str]] = {} + for side, input_ in zip(T.get_args(T.Literal["a", "b"]), (input_a, input_b)): training_path = os.path.dirname(self._image_paths[side][0]) images[side] = [os.path.join(training_path, os.path.basename(pth)) for pth in get_image_paths(input_)] @@ -1028,7 +860,9 @@ def _setup(self, input_a=None, input_b=None, output=None): self._feeder.set_timelapse_feed(images, batchsize) logger.debug("Set up time-lapse") - def output_timelapse(self, timelapse_kwargs): + def output_timelapse(self, timelapse_kwargs: dict[T.Literal["input_a", + "input_b", + "output"], str]) -> None: """ Generate the time-lapse samples and output the created time-lapse to the specified output folder. @@ -1040,10 +874,10 @@ def output_timelapse(self, timelapse_kwargs): """ logger.debug("Ouputting time-lapse") if not self._output_file: - self._setup(**timelapse_kwargs) + self._setup(**T.cast(dict[str, str], timelapse_kwargs)) logger.debug("Getting time-lapse samples") - self._samples.images = self._feeder.compile_timelapse_sample() + self._samples.images = self._feeder.generate_preview(is_timelapse=True) logger.debug("Got time-lapse samples: %s", {side: len(images) for side, images in self._samples.images.items()}) @@ -1056,7 +890,7 @@ def output_timelapse(self, timelapse_kwargs): logger.debug("Created time-lapse: '%s'", filename) -def _stack_images(images): +def _stack_images(images: np.ndarray) -> np.ndarray: """ Stack images evenly for preview. Parameters diff --git a/plugins/train/trainer/original_defaults.py b/plugins/train/trainer/original_defaults.py index be760eff44..1cc45e07b1 100755 --- a/plugins/train/trainer/original_defaults.py +++ b/plugins/train/trainer/original_defaults.py @@ -54,6 +54,20 @@ rounding=2, min_max=(2, 16), group="evaluation"), + mask_opacity=dict( + default=30, + info="The opacity of the mask overlay in the training preview. Lower values are more " + "transparent.", + datatype=int, + rounding=2, + min_max=(0, 100), + group="evaluation"), + mask_color=dict( + default="#ff0000", + choices="colorchooser", + info="The RGB hex color to use for the mask overlay in the training preview.", + datatype=str, + group="evaluation"), zoom_amount=dict( default=5, info="Percentage amount to randomly zoom each training image in and out.", @@ -88,7 +102,7 @@ color_lightness=dict( default=30, info="Percentage amount to randomly alter the lightness of each training image.\n" - "NB: This is ignored if the 'no-flip' option is enabled", + "NB: This is ignored if the 'no-augment-color' option is enabled", datatype=int, rounding=1, min_max=(0, 75), @@ -96,8 +110,8 @@ color_ab=dict( default=8, info="Percentage amount to randomly alter the 'a' and 'b' colors of the L*a*b* color " - "space of each training image.\nNB: This is ignored if the 'no-flip' option is " - "enabled", + "space of each training image.\nNB: This is ignored if the 'no-augment-color' option" + "is enabled", datatype=int, rounding=1, min_max=(0, 50), diff --git a/requirements/_requirements_base.txt b/requirements/_requirements_base.txt index f3f9ff28f6..a2f04a17a1 100644 --- a/requirements/_requirements_base.txt +++ b/requirements/_requirements_base.txt @@ -1,17 +1,13 @@ -tqdm>=4.64 -psutil>=5.8.0 -numpy>=1.18.0 -opencv-python>=4.5.5.0 -pillow>=8.3.1 -scikit-learn>=1.0.2 -fastcluster>=1.2.4 -matplotlib>=3.5.1 -imageio>=2.9.0 -imageio-ffmpeg>=0.4.7 -ffmpy==0.2.3 -# Exclude badly numbered Python2 version of nvidia-ml-py -#nvidia-ml-py>=11.510,<300 -# Pin nvidida-ml-py to <11.515 until we know if bytes->str is an error or permanent change -nvidia-ml-py<11.515 -typing-extensions -pywin32>=228 ; sys_platform == "win32" +tqdm>=4.65 +psutil>=5.9.0 +numexpr>=2.8.7 +numpy>=1.26.0,<2.0.0 +opencv-python>=4.9.0.0 +pillow>=9.4.0,<10.0.0 +scikit-learn>=1.3.0 +fastcluster>=1.2.6 +matplotlib>=3.8.0 +imageio>=2.33.1 +imageio-ffmpeg>=0.4.9 +ffmpy>=0.3.0 +pywin32>=305 ; sys_platform == "win32" diff --git a/requirements/requirements_amd.txt b/requirements/requirements_amd.txt deleted file mode 100644 index 18a4dd6245..0000000000 --- a/requirements/requirements_amd.txt +++ /dev/null @@ -1,5 +0,0 @@ --r _requirements_base.txt -# tf2.2 is last version that tensorboard logging works with old Keras -protobuf>= 3.19.0,<3.20.0 # TF has started pulling in incompatible protobuf -tensorflow>=2.2.0,<2.3.0 -plaidml-keras==0.7.0 diff --git a/requirements/requirements_apple_silicon.txt b/requirements/requirements_apple_silicon.txt index 216b9ce522..5732337ea2 100644 --- a/requirements/requirements_apple_silicon.txt +++ b/requirements/requirements_apple_silicon.txt @@ -1,5 +1,7 @@ -protobuf>= 3.19.0,<3.20.0 # TF has started pulling in incompatible protobuf -tensorflow-macos>=2.8.0,<2.9.0 -tensorflow-deps>=2.8.0,<2.9.0 -tensorflow-metal>=0.4.0,<0.5.0 -libblas # Conda only +-r _requirements_base.txt +tensorflow-macos>=2.10.0,<2.11.0 +tensorflow-deps>=2.10.0,<2.11.0 +tensorflow-metal>=0.6.0,<0.7.0 +# These next 2 should have been installed, but some users complain of errors +decorator +cloudpickle diff --git a/requirements/requirements_cpu.txt b/requirements/requirements_cpu.txt index 2f03386aa7..873e3d3561 100644 --- a/requirements/requirements_cpu.txt +++ b/requirements/requirements_cpu.txt @@ -1,2 +1,2 @@ -r _requirements_base.txt -tensorflow>=2.4.0,<2.9.0 +tensorflow-cpu>=2.10.0,<2.11.0 diff --git a/requirements/requirements_directml.txt b/requirements/requirements_directml.txt new file mode 100644 index 0000000000..d7e0dbc227 --- /dev/null +++ b/requirements/requirements_directml.txt @@ -0,0 +1,4 @@ +-r _requirements_base.txt +tensorflow-cpu>=2.10.0,<2.11.0 +tensorflow-directml-plugin +comtypes diff --git a/requirements/requirements_nvidia.txt b/requirements/requirements_nvidia.txt index fa0364f6d9..45a558911f 100644 --- a/requirements/requirements_nvidia.txt +++ b/requirements/requirements_nvidia.txt @@ -1,3 +1,5 @@ -r _requirements_base.txt -tensorflow-gpu>=2.4.0,<2.9.0 +# Exclude badly numbered Python2 version of nvidia-ml-py +nvidia-ml-py>=12.535,<300 pynvx==1.0.0 ; sys_platform == "darwin" +tensorflow>=2.10.0,<2.11.0 diff --git a/requirements/requirements_rocm.txt b/requirements/requirements_rocm.txt new file mode 100644 index 0000000000..b23ce01590 --- /dev/null +++ b/requirements/requirements_rocm.txt @@ -0,0 +1,2 @@ +-r _requirements_base.txt +tensorflow-rocm>=2.10.0,<2.11.0 diff --git a/scripts/convert.py b/scripts/convert.py index 243a5d1124..7ae3f0eca6 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -1,10 +1,12 @@ #!/usr/bin python3 """ Main entry point to the convert process of FaceSwap """ - +from __future__ import annotations +from dataclasses import dataclass, field import logging import re import os import sys +import typing as T from threading import Event from time import sleep @@ -20,14 +22,47 @@ from lib.image import read_image_meta_batch, ImagesLoader from lib.multithreading import MultiThread, total_cpus from lib.queue_manager import queue_manager -from lib.utils import FaceswapError, get_backend, get_folder, get_image_paths -from plugins.extract.pipeline import Extractor, ExtractMedia +from lib.utils import FaceswapError, get_folder, get_image_paths, handle_deprecated_cliopts +from plugins.extract import ExtractMedia, Extractor from plugins.plugin_loader import PluginLoader -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Callable + from plugins.convert.writer._base import Output + from plugins.train.model._base import ModelBase + from lib.align.aligned_face import CenteringType + from lib.queue_manager import EventQueue + + +logger = logging.getLogger(__name__) + + +@dataclass +class ConvertItem: + """ A single frame with associated objects passing through the convert process. + Parameters + ---------- + input: :class:`~plugins.extract.extract_media.ExtractMedia` + The ExtractMedia object holding the :attr:`filename`, :attr:`image` and attr:`list` of + :class:`~lib.align.DetectedFace` objects loaded from disk + feed_faces: list, Optional + list of :class:`lib.align.AlignedFace` objects for feeding into the model's predict + function + reference_faces: list, Optional + list of :class:`lib.align.AlignedFace` objects at model output sized for using as reference + in the convert functionfor feeding into the model's predict + swapped_faces: :class:`np.ndarray` + The swapped faces returned from the model's predict function + """ + inbound: ExtractMedia + feed_faces: list[AlignedFace] = field(default_factory=list) + reference_faces: list[AlignedFace] = field(default_factory=list) + swapped_faces: np.ndarray = np.array([]) -class Convert(): # pylint:disable=too-few-public-methods + +class Convert(): """ The Faceswap Face Conversion Process. The conversion process is responsible for swapping the faces on source frames with the output @@ -45,24 +80,18 @@ class Convert(): # pylint:disable=too-few-public-methods The arguments to be passed to the convert process as generated from Faceswap's command line arguments """ - def __init__(self, arguments): + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s: (args: %s)", self.__class__.__name__, arguments) - self._args = arguments + self._args = handle_deprecated_cliopts(arguments) - self._patch_threads = None self._images = ImagesLoader(self._args.input_dir, fast_count=True) - self._alignments = Alignments(self._args, False, self._images.is_video) - if self._alignments.version == 1.0: - logger.error("The alignments file format has been updated since the given alignments " - "file was generated. You need to update the file to proceed.") - logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.") - sys.exit(1) - + self._alignments = self._get_alignments() self._opts = OptionalActions(self._args, self._images.file_list, self._alignments) self._add_queues() - self._disk_io = DiskIO(self._alignments, self._images, arguments) - self._predictor = Predict(self._disk_io.load_queue, self._queue_size, arguments) + self._predictor = Predict(self._queue_size, arguments) + self._disk_io = DiskIO(self._alignments, self._images, self._predictor, arguments) + self._predictor.launch(self._disk_io.load_queue) self._validate() get_folder(self._args.output_dir) @@ -74,21 +103,18 @@ def __init__(self, arguments): self._disk_io.pre_encode, arguments, configfile=configfile) - + self._patch_threads = self._get_threads() logger.debug("Initialized %s", self.__class__.__name__) @property - def _queue_size(self): - """ int: Size of the converter queues. 16 for single process otherwise 32 """ - if self._args.singleprocess: - retval = 16 - else: - retval = 32 + def _queue_size(self) -> int: + """ int: Size of the converter queues. 2 for single process otherwise 4 """ + retval = 2 if self._args.singleprocess or self._args.jobs == 1 else 4 logger.debug(retval) return retval @property - def _pool_processes(self): + def _pool_processes(self) -> int: """ int: The number of threads to run in parallel. Based on user options and number of available processors. """ if self._args.singleprocess: @@ -101,7 +127,25 @@ def _pool_processes(self): logger.debug(retval) return retval - def _validate(self): + def _get_alignments(self) -> Alignments: + """ Perform validation checks and legacy updates and return alignemnts object + + Returns + ------- + :class:`~lib.align.alignments.Alignments` + The alignments file for the extract job + """ + retval = Alignments(self._args, False, self._images.is_video) + if retval.version == 1.0: + logger.error("The alignments file format has been updated since the given alignments " + "file was generated. You need to update the file to proceed.") + logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.") + sys.exit(1) + + retval.update_legacy_has_source(os.path.basename(self._args.input_dir)) + return retval + + def _validate(self) -> None: """ Validate the Command Line Options. Ensure that certain cli selections are valid and won't result in an error. Checks: @@ -133,12 +177,12 @@ def _validate(self): if (not self._args.on_the_fly and self._args.mask_type not in ("none", "predicted") and not self._alignments.mask_is_valid(self._args.mask_type)): - msg = ("You have selected the Mask Type `{}` but at least one face does not have this " - "mask stored in the Alignments File.\nYou should generate the required masks " - "with the Mask Tool or set the Mask Type option to an existing Mask Type.\nA " - "summary of existing masks is as follows:\nTotal faces: {}, Masks: " - "{}".format(self._args.mask_type, self._alignments.faces_count, - self._alignments.mask_summary)) + msg = (f"You have selected the Mask Type `{self._args.mask_type}` but at least one " + "face does not have this mask stored in the Alignments File.\nYou should " + "generate the required masks with the Mask Tool or set the Mask Type option to " + "an existing Mask Type.\nA summary of existing masks is as follows:\nTotal " + f"faces: {self._alignments.faces_count}, " + f"Masks: {self._alignments.mask_summary}") raise FaceswapError(msg) if self._args.mask_type == "predicted" and not self._predictor.has_predicted_mask: @@ -154,16 +198,33 @@ def _validate(self): "mask. Selecting first available mask: '%s'", mask_type) self._args.mask_type = mask_type - def _add_queues(self): + def _add_queues(self) -> None: """ Add the queues for in, patch and out. """ logger.debug("Adding queues. Queue size: %s", self._queue_size) for qname in ("convert_in", "convert_out", "patch"): queue_manager.add_queue(qname, self._queue_size) - def process(self): + def _get_threads(self) -> MultiThread: + """ Get the threads for patching the converted faces onto the frames. + + Returns + :class:`lib.multithreading.MultiThread` + The threads that perform the patching of swapped faces onto the output frames + """ + save_queue = queue_manager.get_queue("convert_out") + patch_queue = queue_manager.get_queue("patch") + return MultiThread(self._converter.process, patch_queue, save_queue, + thread_count=self._pool_processes, name="patch") + + def process(self) -> None: """ The entry point for triggering the Conversion Process. Should only be called from :class:`lib.cli.launcher.ScriptExecutor` + + Raises + ------ + FaceswapError + Error raised if the process runs out of memory """ logger.debug("Starting Conversion") # queue_manager.debug_monitor(5) @@ -184,15 +245,10 @@ def process(self): "'singleprocess' flag (-sp) or lowering the number of parallel jobs (-j).") raise FaceswapError(msg) from err - def _convert_images(self): + def _convert_images(self) -> None: """ Start the multi-threaded patching process, monitor all threads for errors and join on completion. """ logger.debug("Converting images") - save_queue = queue_manager.get_queue("convert_out") - patch_queue = queue_manager.get_queue("patch") - self._patch_threads = MultiThread(self._converter.process, patch_queue, save_queue, - thread_count=self._pool_processes, name="patch") - self._patch_threads.start() while True: self._check_thread_error() @@ -206,11 +262,17 @@ def _convert_images(self): self._patch_threads.join() logger.debug("Putting EOF") - save_queue.put("EOF") + queue_manager.get_queue("convert_out").put("EOF") logger.debug("Converted images") - def _check_thread_error(self): - """ Monitor all running threads for errors, and raise accordingly. """ + def _check_thread_error(self) -> None: + """ Monitor all running threads for errors, and raise accordingly. + + Raises + ------ + Error + Re-raises any error encountered within any of the running threads + """ for thread in (self._predictor.thread, self._disk_io.load_thread, self._disk_io.save_thread, @@ -231,14 +293,20 @@ class DiskIO(): The alignments for the input video images: :class:`lib.image.ImagesLoader` The input images + predictor: :class:`Predict` + The object for generating predictions from the model arguments: :class:`argparse.Namespace` The arguments that were passed to the convert process as generated from Faceswap's command line arguments """ - def __init__(self, alignments, images, arguments): - logger.debug("Initializing %s: (alignments: %s, images: %s, arguments: %s)", - self.__class__.__name__, alignments, images, arguments) + def __init__(self, + alignments: Alignments, + images: ImagesLoader, + predictor: Predict, + arguments: Namespace) -> None: + logger.debug("Initializing %s: (alignments: %s, images: %s, predictor: %s, arguments: %s)", + self.__class__.__name__, alignments, images, predictor, arguments) self._alignments = alignments self._images = images self._args = arguments @@ -248,68 +316,74 @@ def __init__(self, alignments, images, arguments): # For frame skipping self._imageidxre = re.compile(r"(\d+)(?!.*\d\.)(?=\.\w+$)") self._frame_ranges = self._get_frame_ranges() - self._writer = self._get_writer() + self._writer = self._get_writer(predictor) # Extractor for on the fly detection self._extractor = self._load_extractor() - self._queues = dict(load=None, save=None) - self._threads = dict(oad=None, save=None) + self._queues: dict[T.Literal["load", "save"], EventQueue] = {} + self._threads: dict[T.Literal["load", "save"], MultiThread] = {} self._init_threads() logger.debug("Initialized %s", self.__class__.__name__) @property - def completion_event(self): + def completion_event(self) -> Event: """ :class:`event.Event`: Event is set when the DiskIO Save task is complete """ return self._completion_event @property - def draw_transparent(self): + def draw_transparent(self) -> bool: """ bool: ``True`` if the selected writer's Draw_transparent configuration item is set otherwise ``False`` """ return self._writer.config.get("draw_transparent", False) @property - def pre_encode(self): + def pre_encode(self) -> Callable[[np.ndarray, T.Any], list[bytes]] | None: """ python function: Selected writer's pre-encode function, if it has one, otherwise ``None`` """ dummy = np.zeros((20, 20, 3), dtype="uint8") test = self._writer.pre_encode(dummy) - retval = None if test is None else self._writer.pre_encode + retval: Callable | None = None if test is None else self._writer.pre_encode logger.debug("Writer pre_encode function: %s", retval) return retval @property - def save_thread(self): + def save_thread(self) -> MultiThread: """ :class:`lib.multithreading.MultiThread`: The thread that is running the image writing operation. """ return self._threads["save"] @property - def load_thread(self): + def load_thread(self) -> MultiThread: """ :class:`lib.multithreading.MultiThread`: The thread that is running the image loading operation. """ return self._threads["load"] @property - def load_queue(self): - """ :class:`queue.Queue()`: The queue that images and detected faces are loaded into. """ + def load_queue(self) -> EventQueue: + """ :class:`~lib.queue_manager.EventQueue`: The queue that images and detected faces are " + "loaded into. """ return self._queues["load"] @property - def _total_count(self): + def _total_count(self) -> int: """ int: The total number of frames to be converted """ if self._frame_ranges and not self._args.keep_unchanged: - retval = sum([fr[1] - fr[0] + 1 for fr in self._frame_ranges]) + retval = sum(fr[1] - fr[0] + 1 for fr in self._frame_ranges) else: retval = self._images.count logger.debug(retval) return retval # Initialization - def _get_writer(self): + def _get_writer(self, predictor: Predict) -> Output: """ Load the selected writer plugin. + Parameters + ---------- + predictor: :class:`Predict` + The object for generating predictions from the model + Returns ------- :mod:`plugins.convert.writer` plugin @@ -323,12 +397,14 @@ def _get_writer(self): args.append(self._args.input_dir) else: args.append(self._args.reference_video) + if self._args.writer == "patch": + args.append(predictor.output_size) logger.debug("Writer args: %s", args) configfile = self._args.configfile if hasattr(self._args, "configfile") else None return PluginLoader.get_converter("writer", self._args.writer)(*args, configfile=configfile) - def _get_frame_ranges(self): + def _get_frame_ranges(self) -> list[tuple[int, int]] | None: """ Obtain the frame ranges that are to be converted. If frame ranges have been specified, then split the command line formatted arguments into @@ -357,7 +433,7 @@ def _get_frame_ranges(self): raise FaceswapError("Frame Ranges specified, but could not determine frame numbering " "from filenames") - retval = list() + retval = [] for rng in self._args.frame_ranges: if "-" not in rng: raise FaceswapError("Frame Ranges not specified in the correct format") @@ -366,7 +442,7 @@ def _get_frame_ranges(self): logger.debug("frame ranges: %s", retval) return retval - def _load_extractor(self): + def _load_extractor(self) -> Extractor | None: """ Load the CV2-DNN Face Extractor Chain. For On-The-Fly conversion we use a CPU based extractor to avoid stacking the GPU. @@ -405,18 +481,18 @@ def _load_extractor(self): logger.debug("Loaded extractor") return extractor - def _init_threads(self): + def _init_threads(self) -> None: """ Initialize queues and threads. Creates the load and save queues and the load and save threads. Starts the threads. """ logger.debug("Initializing DiskIO Threads") - for task in ("load", "save"): + for task in T.get_args(T.Literal["load", "save"]): self._add_queue(task) self._start_thread(task) logger.debug("Initialized DiskIO Threads") - def _add_queue(self, task): + def _add_queue(self, task: T.Literal["load", "save"]) -> None: """ Add the queue to queue_manager and to :attr:`self._queues` for the given task. Parameters @@ -434,7 +510,7 @@ def _add_queue(self, task): self._queues[task] = queue_manager.get_queue(q_name) logger.debug("Added queue for task: '%s'", task) - def _start_thread(self, task): + def _start_thread(self, task: T.Literal["load", "save"]) -> None: """ Create the thread for the given task, add it it :attr:`self._threads` and start it. Parameters @@ -444,14 +520,14 @@ def _start_thread(self, task): """ logger.debug("Starting thread: '%s'", task) args = self._completion_event if task == "save" else None - func = getattr(self, "_{}".format(task)) + func = getattr(self, f"_{task}") io_thread = MultiThread(func, args, thread_count=1) io_thread.start() self._threads[task] = io_thread logger.debug("Started thread: '%s'", task) # Loading tasks - def _load(self, *args): # pylint: disable=unused-argument + def _load(self, *args) -> None: # pylint:disable=unused-argument """ Load frames from disk. In a background thread: @@ -474,23 +550,23 @@ def _load(self, *args): # pylint: disable=unused-argument continue if self._check_skipframe(filename): if self._args.keep_unchanged: - logger.trace("Saving unchanged frame: %s", filename) + logger.trace("Saving unchanged frame: %s", filename) # type:ignore out_file = os.path.join(self._args.output_dir, os.path.basename(filename)) self._queues["save"].put((out_file, image)) else: - logger.trace("Discarding frame: '%s'", filename) + logger.trace("Discarding frame: '%s'", filename) # type:ignore continue detected_faces = self._get_detected_faces(filename, image) - item = dict(filename=filename, image=image, detected_faces=detected_faces) - self._pre_process.do_actions(item) + item = ConvertItem(ExtractMedia(filename, image, detected_faces)) + self._pre_process.do_actions(item.inbound) self._queues["load"].put(item) logger.debug("Putting EOF") self._queues["load"].put("EOF") logger.debug("Load Images: Complete") - def _check_skipframe(self, filename): + def _check_skipframe(self, filename: str) -> bool: """ Check whether a frame is to be skipped. Parameters @@ -504,18 +580,18 @@ def _check_skipframe(self, filename): ``True`` if the frame is to be skipped otherwise ``False`` """ if not self._frame_ranges: - return None + return False indices = self._imageidxre.findall(filename) if not indices: logger.warning("Could not determine frame number. Frame will be converted: '%s'", filename) return False - idx = int(indices[0]) if indices else None + idx = int(indices[0]) skipframe = not any(map(lambda b: b[0] <= idx <= b[1], self._frame_ranges)) - logger.trace("idx: %s, skipframe: %s", idx, skipframe) + logger.trace("idx: %s, skipframe: %s", idx, skipframe) # type: ignore[attr-defined] return skipframe - def _get_detected_faces(self, filename, image): + def _get_detected_faces(self, filename: str, image: np.ndarray) -> list[DetectedFace]: """ Return the detected faces for the given image. If we have an alignments file, then the detected faces are created from that file. If @@ -533,15 +609,15 @@ def _get_detected_faces(self, filename, image): list List of :class:`lib.align.DetectedFace` objects """ - logger.trace("Getting faces for: '%s'", filename) + logger.trace("Getting faces for: '%s'", filename) # type:ignore if not self._extractor: detected_faces = self._alignments_faces(os.path.basename(filename), image) else: detected_faces = self._detect_faces(filename, image) - logger.trace("Got %s faces for: '%s'", len(detected_faces), filename) + logger.trace("Got %s faces for: '%s'", len(detected_faces), filename) # type:ignore return detected_faces - def _alignments_faces(self, frame_name, image): + def _alignments_faces(self, frame_name: str, image: np.ndarray) -> list[DetectedFace]: """ Return detected faces from an alignments file. Parameters @@ -557,10 +633,10 @@ def _alignments_faces(self, frame_name, image): List of :class:`lib.align.DetectedFace` objects """ if not self._check_alignments(frame_name): - return list() + return [] faces = self._alignments.get_faces_in_frame(frame_name) - detected_faces = list() + detected_faces = [] for rawface in faces: face = DetectedFace() @@ -568,7 +644,7 @@ def _alignments_faces(self, frame_name, image): detected_faces.append(face) return detected_faces - def _check_alignments(self, frame_name): + def _check_alignments(self, frame_name: str) -> bool: """ Ensure that we have alignments for the current frame. If we have no alignments for this image, skip it and output a message. @@ -585,11 +661,10 @@ def _check_alignments(self, frame_name): """ have_alignments = self._alignments.frame_exists(frame_name) if not have_alignments: - tqdm.write("No alignment found for {}, " - "skipping".format(frame_name)) + tqdm.write(f"No alignment found for {frame_name}, skipping") return have_alignments - def _detect_faces(self, filename, image): + def _detect_faces(self, filename: str, image: np.ndarray) -> list[DetectedFace]: """ Extract the face from a frame for On-The-Fly conversion. Pulls detected faces out of the Extraction pipeline. @@ -606,12 +681,13 @@ def _detect_faces(self, filename, image): list List of :class:`lib.align.DetectedFace` objects """ + assert self._extractor is not None self._extractor.input_queue.put(ExtractMedia(filename, image)) faces = next(self._extractor.detected_faces()) return faces.detected_faces # Saving tasks - def _save(self, completion_event): + def _save(self, completion_event: Event) -> None: """ Save the converted images. Puts the selected writer into a background thread and feeds it from the output of the @@ -630,7 +706,7 @@ def _save(self, completion_event): if self._queues["save"].shutdown.is_set(): logger.debug("Save Queue: Stop signal received. Terminating") break - item = self._queues["save"].get() + item: tuple[str, np.ndarray | bytes] | T.Literal["EOF"] = self._queues["save"].get() if item == "EOF": logger.debug("EOF Received") break @@ -638,6 +714,7 @@ def _save(self, completion_event): # Write out preview image for the GUI every 10 frames if writing to stream if write_preview and idx % 10 == 0 and not os.path.exists(preview_image): logger.debug("Writing GUI Preview image: '%s'", preview_image) + assert isinstance(image, np.ndarray) cv2.imwrite(preview_image, image) self._writer.write(filename, image) self._writer.close() @@ -650,19 +727,17 @@ class Predict(): Parameters ---------- - in_queue: :class:`queue.Queue` - The queue that contains images and detected faces for feeding the model queue_size: int The maximum size of the input queue arguments: :class:`argparse.Namespace` The arguments that were passed to the convert process as generated from Faceswap's command line arguments """ - def __init__(self, in_queue, queue_size, arguments): - logger.debug("Initializing %s: (args: %s, queue_size: %s, in_queue: %s)", - self.__class__.__name__, arguments, queue_size, in_queue) + def __init__(self, queue_size: int, arguments: Namespace) -> None: + logger.debug("Initializing %s: (args: %s, queue_size: %s)", + self.__class__.__name__, arguments, queue_size) self._args = arguments - self._in_queue = in_queue + self._in_queue: EventQueue | None = None self._out_queue = queue_manager.get_queue("patch") self._serializer = get_serializer("json") self._faces_count = 0 @@ -674,56 +749,58 @@ def __init__(self, in_queue, queue_size, arguments): self._coverage_ratio = self._model.coverage_ratio self._centering = self._model.config["centering"] - self._thread = self._launch_predictor() + self._thread: MultiThread | None = None logger.debug("Initialized %s: (out_queue: %s)", self.__class__.__name__, self._out_queue) @property - def thread(self): + def thread(self) -> MultiThread: """ :class:`~lib.multithreading.MultiThread`: The thread that is running the prediction function from the Faceswap model. """ + assert self._thread is not None return self._thread @property - def in_queue(self): - """ :class:`queue.Queue`: The input queue to the predictor. """ + def in_queue(self) -> EventQueue: + """ :class:`~lib.queue_manager.EventQueue`: The input queue to the predictor. """ + assert self._in_queue is not None return self._in_queue @property - def out_queue(self): - """ :class:`queue.Queue`: The output queue from the predictor. """ + def out_queue(self) -> EventQueue: + """ :class:`~lib.queue_manager.EventQueue`: The output queue from the predictor. """ return self._out_queue @property - def faces_count(self): + def faces_count(self) -> int: """ int: The total number of faces seen by the Predictor. """ return self._faces_count @property - def verify_output(self): + def verify_output(self) -> bool: """ bool: ``True`` if multiple faces have been found in frames, otherwise ``False``. """ return self._verify_output @property - def coverage_ratio(self): + def coverage_ratio(self) -> float: """ float: The coverage ratio that the model was trained at. """ return self._coverage_ratio @property - def centering(self): - """ str: The centering that the model was trained on (`"face"` or `"legacy"`) """ + def centering(self) -> CenteringType: + """ str: The centering that the model was trained on (`"head", "face"` or `"legacy"`) """ return self._centering @property - def has_predicted_mask(self): + def has_predicted_mask(self) -> bool: """ bool: ``True`` if the model was trained to learn a mask, otherwise ``False``. """ return bool(self._model.config.get("learn_mask", False)) @property - def output_size(self): + def output_size(self) -> int: """ int: The size in pixels of the Faceswap model output. """ return self._sizes["output"] - def _get_io_sizes(self): + def _get_io_sizes(self) -> dict[str, int]: """ Obtain the input size and output size of the model. Returns @@ -735,11 +812,11 @@ def _get_io_sizes(self): input_shape = [input_shape] if not isinstance(input_shape, list) else input_shape output_shape = self._model.model.output_shape output_shape = [output_shape] if not isinstance(output_shape, list) else output_shape - retval = dict(input=input_shape[0][1], output=output_shape[-1][1]) + retval = {"input": input_shape[0][1], "output": output_shape[-1][1]} logger.debug(retval) return retval - def _load_model(self): + def _load_model(self) -> ModelBase: """ Load the Faceswap model. Returns @@ -757,7 +834,7 @@ def _load_model(self): logger.debug("Loaded Model") return model - def _get_batchsize(self, queue_size): + def _get_batchsize(self, queue_size: int) -> int: """ Get the batch size for feeding the model. Sets the batch size to 1 if inference is being run on CPU, otherwise the minimum of the @@ -777,15 +854,13 @@ def _get_batchsize(self, queue_size): is_cpu = GPUStats().device_count == 0 batchsize = 1 if is_cpu else self._model.config["convert_batchsize"] batchsize = min(queue_size, batchsize) - logger.debug("Batchsize: %s", batchsize) logger.debug("Got batchsize: %s", batchsize) return batchsize - def _get_model_name(self, model_dir): + def _get_model_name(self, model_dir: str) -> str: """ Return the name of the Faceswap model used. - If a "trainer" option has been selected in the command line arguments, use that value, - otherwise retrieve the name of the model from the model's state file. + Retrieve the name of the model from the model's state file. Parameters ---------- @@ -798,42 +873,36 @@ def _get_model_name(self, model_dir): The name of the Faceswap model being used. """ - if hasattr(self._args, "trainer") and self._args.trainer: - logger.debug("Trainer name provided: '%s'", self._args.trainer) - return self._args.trainer - - statefile = [fname for fname in os.listdir(str(model_dir)) - if fname.endswith("_state.json")] - if len(statefile) != 1: - raise FaceswapError("There should be 1 state file in your model folder. {} were " - "found. Specify a trainer with the '-t', '--trainer' " - "option.".format(len(statefile))) - statefile = os.path.join(str(model_dir), statefile[0]) + statefiles = [fname for fname in os.listdir(str(model_dir)) + if fname.endswith("_state.json")] + if len(statefiles) != 1: + raise FaceswapError("There should be 1 state file in your model folder. " + f"{len(statefiles)} were found.") + statefile = os.path.join(str(model_dir), statefiles[0]) state = self._serializer.load(statefile) trainer = state.get("name", None) if not trainer: - raise FaceswapError("Trainer name could not be read from state file. " - "Specify a trainer with the '-t', '--trainer' option.") + raise FaceswapError("Trainer name could not be read from state file.") logger.debug("Trainer from state file: '%s'", trainer) return trainer - def _launch_predictor(self): + def launch(self, load_queue: EventQueue) -> None: """ Launch the prediction process in a background thread. Starts the prediction thread and returns the thread. - Returns - ------- - :class:`~lib.multithreading.MultiThread` - The started Faceswap model prediction thread. + Parameters + ---------- + load_queue: :class:`~lib.queue_manager.EventQueue` + The queue that contains images and detected faces for feeding the model """ - thread = MultiThread(self._predict_faces, thread_count=1) - thread.start() - return thread + self._in_queue = load_queue + self._thread = MultiThread(self._predict_faces, thread_count=1) + self._thread.start() - def _predict_faces(self): + def _predict_faces(self) -> None: """ Run Prediction on the Faceswap model in a background thread. Reads from the :attr:`self._in_queue`, prepares images for prediction @@ -841,64 +910,77 @@ def _predict_faces(self): """ faces_seen = 0 consecutive_no_faces = 0 - batch = list() - is_amd = get_backend() == "amd" + batch: list[ConvertItem] = [] + assert self._in_queue is not None while True: - item = self._in_queue.get() - if item != "EOF": - logger.trace("Got from queue: '%s'", item["filename"]) - faces_count = len(item["detected_faces"]) - - # Safety measure. If a large stream of frames appear that do not have faces, - # these will stack up into RAM. Keep a count of consecutive frames with no faces. - # If self._batchsize number of frames appear, force the current batch through - # to clear RAM. - consecutive_no_faces = consecutive_no_faces + 1 if faces_count == 0 else 0 - self._faces_count += faces_count - if faces_count > 1: - self._verify_output = True - logger.verbose("Found more than one face in an image! '%s'", - os.path.basename(item["filename"])) - - self.load_aligned(item) - - faces_seen += faces_count - batch.append(item) - - if item != "EOF" and (faces_seen < self._batchsize and - consecutive_no_faces < self._batchsize): - logger.trace("Continuing. Current batchsize: %s, consecutive_no_faces: %s", - faces_seen, consecutive_no_faces) + item: T.Literal["EOF"] | ConvertItem = self._in_queue.get() + if item == "EOF": + logger.debug("EOF Received") + if batch: # Process out any remaining items + self._process_batch(batch, faces_seen) + break + logger.trace("Got from queue: '%s'", item.inbound.filename) # type:ignore + faces_count = len(item.inbound.detected_faces) + + # Safety measure. If a large stream of frames appear that do not have faces, + # these will stack up into RAM. Keep a count of consecutive frames with no faces. + # If self._batchsize number of frames appear, force the current batch through + # to clear RAM. + consecutive_no_faces = consecutive_no_faces + 1 if faces_count == 0 else 0 + self._faces_count += faces_count + if faces_count > 1: + self._verify_output = True + logger.verbose("Found more than one face in an image! '%s'", # type:ignore + os.path.basename(item.inbound.filename)) + + self.load_aligned(item) + faces_seen += faces_count + + batch.append(item) + + if faces_seen < self._batchsize and consecutive_no_faces < self._batchsize: + logger.trace("Continuing. Current batchsize: %s, " # type:ignore + "consecutive_no_faces: %s", faces_seen, consecutive_no_faces) continue - if batch: - logger.trace("Batching to predictor. Frames: %s, Faces: %s", - len(batch), faces_seen) - feed_batch = [feed_face for item in batch - for feed_face in item["feed_faces"]] - if faces_seen != 0: - feed_faces = self._compile_feed_faces(feed_batch) - batch_size = None - if is_amd and feed_faces.shape[0] != self._batchsize: - logger.verbose("Fallback to BS=1") - batch_size = 1 - predicted = self._predict(feed_faces, batch_size) - else: - predicted = list() - - self._queue_out_frames(batch, predicted) + self._process_batch(batch, faces_seen) consecutive_no_faces = 0 faces_seen = 0 - batch = list() - if item == "EOF": - logger.debug("EOF Received") - break + batch = [] + logger.debug("Putting EOF") self._out_queue.put("EOF") logger.debug("Load queue complete") - def load_aligned(self, item): + def _process_batch(self, batch: list[ConvertItem], faces_seen: int): + """ Predict faces on the given batch of images and queue out to patch thread + + Parameters + ---------- + batch: list + List of :class:`ConvertItem` objects for the current batch + faces_seen: int + The number of faces seen in the current batch + + Returns + ------- + :class:`np.narray` + The predicted faces for the current batch + """ + logger.trace("Batching to predictor. Frames: %s, Faces: %s", # type:ignore + len(batch), faces_seen) + feed_batch = [feed_face for item in batch for feed_face in item.feed_faces] + if faces_seen != 0: + feed_faces = self._compile_feed_faces(feed_batch) + batch_size = None + predicted = self._predict(feed_faces, batch_size) + else: + predicted = np.array([]) + + self._queue_out_frames(batch, predicted) + + def load_aligned(self, item: ConvertItem) -> None: """ Load the model's feed faces and the reference output faces. For each detected face in the incoming item, load the feed face and reference face @@ -906,18 +988,15 @@ def load_aligned(self, item): Parameters ---------- - item: dict - The incoming image, list of :class:`~lib.align.DetectedFace` objects and list of - :class:`~lib.align.AlignedFace` objects for the feed face(s) and list of - :class:`~lib.align.AlignedFace` objects for the reference face(s) - + item: :class:`ConvertMedia` + The convert media object, containing the ExctractMedia for the current image """ - logger.trace("Loading aligned faces: '%s'", item["filename"]) + logger.trace("Loading aligned faces: '%s'", item.inbound.filename) # type:ignore feed_faces = [] reference_faces = [] - for detected_face in item["detected_faces"]: + for detected_face in item.inbound.detected_faces: feed_face = AlignedFace(detected_face.landmarks_xy, - image=item["image"], + image=item.inbound.image, centering=self._centering, size=self._sizes["input"], coverage_ratio=self._coverage_ratio, @@ -926,18 +1005,18 @@ def load_aligned(self, item): reference_faces.append(feed_face) else: reference_faces.append(AlignedFace(detected_face.landmarks_xy, - image=item["image"], + image=item.inbound.image, centering=self._centering, size=self._sizes["output"], coverage_ratio=self._coverage_ratio, dtype="float32")) feed_faces.append(feed_face) - item["feed_faces"] = feed_faces - item["reference_faces"] = reference_faces - logger.trace("Loaded aligned faces: '%s'", item["filename"]) + item.feed_faces = feed_faces + item.reference_faces = reference_faces + logger.trace("Loaded aligned faces: '%s'", item.inbound.filename) # type:ignore @staticmethod - def _compile_feed_faces(feed_faces): + def _compile_feed_faces(feed_faces: list[AlignedFace]) -> np.ndarray: """ Compile a batch of faces for feeding into the Predictor. Parameters @@ -950,12 +1029,13 @@ def _compile_feed_faces(feed_faces): :class:`numpy.ndarray` A batch of faces ready for feeding into the Faceswap model. """ - logger.trace("Compiling feed face. Batchsize: %s", len(feed_faces)) - retval = np.stack([feed_face.face[..., :3] for feed_face in feed_faces]) / 255.0 - logger.trace("Compiled Feed faces. Shape: %s", retval.shape) + logger.trace("Compiling feed face. Batchsize: %s", len(feed_faces)) # type:ignore + retval = np.stack([T.cast(np.ndarray, feed_face.face)[..., :3] + for feed_face in feed_faces]) / 255.0 + logger.trace("Compiled Feed faces. Shape: %s", retval.shape) # type:ignore return retval - def _predict(self, feed_faces, batch_size=None): + def _predict(self, feed_faces: np.ndarray, batch_size: int | None = None) -> np.ndarray: """ Run the Faceswap models' prediction function. Parameters @@ -971,32 +1051,33 @@ def _predict(self, feed_faces, batch_size=None): :class:`numpy.ndarray` The swapped faces for the given batch """ - logger.trace("Predicting: Batchsize: %s", len(feed_faces)) + logger.trace("Predicting: Batchsize: %s", len(feed_faces)) # type:ignore if self._model.color_order.lower() == "rgb": feed_faces = feed_faces[..., ::-1] feed = [feed_faces] - logger.trace("Input shape(s): %s", [item.shape for item in feed]) + logger.trace("Input shape(s): %s", [item.shape for item in feed]) # type:ignore - predicted = self._model.model.predict(feed, batch_size=batch_size) - predicted = predicted if isinstance(predicted, list) else [predicted] + inbound = self._model.model.predict(feed, verbose=0, batch_size=batch_size) + predicted: list[np.ndarray] = inbound if isinstance(inbound, list) else [inbound] if self._model.color_order.lower() == "rgb": predicted[0] = predicted[0][..., ::-1] - logger.trace("Output shape(s): %s", [predict.shape for predict in predicted]) + logger.trace("Output shape(s): %s", # type:ignore + [predict.shape for predict in predicted]) # Only take last output(s) if predicted[-1].shape[-1] == 1: # Merge mask to alpha channel - predicted = np.concatenate(predicted[-2:], axis=-1).astype("float32") + retval = np.concatenate(predicted[-2:], axis=-1).astype("float32") else: - predicted = predicted[-1].astype("float32") + retval = predicted[-1].astype("float32") - logger.trace("Final shape: %s", predicted.shape) - return predicted + logger.trace("Final shape: %s", retval.shape) # type:ignore + return retval - def _queue_out_frames(self, batch, swapped_faces): + def _queue_out_frames(self, batch: list[ConvertItem], swapped_faces: np.ndarray) -> None: """ Compile the batch back to original frames and put to the Out Queue. For batching, faces are split away from their frames. This compiles all detected faces @@ -1009,21 +1090,20 @@ def _queue_out_frames(self, batch, swapped_faces): swapped_faces: :class:`numpy.ndarray` The predictions returned from the model's predict function """ - logger.trace("Queueing out batch. Batchsize: %s", len(batch)) + logger.trace("Queueing out batch. Batchsize: %s", len(batch)) # type:ignore pointer = 0 for item in batch: - num_faces = len(item["detected_faces"]) - if num_faces == 0: - item["swapped_faces"] = np.array(list()) - else: - item["swapped_faces"] = swapped_faces[pointer:pointer + num_faces] - - logger.trace("Putting to queue. ('%s', detected_faces: %s, reference_faces: %s, " - "swapped_faces: %s)", item["filename"], len(item["detected_faces"]), - len(item["reference_faces"]), item["swapped_faces"].shape[0]) + num_faces = len(item.inbound.detected_faces) + if num_faces != 0: + item.swapped_faces = swapped_faces[pointer:pointer + num_faces] + + logger.trace("Putting to queue. ('%s', detected_faces: %s, " # type:ignore + "reference_faces: %s, swapped_faces: %s)", item.inbound.filename, + len(item.inbound.detected_faces), len(item.reference_faces), + item.swapped_faces.shape[0]) pointer += num_faces self._out_queue.put(batch) - logger.trace("Queued out batch. Batchsize: %s", len(batch)) + logger.trace("Queued out batch. Batchsize: %s", len(batch)) # type:ignore class OptionalActions(): # pylint:disable=too-few-public-methods @@ -1041,8 +1121,10 @@ class OptionalActions(): # pylint:disable=too-few-public-methods alignments: :class:`lib.align.Alignments` The alignments file for this conversion """ - - def __init__(self, arguments, input_images, alignments): + def __init__(self, + arguments: Namespace, + input_images: list[np.ndarray], + alignments: Alignments) -> None: logger.debug("Initializing %s", self.__class__.__name__) self._args = arguments self._input_images = input_images @@ -1052,7 +1134,7 @@ def __init__(self, arguments, input_images, alignments): logger.debug("Initialized %s", self.__class__.__name__) # SKIP FACES # - def _remove_skipped_faces(self): + def _remove_skipped_faces(self) -> None: """ If the user has specified an input aligned directory, remove any non-matching faces from the alignments file. """ logger.debug("Filtering Faces") @@ -1064,7 +1146,7 @@ def _remove_skipped_faces(self): self._alignments.filter_faces(accept_dict, filter_out=False) logger.info("Faces filtered out: %s", pre_face_count - self._alignments.faces_count) - def _get_face_metadata(self): + def _get_face_metadata(self) -> dict[str, list[int]]: """ Check for the existence of an aligned directory for identifying which faces in the target frames should be swapped. If it exists, scan the folder for face's metadata @@ -1073,12 +1155,12 @@ def _get_face_metadata(self): dict Dictionary of source frame names with a list of associated face indices to be skipped """ - retval = dict() + retval: dict[str, list[int]] = {} input_aligned_dir = self._args.input_aligned_dir if input_aligned_dir is None: - logger.verbose("Aligned directory not specified. All faces listed in the " - "alignments file will be converted") + logger.verbose("Aligned directory not specified. All faces listed in " # type:ignore + "the alignments file will be converted") return retval if not os.path.isdir(input_aligned_dir): logger.warning("Aligned directory not found. All faces listed in the " @@ -1100,13 +1182,13 @@ def _get_face_metadata(self): data = update_legacy_png_header(fullpath, self._alignments) if not data: raise FaceswapError( - "Some of the faces being passed in from '{}' could not be matched to the " - "alignments file '{}'\nPlease double check your sources and try " - "again.".format(input_aligned_dir, self._alignments.file)) + f"Some of the faces being passed in from '{input_aligned_dir}' could not " + f"be matched to the alignments file '{self._alignments.file}'\n" + "Please double check your sources and try again.") meta = data["source"] else: meta = metadata["itxt"]["source"] - retval.setdefault(meta["source_filename"], list()).append(meta["face_index"]) + retval.setdefault(meta["source_filename"], []).append(meta["face_index"]) if not retval: raise FaceswapError("Aligned directory is empty, no faces will be converted!") diff --git a/scripts/extract.py b/scripts/extract.py index 91f4fa1219..da9edf1d15 100644 --- a/scripts/extract.py +++ b/scripts/extract.py @@ -2,29 +2,32 @@ """ Main entry point to the extract process of FaceSwap """ from __future__ import annotations - import logging import os import sys -from typing import TYPE_CHECKING, Optional +import typing as T + +from argparse import Namespace +from multiprocessing import Process +import numpy as np from tqdm import tqdm +from lib.align.alignments import PNGHeaderDict -from lib.image import encode_image, generate_thumbnail, ImagesLoader, ImagesSaver +from lib.image import encode_image, generate_thumbnail, ImagesLoader, ImagesSaver, read_image_meta from lib.multithreading import MultiThread -from lib.utils import get_folder -from plugins.extract.pipeline import Extractor, ExtractMedia +from lib.utils import get_folder, handle_deprecated_cliopts, IMAGE_EXTENSIONS, VIDEO_EXTENSIONS +from plugins.extract import ExtractMedia, Extractor from scripts.fsmedia import Alignments, PostProcess, finalize -if TYPE_CHECKING: - import argparse +if T.TYPE_CHECKING: + from lib.align.alignments import PNGHeaderAlignmentsDict +# tqdm.monitor_interval = 0 # workaround for TqdmSynchronisationWarning # TODO? +logger = logging.getLogger(__name__) -tqdm.monitor_interval = 0 # workaround for TqdmSynchronisationWarning -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - -class Extract(): # pylint:disable=too-few-public-methods +class Extract(): """ The Faceswap Face Extraction Process. The extraction process is responsible for detecting faces in a series of images/video, aligning @@ -42,78 +45,105 @@ class Extract(): # pylint:disable=too-few-public-methods The arguments to be passed to the extraction process as generated from Faceswap's command line arguments """ - def __init__(self, arguments: argparse.Namespace) -> None: + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments) - self._args = arguments - self._output_dir = None if self._args.skip_saving_faces else get_folder( - self._args.output_dir) + self._args = handle_deprecated_cliopts(arguments) + self._input_locations = self._get_input_locations() + self._validate_batchmode() - logger.info("Output Directory: %s", self._args.output_dir) - self._images = ImagesLoader(self._args.input_dir, fast_count=True) - self._alignments = Alignments(self._args, True, self._images.is_video) - - self._existing_count = 0 - self._set_skip_list() - - self._post_process = PostProcess(arguments) configfile = self._args.configfile if hasattr(self._args, "configfile") else None normalization = None if self._args.normalization == "none" else self._args.normalization - maskers = ["components", "extended"] maskers += self._args.masker if self._args.masker else [] + recognition = ("vgg_face2" + if arguments.identity or arguments.filter or arguments.nfilter + else None) self._extractor = Extractor(self._args.detector, self._args.aligner, maskers, + recognition=recognition, configfile=configfile, multiprocess=not self._args.singleprocess, exclude_gpus=self._args.exclude_gpus, rotate_images=self._args.rotate_images, min_size=self._args.min_size, normalize_method=normalization, - re_feed=self._args.re_feed) - self._threads = [] - self._verify_output = False - logger.debug("Initialized %s", self.__class__.__name__) + re_feed=self._args.re_feed, + re_align=self._args.re_align) + self._filter = Filter(self._args.ref_threshold, + self._args.filter, + self._args.nfilter, + self._extractor) + + def _get_input_locations(self) -> list[str]: + """ Obtain the full path to input locations. Will be a list of locations if batch mode is + selected, or a containing a single location if batch mode is not selected. + + Returns + ------- + list: + The list of input location paths + """ + if not self._args.batch_mode or os.path.isfile(self._args.input_dir): + return [self._args.input_dir] # Not batch mode or a single file - @property - def _save_interval(self) -> Optional[int]: - """ int: The number of frames to be processed between each saving of the alignments file if - it has been provided, otherwise ``None`` """ - if hasattr(self._args, "save_interval"): - return self._args.save_interval - return None + retval = [os.path.join(self._args.input_dir, fname) + for fname in os.listdir(self._args.input_dir) + if (os.path.isdir(os.path.join(self._args.input_dir, fname)) # folder images + and any(os.path.splitext(iname)[-1].lower() in IMAGE_EXTENSIONS + for iname in os.listdir(os.path.join(self._args.input_dir, fname)))) + or os.path.splitext(fname)[-1].lower() in VIDEO_EXTENSIONS] # video - @property - def _skip_num(self) -> int: - """ int: Number of frames to skip if extract_every_n has been provided """ - return self._args.extract_every_n if hasattr(self._args, "extract_every_n") else 1 + logger.debug("Input locations: %s", retval) + return retval - def _set_skip_list(self) -> None: - """ Add the skip list to the image loader + def _validate_batchmode(self) -> None: + """ Validate the command line arguments. + + If batch-mode selected and there is only one object to extract from, then batch mode is + disabled + + If processing in batch mode, some of the given arguments may not make sense, in which case + a warning is shown and those options are reset. - Checks against `extract_every_n` and the existence of alignments data (can exist if - `skip_existing` or `skip_existing_faces` has been provided) and compiles a list of frame - indices that should not be processed, providing these to :class:`lib.image.ImagesLoader`. """ - if self._skip_num == 1 and not self._alignments.data: - logger.debug("No frames to be skipped") + if not self._args.batch_mode: return - skip_list = [] - for idx, filename in enumerate(self._images.file_list): - if idx % self._skip_num != 0: - logger.trace("Adding image '%s' to skip list due to extract_every_n = %s", - filename, self._skip_num) - skip_list.append(idx) - # Items may be in the alignments file if skip-existing[-faces] is selected - elif os.path.basename(filename) in self._alignments.data: - self._existing_count += 1 - logger.trace("Removing image: '%s' due to previously existing", filename) - skip_list.append(idx) - if self._existing_count != 0: - logger.info("Skipping %s frames due to skip_existing/skip_existing_faces.", - self._existing_count) - logger.debug("Adding skip list: %s", skip_list) - self._images.add_skip_list(skip_list) + + if os.path.isfile(self._args.input_dir): + logger.warning("Batch mode selected but input is not a folder. Switching to normal " + "mode") + self._args.batch_mode = False + + if not self._input_locations: + logger.error("Batch mode selected, but no valid files found in input location: '%s'. " + "Exiting.", self._args.input_dir) + sys.exit(1) + + if self._args.alignments_path: + logger.warning("Custom alignments path not supported for batch mode. " + "Reverting to default.") + self._args.alignments_path = None + + def _output_for_input(self, input_location: str) -> str: + """ Obtain the path to an output folder for faces for a given input location. + + If not running in batch mode, then the user supplied output location will be returned, + otherwise a sub-folder within the user supplied output location will be returned based on + the input filename + + Parameters + ---------- + input_location: str + The full path to an input video or folder of images + """ + if not self._args.batch_mode: + return self._args.output_dir + + retval = os.path.join(self._args.output_dir, + os.path.splitext(os.path.basename(input_location))[0]) + logger.debug("Returning output: '%s' for input: '%s'", retval, input_location) + return retval def process(self) -> None: """ The entry point for triggering the Extraction Process. @@ -121,17 +151,408 @@ def process(self) -> None: Should only be called from :class:`lib.cli.launcher.ScriptExecutor` """ logger.info('Starting, this may take a while...') - # from lib.queue_manager import queue_manager ; queue_manager.debug_monitor(3) + if self._args.batch_mode: + logger.info("Batch mode selected processing: %s", self._input_locations) + for job_no, location in enumerate(self._input_locations): + if self._args.batch_mode: + logger.info("Processing job %s of %s: '%s'", + job_no + 1, len(self._input_locations), location) + arguments = Namespace(**self._args.__dict__) + arguments.input_dir = location + arguments.output_dir = self._output_for_input(location) + else: + arguments = self._args + extract = _Extract(self._extractor, arguments) + if sys.platform == "linux" and len(self._input_locations) > 1: + # TODO - Running this in a process is hideously hacky. However, there is a memory + # leak in some instances when running in batch mode. Many days have been spent + # trying to track this down to no avail (most likely coming from C-code.) Running + # the extract job inside a process prevents the memory leak in testing. This should + # be replaced if/when the memory leak is found + # Only done for Linux as not reported elsewhere and this new process won't work in + # Windows because it can't fork. + proc = Process(target=extract.process) + proc.start() + proc.join() + else: + extract.process() + self._extractor.reset_phase_index() + + +class Filter(): + """ Obtains and holds face identity embeddings for any filter/nfilter image files + passed in from the command line. + + Parameters + ---------- + filter_files: list or ``None`` + The list of filter file(s) passed in as command line arguments + nfilter_files: list or ``None`` + The list of nfilter file(s) passed in as command line arguments + extractor: :class:`~plugins.extract.pipeline.Extractor` + The extractor pipeline for obtaining face identity from images + """ + def __init__(self, + threshold: float, + filter_files: list[str] | None, + nfilter_files: list[str] | None, + extractor: Extractor) -> None: + logger.debug("Initializing %s: (threshold: %s, filter_files: %s, nfilter_files: %s " + "extractor: %s)", self.__class__.__name__, threshold, filter_files, + nfilter_files, extractor) + self._threshold = threshold + self._filter_files, self._nfilter_files = self._validate_inputs(filter_files, + nfilter_files) + + if not self._filter_files and not self._nfilter_files: + logger.debug("Filter not selected. Exiting %s", self.__class__.__name__) + return + + self._embeddings: list[np.ndarray] = [np.array([]) for _ in self._filter_files] + self._nembeddings: list[np.ndarray] = [np.array([]) for _ in self._nfilter_files] + self._extractor = extractor + + self._get_embeddings() + self._extractor.recognition.add_identity_filters(self.embeddings, + self.n_embeddings, + self._threshold) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def active(self): + """ bool: ``True`` if filter files have been passed in command line arguments. ``False`` if + no filter files have been provided """ + return bool(self._filter_files) or bool(self._nfilter_files) + + @property + def embeddings(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The filter embeddings""" + if self._embeddings and all(np.any(e) for e in self._embeddings): + retval = np.concatenate(self._embeddings, axis=0) + else: + retval = np.array([]) + return retval + + @property + def n_embeddings(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The n-filter embeddings""" + if self._nembeddings and all(np.any(e) for e in self._nembeddings): + retval = np.concatenate(self._nembeddings, axis=0) + else: + retval = np.array([]) + return retval + + @classmethod + def _files_from_folder(cls, input_location: list[str]) -> list[str]: + """ Test whether the input location is a folder and if so, return the list of contained + image files, otherwise return the original input location + + Parameters + --------- + input_files: list + A list of full paths to individual files or to a folder location + + Returns + ------- + bool + Either the original list of files provided, or the image files that exist in the + provided folder location + """ + if not input_location or len(input_location) > 1: + return input_location + + test_folder = input_location[0] + if not os.path.isdir(test_folder): + logger.debug("'%s' is not a folder. Returning original list", test_folder) + return input_location + + retval = [os.path.join(test_folder, fname) + for fname in os.listdir(test_folder) + if os.path.splitext(fname)[-1].lower() in IMAGE_EXTENSIONS] + logger.info("Collected files from folder '%s': %s", test_folder, + [os.path.basename(f) for f in retval]) + return retval + + def _validate_inputs(self, + filter_files: list[str] | None, + nfilter_files: list[str] | None) -> tuple[list[str], list[str]]: + """ Validates that the given filter/nfilter files exist, are image files and are unique + + Parameters + ---------- + filter_files: list or ``None`` + The list of filter file(s) passed in as command line arguments + nfilter_files: list or ``None`` + The list of nfilter file(s) passed in as command line arguments + + Returns + ------- + filter_files: list + List of full paths to filter files + nfilter_files: list + List of full paths to nfilter files + """ + error = False + retval: list[list[str]] = [] + + for files in (filter_files, nfilter_files): + filt_files = [] if files is None else self._files_from_folder(files) + for file in filt_files: + if (not os.path.isfile(file) or + os.path.splitext(file)[-1].lower() not in IMAGE_EXTENSIONS): + logger.warning("Filter file '%s' does not exist or is not an image file", file) + error = True + retval.append(filt_files) + + filters = retval[0] + nfilters = retval[1] + f_fnames = set(os.path.basename(fname) for fname in filters) + n_fnames = set(os.path.basename(fname) for fname in nfilters) + if f_fnames.intersection(n_fnames): + error = True + logger.warning("filter and nfilter filenames should be unique. The following " + "filenames exist in both folders: %s", f_fnames.intersection(n_fnames)) + + if error: + logger.error("There was a problem processing filter files. See the above warnings for " + "details") + sys.exit(1) + logger.debug("filter_files: %s, nfilter_files: %s", retval[0], retval[1]) + + return filters, nfilters + + @classmethod + def _identity_from_extracted(cls, filename) -> tuple[np.ndarray, bool]: + """ Test whether the given image is a faceswap extracted face and contains identity + information. If so, return the identity embedding + + Parameters + ---------- + filename: str + Full path to the image file to load + + Returns + ------- + :class:`numpy.ndarray` + The identity embeddings, if they can be obtained from the image header, otherwise an + empty array + bool + ``True`` if the image is a faceswap extracted image otherwise ``False`` + """ + if os.path.splitext(filename)[-1].lower() != ".png": + logger.debug("'%s' not a png. Returning empty array", filename) + return np.array([]), False + + meta = read_image_meta(filename) + if "itxt" not in meta or "alignments" not in meta["itxt"]: + logger.debug("'%s' does not contain faceswap data. Returning empty array", filename) + return np.array([]), False + + align: "PNGHeaderAlignmentsDict" = meta["itxt"]["alignments"] + if "identity" not in align or "vggface2" not in align["identity"]: + logger.debug("'%s' does not contain identity data. Returning empty array", filename) + return np.array([]), True + + retval = np.array(align["identity"]["vggface2"]) + logger.debug("Obtained identity for '%s'. Shape: %s", filename, retval.shape) + + return retval, True + + def _process_extracted(self, item: ExtractMedia) -> None: + """ Process the output from the extraction pipeline. + + If no face has been detected, or multiple faces are detected for the inclusive filter, + embeddings and filenames are removed from the filter. + + if a single face is detected or multiple faces are detected for the exclusive filter, + embeddings are added to the relevent filter list + + Parameters + ---------- + item: :class:`plugins.extract.Pipeline.ExtracMedia` + The output from the extraction pipeline containing the identity encodings + """ + is_filter = item.filename in self._filter_files + lbl = "filter" if is_filter else "nfilter" + filelist = self._filter_files if is_filter else self._nfilter_files + embeddings = self._embeddings if is_filter else self._nembeddings + identities = np.array([face.identity["vggface2"] for face in item.detected_faces]) + idx = filelist.index(item.filename) + + if len(item.detected_faces) == 0: + logger.warning("No faces detected for %s in file '%s'. Image will not be used", + lbl, os.path.basename(item.filename)) + filelist.pop(idx) + embeddings.pop(idx) + return + + if len(item.detected_faces) == 1: + logger.debug("Adding identity for %s from file '%s'", lbl, item.filename) + embeddings[idx] = identities + return + + if len(item.detected_faces) > 1 and is_filter: + logger.warning("%s faces detected for filter in '%s'. These identies will not be used", + len(item.detected_faces), os.path.basename(item.filename)) + filelist.pop(idx) + embeddings.pop(idx) + return + + if len(item.detected_faces) > 1 and not is_filter: + logger.warning("%s faces detected for nfilter in '%s'. All of these identies will be " + "used", len(item.detected_faces), os.path.basename(item.filename)) + embeddings[idx] = identities + return + + def _identity_from_extractor(self, file_list: list[str], aligned: list[str]) -> None: + """ Obtain the identity embeddings from the extraction pipeline + + Parameters + ---------- + filesile_list: list + List of full path to images to run through the extraction pipeline + aligned: list + List of full path to images that exist in attr:`filelist` that are faceswap aligned + images + """ + logger.info("Extracting faces to obtain identity from images") + logger.debug("Files requiring full extraction: %s", + [fname for fname in file_list if fname not in aligned]) + logger.debug("Aligned files requiring identity info: %s", aligned) + + loader = PipelineLoader(file_list, self._extractor, aligned_filenames=aligned) + loader.launch() + + for phase in range(self._extractor.passes): + is_final = self._extractor.final_pass + detected_faces: dict[str, ExtractMedia] = {} + self._extractor.launch() + desc = "Obtaining reference face Identity" + if self._extractor.passes > 1: + desc = (f"{desc } pass {phase + 1} of {self._extractor.passes}: " + f"{self._extractor.phase_text}") + for extract_media in tqdm(self._extractor.detected_faces(), + total=len(file_list), + file=sys.stdout, + desc=desc): + if is_final: + self._process_extracted(extract_media) + else: + extract_media.remove_image() + # cache extract_media for next run + detected_faces[extract_media.filename] = extract_media + + if not is_final: + logger.debug("Reloading images") + loader.reload(detected_faces) + + self._extractor.reset_phase_index() + + def _get_embeddings(self) -> None: + """ Obtain the embeddings for the given filter lists """ + needs_extraction: list[str] = [] + aligned: list[str] = [] + + for files, embed in zip((self._filter_files, self._nfilter_files), + (self._embeddings, self._nembeddings)): + for idx, file in enumerate(files): + identity, is_aligned = self._identity_from_extracted(file) + if np.any(identity): + logger.debug("Obtained identity from png header: '%s'", file) + embed[idx] = identity[None, ...] + continue + + needs_extraction.append(file) + if is_aligned: + aligned.append(file) + + if needs_extraction: + self._identity_from_extractor(needs_extraction, aligned) + + if not self._nfilter_files and not self._filter_files: + logger.error("No faces were detected from your selected identity filter files") + sys.exit(1) + + logger.debug("Filter: (filenames: %s, shape: %s), nFilter: (filenames: %s, shape: %s)", + [os.path.basename(f) for f in self._filter_files], + self.embeddings.shape, + [os.path.basename(f) for f in self._nfilter_files], + self.n_embeddings.shape) + + +class PipelineLoader(): + """ Handles loading and reloading images into the extraction pipeline. + + Parameters + ---------- + path: str or list of str + Full path to a folder of images or a video file or a list of image files + extractor: :class:`~plugins.extract.pipeline.Extractor` + The extractor pipeline for obtaining face identity from images + aligned_filenames: list, optional + Used for when the loader is used for getting face filter embeddings. List of full path to + image files that exist in :attr:`path` that are aligned faceswap images + """ + def __init__(self, + path: str | list[str], + extractor: Extractor, + aligned_filenames: list[str] | None = None) -> None: + logger.debug("Initializing %s: (path: %s, extractor: %s, aligned_filenames: %s)", + self.__class__.__name__, path, extractor, aligned_filenames) + self._images = ImagesLoader(path, fast_count=True) + self._extractor = extractor + self._threads: list[MultiThread] = [] + self._aligned_filenames = [] if aligned_filenames is None else aligned_filenames + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def is_video(self) -> bool: + """ bool: ``True`` if the input location is a video file, ``False`` if it is a folder of + images """ + return self._images.is_video + + @property + def file_list(self) -> list[str]: + """ list: A full list of files in the source location. If the input is a video + then this is a list of dummy filenames as corresponding to an alignments file """ + return self._images.file_list + + @property + def process_count(self) -> int: + """ int: The number of images or video frames to be processed (IE the total count less + items that are to be skipped from the :attr:`skip_list`)""" + return self._images.process_count + + def add_skip_list(self, skip_list: list[int]) -> None: + """ Add a skip list to the :class:`ImagesLoader` + + Parameters + ---------- + skip_list: list + A list of indices corresponding to the frame indices that should be skipped by the + :func:`load` function. + """ + self._images.add_skip_list(skip_list) + + def launch(self) -> None: + """ Launch the image loading pipeline """ self._threaded_redirector("load") - self._run_extraction() + + def reload(self, detected_faces: dict[str, ExtractMedia]) -> None: + """ Reload images for multiple pipeline passes """ + self._threaded_redirector("reload", (detected_faces, )) + + def check_thread_error(self) -> None: + """ Check if any errors have occurred in the running threads and raise their errors """ + for thread in self._threads: + thread.check_and_raise_error() + + def join(self) -> None: + """ Join all open loader threads """ for thread in self._threads: thread.join() - self._alignments.save() - finalize(self._images.process_count + self._existing_count, - self._alignments.faces_count, - self._verify_output) - def _threaded_redirector(self, task: str, io_args: Optional[tuple] = None) -> None: + def _threaded_redirector(self, task: str, io_args: tuple | None = None) -> None: """ Redirect image input/output tasks to relevant queues in background thread Parameters @@ -142,7 +563,7 @@ def _threaded_redirector(self, task: str, io_args: Optional[tuple] = None) -> No Any arguments that need to be provided to the background function """ logger.debug("Threading task: (Task: '%s')", task) - io_args = tuple() if io_args is None else (io_args, ) + io_args = tuple() if io_args is None else io_args func = getattr(self, f"_{task}") io_thread = MultiThread(func, *io_args, thread_count=1) io_thread.start() @@ -160,7 +581,8 @@ def _load(self) -> None: if load_queue.shutdown.is_set(): logger.debug("Load Queue: Stop signal received. Terminating") break - item = ExtractMedia(filename, image[..., :3]) + is_aligned = filename in self._aligned_filenames + item = ExtractMedia(filename, image[..., :3], is_aligned=is_aligned) load_queue.put(item) load_queue.put("EOF") logger.debug("Load Images: Complete") @@ -174,8 +596,8 @@ def _reload(self, detected_faces: dict[str, ExtractMedia]) -> None: Parameters ---------- detected_faces: dict - Dictionary of :class:`plugins.extract.pipeline.ExtractMedia` with the filename as the - key for repopulating the image attribute. + Dictionary of :class:`~plugins.extract.extract_media.ExtractMedia` with the filename as + the key for repopulating the image attribute. """ logger.debug("Reload Images: Start. Detected Faces Count: %s", len(detected_faces)) load_queue = self._extractor.input_queue @@ -183,7 +605,7 @@ def _reload(self, detected_faces: dict[str, ExtractMedia]) -> None: if load_queue.shutdown.is_set(): logger.debug("Reload Queue: Stop signal received. Terminating") break - logger.trace("Reloading image: '%s'", filename) + logger.trace("Reloading image: '%s'", filename) # type: ignore extract_media = detected_faces.pop(filename, None) if not extract_media: logger.warning("Couldn't find faces for: %s", filename) @@ -193,6 +615,98 @@ def _reload(self, detected_faces: dict[str, ExtractMedia]) -> None: load_queue.put("EOF") logger.debug("Reload Images: Complete") + +class _Extract(): + """ The Actual extraction process. + + This class is called by the parent :class:`Extract` process + + Parameters + ---------- + extractor: :class:`~plugins.extract.pipeline.Extractor` + The extractor pipeline for running extractions + arguments: :class:`argparse.Namespace` + The arguments to be passed to the extraction process as generated from Faceswap's command + line arguments + """ + def __init__(self, + extractor: Extractor, + arguments: Namespace) -> None: + logger.debug("Initializing %s: (extractor: %s, args: %s)", self.__class__.__name__, + extractor, arguments) + self._args = arguments + self._output_dir = None if self._args.skip_saving_faces else get_folder( + self._args.output_dir) + + logger.info("Output Directory: %s", self._output_dir) + self._loader = PipelineLoader(self._args.input_dir, extractor) + + self._alignments = Alignments(self._args, True, self._loader.is_video) + self._extractor = extractor + self._extractor.import_data(self._args.input_dir) + + self._existing_count = 0 + self._set_skip_list() + + self._post_process = PostProcess(arguments) + self._verify_output = False + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def _save_interval(self) -> int | None: + """ int: The number of frames to be processed between each saving of the alignments file if + it has been provided, otherwise ``None`` """ + if hasattr(self._args, "save_interval"): + return self._args.save_interval + return None + + @property + def _skip_num(self) -> int: + """ int: Number of frames to skip if extract_every_n has been provided """ + return self._args.extract_every_n if hasattr(self._args, "extract_every_n") else 1 + + def _set_skip_list(self) -> None: + """ Add the skip list to the image loader + + Checks against `extract_every_n` and the existence of alignments data (can exist if + `skip_existing` or `skip_existing_faces` has been provided) and compiles a list of frame + indices that should not be processed, providing these to :class:`lib.image.ImagesLoader`. + """ + if self._skip_num == 1 and not self._alignments.data: + logger.debug("No frames to be skipped") + return + skip_list = [] + for idx, filename in enumerate(self._loader.file_list): + if idx % self._skip_num != 0: + logger.trace("Adding image '%s' to skip list due to " # type: ignore + "extract_every_n = %s", filename, self._skip_num) + skip_list.append(idx) + # Items may be in the alignments file if skip-existing[-faces] is selected + elif os.path.basename(filename) in self._alignments.data: + self._existing_count += 1 + logger.trace("Removing image: '%s' due to previously existing", # type: ignore + filename) + skip_list.append(idx) + if self._existing_count != 0: + logger.info("Skipping %s frames due to skip_existing/skip_existing_faces.", + self._existing_count) + logger.debug("Adding skip list: %s", skip_list) + self._loader.add_skip_list(skip_list) + + def process(self) -> None: + """ The entry point for triggering the Extraction Process. + + Should only be called from :class:`lib.cli.launcher.ScriptExecutor` + """ + # from lib.queue_manager import queue_manager ; queue_manager.debug_monitor(3) + self._loader.launch() + self._run_extraction() + self._loader.join() + self._alignments.save() + finalize(self._loader.process_count + self._existing_count, + self._alignments.faces_count, + self._verify_output) + def _run_extraction(self) -> None: """ The main Faceswap Extraction process @@ -203,22 +717,19 @@ def _run_extraction(self) -> None: size = self._args.size if hasattr(self._args, "size") else 256 saver = None if self._args.skip_saving_faces else ImagesSaver(self._output_dir, as_bytes=True) - exception = False - for phase in range(self._extractor.passes): - if exception: - break is_final = self._extractor.final_pass - detected_faces = {} + detected_faces: dict[str, ExtractMedia] = {} self._extractor.launch() - self._check_thread_error() + self._loader.check_thread_error() ph_desc = "Extraction" if self._extractor.passes == 1 else self._extractor.phase_text desc = f"Running pass {phase + 1} of {self._extractor.passes}: {ph_desc}" for idx, extract_media in enumerate(tqdm(self._extractor.detected_faces(), - total=self._images.process_count, + total=self._loader.process_count, file=sys.stdout, - desc=desc)): - self._check_thread_error() + desc=desc, + leave=False)): + self._loader.check_thread_error() if is_final: self._output_processing(extract_media, size) self._output_faces(saver, extract_media) @@ -231,15 +742,10 @@ def _run_extraction(self) -> None: if not is_final: logger.debug("Reloading images") - self._threaded_redirector("reload", detected_faces) - if not self._args.skip_saving_faces: + self._loader.reload(detected_faces) + if saver is not None: saver.close() - def _check_thread_error(self) -> None: - """ Check if any errors have occurred in the running threads and their errors """ - for thread in self._threads: - thread.check_and_raise_error() - def _output_processing(self, extract_media: ExtractMedia, size: int) -> None: """ Prepare faces for output @@ -248,7 +754,7 @@ def _output_processing(self, extract_media: ExtractMedia, size: int) -> None: Parameters ---------- - extract_media: :class:`plugins.extract.pipeline.ExtractMedia` + extract_media: :class:`~plugins.extract.extract_media.ExtractMedia` Output from :class:`plugins.extract.pipeline.Extractor` size: int The size that the aligned face should be created at @@ -263,13 +769,13 @@ def _output_processing(self, extract_media: ExtractMedia, size: int) -> None: faces_count = len(extract_media.detected_faces) if faces_count == 0: - logger.verbose("No faces were detected in image: %s", + logger.verbose("No faces were detected in image: %s", # type: ignore os.path.basename(extract_media.filename)) if not self._verify_output and faces_count > 1: self._verify_output = True - def _output_faces(self, saver: ImagesSaver, extract_media: ExtractMedia) -> None: + def _output_faces(self, saver: ImagesSaver | None, extract_media: ExtractMedia) -> None: """ Output faces to save thread Set the face filename based on the frame name and put the face to the @@ -278,29 +784,43 @@ def _output_faces(self, saver: ImagesSaver, extract_media: ExtractMedia) -> None Parameters ---------- - saver: lib.images.ImagesSaver - The background saver for saving the image - extract_media: :class:`~plugins.extract.pipeline.ExtractMedia` + saver: :class:`lib.images.ImagesSaver` or ``None`` + The background saver for saving the image or ``None`` if faces are not to be saved + extract_media: :class:`~plugins.extract.extract_media.ExtractMedia` The output from :class:`~plugins.extract.Pipeline.Extractor` """ - logger.trace("Outputting faces for %s", extract_media.filename) + logger.trace("Outputting faces for %s", extract_media.filename) # type: ignore final_faces = [] filename = os.path.splitext(os.path.basename(extract_media.filename))[0] - extension = ".png" - - for idx, face in enumerate(extract_media.detected_faces): - output_filename = f"{filename}_{idx}{extension}" - meta = dict(alignments=face.to_png_meta(), - source=dict(alignments_version=self._alignments.version, - original_filename=output_filename, - face_index=idx, - source_filename=os.path.basename(extract_media.filename), - source_is_video=self._images.is_video, - source_frame_dims=extract_media.image_size)) - image = encode_image(face.aligned.face, extension, metadata=meta) - - if not self._args.skip_saving_faces: - saver.save(output_filename, image) + + skip_idx = 0 + for face_id, face in enumerate(extract_media.detected_faces): + real_face_id = face_id - skip_idx + output_filename = f"{filename}_{real_face_id}.png" + aligned = face.aligned.face + assert aligned is not None + meta: PNGHeaderDict = { + "alignments": face.to_png_meta(), + "source": {"alignments_version": self._alignments.version, + "original_filename": output_filename, + "face_index": real_face_id, + "source_filename": os.path.basename(extract_media.filename), + "source_is_video": self._loader.is_video, + "source_frame_dims": extract_media.image_size}} + image = encode_image(aligned, ".png", metadata=meta) + + sub_folder = extract_media.sub_folders[face_id] + # Binned faces shouldn't risk filename clash, so just use original id + out_name = output_filename if not sub_folder else f"{filename}_{face_id}.png" + + if saver is not None: + saver.save(out_name, image, sub_folder) + + if sub_folder: # This is a filtered out face being binned + skip_idx += 1 + continue final_faces.append(face.to_alignment()) - self._alignments.data[os.path.basename(extract_media.filename)] = dict(faces=final_faces) + + self._alignments.data[os.path.basename(extract_media.filename)] = {"faces": final_faces, + "video_meta": {}} del extract_media diff --git a/scripts/fsmedia.py b/scripts/fsmedia.py index e721f47fd5..9d1fbdbb4e 100644 --- a/scripts/fsmedia.py +++ b/scripts/fsmedia.py @@ -5,24 +5,32 @@ Holds optional pre/post processing functions for convert and extract. """ - +from __future__ import annotations import logging import os import sys +import typing as T + +from collections.abc import Iterator import cv2 import numpy as np import imageio -from lib.align import Alignments as AlignmentsBase -from lib.face_filter import FaceFilter as FilterFunc +from lib.align import Alignments as AlignmentsBase, get_centered_size from lib.image import count_frames, read_image -from lib.utils import (camel_case_split, get_image_paths, _video_extensions) +from lib.utils import (camel_case_split, get_image_paths, VIDEO_EXTENSIONS) + +if T.TYPE_CHECKING: + from collections.abc import Generator + from argparse import Namespace + from lib.align import AlignedFace + from plugins.extract import ExtractMedia -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -def finalize(images_found, num_faces_detected, verify_output): +def finalize(images_found: int, num_faces_detected: int, verify_output: bool) -> None: """ Output summary statistics at the end of the extract or convert processes. Parameters @@ -45,7 +53,7 @@ def finalize(images_found, num_faces_detected, verify_output): logger.info("Double check your results.") logger.info("-------------------------") - logger.info("Process Succesfully Completed. Shutting Down...") + logger.info("Process Successfully Completed. Shutting Down...") class Alignments(AlignmentsBase): @@ -62,7 +70,10 @@ class Alignments(AlignmentsBase): ``True`` if the input to the process is a video, ``False`` if it is a folder of images. Default: False """ - def __init__(self, arguments, is_extract, input_is_video=False): + def __init__(self, + arguments: Namespace, + is_extract: bool, + input_is_video: bool = False) -> None: logger.debug("Initializing %s: (is_extract: %s, input_is_video: %s)", self.__class__.__name__, is_extract, input_is_video) self._args = arguments @@ -71,7 +82,7 @@ def __init__(self, arguments, is_extract, input_is_video=False): super().__init__(folder, filename=filename) logger.debug("Initialized %s", self.__class__.__name__) - def _set_folder_filename(self, input_is_video): + def _set_folder_filename(self, input_is_video: bool) -> tuple[str, str]: """ Return the folder and the filename for the alignments file. If the input is a video, the alignments file will be stored in the same folder @@ -98,7 +109,7 @@ def _set_folder_filename(self, input_is_video): elif input_is_video: logger.debug("Alignments from Video File: '%s'", self._args.input_dir) folder, filename = os.path.split(self._args.input_dir) - filename = "{}_alignments".format(os.path.splitext(filename)[0]) + filename = f"{os.path.splitext(filename)[0]}_alignments.fsa" else: logger.debug("Alignments from Input Folder: '%s'", self._args.input_dir) folder = str(self._args.input_dir) @@ -106,7 +117,7 @@ def _set_folder_filename(self, input_is_video): logger.debug("Setting Alignments: (folder: '%s' filename: '%s')", folder, filename) return folder, filename - def _load(self): + def _load(self) -> dict[str, T.Any]: """ Override the parent :func:`~lib.align.Alignments._load` to handle skip existing frames and faces on extract. @@ -119,10 +130,10 @@ def _load(self): Any alignments that have already been extracted if skip existing has been selected otherwise an empty dictionary """ - data = dict() + data: dict[str, T.Any] = {} + if not self._is_extract and not self.have_alignments_file: + return data if not self._is_extract: - if not self.have_alignments_file: - return data data = super()._load() return data @@ -146,7 +157,8 @@ def _load(self): logger.debug("Frames with no faces selected for redetection: %s", len(del_keys)) for key in del_keys: if key in data: - logger.trace("Selected for redetection: '%s'", key) + logger.trace("Selected for redetection: '%s'", # type:ignore[attr-defined] + key) del data[key] return data @@ -160,7 +172,7 @@ class Images(): arguments: :class:`argparse.Namespace` The command line arguments that were passed to Faceswap """ - def __init__(self, arguments): + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s", self.__class__.__name__) self._args = arguments self._is_video = self._check_input_folder() @@ -169,22 +181,22 @@ def __init__(self, arguments): logger.debug("Initialized %s", self.__class__.__name__) @property - def is_video(self): + def is_video(self) -> bool: """bool: ``True`` if the input is a video file otherwise ``False``. """ return self._is_video @property - def input_images(self): + def input_images(self) -> str | list[str]: """str or list: Path to the video file if the input is a video otherwise list of image paths. """ return self._input_images @property - def images_found(self): + def images_found(self) -> int: """int: The number of frames that exist in the video file, or the folder of images. """ return self._images_found - def _count_images(self): + def _count_images(self) -> int: """ Get the number of Frames from a video file or folder of images. Returns @@ -198,7 +210,7 @@ def _count_images(self): retval = len(self._input_images) return retval - def _check_input_folder(self): + def _check_input_folder(self) -> bool: """ Check whether the input is a folder or video. Returns @@ -210,7 +222,7 @@ def _check_input_folder(self): logger.error("Input location %s not found.", self._args.input_dir) sys.exit(1) if (os.path.isfile(self._args.input_dir) and - os.path.splitext(self._args.input_dir)[1].lower() in _video_extensions): + os.path.splitext(self._args.input_dir)[1].lower() in VIDEO_EXTENSIONS): logger.info("Input Video: %s", self._args.input_dir) retval = True else: @@ -218,7 +230,7 @@ def _check_input_folder(self): retval = False return retval - def _get_input_images(self): + def _get_input_images(self) -> str | list[str]: """ Return the list of images or path to video file that is to be processed. Returns @@ -233,7 +245,7 @@ def _get_input_images(self): return input_images - def load(self): + def load(self) -> Generator[tuple[str, np.ndarray], None, None]: """ Generator to load frames from a folder of images or from a video file. Yields @@ -247,7 +259,7 @@ def load(self): for filename, image in iterator(): yield filename, image - def _load_disk_frames(self): + def _load_disk_frames(self) -> Generator[tuple[str, np.ndarray], None, None]: """ Generator to load frames from a folder of images. Yields @@ -264,7 +276,7 @@ def _load_disk_frames(self): continue yield filename, image - def _load_video_frames(self): + def _load_video_frames(self) -> Generator[tuple[str, np.ndarray], None, None]: """ Generator to load frames from a video file. Yields @@ -275,17 +287,17 @@ def _load_video_frames(self): A single frame """ logger.debug("Input is video. Capturing frames") - vidname = os.path.splitext(os.path.basename(self._args.input_dir))[0] - reader = imageio.get_reader(self._args.input_dir, "ffmpeg") - for i, frame in enumerate(reader): + vidname, ext = os.path.splitext(os.path.basename(self._args.input_dir)) + reader = imageio.get_reader(self._args.input_dir, "ffmpeg") # type:ignore[arg-type] + for i, frame in enumerate(T.cast(Iterator[np.ndarray], reader)): # Convert to BGR for cv2 compatibility frame = frame[:, :, ::-1] - filename = "{}_{:06d}.png".format(vidname, i + 1) - logger.trace("Loading video frame: '%s'", filename) + filename = f"{vidname}_{i + 1:06d}{ext}" + logger.trace("Loading video frame: '%s'", filename) # type:ignore[attr-defined] yield filename, frame reader.close() - def load_one_image(self, filename): + def load_one_image(self, filename) -> np.ndarray: """ Obtain a single image for the given filename. Parameters @@ -299,19 +311,20 @@ def load_one_image(self, filename): The image for the requested filename, """ - logger.trace("Loading image: '%s'", filename) + logger.trace("Loading image: '%s'", filename) # type:ignore[attr-defined] if self._is_video: if filename.isdigit(): frame_no = filename else: frame_no = os.path.splitext(filename)[0][filename.rfind("_") + 1:] - logger.trace("Extracted frame_no %s from filename '%s'", frame_no, filename) + logger.trace( # type:ignore[attr-defined] + "Extracted frame_no %s from filename '%s'", frame_no, filename) retval = self._load_one_video_frame(int(frame_no)) else: retval = read_image(filename, raise_error=True) return retval - def _load_one_video_frame(self, frame_no): + def _load_one_video_frame(self, frame_no: int) -> np.ndarray: """ Obtain a single frame from a video file. Parameters @@ -324,15 +337,15 @@ def _load_one_video_frame(self, frame_no): :class:`numpy.ndarray` The image for the requested frame index, """ - logger.trace("Loading video frame: %s", frame_no) - reader = imageio.get_reader(self._args.input_dir, "ffmpeg") + logger.trace("Loading video frame: %s", frame_no) # type:ignore[attr-defined] + reader = imageio.get_reader(self._args.input_dir, "ffmpeg") # type:ignore[arg-type] reader.set_image_index(frame_no - 1) - frame = reader.get_next_data()[:, :, ::-1] + frame = reader.get_next_data()[:, :, ::-1] # type:ignore[index] reader.close() return frame -class PostProcess(): # pylint:disable=too-few-public-methods +class PostProcess(): """ Optional pre/post processing tasks for convert and extract. Builds a pipeline of actions that have optionally been requested to be performed @@ -343,13 +356,13 @@ class PostProcess(): # pylint:disable=too-few-public-methods arguments: :class:`argparse.Namespace` The command line arguments that were passed to Faceswap """ - def __init__(self, arguments): + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s", self.__class__.__name__) self._args = arguments self._actions = self._set_actions() logger.debug("Initialized %s", self.__class__.__name__) - def _set_actions(self): + def _set_actions(self) -> list[PostProcessAction]: """ Compile the requested actions to be performed into a list Returns @@ -358,25 +371,25 @@ def _set_actions(self): The list of :class:`PostProcessAction` to be performed """ postprocess_items = self._get_items() - actions = list() + actions: list["PostProcessAction"] = [] for action, options in postprocess_items.items(): - options = dict() if options is None else options + options = {} if options is None else options args = options.get("args", tuple()) - kwargs = options.get("kwargs", dict()) + kwargs = options.get("kwargs", {}) args = args if isinstance(args, tuple) else tuple() - kwargs = kwargs if isinstance(kwargs, dict) else dict() + kwargs = kwargs if isinstance(kwargs, dict) else {} task = globals()[action](*args, **kwargs) if task.valid: logger.debug("Adding Postprocess action: '%s'", task) actions.append(task) - for action in actions: - action_name = camel_case_split(action.__class__.__name__) + for ppaction in actions: + action_name = camel_case_split(ppaction.__class__.__name__) logger.info("Adding post processing item: %s", " ".join(action_name)) return actions - def _get_items(self): + def _get_items(self) -> dict[str, dict[str, tuple | dict] | None]: """ Check the passed in command line arguments for requested actions, For any requested actions, add the item to the actions list along with @@ -388,61 +401,35 @@ def _get_items(self): The name of the action to be performed as the key. Any action specific arguments and keyword arguments as the value. """ - postprocess_items = dict() + postprocess_items: dict[str, dict[str, tuple | dict] | None] = {} # Debug Landmarks if (hasattr(self._args, 'debug_landmarks') and self._args.debug_landmarks): postprocess_items["DebugLandmarks"] = None - # Face Filter post processing - if ((hasattr(self._args, "filter") and self._args.filter is not None) or - (hasattr(self._args, "nfilter") and - self._args.nfilter is not None)): - - if hasattr(self._args, "detector"): - detector = self._args.detector.replace("-", "_").lower() - else: - detector = "cv2_dnn" - if hasattr(self._args, "aligner"): - aligner = self._args.aligner.replace("-", "_").lower() - else: - aligner = "cv2_dnn" - - face_filter = dict(detector=detector, - aligner=aligner, - multiprocess=not self._args.singleprocess) - filter_lists = dict() - if hasattr(self._args, "ref_threshold"): - face_filter["ref_threshold"] = self._args.ref_threshold - for filter_type in ('filter', 'nfilter'): - filter_args = getattr(self._args, filter_type, None) - filter_args = None if not filter_args else filter_args - filter_lists[filter_type] = filter_args - face_filter["filter_lists"] = filter_lists - postprocess_items["FaceFilter"] = {"kwargs": face_filter} - logger.debug("Postprocess Items: %s", postprocess_items) return postprocess_items - def do_actions(self, extract_media): + def do_actions(self, extract_media: ExtractMedia) -> None: """ Perform the requested optional post-processing actions on the given image. Parameters ---------- - extract_media: :class:`~plugins.extract.pipeline.ExtractMedia` - The :class:`~plugins.extract.pipeline.ExtractMedia` object to perform the + extract_media: :class:`~plugins.extract.extract_media.ExtractMedia` + The :class:`~plugins.extract.extract_media.ExtractMedia` object to perform the action on. Returns ------- - :class:`~plugins.extract.pipeline.ExtractMedia` - The original :class:`~plugins.extract.pipeline.ExtractMedia` with any actions applied + :class:`~plugins.extract.extract_media.ExtractMedia` + The original :class:`~plugins.extract.extract_media.ExtractMedia` with any actions + applied """ for action in self._actions: logger.debug("Performing postprocess action: '%s'", action.__class__.__name__) action.process(extract_media) -class PostProcessAction(): # pylint: disable=too-few-public-methods +class PostProcessAction(): """ Parent class for Post Processing Actions. Usable in Extract or Convert or both depending on context. Any post-processing actions should @@ -455,186 +442,177 @@ class PostProcessAction(): # pylint: disable=too-few-public-methods kwargs: dict Varies for specific post process action """ - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: logger.debug("Initializing %s: (args: %s, kwargs: %s)", self.__class__.__name__, args, kwargs) self._valid = True # Set to False if invalid parameters passed in to disable logger.debug("Initialized base class %s", self.__class__.__name__) @property - def valid(self): + def valid(self) -> bool: """bool: ``True`` if the action if the parameters passed in for this action are valid, otherwise ``False`` """ return self._valid - def process(self, extract_media): + def process(self, extract_media: ExtractMedia) -> None: """ Override for specific post processing action Parameters ---------- - extract_media: :class:`~plugins.extract.pipeline.ExtractMedia` - The :class:`~plugins.extract.pipeline.ExtractMedia` object to perform the + extract_media: :class:`~plugins.extract.extract_media.ExtractMedia` + The :class:`~plugins.extract.extract_media.ExtractMedia` object to perform the action on. """ raise NotImplementedError -class DebugLandmarks(PostProcessAction): # pylint: disable=too-few-public-methods +class DebugLandmarks(PostProcessAction): """ Draw debug landmarks on face output. Extract Only """ + def __init__(self, *args, **kwargs) -> None: + super().__init__(self, *args, **kwargs) + self._face_size = 0 + self._legacy_size = 0 + self._font = cv2.FONT_HERSHEY_SIMPLEX + self._font_scale = 0.0 + self._font_pad = 0 - def process(self, extract_media): - """ Draw landmarks on a face. + def _initialize_font(self, size: int) -> None: + """ Set the font scaling sizes on first call Parameters ---------- - extract_media: :class:`~plugins.extract.pipeline.ExtractMedia` - The :class:`~plugins.extract.pipeline.ExtractMedia` object that contains the faces to - draw the landmarks on to - - Returns - ------- - :class:`~plugins.extract.pipeline.ExtractMedia` - The original :class:`~plugins.extract.pipeline.ExtractMedia` with landmarks drawn - onto the face + size: int + The pixel size of the saved aligned face """ - frame = os.path.splitext(os.path.basename(extract_media.filename))[0] - for idx, face in enumerate(extract_media.detected_faces): - logger.trace("Drawing Landmarks. Frame: '%s'. Face: %s", frame, idx) - # Landmarks - for (pos_x, pos_y) in face.aligned.landmarks.astype("int32"): - cv2.circle(face.aligned.face, (pos_x, pos_y), 1, (0, 255, 255), -1) - # Pose - center = tuple(np.int32((face.aligned.size / 2, face.aligned.size / 2))) - points = (face.aligned.pose.xyz_2d * face.aligned.size).astype("int32") - cv2.line(face.aligned.face, center, tuple(points[1]), (0, 255, 0), 1) - cv2.line(face.aligned.face, center, tuple(points[0]), (255, 0, 0), 1) - cv2.line(face.aligned.face, center, tuple(points[2]), (0, 0, 255), 1) - # Face centering - roi = face.aligned.get_cropped_roi("face") - cv2.rectangle(face.aligned.face, tuple(roi[:2]), tuple(roi[2:]), (0, 255, 0), 1) - + self._font_scale = size / 512 + self._font_pad = size // 64 -class FaceFilter(PostProcessAction): - """ Filter in or out faces based on input image(s). Extract or Convert - - Parameters - ----------- - args: tuple - Unused - kwargs: dict - Keyword arguments for face filter: - - * **detector** (`str`) - The detector to use - - * **aligner** (`str`) - The aligner to use - - * **multiprocess** (`bool`) - Whether to run the extraction pipeline in single process \ - mode or not - - * **ref_threshold** (`float`) - The reference threshold for a positive match - - * **filter_lists** (`dict`) - The filter and nfilter image paths - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - logger.info("Extracting and aligning face for Face Filter...") - self._filter = self._load_face_filter(**kwargs) - logger.debug("Initialized %s", self.__class__.__name__) - - def _load_face_filter(self, filter_lists, ref_threshold, aligner, detector, multiprocess): - """ Set up and load the :class:`~lib.face_filter.FaceFilter`. + def _border_text(self, + image: np.ndarray, + text: str, + color: tuple[int, int, int], + position: tuple[int, int]) -> None: + """ Create text on an image with a black border Parameters ---------- - filter_lists: dict - The filter and nfilter image paths - ref_threshold: float - The reference threshold for a positive match - aligner: str - The aligner to use - detector: str - The detector to use - multiprocess: bool - Whether to run the extraction pipeline in single process mode or not - - Returns - ------- - :class:`~lib.face_filter.FaceFilter` - The face filter + image: :class:`numpy.ndarray` + The image to put bordered text on to + text: str + The text to place the image + color: tuple + The color of the text + position: tuple + The (x, y) co-ordinates to place the text """ - if not any(val for val in filter_lists.values()): - return None - - facefilter = None - filter_files = [self._set_face_filter(f_type, filter_lists[f_type]) - for f_type in ("filter", "nfilter")] - - if any(filters for filters in filter_files): - facefilter = FilterFunc(filter_files[0], - filter_files[1], - detector, - aligner, - multiprocess, - ref_threshold) - logger.debug("Face filter: %s", facefilter) - else: - self.valid = False - return facefilter - - @staticmethod - def _set_face_filter(f_type, f_args): - """ Check filter files exist and add the filter file paths to a list. + thickness = 2 + for idx in range(2): + text_color = (0, 0, 0) if idx == 0 else color + cv2.putText(image, + text, + position, + self._font, + self._font_scale, + text_color, + thickness, + lineType=cv2.LINE_AA) + thickness //= 2 + + def _annotate_face_box(self, face: AlignedFace) -> None: + """ Annotate the face extract box and print the original size in pixels + + face: :class:`~lib.align.AlignedFace` + The object containing the aligned face to annotate + """ + assert face.face is not None + color = (0, 255, 0) + roi = face.get_cropped_roi(face.size, self._face_size, "face") + cv2.rectangle(face.face, tuple(roi[:2]), tuple(roi[2:]), color, 1) + + # Size in top right corner + roi_pnts = np.array([[roi[0], roi[1]], + [roi[0], roi[3]], + [roi[2], roi[3]], + [roi[2], roi[1]]]) + orig_roi = face.transform_points(roi_pnts, invert=True) + size = int(round(((orig_roi[1][0] - orig_roi[0][0]) ** 2 + + (orig_roi[1][1] - orig_roi[0][1]) ** 2) ** 0.5)) + text_img = face.face.copy() + text = f"{size}px" + text_size = cv2.getTextSize(text, self._font, self._font_scale, 1)[0] + pos_x = roi[2] - (text_size[0] + self._font_pad) + pos_y = roi[1] + text_size[1] + self._font_pad + + self._border_text(text_img, text, color, (pos_x, pos_y)) + cv2.addWeighted(text_img, 0.75, face.face, 0.25, 0, face.face) + + def _print_stats(self, face: AlignedFace) -> None: + """ Print various metrics on the output face images Parameters ---------- - f_type: {"filter", "nfilter"} - The type of filter to create this list for - f_args: str or list - The filter image(s) to use - - Returns - ------- - list - The confirmed existing paths to filter files to use + face: :class:`~lib.align.AlignedFace` + The loaded aligned face """ - if not f_args: - return list() - - logger.info("%s: %s", f_type.title(), f_args) - filter_files = f_args if isinstance(f_args, list) else [f_args] - filter_files = list(filter(lambda fpath: os.path.exists(fpath), filter_files)) - if not filter_files: - logger.warning("Face %s files were requested, but no files could be found. This " - "filter will not be applied.", f_type) - logger.debug("Face Filter files: %s", filter_files) - return filter_files - - def process(self, extract_media): - """ Filters in or out any wanted or unwanted faces based on command line arguments. + assert face.face is not None + text_image = face.face.copy() + texts = [f"pitch: {face.pose.pitch:.2f}", + f"yaw: {face.pose.yaw:.2f}", + f"roll: {face.pose.roll: .2f}", + f"distance: {face.average_distance:.2f}"] + colors = [(255, 0, 0), (0, 0, 255), (0, 255, 0), (255, 255, 255)] + text_sizes = [cv2.getTextSize(text, self._font, self._font_scale, 1)[0] for text in texts] + + final_y = face.size - text_sizes[-1][1] + pos_y = [(size[1] + self._font_pad) * (idx + 1) + for idx, size in enumerate(text_sizes)][:-1] + [final_y] + pos_x = self._font_pad + + for idx, text in enumerate(texts): + self._border_text(text_image, text, colors[idx], (pos_x, pos_y[idx])) + + # Apply text to face + cv2.addWeighted(text_image, 0.75, face.face, 0.25, 0, face.face) + + def process(self, extract_media: ExtractMedia) -> None: + """ Draw landmarks on a face. Parameters ---------- - extract_media: :class:`~plugins.extract.pipeline.ExtractMedia` - The :class:`~plugins.extract.pipeline.ExtractMedia` object to perform the - face filtering on. - - Returns - ------- - :class:`~plugins.extract.pipeline.ExtractMedia` - The original :class:`~plugins.extract.pipeline.ExtractMedia` with any requested filters - applied + extract_media: :class:`~plugins.extract.extract_media.ExtractMedia` + The :class:`~plugins.extract.extract_media.ExtractMedia` object that contains the faces + to draw the landmarks on to """ - if not self._filter: - return - ret_faces = list() - for idx, detect_face in enumerate(extract_media.detected_faces): - check_item = detect_face["face"] if isinstance(detect_face, dict) else detect_face - if not self._filter.check(extract_media.image, check_item): - logger.verbose("Skipping not recognized face: (Frame: %s Face %s)", - extract_media.filename, idx) - continue - logger.trace("Accepting recognised face. Frame: %s. Face: %s", - extract_media.filename, idx) - ret_faces.append(detect_face) - extract_media.add_detected_faces(ret_faces) + frame = os.path.splitext(os.path.basename(extract_media.filename))[0] + for idx, face in enumerate(extract_media.detected_faces): + if not self._face_size: + self._face_size = get_centered_size(face.aligned.centering, + "face", + face.aligned.size) + logger.debug("set face size: %s", self._face_size) + if not self._legacy_size: + self._legacy_size = get_centered_size(face.aligned.centering, + "legacy", + face.aligned.size) + logger.debug("set legacy size: %s", self._legacy_size) + if not self._font_scale: + self._initialize_font(face.aligned.size) + + logger.trace("Drawing Landmarks. Frame: '%s'. Face: %s", # type:ignore[attr-defined] + frame, idx) + # Landmarks + assert face.aligned.face is not None + for (pos_x, pos_y) in face.aligned.landmarks.astype("int32"): + cv2.circle(face.aligned.face, (pos_x, pos_y), 1, (0, 255, 255), -1) + # Pose + center = (face.aligned.size // 2, face.aligned.size // 2) + points = (face.aligned.pose.xyz_2d * face.aligned.size).astype("int32") + cv2.line(face.aligned.face, center, tuple(points[1]), (0, 255, 0), 1) + cv2.line(face.aligned.face, center, tuple(points[0]), (255, 0, 0), 1) + cv2.line(face.aligned.face, center, tuple(points[2]), (0, 0, 255), 1) + # Face centering + self._annotate_face_box(face.aligned) + # Legacy centering + roi = face.aligned.get_cropped_roi(face.aligned.size, self._legacy_size, "legacy") + cv2.rectangle(face.aligned.face, tuple(roi[:2]), tuple(roi[2:]), (0, 0, 255), 1) + self._print_stats(face.aligned) diff --git a/scripts/gui.py b/scripts/gui.py index da69b9042a..eea446f54b 100644 --- a/scripts/gui.py +++ b/scripts/gui.py @@ -10,7 +10,7 @@ get_images, initialize_images, initialize_config, LastSession, MainMenuBar, preview_trigger, ProcessWrapper, StatusBar) -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class FaceswapGui(tk.Tk): @@ -133,7 +133,7 @@ def rebuild(self): self._last_session.from_dict(session_state) logger.debug("GUI Redrawn") - def close_app(self, *args): # pylint: disable=unused-argument + def close_app(self, *args): # pylint:disable=unused-argument """ Close Python. This is here because the graph animation function continues to run even when tkinter has gone away """ @@ -144,7 +144,7 @@ def close_app(self, *args): # pylint: disable=unused-argument if not self._config.project.confirm_close(): return - if self._config.tk_vars["runningtask"].get(): + if self._config.tk_vars.running_task.get(): self.wrapper.task.terminate() self._last_session.save() @@ -161,7 +161,7 @@ def _confirm_close_on_running_task(self): ------- bool: ``True`` if user confirms close, ``False`` if user cancels close """ - if not self._config.tk_vars["runningtask"].get(): + if not self._config.tk_vars.running_task.get(): logger.debug("No tasks currently running") return True @@ -173,7 +173,7 @@ def _confirm_close_on_running_task(self): return True -class Gui(): # pylint: disable=too-few-public-methods +class Gui(): """ The GUI process. """ def __init__(self, arguments): self.root = FaceswapGui(arguments.debug) diff --git a/scripts/train.py b/scripts/train.py index 0ddc594808..bd455bcb4a 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1,25 +1,37 @@ #!/usr/bin python3 """ Main entry point to the training process of FaceSwap """ - +from __future__ import annotations import logging import os import sys +import typing as T -from threading import Lock from time import sleep +from threading import Event import cv2 +import numpy as np +from lib.gui.utils.image import TRAININGPREVIEW from lib.image import read_image_meta from lib.keypress import KBHit -from lib.multithreading import MultiThread -from lib.utils import (get_folder, get_image_paths, FaceswapError, _image_extensions) +from lib.multithreading import MultiThread, FSThread +from lib.training import Preview, PreviewBuffer, TriggerType +from lib.utils import (get_folder, get_image_paths, handle_deprecated_cliopts, + FaceswapError, IMAGE_EXTENSIONS) from plugins.plugin_loader import PluginLoader -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + import argparse + from collections.abc import Callable + from plugins.train.model._base import ModelBase + from plugins.train.trainer._base import TrainerBase + +logger = logging.getLogger(__name__) -class Train(): # pylint:disable=too-few-public-methods + +class Train(): """ The Faceswap Training Process. The training process is responsible for training a model on a set of source faces and a set of @@ -34,9 +46,10 @@ class Train(): # pylint:disable=too-few-public-methods The arguments to be passed to the training process as generated from Faceswap's command line arguments """ - def __init__(self, arguments): + def __init__(self, arguments: argparse.Namespace) -> None: logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments) - self._args = arguments + self._args = handle_deprecated_cliopts(arguments) + if self._args.summary: # If just outputting summary we don't need to initialize everything return @@ -45,18 +58,16 @@ def __init__(self, arguments): self._timelapse = self._set_timelapse() gui_cache = os.path.join( os.path.realpath(os.path.dirname(sys.argv[0])), "lib", "gui", ".cache") - self._gui_triggers = dict(update=os.path.join(gui_cache, ".preview_trigger"), - mask_toggle=os.path.join(gui_cache, ".preview_mask_toggle")) - self._stop = False - self._save_now = False - self._toggle_preview_mask = False - self._refresh_preview = False - self._preview_buffer = {} - self._lock = Lock() + self._gui_triggers: dict[T.Literal["mask", "refresh"], str] = { + "mask": os.path.join(gui_cache, ".preview_mask_toggle"), + "refresh": os.path.join(gui_cache, ".preview_trigger")} + self._stop: bool = False + self._save_now: bool = False + self._preview = PreviewInterface(self._args.preview) logger.debug("Initialized %s", self.__class__.__name__) - def _get_images(self): + def _get_images(self) -> dict[T.Literal["a", "b"], list[str]]: """ Check the image folders exist and contains valid extracted faces. Obtain image paths. Returns @@ -68,6 +79,7 @@ def _get_images(self): logger.debug("Getting image paths") images = {} for side in ("a", "b"): + side = T.cast(T.Literal["a", "b"], side) image_dir = getattr(self._args, f"input_{side}") if not os.path.isdir(image_dir): logger.error("Error: '%s' does not exist", image_dir) @@ -96,7 +108,7 @@ def _get_images(self): return images @classmethod - def _validate_image_counts(cls, images): + def _validate_image_counts(cls, images: dict[T.Literal["a", "b"], list[str]]) -> None: """ Validate that there are sufficient images to commence training without raising an error. @@ -124,7 +136,7 @@ def _validate_image_counts(cls, images): "Results are likely to be poor.") logger.warning(msg) - def _set_timelapse(self): + def _set_timelapse(self) -> dict[T.Literal["input_a", "input_b", "output"], str]: """ Set time-lapse paths if requested. Returns @@ -136,7 +148,7 @@ def _set_timelapse(self): if (not self._args.timelapse_input_a and not self._args.timelapse_input_b and not self._args.timelapse_output): - return None + return {} if (not self._args.timelapse_input_a or not self._args.timelapse_input_b or not self._args.timelapse_output): @@ -147,6 +159,7 @@ def _set_timelapse(self): timelapse_output = get_folder(self._args.timelapse_output) for side in ("a", "b"): + side = T.cast(T.Literal["a", "b"], side) folder = getattr(self._args, f"timelapse_input_{side}") if folder is not None and not os.path.isdir(folder): raise FaceswapError(f"The Timelapse path '{folder}' does not exist") @@ -156,7 +169,7 @@ def _set_timelapse(self): continue # Time-lapse folder is training folder filenames = [fname for fname in os.listdir(folder) - if os.path.splitext(fname)[-1].lower() in _image_extensions] + if os.path.splitext(fname)[-1].lower() in IMAGE_EXTENSIONS] if not filenames: raise FaceswapError(f"The Timelapse path '{folder}' does not contain any valid " "images") @@ -168,13 +181,14 @@ def _set_timelapse(self): raise FaceswapError(f"All images in the Timelapse folder '{folder}' must exist in " f"the training folder '{training_folder}'") - kwargs = {"input_a": self._args.timelapse_input_a, - "input_b": self._args.timelapse_input_b, - "output": timelapse_output} + TKey = T.Literal["input_a", "input_b", "output"] + kwargs = {T.cast(TKey, "input_a"): self._args.timelapse_input_a, + T.cast(TKey, "input_b"): self._args.timelapse_input_b, + T.cast(TKey, "output"): timelapse_output} logger.debug("Timelapse enabled: %s", kwargs) return kwargs - def process(self): + def process(self) -> None: """ The entry point for triggering the Training Process. Should only be called from :class:`lib.cli.launcher.ScriptExecutor` @@ -190,7 +204,7 @@ def process(self): self._end_thread(thread, err) logger.debug("Completed Training Process") - def _start_thread(self): + def _start_thread(self) -> MultiThread: """ Put the :func:`_training` into a background thread so we can keep control. Returns @@ -204,7 +218,7 @@ def _start_thread(self): logger.debug("Launched Trainer thread") return thread - def _end_thread(self, thread, err): + def _end_thread(self, thread: MultiThread, err: bool) -> None: """ Output message and join thread back to main on termination. Parameters @@ -231,19 +245,22 @@ def _end_thread(self, thread, err): sys.stdout.flush() logger.debug("Ended training thread") - def _training(self): + def _training(self) -> None: """ The training process to be run inside a thread. """ try: - sleep(1) # Let preview instructions flush out to logger + sleep(0.5) # Let preview instructions flush out to logger logger.debug("Commencing Training") logger.info("Loading data, this may take a while...") model = self._load_model() trainer = self._load_trainer(model) + if trainer.exit_early: + self._stop = True + return self._run_training_cycle(model, trainer) except KeyboardInterrupt: try: logger.debug("Keyboard Interrupt Caught. Saving Weights and exiting") - model.save() + model.io.save(is_exit=True) trainer.clear_tensorboard() except KeyboardInterrupt: logger.info("Saving model weights has been cancelled!") @@ -251,7 +268,7 @@ def _training(self): except Exception as err: raise err - def _load_model(self): + def _load_model(self) -> ModelBase: """ Load the model requested for training. Returns @@ -261,7 +278,7 @@ def _load_model(self): """ logger.debug("Loading Model") model_dir = get_folder(self._args.model_dir) - model = PluginLoader.get_model(self._args.trainer)( + model: ModelBase = PluginLoader.get_model(self._args.trainer)( model_dir, self._args, predict=False) @@ -269,7 +286,7 @@ def _load_model(self): logger.debug("Loaded Model") return model - def _load_trainer(self, model): + def _load_trainer(self, model: ModelBase) -> TrainerBase: """ Load the trainer requested for training. Parameters @@ -283,15 +300,15 @@ def _load_trainer(self, model): The requested model trainer plugin """ logger.debug("Loading Trainer") - trainer = PluginLoader.get_trainer(model.trainer) - trainer = trainer(model, - self._images, - self._args.batch_size, - self._args.configfile) + base = PluginLoader.get_trainer(model.trainer) + trainer: TrainerBase = base(model, + self._images, + self._args.batch_size, + self._args.configfile) logger.debug("Loaded Trainer") return trainer - def _run_training_cycle(self, model, trainer): + def _run_training_cycle(self, model: ModelBase, trainer: TrainerBase) -> None: """ Perform the training cycle. Handles the background training, updating previews/time-lapse on each save interval, @@ -305,56 +322,53 @@ def _run_training_cycle(self, model, trainer): The requested model trainer plugin """ logger.debug("Running Training Cycle") + update_preview_images = False if self._args.write_image or self._args.redirect_gui or self._args.preview: - display_func = self._show + display_func: Callable | None = self._show else: display_func = None for iteration in range(1, self._args.iterations + 1): - logger.trace("Training iteration: %s", iteration) + logger.trace("Training iteration: %s", iteration) # type:ignore save_iteration = iteration % self._args.save_interval == 0 or iteration == 1 + gui_triggers = self._process_gui_triggers() - if self._toggle_preview_mask: + if self._preview.should_toggle_mask or gui_triggers["mask"]: trainer.toggle_mask() - self._toggle_preview_mask = False - self._refresh_preview = True + update_preview_images = True - if save_iteration or self._save_now or self._refresh_preview: + if self._preview.should_refresh or gui_triggers["refresh"] or update_preview_images: viewer = display_func + update_preview_images = False else: viewer = None - timelapse = self._timelapse if save_iteration else None + + timelapse = self._timelapse if save_iteration else {} trainer.train_one_step(viewer, timelapse) + + if viewer is not None and not save_iteration: + # Spammy but required by GUI to know to update window + print("") + logger.info("[Preview Updated]") + if self._stop: logger.debug("Stop received. Terminating") break - if self._refresh_preview and viewer is not None: - if self._args.redirect_gui: # Remove any gui trigger files following an update - print("\n") - logger.info("[Preview Updated]") - self._refresh_preview = False - - if save_iteration: - logger.debug("Save Iteration: (iteration: %s", iteration) - model.save() - elif self._save_now: - logger.debug("Save Requested: (iteration: %s", iteration) - model.save() + if save_iteration or self._save_now: + logger.debug("Saving (save_iterations: %s, save_now: %s) Iteration: " + "(iteration: %s)", save_iteration, self._save_now, iteration) + model.io.save(is_exit=False) self._save_now = False + update_preview_images = True + logger.debug("Training cycle complete") - model.save() + model.io.save(is_exit=True) trainer.clear_tensorboard() self._stop = True - def _monitor(self, thread): - """ Monitor the background :func:`_training` thread for key presses and errors. - - Returns - ------- - bool - ``True`` if there has been an error in the background thread otherwise ``False`` - """ + def _output_startup_info(self) -> None: + """ Print the startup information to the console. """ logger.debug("Launching Monitor") logger.info("===================================================") logger.info(" Starting") @@ -362,28 +376,77 @@ def _monitor(self, thread): logger.info(" Using live preview") if sys.stdout.isatty(): logger.info(" Press '%s' to save and quit", - "Stop" if self._args.redirect_gui or self._args.colab else "ENTER") - if not self._args.redirect_gui and not self._args.colab and sys.stdout.isatty(): + "Stop" if self._args.redirect_gui else "ENTER") + if not self._args.redirect_gui and sys.stdout.isatty(): logger.info(" Press 'S' to save model weights immediately") logger.info("===================================================") + def _check_keypress(self, keypress: KBHit) -> bool: + """ Check if a keypress has been detected. + + Parameters + ---------- + keypress: :class:`lib.keypress.KBHit` + The keypress monitor + + Returns + ------- + bool + ``True`` if an exit keypress has been detected otherwise ``False`` + """ + retval = False + if keypress.kbhit(): + console_key = keypress.getch() + if console_key in ("\n", "\r"): + logger.debug("Exit requested") + retval = True + if console_key in ("s", "S"): + logger.info("Save requested") + self._save_now = True + return retval + + def _process_gui_triggers(self) -> dict[T.Literal["mask", "refresh"], bool]: + """ Check whether a file drop has occurred from the GUI to manually update the preview. + + Returns + ------- + dict + The trigger name as key and boolean as value + """ + retval: dict[T.Literal["mask", "refresh"], bool] = {key: False + for key in self._gui_triggers} + if not self._args.redirect_gui: + return retval + + for trigger, filename in self._gui_triggers.items(): + if os.path.isfile(filename): + logger.debug("GUI Trigger received for: '%s'", trigger) + retval[trigger] = True + logger.debug("Removing gui trigger file: %s", filename) + os.remove(filename) + if trigger == "refresh": + print("") # Let log print on different line from loss output + logger.info("Refresh preview requested...") + return retval + + def _monitor(self, thread: MultiThread) -> bool: + """ Monitor the background :func:`_training` thread for key presses and errors. + + Parameters + ---------- + thread: :class:~`lib.multithreading.MultiThread` + The thread containing the training loop + + Returns + ------- + bool + ``True`` if there has been an error in the background thread otherwise ``False`` + """ + self._output_startup_info() keypress = KBHit(is_gui=self._args.redirect_gui) - window_created = False err = False while True: try: - if self._args.preview: - with self._lock: - for name, image in self._preview_buffer.items(): - if not window_created: - self._create_resizable_window(name, image.shape) - cv2.imshow(name, image) # pylint: disable=no-member - if not window_created: - window_created = bool(self._preview_buffer) - cv_key = cv2.waitKey(1000) # pylint: disable=no-member - else: - cv_key = None - if thread.has_error: logger.debug("Thread error detected") err = True @@ -393,103 +456,25 @@ def _monitor(self, thread): break # Preview Monitor - if not self._preview_monitor(cv_key): + if self._preview.should_quit: break + if self._preview.should_save: + self._save_now = True # Console Monitor - if keypress.kbhit(): - console_key = keypress.getch() - if console_key in ("\n", "\r"): - logger.debug("Exit requested") - break - if console_key in ("s", "S"): - logger.info("Save requested") - self._save_now = True - - # GUI Preview trigger update monitor - self._process_gui_triggers() + if self._check_keypress(keypress): + break # Exit requested sleep(1) except KeyboardInterrupt: logger.debug("Keyboard Interrupt received") break + self._preview.shutdown() keypress.set_normal_term() logger.debug("Closed Monitor") return err - @classmethod - def _create_resizable_window(cls, name: str, image_shape: tuple) -> None: - """ Create a resizable OpenCV window to hold the preview image. - - Parameters - ---------- - name: str - The name to display in the window header and for window identification - shape: tuple - The (`rows`, `columns`, `channels`) of the image to be displayed - """ - logger.debug("Creating named window '%s' for image shape %s", name, image_shape) - height, width = image_shape[:2] - cv2.namedWindow(name, cv2.WINDOW_GUI_EXPANDED) - cv2.resizeWindow(name, width, height) - - def _preview_monitor(self, key_press): - """ Monitors keyboard presses on the pop-up OpenCV Preview Window. - - Parameters - ---------- - key_press: str - The key press received from OpenCV or ``None`` if no press received - - Returns - ------- - bool - ``True`` if the process should continue training. ``False`` if an exit has been - requested and process should terminate - """ - if not self._args.preview: - return True - - if key_press == ord("\n") or key_press == ord("\r"): - logger.debug("Exit requested") - return False - - if key_press == ord("s"): - print("\n") - logger.info("Save requested") - self._save_now = True - if key_press == ord("r"): - print("\n") - logger.info("Refresh preview requested") - self._refresh_preview = True - if key_press == ord("m"): - print("\n") - logger.verbose("Toggle mask display requested") - self._toggle_preview_mask = True - - return True - - def _process_gui_triggers(self): - """ Check whether a file drop has occurred from the GUI to manually update the preview. """ - if not self._args.redirect_gui: - return - - parent_flags = dict(mask_toggle="_toggle_preview_mask", update="_refresh_preview") - for trigger in ("mask_toggle", "update"): - filename = self._gui_triggers[trigger] - if os.path.isfile(filename): - logger.debug("GUI Trigger received for: '%s'", trigger) - - logger.debug("Removing gui trigger file: %s", filename) - os.remove(filename) - - if trigger == "update": - print("\n") - logger.info("Refresh preview requested") - - setattr(self, parent_flags[trigger], True) - - def _show(self, image, name=""): + def _show(self, image: np.ndarray, name: str = "") -> None: """ Generate the preview and write preview file output. Handles the output and display of preview images. @@ -500,30 +485,132 @@ def _show(self, image, name=""): The preview image to be displayed and/or written out name: str, optional The name of the image for saving or display purposes. If an empty string is passed - then it will automatically be names. Default: "" + then it will automatically be named. Default: "" """ logger.debug("Updating preview: (name: %s)", name) try: scriptpath = os.path.realpath(os.path.dirname(sys.argv[0])) if self._args.write_image: logger.debug("Saving preview to disk") - img = "training_preview.jpg" + img = "training_preview.png" imgfile = os.path.join(scriptpath, img) - cv2.imwrite(imgfile, image) # pylint: disable=no-member + cv2.imwrite(imgfile, image) # pylint:disable=no-member logger.debug("Saved preview to: '%s'", img) if self._args.redirect_gui: logger.debug("Generating preview for GUI") - img = ".gui_training_preview.jpg" - imgfile = os.path.join(scriptpath, "lib", "gui", - ".cache", "preview", img) - cv2.imwrite(imgfile, image) # pylint: disable=no-member - logger.debug("Generated preview for GUI: '%s'", img) + img = TRAININGPREVIEW + imgfile = os.path.join(scriptpath, "lib", "gui", ".cache", "preview", img) + cv2.imwrite(imgfile, image) # pylint:disable=no-member + logger.debug("Generated preview for GUI: '%s'", imgfile) if self._args.preview: logger.debug("Generating preview for display: '%s'", name) - with self._lock: - self._preview_buffer[name] = image + self._preview.buffer.add_image(name, image) logger.debug("Generated preview for display: '%s'", name) except Exception as err: logging.error("could not preview sample") raise err logger.debug("Updated preview: (name: %s)", name) + + +class PreviewInterface(): + """ Run the preview window in a thread and interface with it + + Parameters + ---------- + use_preview: bool + ``True`` if pop-up preview window has been requested otherwise ``False`` + """ + def __init__(self, use_preview: bool) -> None: + self._active = use_preview + self._triggers: TriggerType = {"toggle_mask": Event(), + "refresh": Event(), + "save": Event(), + "quit": Event(), + "shutdown": Event()} + self._buffer = PreviewBuffer() + self._thread = self._launch_thread() + + @property + def buffer(self) -> PreviewBuffer: + """ :class:`PreviewBuffer`: The thread save preview image object """ + return self._buffer + + @property + def should_toggle_mask(self) -> bool: + """ bool: Check whether the mask should be toggled and return the value. If ``True`` is + returned then resets mask toggle back to ``False`` """ + if not self._active: + return False + retval = self._triggers["toggle_mask"].is_set() + if retval: + logger.debug("Sending toggle mask") + self._triggers["toggle_mask"].clear() + return retval + + @property + def should_refresh(self) -> bool: + """ bool: Check whether the preview should be updated and return the value. If ``True`` is + returned then resets the refresh trigger back to ``False`` """ + if not self._active: + return False + retval = self._triggers["refresh"].is_set() + if retval: + logger.debug("Sending should refresh") + self._triggers["refresh"].clear() + return retval + + @property + def should_save(self) -> bool: + """ bool: Check whether a save request has been made. If ``True`` is returned then save + trigger is set back to ``False`` """ + if not self._active: + return False + retval = self._triggers["save"].is_set() + if retval: + logger.debug("Sending should save") + self._triggers["save"].clear() + return retval + + @property + def should_quit(self) -> bool: + """ bool: Check whether an exit request has been made. ``True`` if an exit request has + been made otherwise ``False``. + + Raises + ------ + Error + Re-raises any error within the preview thread + """ + if self._thread is None: + return False + + self._thread.check_and_raise_error() + + retval = self._triggers["quit"].is_set() + if retval: + logger.debug("Sending should stop") + return retval + + def _launch_thread(self) -> FSThread | None: + """ Launch the preview viewer in it's own thread if preview has been selected + + Returns + ------- + :class:`lib.multithreading.FSThread` or ``None`` + The thread that holds the preview viewer if preview is selected otherwise ``None`` + """ + if not self._active: + return None + thread = FSThread(target=Preview, + name="preview", + args=(self._buffer, ), + kwargs={"triggers": self._triggers}) + thread.start() + return thread + + def shutdown(self) -> None: + """ Send a signal to shutdown the preview window. """ + if not self._active: + return + logger.debug("Sending shutdown to preview viewer") + self._triggers["shutdown"].set() diff --git a/setup.cfg b/setup.cfg index 13459d1465..6427fca936 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,3 +1,57 @@ [flake8] max-line-length = 99 +max-complexity=10 +statistics = True +count = True exclude = .git, __pycache__ +per-file-ignores = + __init__.py:F401 + lib/gpu_stats/directml.py:E402 + +[mypy] +[mypy-comtypes.*] +ignore_missing_imports = True +[mypy-cv2.*] +ignore_missing_imports = True +[mypy-fastcluster.*] +ignore_missing_imports = True +[mypy-ffmpy.*] +ignore_missing_imports = True +[mypy-imageio.*] +ignore_missing_imports = True +[mypy-imageio_ffmpeg.*] +ignore_missing_imports = True +[mypy-keras.*] +ignore_missing_imports = True +[mypy-matplotlib.*] +ignore_missing_imports = True +[mypy-numexpr.*] +ignore_missing_imports = True +[mypy-numpy.*] +ignore_missing_imports = True +[mypy-numpy.core._multiarray_umath.*] +ignore_missing_imports = True +[mypy-pexpect.*] +ignore_missing_imports = True +[mypy-PIL.*] +ignore_missing_imports = True +[mypy-psutil.*] +ignore_missing_imports = True +[mypy-pynvml.*] +ignore_missing_imports = True +[mypy-pynvx.*] +ignore_missing_imports = True +[mypy-pytest.*] +ignore_missing_imports = True +[mypy-scipy.*] +ignore_missing_imports = True +[mypy-sklearn.*] +ignore_missing_imports = True +[mypy-tensorflow.*] +ignore_missing_imports = True +[mypy-tqdm.*] +ignore_missing_imports = True +[mypy-win32console.*] +ignore_missing_imports = True +[mypy-winpty.*] +ignore_missing_imports = True diff --git a/setup.py b/setup.py index 1dc2cd5b83..56c0188a55 100755 --- a/setup.py +++ b/setup.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 """ Install packages for faceswap.py """ +# pylint:disable=too-many-lines -# >>> Environment +import logging import ctypes import json import locale @@ -10,61 +11,87 @@ import os import re import sys -from subprocess import CalledProcessError, run, PIPE, Popen -from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union - -from pkg_resources import parse_requirements, Requirement - -if TYPE_CHECKING: - from logging import Logger - - -INSTALL_FAILED = False +import typing as T +from shutil import which +from subprocess import list2cmdline, PIPE, Popen, run, STDOUT + +from pkg_resources import parse_requirements + +from lib.logger import log_setup + +logger = logging.getLogger(__name__) +backend_type: T.TypeAlias = T.Literal['nvidia', 'apple_silicon', 'directml', 'cpu', 'rocm', "all"] + +_INSTALL_FAILED = False +# Packages that are explicitly required for setup.py +_INSTALLER_REQUIREMENTS: list[tuple[str, str]] = [("pexpect>=4.8.0", "!Windows"), + ("pywinpty==2.0.2", "Windows")] +# Conda packages that are required for a specific backend +# TODO zlib-wapi is required on some Windows installs where cuDNN complains: +# Could not locate zlibwapi.dll. Please make sure it is in your library path! +# This only seems to occur on Anaconda cuDNN not conda-forge +_BACKEND_SPECIFIC_CONDA: dict[backend_type, list[str]] = { + "nvidia": ["cudatoolkit", "cudnn", "zlib-wapi"], + "apple_silicon": ["libblas"]} +# Packages that should only be installed through pip +_FORCE_PIP: dict[backend_type, list[str]] = { + "nvidia": ["tensorflow"], + "all": [ + "tensorflow-cpu", # conda-forge leads to flatbuffer errors because of mixed sources + "imageio-ffmpeg"]} # 17/11/23 Conda forge uses incorrect ffmpeg, so fallback to pip # Revisions of tensorflow GPU and cuda/cudnn requirements. These relate specifically to the # Tensorflow builds available from pypi -TENSORFLOW_REQUIREMENTS = {">=2.4.0,<2.5.0": ["11.0", "8.0"], - ">=2.5.0,<2.9.0": ["11.2", "8.1"]} +_TENSORFLOW_REQUIREMENTS = {">=2.10.0,<2.11.0": [">=11.2,<11.3", ">=8.1,<8.2"]} +# ROCm min/max version requirements for Tensorflow +_TENSORFLOW_ROCM_REQUIREMENTS = {">=2.10.0,<2.11.0": ((5, 2, 0), (5, 4, 0))} +# TODO tensorflow-metal versioning + # Mapping of Python packages to their conda names if different from pip or in non-default channel -CONDA_MAPPING = { - # "opencv-python": ("opencv", "conda-forge"), # Periodic issues with conda-forge opencv +_CONDA_MAPPING: dict[str, tuple[str, str]] = { + "cudatoolkit": ("cudatoolkit", "conda-forge"), + "cudnn": ("cudnn", "conda-forge"), "fastcluster": ("fastcluster", "conda-forge"), - "imageio-ffmpeg": ("imageio-ffmpeg", "conda-forge"), + "ffmpy": ("ffmpy", "conda-forge"), + # "imageio-ffmpeg": ("imageio-ffmpeg", "conda-forge"), + "nvidia-ml-py": ("nvidia-ml-py", "conda-forge"), "tensorflow-deps": ("tensorflow-deps", "apple"), - "libblas": ("libblas", "conda-forge")} + "libblas": ("libblas", "conda-forge"), + "zlib-wapi": ("zlib-wapi", "conda-forge"), + "xorg-libxft": ("xorg-libxft", "conda-forge")} + +# Force output to utf-8 +sys.stdout.reconfigure(encoding="utf-8", errors="replace") # type:ignore[attr-defined] class Environment(): - """ The current install environment """ - def __init__(self, logger: Optional["Logger"] = None, updater: bool = False) -> None: - """ logger will override built in Output() function if passed in - updater indicates that this is being run from update_deps.py - so certain steps can be skipped/output limited """ - self.conda_required_packages: List[Tuple[str, ...]] = [("tk", )] - self.output: Union["Logger", "Output"] = logger if logger else Output() + """ The current install environment + + Parameters + ---------- + updater: bool, Optional + ``True`` if the script is being called by Faceswap's internal updater. ``False`` if full + setup is running. Default: ``False`` + """ + + _backends = (("nvidia", "apple_silicon", "directml", "rocm", "cpu")) + + def __init__(self, updater: bool = False) -> None: self.updater = updater # Flag that setup is being run by installer so steps can be skipped self.is_installer: bool = False - self.cuda_version: str = "" - self.cudnn_version: str = "" - self.enable_amd: bool = False - self.enable_apple_silicon: bool = False + self.backend: backend_type | None = None self.enable_docker: bool = False - self.enable_cuda: bool = False - self.required_packages: List[Tuple[str, Tuple[str, str]]] = [] - self.missing_packages: List[str] = [] - self.conda_missing_packages: List[str] = [] - - self.process_arguments() - self.check_permission() - self.check_system() - self.check_python() - self.output_runtime_info() - self.check_pip() - self.upgrade_pip() - self.set_ld_library_path() - - self.installed_packages = self.get_installed_packages() - self.installed_packages.update(self.get_installed_conda_packages()) + self.cuda_cudnn = ["", ""] + self.rocm_version: tuple[int, ...] = (0, 0, 0) + + self._process_arguments() + self._check_permission() + self._check_system() + self._check_python() + self._output_runtime_info() + self._check_pip() + self._upgrade_pip() + self._set_env_vars() @property def encoding(self) -> str: @@ -72,12 +99,12 @@ def encoding(self) -> str: return locale.getpreferredencoding() @property - def os_version(self) -> Tuple[str, str]: + def os_version(self) -> tuple[str, str]: """ Get OS Version """ return platform.system(), platform.release() @property - def py_version(self) -> Tuple[str, str]: + def py_version(self) -> tuple[str, str]: """ Get Python Version """ return platform.python_version(), platform.architecture()[0] @@ -91,11 +118,21 @@ def is_conda(self) -> bool: def is_admin(self) -> bool: """ Check whether user is admin """ try: - retval = os.getuid() == 0 + retval = os.getuid() == 0 # type: ignore except AttributeError: retval = ctypes.windll.shell32.IsUserAnAdmin() != 0 # type: ignore return retval + @property + def cuda_version(self) -> str: + """ str: The detected globally installed Cuda Version """ + return self.cuda_cudnn[0] + + @property + def cudnn_version(self) -> str: + """ str: The detected globally installed cuDNN Version """ + return self.cuda_cudnn[1] + @property def is_virtualenv(self) -> bool: """ Check whether this is a virtual environment """ @@ -104,115 +141,85 @@ def is_virtualenv(self) -> bool: (hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix)) else: prefix = os.path.dirname(sys.prefix) - retval = (os.path.basename(prefix) == "envs") + retval = os.path.basename(prefix) == "envs" return retval - def process_arguments(self) -> None: + def _process_arguments(self) -> None: """ Process any cli arguments and dummy in cli arguments if calling from updater. """ args = [arg for arg in sys.argv] # pylint:disable=unnecessary-comprehension if self.updater: from lib.utils import get_backend # pylint:disable=import-outside-toplevel args.append(f"--{get_backend()}") + logger.debug(args) for arg in args: if arg == "--installer": self.is_installer = True - if arg == "--nvidia": - self.enable_cuda = True - if arg == "--amd": - self.enable_amd = True - if arg == "--apple-silicon": - self.enable_apple_silicon = True + if not self.backend and (arg.startswith("--") and + arg.replace("--", "") in self._backends): + self.backend = arg.replace("--", "").lower() # type:ignore - def get_required_packages(self) -> None: - """ Load requirements list """ - if self.enable_amd: - suffix = "amd.txt" - elif self.enable_cuda: - suffix = "nvidia.txt" - elif self.enable_apple_silicon: - suffix = "apple_silicon.txt" - else: - suffix = "cpu.txt" - req_files = ["_requirements_base.txt", f"requirements_{suffix}"] - pypath = os.path.dirname(os.path.realpath(__file__)) - requirements = [] - for req_file in req_files: - requirements_file = os.path.join(pypath, "requirements", req_file) - with open(requirements_file, encoding="utf8") as req: - for package in req.readlines(): - package = package.strip() - if package and (not package.startswith(("#", "-r"))): - requirements.append(package) - self.required_packages = [(pkg.name, pkg.specs) - for pkg in parse_requirements(requirements) - if pkg.marker is None or pkg.marker.evaluate()] - - def check_permission(self) -> None: + def _check_permission(self) -> None: """ Check for Admin permissions """ if self.updater: return if self.is_admin: - self.output.info("Running as Root/Admin") + logger.info("Running as Root/Admin") else: - self.output.info("Running without root/admin privileges") + logger.info("Running without root/admin privileges") - def check_system(self) -> None: + def _check_system(self) -> None: """ Check the system """ if not self.updater: - self.output.info("The tool provides tips for installation\n" - "and installs required python packages") - self.output.info(f"Setup in {self.os_version[0]} {self.os_version[1]}") + logger.info("The tool provides tips for installation and installs required python " + "packages") + logger.info("Setup in %s %s", self.os_version[0], self.os_version[1]) if not self.updater and not self.os_version[0] in ["Windows", "Linux", "Darwin"]: - self.output.error(f"Your system {self.os_version[0]} is not supported!") + logger.error("Your system %s is not supported!", self.os_version[0]) sys.exit(1) if self.os_version[0].lower() == "darwin" and platform.machine() == "arm64": - self.enable_apple_silicon = True + self.backend = "apple_silicon" if not self.updater and not self.is_conda: - self.output.error("Setting up Faceswap for Apple Silicon outside of a Conda " - "environment is unsupported") + logger.error("Setting up Faceswap for Apple Silicon outside of a Conda " + "environment is unsupported") sys.exit(1) - def check_python(self) -> None: + def _check_python(self) -> None: """ Check python and virtual environment status """ - self.output.info(f"Installed Python: {self.py_version[0]} {self.py_version[1]}") + logger.info("Installed Python: %s %s", self.py_version[0], self.py_version[1]) if self.updater: return - if not ((3, 7) <= sys.version_info < (3, 10) and self.py_version[1] == "64bit"): - self.output.error("Please run this script with Python version 3.7 to 3.9 " - "64bit and try again.") - sys.exit(1) - if self.enable_amd and sys.version_info >= (3, 9): - self.output.error("The AMD version of Faceswap cannot be installed on versions of " - "Python higher than 3.8") + if not ((3, 10) <= sys.version_info < (3, 11) and self.py_version[1] == "64bit"): + logger.error("Please run this script with Python version 3.10 64bit and try " + "again.") sys.exit(1) - def output_runtime_info(self) -> None: + def _output_runtime_info(self) -> None: """ Output run time info """ if self.is_conda: - self.output.info("Running in Conda") + logger.info("Running in Conda") if self.is_virtualenv: - self.output.info("Running in a Virtual Environment") - self.output.info(f"Encoding: {self.encoding}") + logger.info("Running in a Virtual Environment") + logger.info("Encoding: %s", self.encoding) - def check_pip(self) -> None: + def _check_pip(self) -> None: """ Check installed pip version """ if self.updater: return try: import pip # noqa pylint:disable=unused-import,import-outside-toplevel except ImportError: - self.output.error("Import pip failed. Please Install python3-pip and try again") + logger.error("Import pip failed. Please Install python3-pip and try again") sys.exit(1) - def upgrade_pip(self) -> None: + def _upgrade_pip(self) -> None: """ Upgrade pip to latest version """ if not self.is_conda: # Don't do this with Conda, as we must use Conda version of pip - self.output.info("Upgrading pip...") + logger.info("Upgrading pip...") pipexe = [sys.executable, "-m", "pip"] pipexe.extend(["install", "--no-cache-dir", "-qq", "--upgrade"]) if not self.is_admin and not self.is_virtualenv: @@ -221,24 +228,227 @@ def upgrade_pip(self) -> None: run(pipexe, check=True) import pip # pylint:disable=import-outside-toplevel pip_version = pip.__version__ - self.output.info(f"Installed pip: {pip_version}") + logger.info("Installed pip: %s", pip_version) + + def set_config(self) -> None: + """ Set the backend in the faceswap config file """ + config = {"backend": self.backend} + pypath = os.path.dirname(os.path.realpath(__file__)) + config_file = os.path.join(pypath, "config", ".faceswap") + with open(config_file, "w", encoding="utf8") as cnf: + json.dump(config, cnf) + logger.info("Faceswap config written to: %s", config_file) - def get_installed_packages(self) -> Dict[str, str]: - """ Get currently installed packages """ + def _set_env_vars(self) -> None: + """ There are some foibles under Conda which need to be worked around in different + situations. + + Linux: + Update the LD_LIBRARY_PATH environment variable when activating a conda environment + and revert it when deactivating. + + Notes + ----- + From Tensorflow 2.7, installing Cuda Toolkit from conda-forge and tensorflow from pip + causes tensorflow to not be able to locate shared libs and hence not use the GPU. + We update the environment variable for all instances using Conda as it shouldn't hurt + anything and may help avoid conflicts with globally installed Cuda + """ + if not self.is_conda: + return + + linux_update = self.os_version[0].lower() == "linux" and self.backend == "nvidia" + + if not linux_update: + return + + conda_prefix = os.environ["CONDA_PREFIX"] + activate_folder = os.path.join(conda_prefix, "etc", "conda", "activate.d") + deactivate_folder = os.path.join(conda_prefix, "etc", "conda", "deactivate.d") + os.makedirs(activate_folder, exist_ok=True) + os.makedirs(deactivate_folder, exist_ok=True) + + activate_script = os.path.join(conda_prefix, activate_folder, "env_vars.sh") + deactivate_script = os.path.join(conda_prefix, deactivate_folder, "env_vars.sh") + + if os.path.isfile(activate_script): + # Only create file if it does not already exist. There may be instances where people + # have created their own scripts, but these should be few and far between and those + # people should already know what they are doing. + return + + conda_libs = os.path.join(conda_prefix, "lib") + activate = ["#!/bin/sh\n\n", + "export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\n", + f"export LD_LIBRARY_PATH='{conda_libs}':${{LD_LIBRARY_PATH}}\n"] + deactivate = ["#!/bin/sh\n\n", + "export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\n", + "unset OLD_LD_LIBRARY_PATH\n"] + logger.info("Cuda search path set to '%s'", conda_libs) + + with open(activate_script, "w", encoding="utf8") as afile: + afile.writelines(activate) + with open(deactivate_script, "w", encoding="utf8") as afile: + afile.writelines(deactivate) + + +class Packages(): + """ Holds information about installed and required packages. + Handles updating dependencies based on running platform/backend + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + """ + def __init__(self, environment: Environment) -> None: + self._env = environment + + # Default TK has bad fonts under Linux. There is a better build in Conda-Forge, so set + # channel accordingly + tk_channel = "conda-forge" if self._env.os_version[0].lower() == "linux" else "default" + self._conda_required_packages: list[tuple[list[str] | str, str]] = [("tk", tk_channel), + ("git", "default")] + self._update_backend_specific_conda() + self._installed_packages = self._get_installed_packages() + self._conda_installed_packages = self._get_installed_conda_packages() + self._required_packages: list[tuple[str, list[tuple[str, str]]]] = [] + self._missing_packages: list[tuple[str, list[tuple[str, str]]]] = [] + self._conda_missing_packages: list[tuple[list[str] | str, str]] = [] + + @property + def prerequisites(self) -> list[tuple[str, list[tuple[str, str]]]]: + """ list: Any required packages that the installer needs prior to installing the faceswap + environment on the specific platform that are not already installed """ + all_installed = self._all_installed_packages + candidates = self._format_requirements( + [pkg for pkg, plat in _INSTALLER_REQUIREMENTS + if self._env.os_version[0] == plat or (plat[0] == "!" and + self._env.os_version[0] != plat[1:])]) + retval = [(pkg, spec) for pkg, spec in candidates + if pkg not in all_installed or ( + pkg in all_installed and + not self._validate_spec(spec, all_installed.get(pkg, "")) + )] + return retval + + @property + def packages_need_install(self) -> bool: + """bool: ``True`` if there are packages available that need to be installed """ + return bool(self._missing_packages or self._conda_missing_packages) + + @property + def to_install(self) -> list[tuple[str, list[tuple[str, str]]]]: + """ list: The required packages that need to be installed """ + return self._missing_packages + + @property + def to_install_conda(self) -> list[tuple[list[str] | str, str]]: + """ list: The required conda packages that need to be installed """ + return self._conda_missing_packages + + @property + def _all_installed_packages(self) -> dict[str, str]: + """ dict[str, str]: The package names and version string for all installed packages across + pip and conda """ + return {**self._installed_packages, **self._conda_installed_packages} + + def _update_backend_specific_conda(self) -> None: + """ Add backend specific packages to Conda required packages """ + assert self._env.backend is not None + to_add = _BACKEND_SPECIFIC_CONDA.get(self._env.backend) + if not to_add: + logger.debug("No backend packages to add for '%s'. All optional packages: %s", + self._env.backend, _BACKEND_SPECIFIC_CONDA) + return + + combined_cuda = [] + for pkg in to_add: + pkg, channel = _CONDA_MAPPING.get(pkg, (pkg, "")) + if pkg == "zlib-wapi" and self._env.os_version[0].lower() != "windows": + # TODO move this front and center + continue + if pkg in ("cudatoolkit", "cudnn"): # TODO Handle multiple cuda/cudnn requirements + idx = 0 if pkg == "cudatoolkit" else 1 + pkg = f"{pkg}{list(_TENSORFLOW_REQUIREMENTS.values())[0][idx]}" + + combined_cuda.append(pkg) + continue + + self._conda_required_packages.append((pkg, channel)) + logger.info("Adding conda required package '%s' for backend '%s')", + pkg, self._env.backend) + + if combined_cuda: + self._conda_required_packages.append((combined_cuda, channel)) + logger.info("Adding conda required package '%s' for backend '%s')", + combined_cuda, self._env.backend) + + @classmethod + def _format_requirements(cls, packages: list[str] + ) -> list[tuple[str, list[tuple[str, str]]]]: + """ Parse a list of requirements.txt formatted package strings to a list of pkgresource + formatted requirements """ + return [(package.unsafe_name, package.specs) + for package in parse_requirements(packages) + if package.marker is None or package.marker.evaluate()] + + @classmethod + def _validate_spec(cls, + required: list[tuple[str, str]], + existing: str) -> bool: + """ Validate whether the required specification for a package is met by the installed + version. + + required: list[tuple[str, str]] + The required package version spec to check + existing: str + The version of the installed package + + Returns + ------- + bool + ``True`` if the required specification is met by the existing specification + """ + ops = {"==": operator.eq, ">=": operator.ge, "<=": operator.le, + ">": operator.gt, "<": operator.lt} + if not required: + return True + + return all(ops[spec[0]]([int(s) for s in existing.split(".")], + [int(s) for s in spec[1].split(".")]) + for spec in required) + + def _get_installed_packages(self) -> dict[str, str]: + """ Get currently installed packages and add to :attr:`_installed_packages` + + Returns + ------- + dict[str, str] + The installed package name and version string + """ installed_packages = {} with Popen(f"\"{sys.executable}\" -m pip freeze --local", shell=True, stdout=PIPE) as chk: - installed = chk.communicate()[0].decode(self.encoding).splitlines() + installed = chk.communicate()[0].decode(self._env.encoding, + errors="ignore").splitlines() for pkg in installed: if "==" not in pkg: continue item = pkg.split("==") installed_packages[item[0]] = item[1] + logger.debug(installed_packages) return installed_packages - def get_installed_conda_packages(self) -> Dict[str, str]: - """ Get currently installed conda packages """ - if not self.is_conda: + def _get_installed_conda_packages(self) -> dict[str, str]: + """ Get currently installed conda packages + + Returns + ------- + dict[str, str] + The installed package name and version string + """ + if not self._env.is_conda: return {} chk = os.popen("conda list").read() installed = [re.sub(" +", " ", line.strip()) @@ -247,288 +457,412 @@ def get_installed_conda_packages(self) -> Dict[str, str]: for pkg in installed: item = pkg.split(" ") retval[item[0]] = item[1] + logger.debug(retval) return retval - def update_tf_dep(self) -> None: - """ Update Tensorflow Dependency """ - if self.is_conda or not self.enable_cuda: - # CPU/AMD doesn't need Cuda and Conda handles Cuda and cuDNN so nothing to do here - return + def get_required_packages(self) -> None: + """ Load the requirements from the backend specific requirements list """ + req_files = ["_requirements_base.txt", f"requirements_{self._env.backend}.txt"] + pypath = os.path.dirname(os.path.realpath(__file__)) + requirements = [] + for req_file in req_files: + requirements_file = os.path.join(pypath, "requirements", req_file) + with open(requirements_file, encoding="utf8") as req: + for package in req.readlines(): + package = package.strip() + if package and (not package.startswith(("#", "-r"))): + requirements.append(package) + self._required_packages = self._format_requirements(requirements) + logger.debug(self._required_packages) + + def _update_tf_dep_nvidia(self) -> None: + """ Update the Tensorflow dependency for global Cuda installs """ + if self._env.is_conda: # Conda handles Cuda and cuDNN so nothing to do here + return tf_ver = None - cudnn_inst = self.cudnn_version.split(".") - for key, val in TENSORFLOW_REQUIREMENTS.items(): - cuda_req = val[0] - cudnn_req = val[1].split(".") - if cuda_req == self.cuda_version and (cudnn_req[0] == cudnn_inst[0] and - cudnn_req[1] <= cudnn_inst[1]): + cuda_inst = self._env.cuda_version + cudnn_inst = self._env.cudnn_version + if len(cudnn_inst) == 1: # Sometimes only major version is reported + cudnn_inst = f"{cudnn_inst}.0" + for key, val in _TENSORFLOW_REQUIREMENTS.items(): + cuda_req = next(parse_requirements(f"cuda{val[0]}")).specs + cudnn_req = next(parse_requirements(f"cudnn{val[1]}")).specs + if (self._validate_spec(cuda_req, cuda_inst) + and self._validate_spec(cudnn_req, cudnn_inst)): tf_ver = key break + if tf_ver: # Remove the version of tensorflow in requirements file and add the correct version # that corresponds to the installed Cuda/cuDNN versions - self.required_packages = [pkg for pkg in self.required_packages - if not pkg[0].startswith("tensorflow-gpu")] - tf_ver = f"tensorflow-gpu{tf_ver}" - - tf_ver = f"tensorflow-gpu{tf_ver}" - self.required_packages.append(("tensorflow-gpu", - next(parse_requirements(tf_ver)).specs)) - return - - self.output.warning( - "The minimum Tensorflow requirement is 2.4 \n" - "Tensorflow currently has no official prebuild for your CUDA, cuDNN " - "combination.\nEither install a combination that Tensorflow supports or " - "build and install your own tensorflow-gpu.\r\n" - f"CUDA Version: {self.cuda_version}\r\n" - f"cuDNN Version: {self.cudnn_version}\r\n" + self._required_packages = [pkg for pkg in self._required_packages + if pkg[0] != "tensorflow"] + tf_ver = f"tensorflow{tf_ver}" + self._required_packages.append(("tensorflow", next(parse_requirements(tf_ver)).specs)) + return + + logger.warning( + "The minimum Tensorflow requirement is 2.10 \n" + "Tensorflow currently has no official prebuild for your CUDA, cuDNN combination.\n" + "Either install a combination that Tensorflow supports or build and install your own " + "tensorflow.\r\n" + "CUDA Version: %s\r\n" + "cuDNN Version: %s\r\n" "Help:\n" "Building Tensorflow: https://www.tensorflow.org/install/install_sources\r\n" "Tensorflow supported versions: " - "https://www.tensorflow.org/install/source#tested_build_configurations") + "https://www.tensorflow.org/install/source#tested_build_configurations", + self._env.cuda_version, self._env.cudnn_version) - custom_tf = input("Location of custom tensorflow-gpu wheel (leave " - "blank to manually install): ") + custom_tf = input("Location of custom tensorflow wheel (leave blank to manually " + "install): ") if not custom_tf: return custom_tf = os.path.realpath(os.path.expanduser(custom_tf)) + global _INSTALL_FAILED # pylint:disable=global-statement if not os.path.isfile(custom_tf): - self.output.error(f"{custom_tf} not found") + logger.error("%s not found", custom_tf) + _INSTALL_FAILED = True elif os.path.splitext(custom_tf)[1] != ".whl": - self.output.error(f"{custom_tf} is not a valid pip wheel") + logger.error("%s is not a valid pip wheel", custom_tf) + _INSTALL_FAILED = True elif custom_tf: - self.required_packages.append((custom_tf, (custom_tf, ""))) + self._required_packages.append((custom_tf, [(custom_tf, "")])) - def set_config(self) -> None: - """ Set the backend in the faceswap config file """ - if self.enable_amd: - backend = "amd" - elif self.enable_cuda: - backend = "nvidia" - elif self.enable_apple_silicon: - backend = "apple_silicon" - else: - backend = "cpu" - config = {"backend": backend} - pypath = os.path.dirname(os.path.realpath(__file__)) - config_file = os.path.join(pypath, "config", ".faceswap") - with open(config_file, "w", encoding="utf8") as cnf: - json.dump(config, cnf) - self.output.info(f"Faceswap config written to: {config_file}") + def _update_tf_dep_rocm(self) -> None: + """ Update the Tensorflow dependency for global ROCm installs """ + if not any(self._env.rocm_version): # ROCm was not found and the install will be aborted + return - def set_ld_library_path(self) -> None: - """ Update the LD_LIBRARY_PATH environment variable when activating a conda environment - and revert it when deactivating. Linux/conda only + global _INSTALL_FAILED # pylint:disable=global-statement + candidates = [key for key, val in _TENSORFLOW_ROCM_REQUIREMENTS.items() + if val[0] <= self._env.rocm_version <= val[1]] - Notes - ----- - From Tensorflow 2.7, installing Cuda Toolkit from conda-forge and tensorflow from pip - causes tensorflow to not be able to locate shared libs and hence not use the GPU. - We update the environment variable for all instances using Conda as it shouldn't hurt - anything and may help avoid conflicts with globally installed Cuda + if not candidates: + _INSTALL_FAILED = True + logger.error("No matching Tensorflow candidates found for ROCm %s in %s", + ".".join(str(v) for v in self._env.rocm_version), + _TENSORFLOW_ROCM_REQUIREMENTS) + return + + # set tf_ver to the minimum and maximum compatible range + tf_ver = f"{candidates[0].split(',')[0]},{candidates[-1].split(',')[-1]}" + # Remove the version of tensorflow-rocm in requirements file and add the correct version + # that corresponds to the installed ROCm version + self._required_packages = [pkg for pkg in self._required_packages + if not pkg[0].startswith("tensorflow-rocm")] + tf_ver = f"tensorflow-rocm{tf_ver}" + self._required_packages.append(("tensorflow-rocm", + next(parse_requirements(tf_ver)).specs)) + + def update_tf_dep(self) -> None: + """ Update Tensorflow Dependency. + + Selects a compatible version of Tensorflow for a globally installed GPU library """ - if not self.is_conda or not self.enable_cuda or self.os_version[0].lower() != "linux": + if self._env.backend == "nvidia": + self._update_tf_dep_nvidia() + if self._env.backend == "rocm": + self._update_tf_dep_rocm() + + def _check_conda_missing_dependencies(self) -> None: + """ Check for conda missing dependencies and add to :attr:`_conda_missing_packages` """ + if not self._env.is_conda: return + for pkg in self._conda_required_packages: + reqs = next(parse_requirements(pkg[0])) # TODO Handle '=' vs '==' for conda + key = reqs.unsafe_name + specs = reqs.specs + + if pkg[0] == "tk" and self._env.os_version[0].lower() == "linux": + # Default tk has bad fonts under Linux. We pull in an explicit build from + # Conda-Forge that is compiled with better fonts. + # Ref: https://github.com/ContinuumIO/anaconda-issues/issues/6833 + newpkg = (f"{pkg[0]}=*=xft_*", pkg[1]) # Swap out package for explicit XFT version + self._conda_missing_packages.append(newpkg) + # We also need to bring in xorg-libxft incase libXft does not exist on host system + self._conda_missing_packages.append(_CONDA_MAPPING["xorg-libxft"]) + continue - conda_prefix = os.environ["CONDA_PREFIX"] - activate_folder = os.path.join(conda_prefix, "etc", "conda", "activate.d") - deactivate_folder = os.path.join(conda_prefix, "etc", "conda", "deactivate.d") + if key not in self._conda_installed_packages: + self._conda_missing_packages.append(pkg) + continue - os.makedirs(activate_folder, exist_ok=True) - os.makedirs(deactivate_folder, exist_ok=True) + if not self._validate_spec(specs, self._conda_installed_packages[key]): + self._conda_missing_packages.append(pkg) + logger.debug(self._conda_missing_packages) - activate_script = os.path.join(conda_prefix, activate_folder, "env_vars.sh") - deactivate_script = os.path.join(conda_prefix, deactivate_folder, "env_vars.sh") + def check_missing_dependencies(self) -> None: + """ Check for missing dependencies and add to :attr:`_missing_packages` """ + for key, specs in self._required_packages: - if os.path.isfile(activate_script): - # Only create file if it does not already exist. There may be instances where people - # have created their own scripts, but these should be few and far between and those - # people should already know what they are doing. - return + if self._env.is_conda: # Get Conda alias for Key + key = _CONDA_MAPPING.get(key, (key, None))[0] - conda_libs = os.path.join(conda_prefix, "lib") - shebang = "#!/bin/sh\n\n" + if key not in self._all_installed_packages: + # Add not installed packages to missing packages list + self._missing_packages.append((key, specs)) + continue - with open(activate_script, "w", encoding="utf8") as afile: - afile.write(f"{shebang}") - afile.write("export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\n") - afile.write(f"export LD_LIBRARY_PATH='{conda_libs}':${{LD_LIBRARY_PATH}}\n") + if not self._validate_spec(specs, self._all_installed_packages.get(key, "")): + self._missing_packages.append((key, specs)) - with open(deactivate_script, "w", encoding="utf8") as afile: - afile.write(f"{shebang}") - afile.write("export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\n") - afile.write("unset OLD_LD_LIBRARY_PATH\n") + logger.debug(self._missing_packages) + self._check_conda_missing_dependencies() - self.output.info(f"Cuda search path set to '{conda_libs}'") +class Checks(): # pylint:disable=too-few-public-methods + """ Pre-installation checks -class Output(): - """ Format and display output """ - def __init__(self) -> None: - self.red: str = "\033[31m" - self.green: str = "\033[32m" - self.yellow: str = "\033[33m" - self.default_color: str = "\033[0m" - self.term_support_color: bool = platform.system().lower() in ("linux", "darwin") - - @staticmethod - def __indent_text_block(text: str) -> str: - """ Indent a text block """ - lines = text.splitlines() - if len(lines) > 1: - out = lines[0] + "\r\n" - for i in range(1, len(lines)-1): - out = out + " " + lines[i] + "\r\n" - out = out + " " + lines[-1] - return out - return text - - def info(self, text: str) -> None: - """ Format INFO Text """ - trm = "INFO " - if self.term_support_color: - trm = f"{self.green}INFO {self.default_color} " - print(trm + self.__indent_text_block(text)) - - def warning(self, text: str) -> None: - """ Format WARNING Text """ - trm = "WARNING " - if self.term_support_color: - trm = f"{self.yellow}WARNING{self.default_color} " - print(trm + self.__indent_text_block(text)) - - def error(self, text: str) -> None: - """ Format ERROR Text """ - global INSTALL_FAILED # pylint:disable=global-statement - trm = "ERROR " - if self.term_support_color: - trm = f"{self.red}ERROR {self.default_color} " - print(trm + self.__indent_text_block(text)) - INSTALL_FAILED = True - - -class Checks(): - """ Pre-installation checks """ + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + """ def __init__(self, environment: Environment) -> None: - self.env: Environment = environment - self.output: Output = Output() - self.tips: Tips = Tips() + self._env: Environment = environment + self._tips: Tips = Tips() # Checks not required for installer - if self.env.is_installer: + if self._env.is_installer: return # Checks not required for Apple Silicon - if self.env.enable_apple_silicon: - return - - # Ask AMD/Docker/Cuda - self.amd_ask_enable() - if not self.env.enable_amd: - self.docker_ask_enable() - self.cuda_ask_enable() - if self.env.os_version[0] != "Linux" and self.env.enable_docker and self.env.enable_cuda: - self.docker_confirm() - if self.env.enable_docker: - self.docker_tips() - self.env.set_config() - sys.exit(0) - - # Check for CUDA and cuDNN - if self.env.enable_cuda and self.env.is_conda: - self.output.info("Skipping Cuda/cuDNN checks for Conda install") - elif self.env.enable_cuda and self.env.os_version[0] in ("Linux", "Windows"): - check = CudaCheck() - if check.cuda_version: - self.env.cuda_version = check.cuda_version - self.output.info("CUDA version: " + self.env.cuda_version) - else: - self.output.error("CUDA not found. Install and try again.\n" - "Recommended version: CUDA 10.1 cuDNN 7.6\n" - "CUDA: https://developer.nvidia.com/cuda-downloads\n" - "cuDNN: https://developer.nvidia.com/rdp/cudnn-download") - return + if self._env.backend == "apple_silicon": + return + self._user_input() + self._check_cuda() + self._check_rocm() + if self._env.os_version[0] == "Windows": + self._tips.pip() + + def _rocm_ask_enable(self) -> None: + """ Set backend to 'rocm' if OS is Linux and ROCm support required """ + if self._env.os_version[0] != "Linux": + return + logger.info("ROCm support:\r\nIf you are using an AMD GPU, then select 'yes'." + "\r\nCPU/non-AMD GPU users should answer 'no'.\r\n") + i = input("Enable ROCm Support? [y/N] ") + if i in ("Y", "y"): + logger.info("ROCm Support Enabled") + self._env.backend = "rocm" - if check.cudnn_version: - self.env.cudnn_version = ".".join(check.cudnn_version.split(".")[:2]) - self.output.info(f"cuDNN version: {self.env.cudnn_version}") - else: - self.output.error("cuDNN not found. See " - "https://github.com/deepfakes/faceswap/blob/master/INSTALL.md#" - "cudnn for instructions") - return - elif self.env.enable_cuda and self.env.os_version[0] not in ("Linux", "Windows"): - self.tips.macos() - self.output.warning("Cannot find CUDA on macOS") - self.env.cuda_version = input("Manually specify CUDA version: ") - - self.env.update_tf_dep() - if self.env.os_version[0] == "Windows": - self.tips.pip() - - def amd_ask_enable(self) -> None: - """ Enable or disable Plaidml for AMD""" - self.output.info("AMD Support: AMD GPU support is currently limited.\r\n" - "Nvidia Users MUST answer 'no' to this option.") - i = input("Enable AMD Support? [y/N] ") + def _directml_ask_enable(self) -> None: + """ Set backend to 'directml' if OS is Windows and DirectML support required """ + if self._env.os_version[0] != "Windows": + return + logger.info("DirectML support:\r\nIf you are using an AMD or Intel GPU, then select 'yes'." + "\r\nNvidia users should answer 'no'.") + i = input("Enable DirectML Support? [y/N] ") if i in ("Y", "y"): - self.output.info("AMD Support Enabled") - self.env.enable_amd = True - else: - self.output.info("AMD Support Disabled") - self.env.enable_amd = False + logger.info("DirectML Support Enabled") + self._env.backend = "directml" + + def _user_input(self) -> None: + """ Get user input for AMD/DirectML/ROCm/Cuda/Docker """ + self._directml_ask_enable() + self._rocm_ask_enable() + if not self._env.backend: + self._docker_ask_enable() + self._cuda_ask_enable() + if self._env.os_version[0] != "Linux" and (self._env.enable_docker + and self._env.backend == "nvidia"): + self._docker_confirm() + if self._env.enable_docker: + self._docker_tips() + self._env.set_config() + sys.exit(0) - def docker_ask_enable(self) -> None: + def _docker_ask_enable(self) -> None: """ Enable or disable Docker """ i = input("Enable Docker? [y/N] ") if i in ("Y", "y"): - self.output.info("Docker Enabled") - self.env.enable_docker = True + logger.info("Docker Enabled") + self._env.enable_docker = True else: - self.output.info("Docker Disabled") - self.env.enable_docker = False + logger.info("Docker Disabled") + self._env.enable_docker = False - def docker_confirm(self) -> None: + def _docker_confirm(self) -> None: """ Warn if nvidia-docker on non-Linux system """ - self.output.warning("Nvidia-Docker is only supported on Linux.\r\n" - "Only CPU is supported in Docker for your system") - self.docker_ask_enable() - if self.env.enable_docker: - self.output.warning("CUDA Disabled") - self.env.enable_cuda = False - - def docker_tips(self) -> None: + logger.warning("Nvidia-Docker is only supported on Linux.\r\n" + "Only CPU is supported in Docker for your system") + self._docker_ask_enable() + if self._env.enable_docker: + logger.warning("CUDA Disabled") + self._env.backend = "cpu" + + def _docker_tips(self) -> None: """ Provide tips for Docker use """ - if not self.env.enable_cuda: - self.tips.docker_no_cuda() + if self._env.backend != "nvidia": + self._tips.docker_no_cuda() else: - self.tips.docker_cuda() + self._tips.docker_cuda() - def cuda_ask_enable(self) -> None: + def _cuda_ask_enable(self) -> None: """ Enable or disable CUDA """ i = input("Enable CUDA? [Y/n] ") if i in ("", "Y", "y"): - self.output.info("CUDA Enabled") - self.env.enable_cuda = True + logger.info("CUDA Enabled") + self._env.backend = "nvidia" + + def _check_cuda(self) -> None: + """ Check for Cuda and cuDNN Locations. """ + if self._env.backend != "nvidia": + logger.debug("Skipping Cuda checks as not enabled") + return + + if self._env.is_conda: + logger.info("Skipping Cuda/cuDNN checks for Conda install") + return + + if self._env.os_version[0] in ("Linux", "Windows"): + global _INSTALL_FAILED # pylint:disable=global-statement + check = CudaCheck() + if check.cuda_version: + self._env.cuda_cudnn[0] = check.cuda_version + logger.info("CUDA version: %s", self._env.cuda_version) + else: + logger.error("CUDA not found. Install and try again.\n" + "Recommended version: CUDA 10.1 cuDNN 7.6\n" + "CUDA: https://developer.nvidia.com/cuda-downloads\n" + "cuDNN: https://developer.nvidia.com/rdp/cudnn-download") + _INSTALL_FAILED = True + return + + if check.cudnn_version: + self._env.cuda_cudnn[1] = ".".join(check.cudnn_version.split(".")[:2]) + logger.info("cuDNN version: %s", self._env.cudnn_version) + else: + logger.error("cuDNN not found. See " + "https://github.com/deepfakes/faceswap/blob/master/INSTALL.md#" + "cudnn for instructions") + _INSTALL_FAILED = True + return + + # If we get here we're on MacOS + self._tips.macos() + logger.warning("Cannot find CUDA on macOS") + self._env.cuda_cudnn[0] = input("Manually specify CUDA version: ") + + def _check_rocm(self) -> None: + """ Check for ROCm version """ + if self._env.backend != "rocm" or self._env.os_version[0] != "Linux": + logger.info("Skipping ROCm checks as not enabled") + return + + global _INSTALL_FAILED # pylint:disable=global-statement + check = ROCmCheck() + + str_min = ".".join(str(v) for v in check.version_min) + str_max = ".".join(str(v) for v in check.version_max) + + if check.is_valid: + self._env.rocm_version = check.rocm_version + logger.info("ROCm version: %s", ".".join(str(v) for v in self._env.rocm_version)) else: - self.output.info("CUDA Disabled") - self.env.enable_cuda = False + if check.rocm_version: + msg = f"Incompatible ROCm version: {'.'.join(str(v) for v in check.rocm_version)}" + else: + msg = "ROCm not found" + logger.error("%s.\n" + "A compatible version of ROCm must be installed to proceed.\n" + "ROCm versions between %s and %s are supported.\n" + "ROCm install guide: https://docs.amd.com/bundle/ROCm_Installation_Guide" + "v5.0/page/Overview_of_ROCm_Installation_Methods.html", + msg, + str_min, + str_max) + _INSTALL_FAILED = True + + +def _check_ld_config(lib: str) -> str: + """ Locate a library in ldconfig + + Parameters + ---------- + lib: str The library to locate + + Returns + ------- + str + The library from ldconfig, or empty string if not found + """ + retval = "" + ldconfig = which("ldconfig") + if not ldconfig: + return retval + + retval = next((line.decode("utf-8", errors="replace").strip() + for line in run([ldconfig, "-p"], + capture_output=True, + check=False).stdout.splitlines() + if lib.encode("utf-8") in line), "") + + if retval or (not retval and not os.environ.get("LD_LIBRARY_PATH")): + return retval + + for path in os.environ["LD_LIBRARY_PATH"].split(":"): + if not path or not os.path.exists(path): + continue + + retval = next((fname.strip() for fname in reversed(os.listdir(path)) + if lib in fname), "") + if retval: + break + + return retval + + +class ROCmCheck(): # pylint:disable=too-few-public-methods + """ Find the location of system installed ROCm on Linux """ + def __init__(self) -> None: + self.version_min = min(v[0] for v in _TENSORFLOW_ROCM_REQUIREMENTS.values()) + self.version_max = max(v[1] for v in _TENSORFLOW_ROCM_REQUIREMENTS.values()) + self.rocm_version: tuple[int, ...] = (0, 0, 0) + if platform.system() == "Linux": + self._rocm_check() + + @property + def is_valid(self): + """ bool: `True` if ROCm has been detected and is between the minimum and maximum + compatible versions otherwise ``False`` """ + return self.version_min <= self.rocm_version <= self.version_max + + def _rocm_check(self) -> None: + """ Attempt to locate the installed ROCm version from the dynamic link loader. If not found + with ldconfig then attempt to find it in LD_LIBRARY_PATH. If found, set the + :attr:`rocm_version` to the discovered version + """ + chk = _check_ld_config("librocm-core.so.") + if not chk: + return + + rocm_vers = chk.strip() + version = re.search(r"rocm\-(\d+\.\d+\.\d+)", rocm_vers) + if version is None: + return + try: + self.rocm_version = tuple(int(v) for v in version.groups()[0].split(".")) + except ValueError: + return class CudaCheck(): # pylint:disable=too-few-public-methods """ Find the location of system installed Cuda and cuDNN on Windows and Linux. """ def __init__(self) -> None: - self.cuda_path: Optional[str] = None - self.cuda_version: Optional[str] = None - self.cudnn_version: Optional[str] = None + self.cuda_path: str | None = None + self.cuda_version: str | None = None + self.cudnn_version: str | None = None self._os: str = platform.system().lower() - self._cuda_keys: List[str] = [key + self._cuda_keys: list[str] = [key for key in os.environ if key.lower().startswith("cuda_path_v")] - self._cudnn_header_files: List[str] = ["cudnn_version.h", "cudnn.h"] - + self._cudnn_header_files: list[str] = ["cudnn_version.h", "cudnn.h"] + logger.debug("cuda keys: %s, cudnn header files: %s", + self._cuda_keys, self._cudnn_header_files) if self._os in ("windows", "linux"): self._cuda_check() self._cudnn_check() @@ -544,11 +878,10 @@ def _cuda_check(self) -> None: stdout, stderr = chk.communicate() if not stderr: version = re.search(r".*release (?P\d+\.\d+)", - stdout.decode(locale.getpreferredencoding())) + stdout.decode(locale.getpreferredencoding(), errors="ignore")) if version is not None: self.cuda_version = version.groupdict().get("cuda", None) - locate = "where" if self._os == "windows" else "which" - path = os.popen(f"{locate} nvcc").read() + path = which("nvcc") if path: path = path.split("\n")[0] # Split multiple entries and take first found while True: # Get Cuda root folder @@ -560,23 +893,20 @@ def _cuda_check(self) -> None: # Failed to load nvcc, manual check getattr(self, f"_cuda_check_{self._os}")() + logger.debug("Cuda Version: %s, Cuda Path: %s", self.cuda_version, self.cuda_path) def _cuda_check_linux(self) -> None: """ For Linux check the dynamic link loader for libcudart. If not found with ldconfig then attempt to find it in LD_LIBRARY_PATH. """ - chk = os.popen("ldconfig -p | grep -P \"libcudart.so.\\d+.\\d+\" | head -n 1").read() - if not chk and os.environ.get("LD_LIBRARY_PATH"): - for path in os.environ["LD_LIBRARY_PATH"].split(":"): - chk = os.popen(f"ls {path} | grep -P -o \"libcudart.so.\\d+.\\d+\" | " - "head -n 1").read() - if chk: - break + chk = _check_ld_config("libcudart.so.") if not chk: # Cuda not found return cudavers = chk.strip().replace("libcudart.so.", "") - self.cuda_version = cudavers[:cudavers.find(" ")] - self.cuda_path = chk[chk.find("=>") + 3:chk.find("targets") - 1] + self.cuda_version = cudavers[:cudavers.find(" ")] if " " in cudavers else cudavers + cuda_path = chk[chk.find("=>") + 3:chk.find("targets") - 1] + if os.path.exists(cuda_path): + self.cuda_path = cuda_path def _cuda_check_windows(self) -> None: """ Check Windows CUDA Version and path from Environment Variables""" @@ -585,12 +915,14 @@ def _cuda_check_windows(self) -> None: self.cuda_version = self._cuda_keys[0].lower().replace("cuda_path_v", "").replace("_", ".") self.cuda_path = os.environ[self._cuda_keys[0][0]] - def _cudnn_check(self): - """ Check Linux or Windows cuDNN Version from cudnn.h and add to :attr:`cudnn_version`. """ + def _cudnn_check_files(self) -> bool: + """ Check header files for cuDNN version """ cudnn_checkfiles = getattr(self, f"_get_checkfiles_{self._os}")() cudnn_checkfile = next((hdr for hdr in cudnn_checkfiles if os.path.isfile(hdr)), None) + logger.debug("cudnn checkfiles: %s", cudnn_checkfile) if not cudnn_checkfile: - return + return False + found = 0 with open(cudnn_checkfile, "r", encoding="utf8") as ofile: for line in ofile: @@ -605,11 +937,31 @@ def _cudnn_check(self): found += 1 if found == 3: break - if found != 3: # Full version could not be determined - return + if found != 3: # Full version not determined + return False + self.cudnn_version = ".".join([str(major), str(minor), str(patchlevel)]) + logger.debug("cudnn version: %s", self.cudnn_version) + return True - def _get_checkfiles_linux(self) -> List[str]: + def _cudnn_check(self) -> None: + """ Check Linux or Windows cuDNN Version from cudnn.h and add to :attr:`cudnn_version`. """ + if self._cudnn_check_files(): + return + if self._os == "windows": + return + + chk = _check_ld_config("libcudnn.so.") + if not chk: + return + cudnnvers = chk.strip().replace("libcudnn.so.", "").split()[0] + if not cudnnvers: + return + + self.cudnn_version = cudnnvers + logger.debug("cudnn version: %s", self.cudnn_version) + + def _get_checkfiles_linux(self) -> list[str]: """ Return the the files to check for cuDNN locations for Linux by querying the dynamic link loader. @@ -618,7 +970,7 @@ def _get_checkfiles_linux(self) -> List[str]: list List of header file locations to scan for cuDNN versions """ - chk = os.popen("ldconfig -p | grep -P \"libcudnn.so.\\d+\" | head -n 1").read() + chk = _check_ld_config("libcudnn.so.") chk = chk.strip().replace("libcudnn.so.", "") if not chk: return [] @@ -631,7 +983,7 @@ def _get_checkfiles_linux(self) -> List[str]: cudnn_checkfiles = [os.path.join(cudnn_path, header) for header in header_files] return cudnn_checkfiles - def _get_checkfiles_windows(self) -> List[str]: + def _get_checkfiles_windows(self) -> list[str]: """ Return the check-file locations for Windows. Just looks inside the include folder of the discovered :attr:`cuda_path` @@ -641,273 +993,710 @@ def _get_checkfiles_windows(self) -> List[str]: List of header file locations to scan for cuDNN versions """ # TODO A more reliable way of getting the windows location - if not self.cuda_path: + if not self.cuda_path or not os.path.exists(self.cuda_path): return [] scandir = os.path.join(self.cuda_path, "include") cudnn_checkfiles = [os.path.join(scandir, header) for header in self._cudnn_header_files] return cudnn_checkfiles -class Install(): - """ Install the requirements """ - def __init__(self, environment: Environment): - self._operators = {"==": operator.eq, - ">=": operator.ge, - "<=": operator.le, - ">": operator.gt, - "<": operator.lt} - self.output = environment.output - self.env = environment - - if not self.env.is_installer and not self.env.updater: - self.ask_continue() - self.env.get_required_packages() - self.check_missing_dep() - self.check_conda_missing_dep() - if (self.env.updater and - not self.env.missing_packages and not self.env.conda_missing_packages): - self.output.info("All Dependencies are up to date") - return - self.install_missing_dep() - if self.env.updater: - return - self.output.info("All python3 dependencies are met.\r\nYou are good to go.\r\n\r\n" - "Enter: 'python faceswap.py -h' to see the options\r\n" - " 'python faceswap.py gui' to launch the GUI") - - def ask_continue(self): +class Install(): # pylint:disable=too-few-public-methods + """ Handles installation of Faceswap requirements + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + is_gui: bool, Optional + ``True`` if the caller is the Faceswap GUI. Used to prevent output of progress bars + which get scrambled in the GUI + """ + def __init__(self, environment: Environment, is_gui: bool = False) -> None: + self._env = environment + self._packages = Packages(environment) + self._is_gui = is_gui + + if self._env.os_version[0] == "Windows": + self._installer: type[Installer] = WinPTYInstaller + else: + self._installer = PexpectInstaller + + if not self._env.is_installer and not self._env.updater: + self._ask_continue() + + self._packages.get_required_packages() + self._packages.update_tf_dep() + self._packages.check_missing_dependencies() + + if self._env.updater and not self._packages.packages_need_install: + logger.info("All Dependencies are up to date") + return + + logger.info("Installing Required Python Packages. This may take some time...") + self._install_setup_packages() + self._install_missing_dep() + if self._env.updater: + return + if not _INSTALL_FAILED: + logger.info("All python3 dependencies are met.\r\nYou are good to go.\r\n\r\n" + "Enter: 'python faceswap.py -h' to see the options\r\n" + " 'python faceswap.py gui' to launch the GUI") + else: + logger.error("Some packages failed to install. This may be a temporary error which " + "might be fixed by re-running this script. Otherwise please install " + "these packages manually.") + sys.exit(1) + + def _ask_continue(self) -> None: """ Ask Continue with Install """ - inp = input("Please ensure your System Dependencies are met. Continue? [y/N] ") + text = "Please ensure your System Dependencies are met" + if self._env.backend == "rocm": + text += ("\r\nROCm users: Please ensure that your AMD GPU is supported by the " + "installed ROCm version before proceeding.") + text += "\r\nContinue? [y/N] " + inp = input(text) if inp in ("", "N", "n"): - self.output.error("Please install system dependencies to continue") + logger.error("Please install system dependencies to continue") sys.exit(1) - def check_missing_dep(self): - """ Check for missing dependencies """ - for key, specs in self.env.required_packages: + @classmethod + def _format_package(cls, package: str, version: list[tuple[str, str]]) -> str: + """ Format a parsed requirement package and version string to a format that can be used by + the installer. - if self.env.is_conda: # Get Conda alias for Key - key = CONDA_MAPPING.get(key, (key, None))[0] + Parameters + ---------- + package: str + The package name + version: list + The parsed requirement version strings - if key not in self.env.installed_packages: - # Add not installed packages to missing packages list - self.env.missing_packages.append((key, specs)) - continue - - installed_vers = self.env.installed_packages.get(key, "") + Returns + ------- + str + The formatted full package and version string + """ + return f"{package}{','.join(''.join(spec) for spec in version)}" - if specs and not all(self._operators[spec[0]](installed_vers, spec[1]) - for spec in specs): - self.env.missing_packages.append((key, specs)) + def _install_setup_packages(self) -> None: + """ Install any packages that are required for the setup.py installer to work. This + includes the pexpect package if it is not already installed. - def check_conda_missing_dep(self): - """ Check for conda missing dependencies """ - if not self.env.is_conda: - return - for pkg in self.env.conda_required_packages: - key = pkg[0].split("==")[0] - if key not in self.env.installed_packages: - self.env.conda_missing_packages.append(pkg) - continue - if len(pkg[0].split("==")) > 1: - if pkg[0].split("==")[1] != self.env.installed_conda_packages.get(key): - self.env.conda_missing_packages.append(pkg) - continue + Subprocess is used as we do not currently have pexpect + """ + for pkg in self._packages.prerequisites: + pkg_str = self._format_package(*pkg) + if self._env.is_conda: + cmd = ["conda", "install", "-y"] + if any(char in pkg_str for char in (" ", "<", ">", "*", "|")): + pkg_str = f"\"{pkg_str}\"" + else: + cmd = [sys.executable, "-m", "pip", "install", "--no-cache-dir"] + if self._env.is_admin: + cmd.append("--user") + cmd.append(pkg_str) + + clean_pkg = pkg_str.replace("\"", "") + installer = SubProcInstaller(self._env, clean_pkg, cmd, self._is_gui) + if installer() != 0: + logger.error("Unable to install package: %s. Process aborted", clean_pkg) + sys.exit(1) - def install_missing_dep(self): - """ Install missing dependencies """ - # Install conda packages first - if self.env.conda_missing_packages: - self.install_conda_packages() - if self.env.missing_packages: - self.install_python_packages() + def _install_conda_packages(self) -> None: + """ Install required conda packages """ + logger.info("Installing Required Conda Packages. This may take some time...") + for pkg in self._packages.to_install_conda: + channel = "" if len(pkg) != 2 else pkg[1] + self._from_conda(pkg[0], channel=channel, conda_only=True) - def install_python_packages(self): + def _install_python_packages(self) -> None: """ Install required pip packages """ - self.output.info("Installing Required Python Packages. This may take some time...") conda_only = False - for pkg, version in self.env.missing_packages: - if self.env.is_conda: - pkg = CONDA_MAPPING.get(pkg, (pkg, None)) - channel = None if len(pkg) != 2 else pkg[1] - pkg = pkg[0] - if version: - pkg = f"{pkg}{','.join(''.join(spec) for spec in version)}" - if self.env.is_conda: - if pkg.startswith("tensorflow-gpu"): - # From TF 2.4 onwards, Anaconda Tensorflow becomes a mess. The version of 2.5 - # installed by Anaconda is compiled against an incorrect numpy version which - # breaks Tensorflow. Coupled with this the versions of cudatoolkit and cudnn - # available in the default Anaconda channel are not compatible with the - # official PyPi versions of Tensorflow. With this in mind we will pull in the - # required Cuda/cuDNN from conda-forge, and install Tensorflow with pip - # TODO Revert to Conda if they get their act together - - # Rewrite tensorflow requirement to versions from highest available cuda/cudnn - highest_cuda = sorted(TENSORFLOW_REQUIREMENTS.values())[-1] - compat_tf = next(k for k, v in TENSORFLOW_REQUIREMENTS.items() - if v == highest_cuda) - pkg = f"tensorflow-gpu{compat_tf}" - conda_only = True - - verbose = pkg.startswith("tensorflow") or self.env.updater - if self.conda_installer(pkg, - verbose=verbose, channel=channel, conda_only=conda_only): + assert self._env.backend is not None + for pkg, version in self._packages.to_install: + if self._env.is_conda: + mapping = _CONDA_MAPPING.get(pkg, (pkg, "")) + channel = "" if mapping[1] is None else mapping[1] + pkg = mapping[0] + pip_only = pkg in _FORCE_PIP.get(self._env.backend, []) or pkg in _FORCE_PIP["all"] + pkg = self._format_package(pkg, version) if version else pkg + if self._env.is_conda and not pip_only: + if self._from_conda(pkg, channel=channel, conda_only=conda_only): continue - self.pip_installer(pkg) + self._from_pip(pkg) - def install_conda_packages(self): - """ Install required conda packages """ - self.output.info("Installing Required Conda Packages. This may take some time...") - for pkg in self.env.conda_missing_packages: - channel = None if len(pkg) != 2 else pkg[1] - self.conda_installer(pkg[0], channel=channel, conda_only=True) + def _install_missing_dep(self) -> None: + """ Install missing dependencies """ + self._install_conda_packages() # Install conda packages first + self._install_python_packages() + + def _from_conda(self, + package: list[str] | str, + channel: str = "", + conda_only: bool = False) -> bool: + """ Install a conda package + + Parameters + ---------- + package: list[str] | str + The full formatted package(s), with version(s), to be installed + channel: str, optional + The Conda channel to install from. Select empty string for default channel. + Default: ``""`` (empty string) + conda_only: bool, optional + ``True`` if the package is only available in Conda. Default: ``False`` - def conda_installer(self, package, channel=None, verbose=False, conda_only=False): - """ Install a conda package """ + Returns + ------- + bool + ``True`` if the package was succesfully installed otherwise ``False`` + """ # Packages with special characters need to be enclosed in double quotes success = True condaexe = ["conda", "install", "-y"] - if not verbose or self.env.updater: - condaexe.append("-q") if channel: condaexe.extend(["-c", channel]) - if package.startswith("tensorflow-gpu"): - # Here we will install the cuda/cudnn toolkits, currently only available from - # conda-forge, but fail tensorflow itself so that it can be handled by pip. - specs = Requirement.parse(package).specs - for key, val in TENSORFLOW_REQUIREMENTS.items(): - req_specs = Requirement.parse("foobar" + key).specs - if all(item in req_specs for item in specs): - cuda, cudnn = val - break - condaexe.extend(["-c", "conda-forge", f"cudatoolkit={cuda}", f"cudnn={cudnn}"]) - package = "Cuda Toolkit" - success = False - - if package != "Cuda Toolkit": - if any(char in package for char in (" ", "<", ">", "*", "|")): - package = f"\"{package}\"" - condaexe.append(package) - - clean_pkg = package.replace("\"", "") - self.output.info(f"Installing {clean_pkg}") - shell = self.env.os_version[0] == "Windows" - try: - if verbose: - run(condaexe, check=True, shell=shell) - else: - with open(os.devnull, "w", encoding="utf8") as devnull: - run(condaexe, stdout=devnull, stderr=devnull, check=True, shell=shell) - except CalledProcessError: - if not conda_only: - self.output.info(f"{package} not available in Conda. Installing with pip") - else: - self.output.warning(f"Couldn't install {package} with Conda. " - "Please install this package manually") - success = False + pkgs = package if isinstance(package, list) else [package] + + for i, pkg in enumerate(pkgs): + if any(char in pkg for char in (" ", "<", ">", "*", "|")): + pkgs[i] = f"\"{pkg}\"" + condaexe.extend(pkgs) + + clean_pkg = " ".join([p.replace("\"", "") for p in pkgs]) + installer = self._installer(self._env, clean_pkg, condaexe, self._is_gui) + retcode = installer() + + if retcode != 0 and not conda_only: + logger.info("%s not available in Conda. Installing with pip", package) + elif retcode != 0: + logger.warning("Couldn't install %s with Conda. Please install this package " + "manually", package) + success = retcode == 0 and success return success - def pip_installer(self, package): - """ Install a pip package """ - pipexe = [sys.executable, "-m", "pip"] - # hide info/warning and fix cache hang - pipexe.extend(["install", "--no-cache-dir"]) - if not self.env.updater and not package.startswith("tensorflow"): - pipexe.append("-qq") + def _from_pip(self, package: str) -> None: + """ Install a pip package + + Parameters + ---------- + package: str + The full formatted package, with version, to be installed + """ + pipexe = [sys.executable, "-u", "-m", "pip", "install", "--no-cache-dir"] # install as user to solve perm restriction - if not self.env.is_admin and not self.env.is_virtualenv: + if not self._env.is_admin and not self._env.is_virtualenv: pipexe.append("--user") - msg = f"Installing {package}" - self.output.info(msg) pipexe.append(package) + + installer = self._installer(self._env, package, pipexe, self._is_gui) + if installer() != 0: + logger.warning("Couldn't install %s with pip. Please install this package manually", + package) + global _INSTALL_FAILED # pylint:disable=global-statement + _INSTALL_FAILED = True + + +class ProgressBar(): + """ Simple progress bar using STDLib for intercepting Conda installs and keeping the + terminal from getting jumbled """ + def __init__(self): + self._width_desc = 21 + self._width_size = 9 + self._width_bar = 35 + self._width_pct = 4 + self._marker = "█" + + self._cursor_visible = True + self._current_pos = 0 + self._bars = [] + + @classmethod + def _display_cursor(cls, visible: bool) -> None: + """ Sends ANSI code to display or hide the cursor + + Parameters + ---------- + visible: bool + ``True`` to display the cursor. ``False`` to hide the cursor + """ + code = "\x1b[?25h" if visible else "\x1b[?25l" + print(code, end="\r") + + def _format_bar(self, description: str, size: str, percent: int) -> str: + """ Format the progress bar for display + + Parameters + ---------- + description: str + The description to display for the progress bar + size: str + The size of the download, including units + percent: int + The percentage progress of the bar + """ + size = size[:self._width_size].ljust(self._width_size) + bar_len = int(self._width_bar * (percent / 100)) + progress = f"{self._marker * bar_len}"[:self._width_bar].ljust(self._width_bar) + pct = f"{percent}%"[:self._width_pct].rjust(self._width_pct) + return f" {description}| {size} | {progress} | {pct}" + + def _move_cursor(self, position: int) -> str: + """ Generate ANSI code for moving the cursor to the given progress bar's position + + Parameters + ---------- + position: int + The progress bar position to move to + + Returns + ------- + str + The ansi code to move to the given position + """ + move = position - self._current_pos + retval = "\x1b[A" if move < 0 else "\x1b[B" if move > 0 else "" + retval *= abs(move) + return retval + + def __call__(self, description: str, size: str, percent: int) -> None: + """ Create or update a progress bar + + Parameters + ---------- + description: str + The description to display for the progress bar + size: str + The size of the download, including units + percent: int + The percentage progress of the bar + """ + if self._cursor_visible: + self._display_cursor(visible=False) + + desc = description[:self._width_desc].ljust(self._width_desc) + if desc not in self._bars: + self._bars.append(desc) + + position = self._bars.index(desc) + pbar = self._format_bar(desc, size, percent) + + output = f"{self._move_cursor(position)} {pbar}" + + print(output) + self._current_pos = position + 1 + + def close(self) -> None: + """ Reset all progress bars and re-enable the cursor """ + print(self._move_cursor(len(self._bars)), end="\r") + self._display_cursor(True) + self._cursor_visible = True + self._current_pos = 0 + self._bars = [] + + +class Installer(): + """ Parent class for package installers. + + PyWinPty is used for Windows, Pexpect is used for Linux, as these can provide us with realtime + output. + + Subprocess is used as a fallback if any of the above fail, but this caches output, so it can + look like the process has hung to the end user + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + package: str + The package name that is being installed + command: list + The command to run + is_gui: bool + ``True`` if the process is being called from the Faceswap GUI + """ + def __init__(self, + environment: Environment, + package: str, + command: list[str], + is_gui: bool) -> None: + logger.info("Installing %s", package) + logger.debug("argv: %s", command) + self._env = environment + self._package = package + self._command = command + self._is_conda = "conda" in command + self._is_gui = is_gui + + self._progess_bar = ProgressBar() + self._re_conda = re.compile( + rb"(?P^\S+)\s+\|\s+(?P\d+\.?\d*\s\w+).*\|\s+(?P\d+%)") + self._re_pip_pkg = re.compile(rb"^\s*Downloading\s(?P\w+-.+?)-") + self._re_pip = re.compile(rb"(?P\d+\.?\d*)/(?P\d+\.?\d*\s\w+)") + self._pip_pkg = "" + self._seen_lines: set[str] = set() + + def __call__(self) -> int: + """ Call the subclassed call function + + Returns + ------- + int + The return code of the package install process + """ try: - run(pipexe, check=True) - except CalledProcessError: - self.output.warning(f"Couldn't install {package} with pip. " - "Please install this package manually") + returncode = self.call() + except Exception as err: # pylint:disable=broad-except + logger.debug("Failed to install with %s. Falling back to subprocess. Error: %s", + self.__class__.__name__, str(err)) + self._progess_bar.close() + returncode = SubProcInstaller(self._env, self._package, self._command, self._is_gui)() + + logger.debug("Package: %s, returncode: %s", self._package, returncode) + self._progess_bar.close() + return returncode + + def call(self) -> int: + """ Override for package installer specific logic. + + Returns + ------- + int + The return code of the package install process + """ + raise NotImplementedError() + + def _print_conda(self, text: bytes) -> None: + """ Output progress for Conda installs + + Parameters + ---------- + text: bytes + The text to print + """ + data = self._re_conda.match(text) + if not data: + return + lib = data.groupdict()["lib"].decode("utf-8", errors="replace") + size = data.groupdict()["tot"].decode("utf-8", errors="replace") + progress = int(data.groupdict()["prg"].decode("utf-8", errors="replace")[:-1]) + self._progess_bar(lib, size, progress) + + def _print_pip(self, text: bytes) -> None: + """ Output progress for Pip installs + + Parameters + ---------- + text: bytes + The text to print + """ + pkg = self._re_pip_pkg.match(text) + if pkg: + logger.debug("Collected pip package '%s'", pkg) + self._pip_pkg = pkg.groupdict()["lib"].decode("utf-8", errors="replace") + return + data = self._re_pip.search(text) + if not data: + return + done = float(data.groupdict()["done"].decode("utf-8", errors="replace")) + size = data.groupdict()["tot"].decode("utf-8", errors="replace") + progress = int(round(done / float(size.split()[0]) * 100, 0)) + self._progess_bar(self._pip_pkg, size, progress) + + def _non_gui_print(self, text: bytes) -> None: + """ Print output to console if not running in the GUI + + Parameters + ---------- + text: bytes + The text to print + """ + if self._is_gui: + return + if self._is_conda: + self._print_conda(text) + else: + self._print_pip(text) + + def _seen_line_log(self, text: str) -> None: + """ Output gets spammed to the log file when conda is waiting/processing. Only log each + unique line once. + + Parameters + ---------- + text: str + The text to log + """ + if text in self._seen_lines: + return + logger.debug(text) + self._seen_lines.add(text) + + +class PexpectInstaller(Installer): # pylint:disable=too-few-public-methods + """ Package installer for Linux/macOS using Pexpect + + Uses Pexpect for installing packages allowing access to realtime feedback + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + package: str + The package name that is being installed + command: list + The command to run + is_gui: bool + ``True`` if the process is being called from the Faceswap GUI + """ + def call(self) -> int: + """ Install a package using the Pexpect module + + Returns + ------- + int + The return code of the package install process + """ + import pexpect # pylint:disable=import-outside-toplevel,import-error + proc = pexpect.spawn(" ".join(self._command), timeout=None) + while True: + try: + proc.expect([b"\r\n", b"\r"]) + line: bytes = proc.before + self._seen_line_log(line.decode("utf-8", errors="replace").rstrip()) + self._non_gui_print(line) + except pexpect.EOF: + break + proc.close() + return proc.exitstatus + + +class WinPTYInstaller(Installer): # pylint:disable=too-few-public-methods + """ Package installer for Windows using WinPTY + + Spawns a pseudo PTY for installing packages allowing access to realtime feedback + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + package: str + The package name that is being installed + command: list + The command to run + is_gui: bool + ``True`` if the process is being called from the Faceswap GUI + """ + def __init__(self, + environment: Environment, + package: str, + command: list[str], + is_gui: bool) -> None: + super().__init__(environment, package, command, is_gui) + self._cmd = which(command[0], path=os.environ.get('PATH', os.defpath)) + self._cmdline = list2cmdline(command) + logger.debug("cmd: '%s', cmdline: '%s'", self._cmd, self._cmdline) + + self._pbar = re.compile(r"(?:eta\s[\d\W]+)|(?:\s+\|\s+\d+%)\Z") + self._eof = False + self._read_bytes = 1024 + + self._lines: list[str] = [] + self._out = "" + + def _read_from_pty(self, proc: T.Any, winpty_error: T.Any) -> None: + """ Read :attr:`_num_bytes` from WinPTY. If there is an error reading, recursively halve + the number of bytes read until we get a succesful read. If we get down to 1 byte without a + succesful read, assume we are at EOF. + + Parameters + ---------- + proc: :class:`winpty.PTY` + The WinPTY process + winpty_error: :class:`winpty.WinptyError` + The winpty error exception. Passed in as WinPTY is not in global scope + """ + try: + from_pty = proc.read(self._read_bytes) + except winpty_error: + # TODO Reinsert this check + # The error message "pipe has been ended" is language specific so this check + # fails on non english systems. For now we just swallow all errors until no + # bytes are left to read and then check the return code + # if any(val in str(err) for val in ["EOF", "pipe has been ended"]): + # # Get remaining bytes. On a comms error, the buffer remains unread so keep + # # halving buffer amount until down to 1 when we know we have everything + # if self._read_bytes == 1: + # self._eof = True + # from_pty = "" + # self._read_bytes //= 2 + # else: + # raise + + # Get remaining bytes. On a comms error, the buffer remains unread so keep + # halving buffer amount until down to 1 when we know we have everything + if self._read_bytes == 1: + self._eof = True + from_pty = "" + self._read_bytes //= 2 + + self._out += from_pty + + def _out_to_lines(self) -> None: + """ Process the winpty output into separate lines. Roll over any semi-consumed lines to the + next proc call. """ + if "\n" not in self._out: + return + + self._lines.extend(self._out.split("\n")) + + if self._out.endswith("\n") or self._eof: # Ends on newline or is EOF + self._out = "" + else: # roll over semi-consumed line to next read + self._out = self._lines[-1] + self._lines = self._lines[:-1] + + def call(self) -> int: + """ Install a package using the PyWinPTY module + + Returns + ------- + int + The return code of the package install process + """ + import winpty # pylint:disable=import-outside-toplevel,import-error + # For some reason with WinPTY we need to pass in the full command. Probably a bug + proc = winpty.PTY( + 100, + 24, + backend=winpty.enums.Backend.WinPTY, # ConPTY hangs and has lots of Ansi Escapes + agent_config=winpty.enums.AgentConfig.WINPTY_FLAG_PLAIN_OUTPUT) # Strip all Ansi + + if not proc.spawn(self._cmd, cmdline=self._cmdline): + del proc + raise RuntimeError("Failed to spawn winpty") + + while True: + self._read_from_pty(proc, winpty.WinptyError) + self._out_to_lines() + for line in self._lines: + self._seen_line_log(line.rstrip()) + self._non_gui_print(line.encode("utf-8", errors="replace")) + self._lines = [] + + if self._eof: + returncode = proc.get_exitstatus() + break + + del proc + return returncode + + +class SubProcInstaller(Installer): + """ The fallback package installer if either of the OS specific installers fail. + + Uses the python Subprocess module to install packages. Feedback does not return in realtime + so the process can look like it has hung to the end user + + Parameters + ---------- + environment: :class:`Environment` + Environment class holding information about the running system + package: str + The package name that is being installed + command: list + The command to run + is_gui: bool + ``True`` if the process is being called from the Faceswap GUI + """ + def __init__(self, + environment: Environment, + package: str, + command: list[str], + is_gui: bool) -> None: + super().__init__(environment, package, command, is_gui) + self._shell = self._env.os_version[0] == "Windows" and command[0] == "conda" + + def __call__(self) -> int: + """ Override default call function so we don't recursively call ourselves on failure. """ + returncode = self.call() + logger.debug("Package: %s, returncode: %s", self._package, returncode) + return returncode + + def call(self) -> int: + """ Install a package using the Subprocess module + + Returns + ------- + int + The return code of the package install process + """ + with Popen(self._command, + bufsize=0, stdout=PIPE, stderr=STDOUT, shell=self._shell) as proc: + while True: + if proc.stdout is not None: + lines = proc.stdout.readline() + returncode = proc.poll() + if lines == b"" and returncode is not None: + break + + for line in lines.split(b"\r"): + self._seen_line_log(line.decode("utf-8", errors="replace").rstrip()) + self._non_gui_print(line) + + return returncode class Tips(): """ Display installation Tips """ - def __init__(self): - self.output = Output() - - def docker_no_cuda(self): + @classmethod + def docker_no_cuda(cls) -> None: """ Output Tips for Docker without Cuda """ - - path = os.path.dirname(os.path.realpath(__file__)) - self.output.info( - "1. Install Docker\n" - "https://www.docker.com/community-edition\n\n" - "2. Build Docker Image For Faceswap\n" - "docker build -t deepfakes-cpu -f Dockerfile.cpu .\n\n" - "3. Mount faceswap volume and Run it\n" - "# without GUI\n" - "docker run -tid -p 8888:8888 \\ \n" - "\t--hostname deepfakes-cpu --name deepfakes-cpu \\ \n" - f"\t-v {path}:/srv \\ \n" - "\tdeepfakes-cpu\n\n" - "# with gui. tools.py gui working.\n" - "## enable local access to X11 server\n" - "xhost +local:\n" - "## create container\n" - "nvidia-docker run -tid -p 8888:8888 \\ \n" - "\t--hostname deepfakes-cpu --name deepfakes-cpu \\ \n" - f"\t-v {path}:/srv \\ \n" - "\t-v /tmp/.X11-unix:/tmp/.X11-unix \\ \n" - "\t-e DISPLAY=unix$DISPLAY \\ \n" - "\t-e AUDIO_GID=`getent group audio | cut -d: -f3` \\ \n" - "\t-e VIDEO_GID=`getent group video | cut -d: -f3` \\ \n" - "\t-e GID=`id -g` \\ \n" - "\t-e UID=`id -u` \\ \n" - "\tdeepfakes-cpu \n\n" - "4. Open a new terminal to run faceswap.py in /srv\n" - "docker exec -it deepfakes-cpu bash") - self.output.info("That's all you need to do with a docker. Have fun.") - - def docker_cuda(self): - """ Output Tips for Docker wit Cuda""" - - path = os.path.dirname(os.path.realpath(__file__)) - self.output.info( - "1. Install Docker\n" - "https://www.docker.com/community-edition\n\n" - "2. Install latest CUDA\n" - "CUDA: https://developer.nvidia.com/cuda-downloads\n\n" - "3. Install Nvidia-Docker & Restart Docker Service\n" - "https://github.com/NVIDIA/nvidia-docker\n\n" - "4. Build Docker Image For Faceswap\n" - "docker build -t deepfakes-gpu -f Dockerfile.gpu .\n\n" - "5. Mount faceswap volume and Run it\n" - "# without gui \n" - "docker run -tid -p 8888:8888 \\ \n" - "\t--hostname deepfakes-gpu --name deepfakes-gpu \\ \n" - f"\t-v {path}:/srv \\ \n" - "\tdeepfakes-gpu\n\n" - "# with gui.\n" - "## enable local access to X11 server\n" - "xhost +local:\n" - "## enable nvidia device if working under bumblebee\n" - "echo ON > /proc/acpi/bbswitch\n" - "## create container\n" - "nvidia-docker run -tid -p 8888:8888 \\ \n" - "\t--hostname deepfakes-gpu --name deepfakes-gpu \\ \n" - f"\t-v {path}:/srv \\ \n" - "\t-v /tmp/.X11-unix:/tmp/.X11-unix \\ \n" - "\t-e DISPLAY=unix$DISPLAY \\ \n" - "\t-e AUDIO_GID=`getent group audio | cut -d: -f3` \\ \n" - "\t-e VIDEO_GID=`getent group video | cut -d: -f3` \\ \n" - "\t-e GID=`id -g` \\ \n" - "\t-e UID=`id -u` \\ \n" - "\tdeepfakes-gpu\n\n" - "6. Open a new terminal to interact with the project\n" - "docker exec deepfakes-gpu python /srv/faceswap.py gui\n") - - def macos(self): + logger.info( + "1. Install Docker from: https://www.docker.com/get-started\n\n" + "2. Enter the Faceswap folder and build the Docker Image For Faceswap:\n" + " docker build -t faceswap-cpu -f Dockerfile.cpu .\n\n" + "3. Launch and enter the Faceswap container:\n" + " a. Headless:\n" + " docker run --rm -it -v ./:/srv faceswap-cpu\n\n" + " b. GUI:\n" + " xhost +local: && \\ \n" + " docker run --rm -it \\ \n" + " -v ./:/srv \\ \n" + " -v /tmp/.X11-unix:/tmp/.X11-unix \\ \n" + " -e DISPLAY=${DISPLAY} \\ \n" + " faceswap-cpu \n") + logger.info("That's all you need to do with docker. Have fun.") + + @classmethod + def docker_cuda(cls) -> None: + """ Output Tips for Docker with Cuda""" + logger.info( + "1. Install Docker from: https://www.docker.com/get-started\n\n" + "2. Install latest CUDA 11 and cuDNN 8 from: https://developer.nvidia.com/cuda-" + "downloads\n\n" + "3. Install the the Nvidia Container Toolkit from https://docs.nvidia.com/datacenter/" + "cloud-native/container-toolkit/latest/install-guide\n\n" + "4. Restart Docker Service\n\n" + "5. Enter the Faceswap folder and build the Docker Image For Faceswap:\n" + " docker build -t faceswap-gpu -f Dockerfile.gpu .\n\n" + "6. Launch and enter the Faceswap container:\n" + " a. Headless:\n" + " docker run --runtime=nvidia --rm -it -v ./:/srv faceswap-gpu\n\n" + " b. GUI:\n" + " xhost +local: && \\ \n" + " docker run --runtime=nvidia --rm -it \\ \n" + " -v ./:/srv \\ \n" + " -v /tmp/.X11-unix:/tmp/.X11-unix \\ \n" + " -e DISPLAY=${DISPLAY} \\ \n" + " faceswap-gpu \n") + logger.info("That's all you need to do with docker. Have fun.") + + @classmethod + def macos(cls) -> None: """ Output Tips for macOS""" - self.output.info( + logger.info( "setup.py does not directly support macOS. The following tips should help:\n\n" "1. Install system dependencies:\n" "XCode from the Apple Store\n" @@ -922,17 +1711,21 @@ def macos(self): "CUDA: https://developer.nvidia.com/cuda-downloads" "cuDNN: https://developer.nvidia.com/rdp/cudnn-download\n\n") - def pip(self): + @classmethod + def pip(cls) -> None: """ Pip Tips """ - self.output.info("1. Install PIP requirements\n" - "You may want to execute `chcp 65001` in cmd line\n" - "to fix Unicode issues on Windows when installing dependencies") + logger.info("1. Install PIP requirements\n" + "You may want to execute `chcp 65001` in cmd line\n" + "to fix Unicode issues on Windows when installing dependencies") if __name__ == "__main__": + logfile = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "faceswap_setup.log") + log_setup("INFO", logfile, "setup") + logger.debug("Setup called with args: %s", sys.argv) ENV = Environment() Checks(ENV) ENV.set_config() - if INSTALL_FAILED: + if _INSTALL_FAILED: sys.exit(1) Install(ENV) diff --git a/tools/restore/__init__.py b/tests/lib/gpu_stats/__init__.py similarity index 100% rename from tools/restore/__init__.py rename to tests/lib/gpu_stats/__init__.py diff --git a/tests/lib/gpu_stats/_base_test.py b/tests/lib/gpu_stats/_base_test.py new file mode 100644 index 0000000000..225d11f48b --- /dev/null +++ b/tests/lib/gpu_stats/_base_test.py @@ -0,0 +1,162 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`lib.gpu_stats._base` """ +import typing as T + +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest +import pytest_mock + +# pylint:disable=protected-access +from lib.gpu_stats import _base +from lib.gpu_stats._base import BiggestGPUInfo, GPUInfo, _GPUStats, set_exclude_devices +from lib.utils import get_backend + + +def test_set_exclude_devices(monkeypatch: pytest.MonkeyPatch) -> None: + """ Test that :func:`~lib.gpu_stats._base.set_exclude_devices` adds devices + + Parameters + ---------- + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching _EXCLUDE_DEVICES + """ + monkeypatch.setattr(_base, "_EXCLUDE_DEVICES", []) + assert not _base._EXCLUDE_DEVICES + set_exclude_devices([0, 1]) + assert _base._EXCLUDE_DEVICES == [0, 1] + + +@dataclass +class _DummyData: + """ Dummy data for initializing and testing :class:`~lib.gpu_stats._base._GPUStats` """ + device_count = 2 + active_devices = [0, 1] + handles = [0, 1] + driver = "test_driver" + device_names = ['test_device_0', 'test_device_1'] + vram = [1024, 2048] + free_vram = [512, 1024] + + +@pytest.fixture(name="gpu_stats_instance") +def fixture__gpu_stats_instance(mocker: pytest_mock.MockerFixture) -> _GPUStats: + """ Create a fixture of the :class:`~lib.gpu_stats._base._GPUStats` object + + Parameters + ---------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + """ + mocker.patch.object(_GPUStats, '_initialize') + mocker.patch.object(_GPUStats, '_shutdown') + mocker.patch.object(_GPUStats, '_get_device_count', return_value=_DummyData.device_count) + mocker.patch.object(_GPUStats, '_get_active_devices', return_value=_DummyData.active_devices) + mocker.patch.object(_GPUStats, '_get_handles', return_value=_DummyData.handles) + mocker.patch.object(_GPUStats, '_get_driver', return_value=_DummyData.driver) + mocker.patch.object(_GPUStats, '_get_device_names', return_value=_DummyData.device_names) + mocker.patch.object(_GPUStats, '_get_vram', return_value=_DummyData.vram) + mocker.patch.object(_GPUStats, '_get_free_vram', return_value=_DummyData.free_vram) + gpu_stats = _GPUStats() + return gpu_stats + + +def test__gpu_stats_init_(gpu_stats_instance: _GPUStats) -> None: + """ Test that the base :class:`~lib.gpu_stats._base._GPUStats` class initializes correctly + + Parameters + ---------- + gpu_stats_instance: :class:`_GPUStats` + Fixture instance of the _GPUStats base class + """ + # Ensure that the object is initialized and shutdown correctly + assert gpu_stats_instance._is_initialized is False + assert T.cast(MagicMock, gpu_stats_instance._initialize).call_count == 1 + assert T.cast(MagicMock, gpu_stats_instance._shutdown).call_count == 1 + + # Ensure that the object correctly gets and stores the device count, active devices, + # handles, driver, device names, and VRAM information + assert gpu_stats_instance.device_count == _DummyData.device_count + assert gpu_stats_instance._active_devices == _DummyData.active_devices + assert gpu_stats_instance._handles == _DummyData.handles + assert gpu_stats_instance._driver == _DummyData.driver + assert gpu_stats_instance._device_names == _DummyData.device_names + assert gpu_stats_instance._vram == _DummyData.vram + + +def test__gpu_stats_properties(gpu_stats_instance: _GPUStats) -> None: + """ Test that the :class:`~lib.gpu_stats._base._GPUStats` properties are set and formatted + correctly. + + Parameters + ---------- + gpu_stats_instance: :class:`_GPUStats` + Fixture instance of the _GPUStats base class + """ + assert gpu_stats_instance.cli_devices == ['0: test_device_0', '1: test_device_1'] + assert gpu_stats_instance.sys_info == GPUInfo(vram=_DummyData.vram, + vram_free=_DummyData.free_vram, + driver=_DummyData.driver, + devices=_DummyData.device_names, + devices_active=_DummyData.active_devices) + + +def test__gpu_stats_get_card_most_free(mocker: pytest_mock.MockerFixture, + gpu_stats_instance: _GPUStats) -> None: + """ Confirm that :func:`ib.gpu_stats._base._GPUStats.get_card_most_free` functions + correctly + + Parameters + ---------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + gpu_stats_instance: :class:`_GPUStats` + Fixture instance of the _GPUStats base class + """ + assert gpu_stats_instance.get_card_most_free() == BiggestGPUInfo(card_id=1, + device='test_device_1', + free=1024, + total=2048) + mocker.patch.object(_GPUStats, '_get_active_devices', return_value=[]) + gpu_stats = _GPUStats() + assert gpu_stats.get_card_most_free() == BiggestGPUInfo(card_id=-1, + device='No GPU devices found', + free=2048, + total=2048) + + +def test__gpu_stats_exclude_all_devices(gpu_stats_instance: _GPUStats) -> None: + """ Ensure that the object correctly returns whether all devices are excluded + + Parameters + ---------- + gpu_stats_instance: :class:`_GPUStats` + Fixture instance of the _GPUStats base class + """ + assert gpu_stats_instance.exclude_all_devices is False + set_exclude_devices([0, 1]) + assert gpu_stats_instance.exclude_all_devices is True + + +def test__gpu_stats_no_active_devices( + caplog: pytest.LogCaptureFixture, + gpu_stats_instance: _GPUStats, # pylint:disable=unused-argument + mocker: pytest_mock.MockerFixture) -> None: + """ Ensure that no active GPUs raises a warning when not in CPU mode + + Parameters + ---------- + caplog: :class:`pytest.LogCaptureFixture` + Pytest's log capturing fixture + gpu_stats_instance: :class:`_GPUStats` + Fixture instance of the _GPUStats base class + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + """ + if get_backend() == "cpu": + return + caplog.set_level("WARNING") + mocker.patch.object(_GPUStats, '_get_active_devices', return_value=[]) + _GPUStats() + assert "No GPU detected" in caplog.messages diff --git a/plugins/extract/detect/.cache/.keep b/tests/lib/gui/__init__.py similarity index 100% rename from plugins/extract/detect/.cache/.keep rename to tests/lib/gui/__init__.py diff --git a/plugins/extract/mask/.cache/.keep b/tests/lib/gui/stats/__init__.py similarity index 100% rename from plugins/extract/mask/.cache/.keep rename to tests/lib/gui/stats/__init__.py diff --git a/tests/lib/gui/stats/event_reader_test.py b/tests/lib/gui/stats/event_reader_test.py new file mode 100644 index 0000000000..216790550b --- /dev/null +++ b/tests/lib/gui/stats/event_reader_test.py @@ -0,0 +1,740 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`lib.gui.stats.event_reader` """ +# pylint:disable=protected-access +from __future__ import annotations +import json +import os +import typing as T + +from shutil import rmtree +from time import time +from unittest.mock import MagicMock + +import numpy as np +import pytest +import pytest_mock + +import tensorflow as tf +from tensorflow.core.util import event_pb2 # pylint:disable=no-name-in-module + +from lib.gui.analysis.event_reader import (_Cache, _CacheData, _EventParser, + _LogFiles, EventData, TensorBoardLogs) + +if T.TYPE_CHECKING: + from collections.abc import Iterator + + +def test__logfiles(tmp_path: str): + """ Test the _LogFiles class operates correctly + + Parameters + ---------- + tmp_path: :class:`pathlib.Path` + """ + # dummy logfiles + junk data + sess_1 = os.path.join(tmp_path, "session_1", "train") + sess_2 = os.path.join(tmp_path, "session_2", "train") + os.makedirs(sess_1) + os.makedirs(sess_2) + + test_log_1 = os.path.join(sess_1, "events.out.tfevents.123.456.v2") + test_log_2 = os.path.join(sess_2, "events.out.tfevents.789.012.v2") + test_log_junk = os.path.join(sess_2, "test_file.txt") + + for fname in (test_log_1, test_log_2, test_log_junk): + with open(fname, "a", encoding="utf-8"): + pass + + log_files = _LogFiles(tmp_path) + # Test all correct + assert isinstance(log_files._filenames, dict) + assert len(log_files._filenames) == 2 + assert log_files._filenames == {1: test_log_1, 2: test_log_2} + + assert log_files.session_ids == [1, 2] + + assert log_files.get(1) == test_log_1 + assert log_files.get(2) == test_log_2 + + # Remove a file, refresh and check again + rmtree(sess_1) + log_files.refresh() + assert log_files._filenames == {2: test_log_2} + assert log_files.get(2) == test_log_2 + assert log_files.get(3) == "" + + +def test__cachedata(): + """ Test the _CacheData class operates correctly """ + labels = ["label_a", "label_b"] + timestamps = np.array([1.23, 4.56], dtype="float64") + loss = np.array([[2.34, 5.67], [3.45, 6.78]], dtype="float32") + + # Initial test + cache = _CacheData(labels, timestamps, loss) + assert cache.labels == labels + assert cache._timestamps_shape == timestamps.shape + assert cache._loss_shape == loss.shape + np.testing.assert_array_equal(cache.timestamps, timestamps) + np.testing.assert_array_equal(cache.loss, loss) + + # Add data test + new_timestamps = np.array([2.34, 6.78], dtype="float64") + new_loss = np.array([[3.45, 7.89], [8.90, 1.23]], dtype="float32") + + expected_timestamps = np.concatenate([timestamps, new_timestamps]) + expected_loss = np.concatenate([loss, new_loss]) + + cache.add_live_data(new_timestamps, new_loss) + assert cache.labels == labels + assert cache._timestamps_shape == expected_timestamps.shape + assert cache._loss_shape == expected_loss.shape + np.testing.assert_array_equal(cache.timestamps, expected_timestamps) + np.testing.assert_array_equal(cache.loss, expected_loss) + + +# _Cache tests +class Test_Cache: # pylint:disable=invalid-name + """ Test that :class:`lib.gui.analysis.event_reader._Cache` works correctly """ + @staticmethod + def test_init() -> None: + """ Test __init__ """ + cache = _Cache() + assert isinstance(cache._data, dict) + assert isinstance(cache._carry_over, dict) + assert isinstance(cache._loss_labels, list) + assert not cache._data + assert not cache._carry_over + assert not cache._loss_labels + + @staticmethod + def test_is_cached() -> None: + """ Test is_cached function works """ + cache = _Cache() + + data = _CacheData(["test_1", "test_2"], + np.array([1.23, ], dtype="float64"), + np.array([[2.34, ], [4.56]], dtype="float32")) + cache._data[1] = data + assert cache.is_cached(1) + assert not cache.is_cached(2) + + @staticmethod + def test_cache_data(mocker: pytest_mock.MockerFixture) -> None: + """ Test cache_data function works + + Parameters + ---------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking full_info called from _SysInfo + """ + cache = _Cache() + + session_id = 1 + data = {1: EventData(4., [1., 2.]), 2: EventData(5., [3., 4.])} + labels = ['label1', 'label2'] + is_live = False + + cache.cache_data(session_id, data, labels, is_live) + assert cache._loss_labels == labels + assert cache.is_cached(session_id) + np.testing.assert_array_equal(cache._data[session_id].timestamps, np.array([4., 5.])) + np.testing.assert_array_equal(cache._data[session_id].loss, np.array([[1., 2.], [3., 4.]])) + + add_live = mocker.patch("lib.gui.analysis.event_reader._Cache._add_latest_live") + is_live = True + cache.cache_data(session_id, data, labels, is_live) + assert add_live.called + + @staticmethod + def test__to_numpy() -> None: + """ Test _to_numpy function works """ + cache = _Cache() + cache._loss_labels = ['label1', 'label2'] + data = {1: EventData(4., [1., 2.]), 2: EventData(5., [3., 4.])} + + # Non-live + is_live = False + times, loss = cache._to_numpy(data, is_live) + np.testing.assert_array_equal(times, np.array([4., 5.])) + np.testing.assert_array_equal(loss, np.array([[1., 2.], [3., 4.]])) + + # Correctly collected live + is_live = True + times, loss = cache._to_numpy(data, is_live) + np.testing.assert_array_equal(times, np.array([4., 5.])) + np.testing.assert_array_equal(loss, np.array([[1., 2.], [3., 4.]])) + + # Incorrectly collected live + live_data = {1: EventData(4., [1., 2.]), + 2: EventData(5., [3.]), + 3: EventData(6., [4., 5., 6.])} + times, loss = cache._to_numpy(live_data, is_live) + np.testing.assert_array_equal(times, np.array([4.])) + np.testing.assert_array_equal(loss, np.array([[1., 2.]])) + + @staticmethod + def test__collect_carry_over() -> None: + """ Test _collect_carry_over function works """ + data = {1: EventData(3., [4., 5.]), 2: EventData(6., [7., 8.])} + carry_over = {1: EventData(3., [2., 3.])} + expected = {1: EventData(3., [2., 3., 4., 5.]), 2: EventData(6., [7., 8.])} + + cache = _Cache() + cache._carry_over = carry_over + cache._collect_carry_over(data) + assert data == expected + + @staticmethod + def test__process_data() -> None: + """ Test _process_data function works """ + cache = _Cache() + cache._loss_labels = ['label1', 'label2'] + + data = {1: EventData(4., [5., 6.]), + 2: EventData(5., [7., 8.]), + 3: EventData(6., [9.])} + is_live = False + expected_timestamps = np.array([4., 5.]) + expected_loss = np.array([[5., 6.], [7., 8.]]) + expected_carry_over = {3: EventData(6., [9.])} + + timestamps, loss = cache._process_data(data, is_live) + np.testing.assert_array_equal(timestamps, expected_timestamps) + np.testing.assert_array_equal(loss, expected_loss) + assert not cache._carry_over + + is_live = True + timestamps, loss = cache._process_data(data, is_live) + np.testing.assert_array_equal(timestamps, expected_timestamps) + np.testing.assert_array_equal(loss, expected_loss) + assert cache._carry_over == expected_carry_over + + @staticmethod + def test__add_latest_live() -> None: + """ Test _add_latest_live function works """ + session_id = 1 + labels = ['label1', 'label2'] + data = {1: EventData(3., [5., 6.]), 2: EventData(4., [7., 8.])} + new_timestamp = np.array([5.], dtype="float64") + new_loss = np.array([[8., 9.]], dtype="float32") + expected_timestamps = np.array([3., 4., 5.]) + expected_loss = np.array([[5., 6.], [7., 8.], [8., 9.]]) + + cache = _Cache() + cache.cache_data(session_id, data, labels) # Initial data + cache._add_latest_live(session_id, new_loss, new_timestamp) + + assert cache.is_cached(session_id) + assert cache._loss_labels == labels + np.testing.assert_array_equal(cache._data[session_id].timestamps, expected_timestamps) + np.testing.assert_array_equal(cache._data[session_id].loss, expected_loss) + + @staticmethod + def test_get_data() -> None: + """ Test get_data function works """ + session_id = 1 + + cache = _Cache() + assert cache.get_data(session_id, "loss") is None + assert cache.get_data(session_id, "timestamps") is None + + labels = ['label1', 'label2'] + data = {1: EventData(3., [5., 6.]), 2: EventData(4., [7., 8.])} + expected_timestamps = np.array([3., 4.]) + expected_loss = np.array([[5., 6.], [7., 8.]]) + + cache.cache_data(session_id, data, labels, is_live=False) + get_timestamps = cache.get_data(session_id, "timestamps") + get_loss = cache.get_data(session_id, "loss") + + assert isinstance(get_timestamps, dict) + assert len(get_timestamps) == 1 + assert list(get_timestamps) == [session_id] + result = get_timestamps[session_id] + assert list(result) == ["timestamps"] + np.testing.assert_array_equal(result["timestamps"], expected_timestamps) + + assert isinstance(get_loss, dict) + assert len(get_loss) == 1 + assert list(get_loss) == [session_id] + result = get_loss[session_id] + assert list(result) == ["loss", "labels"] + np.testing.assert_array_equal(result["loss"], expected_loss) + + +# TensorBoardLogs +class TestTensorBoardLogs: + """ Test that :class:`lib.gui.analysis.event_reader.TensorBoardLogs` works correctly """ + + @pytest.fixture(name="tensorboardlogs_instance") + def tensorboardlogs_fixture(self, + tmp_path: str, + request: pytest.FixtureRequest) -> TensorBoardLogs: + """ Pytest fixture for :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + + Parameters + ---------- + tmp_path: :class:`pathlib.Path` + Temporary folder for dummy data + + Returns + ------- + :class::class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance for testing + """ + sess_1 = os.path.join(tmp_path, "session_1", "train") + sess_2 = os.path.join(tmp_path, "session_2", "train") + os.makedirs(sess_1) + os.makedirs(sess_2) + + test_log_1 = os.path.join(sess_1, "events.out.tfevents.123.456.v2") + test_log_2 = os.path.join(sess_2, "events.out.tfevents.789.012.v2") + + for fname in (test_log_1, test_log_2): + with open(fname, "a", encoding="utf-8"): + pass + + tblogs_instance = TensorBoardLogs(tmp_path, False) + + def teardown(): + rmtree(tmp_path) + + request.addfinalizer(teardown) + return tblogs_instance + + @staticmethod + def test_init(tensorboardlogs_instance: TensorBoardLogs) -> None: + """ Test __init__ works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + """ + tb_logs = tensorboardlogs_instance + assert isinstance(tb_logs._log_files, _LogFiles) + assert isinstance(tb_logs._cache, _Cache) + assert not tb_logs._is_training + + is_training = True + folder = tb_logs._log_files._logs_folder + tb_logs = TensorBoardLogs(folder, is_training) + assert tb_logs._is_training + + @staticmethod + def test_session_ids(tensorboardlogs_instance: TensorBoardLogs) -> None: + """ Test session_ids property works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + """ + tb_logs = tensorboardlogs_instance + assert tb_logs.session_ids == [1, 2] + + @staticmethod + def test_set_training(tensorboardlogs_instance: TensorBoardLogs) -> None: + """ Test set_training works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + """ + tb_logs = tensorboardlogs_instance + assert not tb_logs._is_training + assert tb_logs._training_iterator is None + tb_logs.set_training(True) + assert tb_logs._is_training + assert tb_logs._training_iterator is not None + tb_logs.set_training(False) + assert not tb_logs._is_training + assert tb_logs._training_iterator is None + + @staticmethod + def test__cache_data(tensorboardlogs_instance: TensorBoardLogs, + mocker: pytest_mock.MockerFixture) -> None: + """ Test _cache_data works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking event parser caching is called + """ + tb_logs = tensorboardlogs_instance + session_id = 1 + cacher = mocker.patch("lib.gui.analysis.event_reader._EventParser.cache_events") + tb_logs._cache_data(session_id) + assert cacher.called + cacher.reset_mock() + + tb_logs.set_training(True) + tb_logs._cache_data(session_id) + assert cacher.called + + @staticmethod + def test__check_cache(tensorboardlogs_instance: TensorBoardLogs, + mocker: pytest_mock.MockerFixture) -> None: + """ Test _check_cache works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _cache_data is called + """ + is_cached = mocker.patch("lib.gui.analysis.event_reader._Cache.is_cached") + cache_data = mocker.patch("lib.gui.analysis.event_reader.TensorBoardLogs._cache_data") + tb_logs = tensorboardlogs_instance + + # Session ID not training + is_cached.return_value = False + tb_logs._check_cache(1) + assert is_cached.called + assert cache_data.called + is_cached.reset_mock() + cache_data.reset_mock() + + is_cached.return_value = True + tb_logs._check_cache(1) + assert is_cached.called + assert not cache_data.called + is_cached.reset_mock() + cache_data.reset_mock() + + # Session ID and training + tb_logs.set_training(True) + tb_logs._check_cache(1) + assert not cache_data.called + cache_data.reset_mock() + + tb_logs._check_cache(2) + assert cache_data.called + cache_data.reset_mock() + + # No session id + tb_logs.set_training(False) + is_cached.return_value = False + + tb_logs._check_cache(None) + assert is_cached.called + assert cache_data.called + is_cached.reset_mock() + cache_data.reset_mock() + + is_cached.return_value = True + tb_logs._check_cache(None) + assert is_cached.called + assert not cache_data.called + is_cached.reset_mock() + cache_data.reset_mock() + + @staticmethod + def test_get_loss(tensorboardlogs_instance: TensorBoardLogs, + mocker: pytest_mock.MockerFixture) -> None: + """ Test get_loss works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _cache_data is called + """ + tb_logs = tensorboardlogs_instance + + with pytest.raises(tf.errors.NotFoundError): # Invalid session id + tb_logs.get_loss(3) + + check_cache = mocker.patch("lib.gui.analysis.event_reader.TensorBoardLogs._check_cache") + get_data = mocker.patch("lib.gui.analysis.event_reader._Cache.get_data") + get_data.return_value = None + + assert isinstance(tb_logs.get_loss(None), dict) + assert check_cache.call_count == 2 + assert get_data.call_count == 2 + check_cache.reset_mock() + get_data.reset_mock() + + assert isinstance(tb_logs.get_loss(1), dict) + assert check_cache.call_count == 1 + assert get_data.call_count == 1 + check_cache.reset_mock() + get_data.reset_mock() + + @staticmethod + def test_get_timestamps(tensorboardlogs_instance: TensorBoardLogs, + mocker: pytest_mock.MockerFixture) -> None: + """ Test get_timestamps works correctly + + Parameters + ---------- + tensorboadlogs_instance: :class:`lib.gui.analysis.event_reader.TensorBoardLogs` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _cache_data is called + """ + tb_logs = tensorboardlogs_instance + with pytest.raises(tf.errors.NotFoundError): # invalid session_id + tb_logs.get_timestamps(3) + + check_cache = mocker.patch("lib.gui.analysis.event_reader.TensorBoardLogs._check_cache") + get_data = mocker.patch("lib.gui.analysis.event_reader._Cache.get_data") + get_data.return_value = None + + assert isinstance(tb_logs.get_timestamps(None), dict) + assert check_cache.call_count == 2 + assert get_data.call_count == 2 + check_cache.reset_mock() + get_data.reset_mock() + + assert isinstance(tb_logs.get_timestamps(1), dict) + assert check_cache.call_count == 1 + assert get_data.call_count == 1 + check_cache.reset_mock() + get_data.reset_mock() + + +# EventParser +class Test_EventParser: # pylint:disable=invalid-name + """ Test that :class:`lib.gui.analysis.event_reader.TensorBoardLogs` works correctly """ + def _create_example_event(self, + step: int, + loss_value: float, + timestamp: float, + serialize: bool = True) -> bytes: + """ Generate a test TensorBoard event + + Parameters + ---------- + step: int + The step value to use + loss_value: float + The loss value to store + timestamp: float + The timestamp to store + serialize: bool, optional + ``True`` to serialize the event to bytes, ``False`` to return the Event object + """ + tags = {0: "keras", 1: "batch_total", 2: "batch_face_a", 3: "batch_face_b"} + event = event_pb2.Event(step=step) + event.summary.value.add(tag=tags[step], # pylint:disable=no-member + simple_value=loss_value) + event.wall_time = timestamp + retval = event.SerializeToString() if serialize else event + return retval + + @pytest.fixture(name="mock_iterator") + def iterator(self) -> Iterator[bytes]: + """ Dummy iterator for generating test events + + Yields + ------ + bytes + A serialized test Tensorboard Event + """ + return iter([self._create_example_event(i, 1 + (i / 10), time()) for i in range(4)]) + + @pytest.fixture(name="mock_cache") + def mock_cache(self): + """ Dummy :class:`_Cache` for testing""" + class _CacheMock: + def __init__(self): + self.data = {} + self._loss_labels = [] + + def is_cached(self, session_id): + """ Dummy is_cached method""" + return session_id in self.data + + def cache_data(self, session_id, data, labels, + is_live=False): # pylint:disable=unused-argument + """ Dummy cache_data method""" + self.data[session_id] = {'data': data, 'labels': labels} + + return _CacheMock() + + @pytest.fixture(name="event_parser_instance") + def event_parser_fixture(self, + mock_iterator: Iterator[bytes], + mock_cache: _Cache) -> _EventParser: + """ Pytest fixture for :class:`lib.gui.analysis.event_reader._EventParser` + + Parameters + ---------- + mock_iterator: Iterator[bytes] + Dummy iterator for generating TF Event data + mock_cache: :class:'_CacheMock' + Dummy _Cache object + + Returns + ------- + :class::class:`lib.gui.analysis.event_reader._EventParser` + The class instance for testing + """ + event_parser = _EventParser(mock_iterator, mock_cache, live_data=False) + return event_parser + + def test__init_(self, + event_parser_instance: _EventParser, + mock_iterator: Iterator[bytes], + mock_cache: _Cache) -> None: + """ Test __init__ works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + mock_iterator: Iterator[bytes] + Dummy iterator for generating TF Event data + mock_cache: :class:'_CacheMock' + Dummy _Cache object + """ + event_parse = event_parser_instance + assert not hasattr(event_parse._iterator, "__name__") + evp_live = _EventParser(mock_iterator, mock_cache, live_data=True) + assert evp_live._iterator.__name__ == "_get_latest_live" # type:ignore[attr-defined] + + def test__get_latest_live(self, event_parser_instance: _EventParser) -> None: + """ Test _get_latest_live works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + """ + event_parse = event_parser_instance + test = list(event_parse._get_latest_live(event_parse._iterator)) + assert len(test) == 4 + + def test_cache_events(self, + event_parser_instance: _EventParser, + mocker: pytest_mock.MockerFixture, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test cache_events works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for capturing method calls + monkeypatch: :class:`pytest.MonkeyPatch` + For patching different iterators for testing output + """ + monkeypatch.setattr("lib.utils._FS_BACKEND", "cpu") + + event_parse = event_parser_instance + event_parse._parse_outputs = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + event_parse._process_event = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + event_parse._cache.cache_data = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + + # keras model + monkeypatch.setattr(event_parse, + "_iterator", + iter([self._create_example_event(0, 1., time())])) + event_parse.cache_events(1) + assert event_parse._parse_outputs.called + assert not event_parse._process_event.called + assert event_parse._cache.cache_data.called + event_parse._parse_outputs.reset_mock() + event_parse._process_event.reset_mock() + event_parse._cache.cache_data.reset_mock() + + # Batch item + monkeypatch.setattr(event_parse, + "_iterator", + iter([self._create_example_event(1, 1., time())])) + event_parse.cache_events(1) + assert not event_parse._parse_outputs.called + assert event_parse._process_event.called + assert event_parse._cache.cache_data.called + event_parse._parse_outputs.reset_mock() + event_parse._process_event.reset_mock() + event_parse._cache.cache_data.reset_mock() + + # No summary value + monkeypatch.setattr(event_parse, + "_iterator", + iter([event_pb2.Event(step=1).SerializeToString()])) + assert not event_parse._parse_outputs.called + assert not event_parse._process_event.called + assert not event_parse._cache.cache_data.called + event_parse._parse_outputs.reset_mock() + event_parse._process_event.reset_mock() + event_parse._cache.cache_data.reset_mock() + + def test__parse_outputs(self, + event_parser_instance: _EventParser, + mocker: pytest_mock.MockerFixture) -> None: + """ Test _parse_outputs works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for event object + """ + event_parse = event_parser_instance + model = {"config": {"layers": [{"name": "decoder_a", + "config": {"output_layers": [["face_out_a", 0, 0]]}}, + {"name": "decoder_b", + "config": {"output_layers": [["face_out_b", 0, 0]]}}], + "output_layers": [["decoder_a", 1, 0], ["decoder_b", 1, 0]]}} + data = json.dumps(model).encode("utf-8") + + event = mocker.MagicMock() + event.summary.value.__getitem__ = lambda self, x: event + event.tensor.string_val.__getitem__ = lambda self, x: data + + assert not event_parse._loss_labels + event_parse._parse_outputs(event) + assert event_parse._loss_labels == ["face_out_a", "face_out_b"] + + def test__get_outputs(self, event_parser_instance: _EventParser) -> None: + """ Test _get_outputs works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + """ + outputs = [["decoder_a", 1, 0], ["decoder_b", 1, 0]] + model_config = {"output_layers": outputs} + + expected = np.array([[out] for out in outputs]) + actual = event_parser_instance._get_outputs(model_config) + assert isinstance(actual, np.ndarray) + assert actual.shape == (2, 1, 3) + np.testing.assert_equal(expected, actual) + + def test__process_event(self, event_parser_instance: _EventParser) -> None: + """ Test _process_event works correctly + + Parameters + ---------- + event_parser_instance: :class:`lib.gui.analysis.event_reader._EventParser` + The class instance to test + """ + event_parse = event_parser_instance + event_data = EventData() + assert not event_data.timestamp + assert not event_data.loss + timestamp = time() + loss = [1.1, 2.2] + event = self._create_example_event(1, 1.0, timestamp, serialize=False) # batch_total + event_parse._process_event(event, event_data) + event = self._create_example_event(2, loss[0], time(), serialize=False) # face A + event_parse._process_event(event, event_data) + event = self._create_example_event(3, loss[1], time(), serialize=False) # face B + event_parse._process_event(event, event_data) + + # Original timestamp and both loss values collected + assert event_data.timestamp == timestamp + np.testing.assert_almost_equal(event_data.loss, loss) # float rounding diff --git a/tests/lib/model/initializers_test.py b/tests/lib/model/initializers_test.py index 2e4921586c..ff52a3369c 100644 --- a/tests/lib/model/initializers_test.py +++ b/tests/lib/model/initializers_test.py @@ -7,18 +7,12 @@ import pytest import numpy as np +from tensorflow.keras import backend as K # pylint:disable=import-error +from tensorflow.keras import initializers as k_initializers # noqa:E501 # pylint:disable=import-error + from lib.model import initializers from lib.utils import get_backend -if get_backend() == "amd": - from keras import backend as K - from keras import initializers as k_initializers -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import backend as K # pylint:disable=import-error - from tensorflow.keras import initializers as k_initializers # pylint:disable=import-error - - CONV_SHAPE = (3, 3, 256, 2048) CONV_ID = get_backend().upper() @@ -49,8 +43,11 @@ def test_icnr(tensor_shape): """ fan_in, _ = initializers.compute_fans(tensor_shape) std = np.sqrt(2. / fan_in) - _runner(initializers.ICNR(initializer=k_initializers.he_uniform(), scale=2), tensor_shape, - target_mean=0, target_std=std) + _runner(initializers.ICNR(initializer=k_initializers.he_uniform(), # pylint:disable=no-member + scale=2), + tensor_shape, + target_mean=0, + target_std=std) @pytest.mark.parametrize('tensor_shape', [CONV_SHAPE], ids=[CONV_ID]) diff --git a/tests/lib/model/layers_test.py b/tests/lib/model/layers_test.py index b6c8fb9286..7b5ec6ebd7 100644 --- a/tests/lib/model/layers_test.py +++ b/tests/lib/model/layers_test.py @@ -10,22 +10,18 @@ from numpy.testing import assert_allclose +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import Input, Model, backend as K # pylint:disable=import-error + from lib.model import layers from lib.utils import get_backend from tests.utils import has_arg -if get_backend() == "amd": - from keras import Input, Model, backend as K -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import Input, Model, backend as K # pylint:disable=import-error - - CONV_SHAPE = (3, 3, 256, 2048) CONV_ID = get_backend().upper() -def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, +def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, # noqa:C901 input_data=None, expected_output=None, expected_output_dtype=None, fixed_batch_size=False): """Test routine for a layer with a single input tensor @@ -40,7 +36,7 @@ def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, for i, var_e in enumerate(input_data_shape): if var_e is None: input_data_shape[i] = np.random.randint(1, 4) - input_data = (10 * np.random.random(input_data_shape)) + input_data = 10 * np.random.random(input_data_shape) input_data = input_data.astype(input_dtype) else: if input_shape is None: @@ -71,7 +67,7 @@ def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, # check with the functional API model = Model(inp, outp) - actual_output = model.predict(input_data) + actual_output = model.predict(input_data, verbose=0) actual_output_shape = actual_output.shape for expected_dim, actual_dim in zip(expected_output_shape, actual_output_shape): @@ -87,7 +83,7 @@ def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, if model.weights: weights = model.get_weights() recovered_model.set_weights(weights) - _output = recovered_model.predict(input_data) + _output = recovered_model.predict(input_data, verbose=0) assert_allclose(_output, actual_output, rtol=1e-3) # test training mode (e.g. useful when the layer has a @@ -105,17 +101,40 @@ def layer_test(layer_cls, kwargs={}, input_shape=None, input_dtype=None, return actual_output +@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) +def test_global_min_pooling_2d(dummy): # pylint:disable=unused-argument + """ Global Min Pooling 2D layer test """ + layer_test(layers.GlobalMinPooling2D, input_shape=(2, 4, 4, 1024)) + + +@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) +def test_global_std_pooling_2d(dummy): # pylint:disable=unused-argument + """ Global Standard Deviation Pooling 2D layer test """ + layer_test(layers.GlobalStdDevPooling2D, input_shape=(2, 4, 4, 1024)) + + +@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) +def test_k_resize_images(dummy): # pylint:disable=unused-argument + """ Global Standard Deviation Pooling 2D layer test """ + layer_test(layers.KResizeImages, input_shape=(2, 4, 4, 1024)) + + +@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) +def test_l2_normalize(dummy): # pylint:disable=unused-argument + """ L2 Normalize layer test """ + layer_test(layers.L2_normalize, kwargs={"axis": 1}, input_shape=(2, 4, 4, 1024)) + + @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) def test_pixel_shuffler(dummy): # pylint:disable=unused-argument """ Pixel Shuffler layer test """ layer_test(layers.PixelShuffler, input_shape=(2, 4, 4, 1024)) -@pytest.mark.skipif(get_backend() == "amd", reason="amd does not support this layer") @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) -def test_subpixel_upscaling(dummy): # pylint:disable=unused-argument - """ Sub Pixel up-scaling layer test """ - layer_test(layers.SubPixelUpscaling, input_shape=(2, 4, 4, 1024)) +def test_quick_gelu(dummy): # pylint:disable=unused-argument + """ Global Standard Deviation Pooling 2D layer test """ + layer_test(layers.QuickGELU, input_shape=(2, 4, 4, 1024)) @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) @@ -125,18 +144,12 @@ def test_reflection_padding_2d(dummy): # pylint:disable=unused-argument @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) -def test_global_min_pooling_2d(dummy): # pylint:disable=unused-argument - """ Global Min Pooling 2D layer test """ - layer_test(layers.GlobalMinPooling2D, input_shape=(2, 4, 4, 1024)) - - -@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) -def test_global_std_pooling_2d(dummy): # pylint:disable=unused-argument - """ Global Standard Deviation Pooling 2D layer test """ - layer_test(layers.GlobalStdDevPooling2D, input_shape=(2, 4, 4, 1024)) +def test_subpixel_upscaling(dummy): # pylint:disable=unused-argument + """ Sub Pixel up-scaling layer test """ + layer_test(layers.SubPixelUpscaling, input_shape=(2, 4, 4, 1024)) @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) -def test_l2_normalize(dummy): # pylint:disable=unused-argument - """ L2 Normalize layer test """ - layer_test(layers.L2_normalize, kwargs={"axis": 1}, input_shape=(2, 4, 4, 1024)) +def test_swish(dummy): # pylint:disable=unused-argument + """ Sub Pixel up-scaling layer test """ + layer_test(layers.Swish, input_shape=(2, 4, 4, 1024)) diff --git a/tests/lib/model/losses_test.py b/tests/lib/model/losses_test.py index 7c83b8755e..ae59b38e10 100644 --- a/tests/lib/model/losses_test.py +++ b/tests/lib/model/losses_test.py @@ -7,15 +7,13 @@ import pytest import numpy as np +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import backend as K, losses as k_losses # noqa:E501 # pylint:disable=import-error + + from lib.model import losses from lib.utils import get_backend -if get_backend() == "amd": - from keras import backend as K, losses as k_losses -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import backend as K, losses as k_losses # pylint:disable=import-error - _PARAMS = [(losses.GeneralizedLoss(), (2, 16, 16)), (losses.GradientLoss(), (2, 16, 16)), # TODO Make sure these output dimensions are correct @@ -29,44 +27,38 @@ @pytest.mark.parametrize(["loss_func", "output_shape"], _PARAMS, ids=_IDS) def test_loss_output(loss_func, output_shape): """ Basic shape tests for loss functions. """ - if get_backend() == "amd" and isinstance(loss_func, losses.GMSDLoss): - pytest.skip("GMSD Loss is not currently compatible with PlaidML") y_a = K.variable(np.random.random((2, 16, 16, 3))) y_b = K.variable(np.random.random((2, 16, 16, 3))) objective_output = loss_func(y_a, y_b) - if get_backend() == "amd": - assert K.eval(objective_output).shape == output_shape - else: - output = objective_output.numpy() - assert output.dtype == "float32" and not np.isnan(output) + output = objective_output.numpy() + assert output.dtype == "float32" and not np.any(np.isnan(output)) -_LWPARAMS = [losses.GeneralizedLoss(), losses.GradientLoss(), losses.GMSDLoss(), - losses.LInfNorm(), k_losses.mean_absolute_error, k_losses.mean_squared_error, - k_losses.logcosh, losses.DSSIMObjective(), losses.MSSSIMLoss()] -_LWIDS = ["GeneralizedLoss", "GradientLoss", "GMSDLoss", "LInfNorm", "mae", "mse", "logcosh", - "DSSIMObjective", "MS-SSIM"] +_LWPARAMS = [losses.DSSIMObjective(), + losses.FocalFrequencyLoss(), + losses.GeneralizedLoss(), + losses.GMSDLoss(), + losses.GradientLoss(), + losses.LaplacianPyramidLoss(), + losses.LDRFLIPLoss(), + losses.LInfNorm(), + k_losses.logcosh, # pylint:disable=no-member + k_losses.mean_absolute_error, + k_losses.mean_squared_error, + losses.MSSIMLoss()] +_LWIDS = ["DSSIMObjective", "FocalFrequencyLoss", "GeneralizedLoss", "GMSDLoss", "GradientLoss", + "LaplacianPyramidLoss", "LInfNorm", "LDRFlipLoss", "logcosh", "mae", "mse", "MS-SSIM"] _LWIDS = [f"{loss}[{get_backend().upper()}]" for loss in _LWIDS] @pytest.mark.parametrize("loss_func", _LWPARAMS, ids=_LWIDS) def test_loss_wrapper(loss_func): """ Test penalized loss wrapper works as expected """ - if get_backend() == "amd": - if isinstance(loss_func, losses.GMSDLoss): - pytest.skip("GMSD Loss is not currently compatible with PlaidML") - if isinstance(loss_func, losses.MSSSIMLoss): - pytest.skip("MS-SSIM Loss is not currently compatible with PlaidML") - if hasattr(loss_func, "__name__") and loss_func.__name__ == "logcosh": - pytest.skip("LogCosh Loss is not currently compatible with PlaidML") - y_a = K.variable(np.random.random((2, 16, 16, 4))) - y_b = K.variable(np.random.random((2, 16, 16, 3))) + y_a = K.variable(np.random.random((2, 64, 64, 4))) + y_b = K.variable(np.random.random((2, 64, 64, 3))) p_loss = losses.LossWrapper() p_loss.add_loss(loss_func, 1.0, -1) p_loss.add_loss(k_losses.mean_squared_error, 2.0, 3) output = p_loss(y_a, y_b) - if get_backend() == "amd": - assert K.dtype(output) == "float32" and K.eval(output).shape == (2, ) - else: - output = output.numpy() - assert output.dtype == "float32" and not np.isnan(output) + output = output.numpy() + assert output.dtype == "float32" and not np.any(np.isnan(output)) diff --git a/tests/lib/model/nn_blocks_test.py b/tests/lib/model/nn_blocks_test.py index d775be8b0e..4793b95ea7 100644 --- a/tests/lib/model/nn_blocks_test.py +++ b/tests/lib/model/nn_blocks_test.py @@ -9,25 +9,17 @@ import pytest import numpy as np - from numpy.testing import assert_allclose +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import Input, Model, backend as K # pylint:disable=import-error + from lib.model import nn_blocks from lib.utils import get_backend -if get_backend() == "amd": - from keras import Input, Model, backend as K -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import Input, Model, backend as K # pylint:disable=import-error - def block_test(layer_func, kwargs={}, input_shape=None): - """Test routine for faceswap neural network blocks. - - Tests are simple and are to ensure that the blocks compile on both tensorflow - and plaidml backends - """ + """Test routine for faceswap neural network blocks. """ # generate input data assert input_shape input_dtype = K.floatx() @@ -35,7 +27,7 @@ def block_test(layer_func, kwargs={}, input_shape=None): for i, var_e in enumerate(input_data_shape): if var_e is None: input_data_shape[i] = np.random.randint(1, 4) - input_data = (10 * np.random.random(input_data_shape)) + input_data = 10 * np.random.random(input_data_shape) input_data = input_data.astype(input_dtype) expected_output_dtype = input_dtype @@ -47,7 +39,7 @@ def block_test(layer_func, kwargs={}, input_shape=None): # check with the functional API model = Model(inp, outp) - actual_output = model.predict(input_data) + actual_output = model.predict(input_data, verbose=0) # test serialization, weight setting at model level model_config = model.get_config() @@ -55,7 +47,7 @@ def block_test(layer_func, kwargs={}, input_shape=None): if model.weights: weights = model.get_weights() recovered_model.set_weights(weights) - _output = recovered_model.predict(input_data) + _output = recovered_model.predict(input_data, verbose=0) assert_allclose(_output, actual_output, rtol=1e-3) # for further checks in the caller function @@ -64,16 +56,16 @@ def block_test(layer_func, kwargs={}, input_shape=None): _PARAMS = ["use_icnr_init", "use_convaware_init", "use_reflect_padding"] _VALUES = list(product([True, False], repeat=len(_PARAMS))) -_IDS = ["{}[{}]".format("|".join([_PARAMS[idx] for idx, b in enumerate(v) if b]), - get_backend().upper()) for v in _VALUES] +_IDS = [f"{'|'.join([_PARAMS[idx] for idx, b in enumerate(v) if b])}[{get_backend().upper()}]" + for v in _VALUES] @pytest.mark.parametrize(_PARAMS, _VALUES, ids=_IDS) def test_blocks(use_icnr_init, use_convaware_init, use_reflect_padding): """ Test for all blocks contained within the NNBlocks Class """ - config = dict(icnr_init=use_icnr_init, - conv_aware_init=use_convaware_init, - reflect_padding=use_reflect_padding) + config = {"icnr_init": use_icnr_init, + "conv_aware_init": use_convaware_init, + "reflect_padding": use_reflect_padding} nn_blocks.set_config(config) block_test(nn_blocks.Conv2DOutput(64, 3), input_shape=(2, 8, 8, 32)) block_test(nn_blocks.Conv2DBlock(64), input_shape=(2, 8, 8, 32)) diff --git a/tests/lib/model/normalization_test.py b/tests/lib/model/normalization_test.py index 7f3b1fb7f6..c447c6e480 100644 --- a/tests/lib/model/normalization_test.py +++ b/tests/lib/model/normalization_test.py @@ -8,17 +8,13 @@ import numpy as np import pytest +from tensorflow.keras import regularizers, models, layers # noqa:E501 # pylint:disable=import-error + from lib.model import normalization from lib.utils import get_backend from tests.lib.model.layers_test import layer_test -if get_backend() == "amd": - from keras import regularizers, models, layers -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import regularizers, models, layers # pylint:disable=import-error - @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) def test_instance_normalization(dummy): # pylint:disable=unused-argument @@ -69,8 +65,8 @@ def test_group_normalization(dummy): # pylint:disable=unused-argument _PARAMS = ["center", "scale"] _VALUES = list(product([True, False], repeat=len(_PARAMS))) -_IDS = ["{}[{}]".format("|".join([_PARAMS[idx] for idx, b in enumerate(v) if b]), - get_backend().upper()) for v in _VALUES] +_IDS = [f"{'|'.join([_PARAMS[idx] for idx, b in enumerate(v) if b])}[{get_backend().upper()}]" + for v in _VALUES] @pytest.mark.parametrize(_PARAMS, _VALUES, ids=_IDS) @@ -86,7 +82,7 @@ def test_adain_normalization(center, scale): model = models.Model(inputs, norm(inputs)) data = [10 * np.random.random(shape) for shape in shapes] - actual_output = model.predict(data) + actual_output = model.predict(data, verbose=0) actual_output_shape = actual_output.shape for expected_dim, actual_dim in zip(expected_output_shape, @@ -95,16 +91,8 @@ def test_adain_normalization(center, scale): assert expected_dim == actual_dim -@pytest.mark.parametrize(_PARAMS, _VALUES, ids=_IDS) -def test_layer_normalization(center, scale): - """ Basic test for layer normalization. """ - layer_test(normalization.LayerNormalization, - kwargs={"center": center, "scale": scale}, - input_shape=(4, 512)) - - _PARAMS = ["partial", "bias"] -_VALUES = [(0.0, False), (0.25, False), (0.5, True), (0.75, False), (1.0, True)] +_VALUES = [(0.0, False), (0.25, False), (0.5, True), (0.75, False), (1.0, True)] # type:ignore _IDS = [f"partial={v[0]}|bias={v[1]}[{get_backend().upper()}]" for v in _VALUES] diff --git a/tests/lib/model/optimizers_test.py b/tests/lib/model/optimizers_test.py index b85581c43e..34d5335824 100644 --- a/tests/lib/model/optimizers_test.py +++ b/tests/lib/model/optimizers_test.py @@ -7,22 +7,16 @@ import numpy as np from numpy.testing import assert_allclose +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow.keras import optimizers as k_optimizers # pylint:disable=import-error +from tensorflow.keras.layers import Dense, Activation # pylint:disable=import-error +from tensorflow.keras.models import Sequential # pylint:disable=import-error from lib.model import optimizers from lib.utils import get_backend from tests.utils import generate_test_data, to_categorical -if get_backend() == "amd": - from keras import optimizers as k_optimizers - from keras.layers import Dense, Activation - from keras.models import Sequential -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow.keras import optimizers as k_optimizers # pylint:disable=import-error - from tensorflow.keras.layers import Dense, Activation # noqa pylint:disable=import-error,no-name-in-module - from tensorflow.keras.models import Sequential # pylint:disable=import-error,no-name-in-module - def get_test_data(): """ Obtain randomized test data for training """ @@ -49,8 +43,7 @@ def _test_optimizer(optimizer, target=0.75): metrics=["accuracy"]) history = model.fit(x_train, y_train, epochs=2, batch_size=16, verbose=0) - accuracy = "acc" if get_backend() == "amd" else "accuracy" - assert history.history[accuracy][-1] >= target + assert history.history["accuracy"][-1] >= target config = k_optimizers.serialize(optimizer) optim = k_optimizers.deserialize(config) new_config = k_optimizers.serialize(optim) @@ -59,9 +52,6 @@ def _test_optimizer(optimizer, target=0.75): assert config == new_config # Test constraints. - if get_backend() == "amd": - # NB: PlaidML does not support constraints, so this test skipped for AMD backends - return model = Sequential() dense = Dense(10, input_shape=(x_train.shape[1],), @@ -80,14 +70,7 @@ def _test_optimizer(optimizer, target=0.75): assert_allclose(bias, 2.) -@pytest.mark.parametrize("dummy", [None], ids=[get_backend().upper()]) -def test_adam(dummy): # pylint:disable=unused-argument - """ Test for custom Adam optimizer """ - _test_optimizer(k_optimizers.Adam(), target=0.45) - _test_optimizer(k_optimizers.Adam(decay=1e-3), target=0.45) - - @pytest.mark.parametrize("dummy", [None], ids=[get_backend().upper()]) def test_adabelief(dummy): # pylint:disable=unused-argument """ Test for custom Adam optimizer """ - _test_optimizer(optimizers.AdaBelief(), target=0.45) + _test_optimizer(optimizers.AdaBelief(), target=0.20) diff --git a/tests/lib/sysinfo_test.py b/tests/lib/sysinfo_test.py new file mode 100644 index 0000000000..215e8c9f46 --- /dev/null +++ b/tests/lib/sysinfo_test.py @@ -0,0 +1,438 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`lib.sysinfo` """ + +import locale +import os +import platform +import sys +import typing as T + +from collections import namedtuple +from io import StringIO +from unittest.mock import MagicMock + +import pytest +import pytest_mock + +from lib.gpu_stats import GPUInfo +from lib.sysinfo import _Configs, _State, _SysInfo, CudaCheck, get_sysinfo + +# pylint:disable=protected-access + + +# _SysInfo +@pytest.fixture(name="sys_info_instance") +def sys_info_fixture() -> _SysInfo: + """ Single :class:~`lib.utils._SysInfo` object for tests + + Returns + ------- + :class:`~lib.utils.sysinfo._SysInfo` + The class instance for testing + """ + return _SysInfo() + + +def test_init(sys_info_instance: _SysInfo) -> None: + """ Test :class:`~lib.utils.sysinfo._SysInfo` __init__ and attributes + + Parameters + ---------- + sys_info_instance: :class:`~lib.utils.sysinfo._SysInfo` + The class instance to test + """ + assert isinstance(sys_info_instance, _SysInfo) + + assert hasattr(sys_info_instance, "_state_file") + assert isinstance(sys_info_instance._state_file, str) + + assert hasattr(sys_info_instance, "_configs") + assert isinstance(sys_info_instance._configs, str) + + assert hasattr(sys_info_instance, "_system") + assert isinstance(sys_info_instance._system, dict) + assert sys_info_instance._system == {"platform": platform.platform(), + "system": platform.system().lower(), + "machine": platform.machine(), + "release": platform.release(), + "processor": platform.processor(), + "cpu_count": os.cpu_count()} + + assert hasattr(sys_info_instance, "_python") + assert isinstance(sys_info_instance._python, dict) + assert sys_info_instance._python == {"implementation": platform.python_implementation(), + "version": platform.python_version()} + + assert hasattr(sys_info_instance, "_gpu") + assert isinstance(sys_info_instance._gpu, GPUInfo) + + assert hasattr(sys_info_instance, "_cuda_check") + assert isinstance(sys_info_instance._cuda_check, CudaCheck) + + +def test_properties(sys_info_instance: _SysInfo) -> None: + """ Test :class:`~lib.utils.sysinfo._SysInfo` properties + + Parameters + ---------- + sys_info_instance: :class:`~lib.utils.sysinfo._SysInfo` + The class instance to test + """ + assert hasattr(sys_info_instance, "_encoding") + assert isinstance(sys_info_instance._encoding, str) + assert sys_info_instance._encoding == locale.getpreferredencoding() + + assert hasattr(sys_info_instance, "_is_conda") + assert isinstance(sys_info_instance._is_conda, bool) + assert sys_info_instance._is_conda == ("conda" in sys.version.lower() or + os.path.exists(os.path.join(sys.prefix, "conda-meta"))) + + assert hasattr(sys_info_instance, "_is_linux") + assert isinstance(sys_info_instance._is_linux, bool) + if platform.system().lower() == "linux": + assert sys_info_instance._is_linux and sys_info_instance._system["system"] == "linux" + assert not sys_info_instance._is_macos + assert not sys_info_instance._is_windows + + assert hasattr(sys_info_instance, "_is_macos") + assert isinstance(sys_info_instance._is_macos, bool) + if platform.system().lower() == "darwin": + assert sys_info_instance._is_macos and sys_info_instance._system["system"] == "darwin" + assert not sys_info_instance._is_linux + assert not sys_info_instance._is_windows + + assert hasattr(sys_info_instance, "_is_windows") + assert isinstance(sys_info_instance._is_windows, bool) + if platform.system().lower() == "windows": + assert sys_info_instance._is_windows and sys_info_instance._system["system"] == "windows" + assert not sys_info_instance._is_linux + assert not sys_info_instance._is_macos + + assert hasattr(sys_info_instance, "_is_virtual_env") + assert isinstance(sys_info_instance._is_virtual_env, bool) + + assert hasattr(sys_info_instance, "_ram_free") + assert isinstance(sys_info_instance._ram_free, int) + + assert hasattr(sys_info_instance, "_ram_total") + assert isinstance(sys_info_instance._ram_total, int) + + assert hasattr(sys_info_instance, "_ram_available") + assert isinstance(sys_info_instance._ram_available, int) + + assert hasattr(sys_info_instance, "_ram_used") + assert isinstance(sys_info_instance._ram_used, int) + + assert hasattr(sys_info_instance, "_fs_command") + assert isinstance(sys_info_instance._fs_command, str) + + assert hasattr(sys_info_instance, "_installed_pip") + assert isinstance(sys_info_instance._installed_pip, str) + + assert hasattr(sys_info_instance, "_installed_conda") + assert isinstance(sys_info_instance._installed_conda, str) + + assert hasattr(sys_info_instance, "_conda_version") + assert isinstance(sys_info_instance._conda_version, str) + + +def test_full_info(sys_info_instance: _SysInfo) -> None: + """ Test the sys_info method of :class:`~lib.utils.sysinfo._SysInfo` returns as expected + + Parameters + ---------- + sys_info_instance: :class:`~lib.utils.sysinfo._SysInfo` + The class instance to test + """ + assert hasattr(sys_info_instance, "full_info") + sys_info = sys_info_instance.full_info() + assert isinstance(sys_info, str) + assert "backend:" in sys_info + assert "os_platform:" in sys_info + assert "os_machine:" in sys_info + assert "os_release:" in sys_info + assert "py_conda_version:" in sys_info + assert "py_implementation:" in sys_info + assert "py_version:" in sys_info + assert "py_command:" in sys_info + assert "py_virtual_env:" in sys_info + assert "sys_cores:" in sys_info + assert "sys_processor:" in sys_info + assert "sys_ram:" in sys_info + assert "encoding:" in sys_info + assert "git_branch:" in sys_info + assert "git_commits:" in sys_info + assert "gpu_cuda:" in sys_info + assert "gpu_cudnn:" in sys_info + assert "gpu_driver:" in sys_info + assert "gpu_devices:" in sys_info + assert "gpu_vram:" in sys_info + assert "gpu_devices_active:" in sys_info + + +def test__format_ram(sys_info_instance: _SysInfo, monkeypatch: pytest.MonkeyPatch) -> None: + """ Test the _format_ram method of :class:`~lib.utils.sysinfo._SysInfo` returns as expected + + Parameters + ---------- + sys_info_instance: :class:`~lib.utils.sysinfo._SysInfo` + The class instance to test + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching psutil.virtual_memory to be consistent + """ + assert hasattr(sys_info_instance, "_format_ram") + svmem = namedtuple("svmem", ["available", "free", "total", "used"]) + data = svmem(12345678, 1234567, 123456789, 123456) + monkeypatch.setattr("psutil.virtual_memory", lambda *args, **kwargs: data) + ram_info = sys_info_instance._format_ram() + + assert isinstance(ram_info, str) + assert ram_info == "Total: 117MB, Available: 11MB, Used: 0MB, Free: 1MB" + + +# get_sys_info +def test_get_sys_info(mocker: pytest_mock.MockerFixture) -> None: + """ Thest that the :func:`~lib.utils.sysinfo.get_sysinfo` function executes correctly + + Parameters + ---------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking full_info called from _SysInfo + """ + sys_info = get_sysinfo() + assert isinstance(sys_info, str) + full_info = mocker.patch("lib.sysinfo._SysInfo.full_info") + get_sysinfo() + assert full_info.called + + +# _Configs +@pytest.fixture(name="configs_instance") +def configs_fixture(): + """ Pytest fixture for :class:`~lib.utils.sysinfo._Configs` + + Returns + ------- + :class:`~lib.utils.sysinfo._Configs` + The class instance for testing + """ + return _Configs() + + +def test__configs__init__(configs_instance: _Configs) -> None: + """ Test __init__ and attributes for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + """ + assert hasattr(configs_instance, "config_dir") + assert isinstance(configs_instance.config_dir, str) + assert hasattr(configs_instance, "configs") + assert isinstance(configs_instance.configs, str) + + +def test__configs__get_configs(configs_instance: _Configs) -> None: + """ Test __init__ and attributes for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + """ + assert hasattr(configs_instance, "_get_configs") + assert isinstance(configs_instance._get_configs(), str) + + +def test__configs__parse_configs(configs_instance: _Configs, + mocker: pytest_mock.MockerFixture) -> None: + """ Test _parse_configs function for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + """ + assert hasattr(configs_instance, "_parse_configs") + assert isinstance(configs_instance._parse_configs([]), str) + configs_instance._parse_ini = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + configs_instance._parse_json = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + configs_instance._parse_configs(config_files=["test.ini", ".faceswap"]) + assert configs_instance._parse_ini.called + assert configs_instance._parse_json.called + + +def test__configs__parse_ini(configs_instance: _Configs, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test _parse_ini function for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`builtins.open` to dummy in ini file + """ + assert hasattr(configs_instance, "_parse_ini") + + file = ("[test.ini_header]\n" + "# Test Header\n\n" + "param = value") + monkeypatch.setattr("builtins.open", lambda *args, **kwargs: StringIO(file)) + + converted = configs_instance._parse_ini("test.ini") + assert isinstance(converted, str) + assert converted == ("\n[test.ini_header]\n" + "param: value\n") + + +def test__configs__parse_json(configs_instance: _Configs, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test _parse_json function for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`builtins.open` to dummy in json file + + """ + assert hasattr(configs_instance, "_parse_json") + file = '{"test": "param"}' + monkeypatch.setattr("builtins.open", lambda *args, **kwargs: StringIO(file)) + + converted = configs_instance._parse_json(".file") + assert isinstance(converted, str) + assert converted == ("test: param\n") + + +def test__configs__format_text(configs_instance: _Configs) -> None: + """ Test _format_text function for :class:`~lib.utils.sysinfo._Configs` + + Parameters + ---------- + configs_instance: :class:`~lib.utils.sysinfo._Configs` + The class instance to test + """ + assert hasattr(configs_instance, "_format_text") + key, val = " test_key ", "test_val " + formatted = configs_instance._format_text(key, val) + assert isinstance(formatted, str) + assert formatted == "test_key: test_val\n" + + +# _State +@pytest.fixture(name="state_instance") +def state_fixture(): + """ Pytest fixture for :class:`~lib.utils.sysinfo._State` + + Returns + ------- + :class:`~lib.utils.sysinfo._State` + The class instance for testing + """ + return _State() + + +def test__state__init__(state_instance: _State) -> None: + """ Test __init__ and attributes for :class:`~lib.utils.sysinfo._State` + + Parameters + ---------- + state_instance: :class:`~lib.utils.sysinfo._State` + The class instance to test + """ + assert hasattr(state_instance, '_model_dir') + assert state_instance._model_dir is None + assert hasattr(state_instance, '_trainer') + assert state_instance._trainer is None + assert hasattr(state_instance, 'state_file') + assert isinstance(state_instance.state_file, str) + + +def test__state__is_training(state_instance: _State, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test _is_training function for :class:`~lib.utils.sysinfo._State` + + Parameters + ---------- + state_instance: :class:`~lib.utils.sysinfo._State` + The class instance to test + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`sys.argv` to dummy in commandline args + + """ + assert hasattr(state_instance, '_is_training') + assert isinstance(state_instance._is_training, bool) + assert not state_instance._is_training + monkeypatch.setattr("sys.argv", ["faceswap.py", "train"]) + assert state_instance._is_training + monkeypatch.setattr("sys.argv", ["faceswap.py", "extract"]) + assert not state_instance._is_training + + +def test__state__get_arg(state_instance: _State, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test _get_arg function for :class:`~lib.utils.sysinfo._State` + + Parameters + ---------- + state_instance: :class:`~lib.utils.sysinfo._State` + The class instance to test + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`sys.argv` to dummy in commandline args + :func:`builtins.input` + """ + assert hasattr(state_instance, '_get_arg') + assert state_instance._get_arg("-t", "--test_arg") is None + monkeypatch.setattr("sys.argv", ["test", "command", "-t", "test_option"]) + assert state_instance._get_arg("-t", "--test_arg") == "test_option" + + +def test__state__get_state_file(state_instance: _State, + mocker: pytest_mock.MockerFixture, + monkeypatch: pytest.MonkeyPatch) -> None: + """ Test _get_state_file function for :class:`~lib.utils.sysinfo._State` + + Parameters + ---------- + state_instance: :class:`~lib.utils.sysinfo._State` + The class instance to test + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`sys.argv` to dummy in commandline args + :func:`builtins.input` +` """ + assert hasattr(state_instance, '_get_state_file') + assert isinstance(state_instance._get_state_file(), str) + + mock_is_training = mocker.patch("lib.sysinfo._State._is_training") + + # Not training or missing training arguments + mock_is_training.return_value = False + assert state_instance._get_state_file() == "" + mock_is_training.return_value = False + + monkeypatch.setattr(state_instance, "_model_dir", None) + assert state_instance._get_state_file() == "" + monkeypatch.setattr(state_instance, "_model_dir", "test_dir") + + monkeypatch.setattr(state_instance, "_trainer", None) + assert state_instance._get_state_file() == "" + monkeypatch.setattr(state_instance, "_trainer", "test_trainer") + + # Training but file not found + assert state_instance._get_state_file() == "" + + # State file is just a json dump + file = ('{\n' + ' "test": "json",\n' + '}') + monkeypatch.setattr("os.path.isfile", lambda *args, **kwargs: True) + monkeypatch.setattr("builtins.open", lambda *args, **kwargs: StringIO(file)) + assert state_instance._get_state_file().endswith(file) diff --git a/tests/lib/utils_test.py b/tests/lib/utils_test.py new file mode 100644 index 0000000000..34f9be5f4f --- /dev/null +++ b/tests/lib/utils_test.py @@ -0,0 +1,622 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`lib.utils` """ +import os +import platform +import time +import typing as T +import warnings +import zipfile + +from io import StringIO +from socket import timeout as socket_timeout, error as socket_error +from shutil import rmtree +from unittest.mock import MagicMock +from urllib import error as urlliberror + +import pytest +import pytest_mock + +from lib import utils +from lib.utils import ( + _Backend, camel_case_split, convert_to_secs, DebugTimes, deprecation_warning, FaceswapError, + full_path_split, get_backend, get_dpi, get_folder, get_image_paths, get_tf_version, GetModel, + safe_shutdown, set_backend, set_system_verbosity) + +from lib.logger import log_setup +# Need to setup logging to avoid trace/verbose errors +log_setup("DEBUG", "pytest_utils.log", "PyTest, False") + + +# pylint:disable=protected-access + + +# Backend tests +def test_set_backend(monkeypatch: pytest.MonkeyPatch) -> None: + """ Test the :func:`~lib.utils.set_backend` function + + Parameters + ---------- + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching _FS_BACKEND + """ + monkeypatch.setattr(utils, "_FS_BACKEND", "cpu") # _FS_BACKEND already defined + set_backend("directml") + assert utils._FS_BACKEND == "directml" + monkeypatch.delattr(utils, "_FS_BACKEND") # _FS_BACKEND is not already defined + set_backend("rocm") + assert utils._FS_BACKEND == "rocm" + + +def test_get_backend(monkeypatch: pytest.MonkeyPatch) -> None: + """ Test the :func:`~lib.utils.get_backend` function + + Parameters + ---------- + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching _FS_BACKEND + """ + monkeypatch.setattr(utils, "_FS_BACKEND", "apple-silicon") + assert get_backend() == "apple-silicon" + + +def test__backend(monkeypatch: pytest.MonkeyPatch) -> None: + """ Test the :class:`~lib.utils._Backend` class + + Parameters + ---------- + monkeypatch: :class:`pytest.MonkeyPatch` + Monkey patching :func:`os.environ`, :func:`os.path.isfile`, :func:`builtins.open` and + :func:`builtins.input` + """ + monkeypatch.setattr("os.environ", {"FACESWAP_BACKEND": "nvidia"}) # Environment variable set + backend = _Backend() + assert backend.backend == "nvidia" + + monkeypatch.setattr("os.environ", {}) # Environment variable not set, dummy in config file + monkeypatch.setattr("os.path.isfile", lambda x: True) + monkeypatch.setattr("builtins.open", lambda *args, **kwargs: StringIO('{"backend": "cpu"}')) + backend = _Backend() + assert backend.backend == "cpu" + + monkeypatch.setattr("os.path.isfile", lambda x: False) # no config file, dummy in user input + monkeypatch.setattr("builtins.input", lambda x: "3") + backend = _Backend() + assert backend._configure_backend() == "nvidia" + + +# Folder and path utils +def test_get_folder(tmp_path: str) -> None: + """ Unit test for :func:`~lib.utils.get_folder` + + Parameters + ---------- + tmp_path: str + pytest temporary path to generate folders + """ + # New folder + path = os.path.join(tmp_path, "test_new_folder") + expected_output = path + assert not os.path.isdir(path) + assert get_folder(path) == expected_output + assert os.path.isdir(path) + + # Test not creating a new folder when it already exists + path = os.path.join(tmp_path, "test_new_folder") + expected_output = path + assert os.path.isdir(path) + stats = os.stat(path) + assert get_folder(path) == expected_output + assert os.path.isdir(path) + assert stats == os.stat(path) + + # Test not creating a new folder when make_folder is False + path = os.path.join(tmp_path, "test_no_folder") + expected_output = "" + assert get_folder(path, make_folder=False) == expected_output + assert not os.path.isdir(path) + + +def test_get_image_paths(tmp_path: str) -> None: + """ Unit test for :func:`~lib.utils.test_get_image_paths` + + Parameters + ---------- + tmp_path: str + pytest temporary path to generate folders + """ + # Test getting image paths from a folder with no images + test_folder = os.path.join(tmp_path, "test_image_folder") + os.makedirs(test_folder) + assert not get_image_paths(test_folder) + + # Populate 2 different image files and 1 text file + test_jpg_path = os.path.join(test_folder, "test_image.jpg") + test_png_path = os.path.join(test_folder, "test_image.png") + test_txt_path = os.path.join(test_folder, "test_file.txt") + for fname in (test_jpg_path, test_png_path, test_txt_path): + with open(fname, "a", encoding="utf-8"): + pass + + # Test getting any image paths from a folder with images and random files + exists = [os.path.join(test_folder, img) + for img in os.listdir(test_folder) if os.path.splitext(img)[-1] != ".txt"] + assert sorted(get_image_paths(test_folder)) == sorted(exists) + + # Test getting image paths from a folder with images with a specific extension + exists = [os.path.join(test_folder, img) + for img in os.listdir(test_folder) if os.path.splitext(img)[-1] == ".png"] + assert sorted(get_image_paths(test_folder, extension=".png")) == sorted(exists) + + +_PARAMS = [("/path/to/file.txt", ["/", "path", "to", "file.txt"]), # Absolute + ("/path/to/directory/", ["/", "path", "to", "directory"]), + ("/path/to/directory", ["/", "path", "to", "directory"]), + ("path/to/file.txt", ["path", "to", "file.txt"]), # Relative + ("path/to/directory/", ["path", "to", "directory"]), + ("path/to/directory", ["path", "to", "directory"]), + ("", []), # Edge cases + ("/", ["/"]), + (".", ["."]), + ("..", [".."])] + + +@pytest.mark.parametrize("path,result", _PARAMS, ids=[f'"{p[0]}"' for p in _PARAMS]) +def test_full_path_split(path: str, result: list[str]) -> None: + """ Test the :func:`~lib.utils.full_path_split` function works correctly + + Parameters + ---------- + path: str + The path to test + result: list + The expected result from the path + """ + split = full_path_split(path) + assert isinstance(split, list) + assert split == result + + +_PARAMS = [("camelCase", ["camel", "Case"]), + ("camelCaseTest", ["camel", "Case", "Test"]), + ("camelCaseTestCase", ["camel", "Case", "Test", "Case"]), + ("CamelCase", ["Camel", "Case"]), + ("CamelCaseTest", ["Camel", "Case", "Test"]), + ("CamelCaseTestCase", ["Camel", "Case", "Test", "Case"]), + ("CAmelCASETestCase", ["C", "Amel", "CASE", "Test", "Case"]), + ("camelcasetestcase", ["camelcasetestcase"]), + ("CAMELCASETESTCASE", ["CAMELCASETESTCASE"]), + ("", [])] + + +@pytest.mark.parametrize("text, result", _PARAMS, ids=[f'"{p[0]}"' for p in _PARAMS]) +def test_camel_case_split(text: str, result: list[str]) -> None: + """ Test the :func:`~lib.utils.camel_case_spli` function works correctly + + Parameters + ---------- + text: str + The camel case text to test + result: list + The expected result from the path + """ + split = camel_case_split(text) + assert isinstance(split, list) + assert split == result + + +# General utils +def test_get_tf_version() -> None: + """ Test the :func:`~lib.utils.get_tf_version` function version returns correctly in range """ + tf_version = get_tf_version() + assert (2, 10) <= tf_version < (2, 11) + + +def test_get_dpi() -> None: + """ Test the :func:`~lib.utils.get_dpi` function version returns correctly in a sane + range """ + dpi = get_dpi() + assert isinstance(dpi, float) or dpi is None + if dpi is None: # No display detected + return + assert dpi > 0 + assert dpi < 600.0 + + +_SECPARAMS = [((1, ), 1), # 1 argument + ((10, ), 10), + ((0, 1), 1), + ((0, 60), 60), # 2 arguments + ((1, 0), 60), + ((1, 1), 61), + ((0, 0, 1), 1), + ((0, 0, 60), 60), # 3 arguments + ((0, 1, 0), 60), + ((1, 0, 0), 3600), + ((1, 1, 1), 3661)] + + +@pytest.mark.parametrize("args,result", _SECPARAMS, ids=[str(p[0]) for p in _SECPARAMS]) +def test_convert_to_secs(args: tuple[int, ...], result: int) -> None: + """ Test the :func:`~lib.utils.convert_to_secs` function works correctly + + Parameters + ---------- + args: tuple + Tuple of 1, 2 or 3 integers to pass to the function + result: int + The expected results for the args tuple + """ + secs = convert_to_secs(*args) + assert isinstance(secs, int) + assert secs == result + + +@pytest.mark.parametrize("log_level", ["DEBUG", "INFO", "WARNING", "ERROR"]) +def test_set_system_verbosity(log_level: str) -> None: + """ Test the :func:`~lib.utils.set_system_verbosity` function works correctly + + Parameters + ---------- + log_level: str + The logging loglevel in upper text format + """ + # Set TF Env Variable + tf_set_level = "0" if log_level == "DEBUG" else "3" + set_system_verbosity(log_level) + tf_get_level = os.environ["TF_CPP_MIN_LOG_LEVEL"] + assert tf_get_level == tf_set_level + warn_filters = [filt for filt in warnings.filters + if filt[0] == "ignore" + and filt[2] in (FutureWarning, DeprecationWarning, UserWarning)] + # Python Warnings + # DeprecationWarning is already ignored by default, so there should be 1 warning for debug + # warning. 3 for the rest + num_warnings = 1 if log_level == "DEBUG" else 3 + warn_count = len(warn_filters) + assert warn_count == num_warnings + + +@pytest.mark.parametrize("additional_info", [None, "additional information"]) +def test_deprecation_warning(caplog: pytest.LogCaptureFixture, additional_info: str) -> None: + """ Test the :func:`~lib.utils.deprecation_warning` function works correctly + + Parameters + ---------- + caplog: :class:`pytest.LogCaptureFixture` + Pytest's log capturing fixture + additional_info: str + Additional information to pass to the warning function + """ + func_name = "function_name" + test = f"{func_name} has been deprecated and will be removed from a future update." + if additional_info: + test = f"{test} {additional_info}" + deprecation_warning(func_name, additional_info=additional_info) + assert test in caplog.text + + +@pytest.mark.parametrize("got_error", [True, False]) +def test_safe_shutdown(caplog: pytest.LogCaptureFixture, got_error: bool) -> None: + """ Test the :func:`~lib.utils.safe_shutdown` function works correctly + + Parameters + ---------- + caplog: :class:`pytest.LogCaptureFixture` + Pytest's log capturing fixture + got_error: bool + The got_error parameter to pass to safe_shutdown + """ + caplog.set_level("DEBUG") + with pytest.raises(SystemExit) as wrapped_exit: + safe_shutdown(got_error=got_error) + + exit_value = 1 if got_error else 0 + assert wrapped_exit.typename == "SystemExit" + assert wrapped_exit.value.code == exit_value + assert "Safely shutting down" in caplog.messages + assert "Cleanup complete. Shutting down queue manager and exiting" in caplog.messages + + +def test_faceswap_error(): + """ Test the :class:`~lib.utils.FaceswapError` raises correctly """ + with pytest.raises(Exception): + raise FaceswapError + + +# GetModel class +@pytest.fixture(name="get_model_instance") +def fixture_get_model_instance(monkeypatch: pytest.MonkeyPatch, + tmp_path: pytest.TempdirFactory, + request: pytest.FixtureRequest) -> GetModel: + """ Create a fixture of the :class:`~lib.utils.GetModel` object, prevent _get() from running at + __init__ and point the cache_dir at our local test folder """ + cache_dir = os.path.join(str(tmp_path), "get_model") + os.mkdir(cache_dir) + + model_filename = "test_model_file_v1.h5" + git_model_id = 123 + + original_get = GetModel._get + # Patch out _get() so it is not called from __init__() + monkeypatch.setattr(utils.GetModel, "_get", lambda x: None) + model_instance = GetModel(model_filename, git_model_id) + # Reinsert _get() so we can test it + monkeypatch.setattr(model_instance, "_get", original_get) + model_instance._cache_dir = cache_dir + + def teardown(): + rmtree(cache_dir) + + request.addfinalizer(teardown) + return model_instance + + +_INPUT = ("test_model_file_v3.h5", + ["test_multi_model_file_v1.1.npy", "test_multi_model_file_v1.2.npy"]) +_EXPECTED = ((["test_model_file_v3.h5"], "test_model_file_v3", "test_model_file", 3), + (["test_multi_model_file_v1.1.npy", "test_multi_model_file_v1.2.npy"], + "test_multi_model_file_v1", "test_multi_model_file", 1)) + + +@pytest.mark.parametrize("filename,results", zip(_INPUT, _EXPECTED), ids=[str(i) for i in _INPUT]) +def test_get_model_model_filename_input( + get_model_instance: GetModel, # pylint:disable=unused-argument + filename: str | list[str], + results: str | list[str]) -> None: + """ Test :class:`~lib.utils.GetModel` filename parsing works + + Parameters + --------- + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + filename: list or str + The test filenames + results: tuple + The expected results for :attr:`_model_filename`, :attr:`_model_full_name`, + :attr:`_model_name`, :attr:`_model_version` respectively + """ + model = GetModel(filename, 123) + assert model._model_filename == results[0] + assert model._model_full_name == results[1] + assert model._model_name == results[2] + assert model._model_version == results[3] + + +def test_get_model_attributes(get_model_instance: GetModel) -> None: + """ Test :class:`~lib.utils.GetModel` private attributes set correctly + + Parameters + --------- + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + """ + model = get_model_instance + assert model._git_model_id == 123 + assert model._url_base == ("https://github.com/deepfakes-models/faceswap-models" + "/releases/download") + assert model._chunk_size == 1024 + assert model._retries == 6 + + +def test_get_model_properties(get_model_instance: GetModel) -> None: + """ Test :class:`~lib.utils.GetModel` calculated attributes return correctly + + Parameters + --------- + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + """ + model = get_model_instance + assert model.model_path == os.path.join(model._cache_dir, "test_model_file_v1.h5") + assert model._model_zip_path == os.path.join(model._cache_dir, "test_model_file_v1.zip") + assert not model._model_exists + assert model._url_download == ("https://github.com/deepfakes-models/faceswap-models/releases/" + "download/v123.1/test_model_file_v1.zip") + assert model._url_partial_size == 0 + + +@pytest.mark.parametrize("model_exists", (True, False)) +def test_get_model__get(mocker: pytest_mock.MockerFixture, + get_model_instance: GetModel, + model_exists: bool) -> None: + """ Test :func:`~lib.utils.GetModel._get` executes logic correctly + + Parameters + --------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + model_exists: bool + For testing the function when a model exists and when it does not + """ + model = get_model_instance + model._download_model = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + model._unzip_model = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + os_remove = mocker.patch("os.remove") + + if model_exists: # Dummy in a model file + assert isinstance(model.model_path, str) + with open(model.model_path, "a", encoding="utf-8"): + pass + + model._get(model) # type:ignore + + assert (model_exists and not model._download_model.called) or ( + not model_exists and model._download_model.called) + assert (model_exists and not model._unzip_model.called) or ( + not model_exists and model._unzip_model.called) + assert model_exists or not (model_exists and os_remove.called) + os_remove.reset_mock() + + +_DLPARAMS = [(None, None), + (socket_error, ()), + (socket_timeout, ()), + (urlliberror.URLError, ("test_reason", )), + (urlliberror.HTTPError, ("test_uri", 400, "", "", 0))] + + +@pytest.mark.parametrize("error_type,error_args", _DLPARAMS, ids=[str(p[0]) for p in _DLPARAMS]) +def test_get_model__download_model(mocker: pytest_mock.MockerFixture, + get_model_instance: GetModel, + error_type: T.Any, + error_args: tuple[str | int, ...]) -> None: + """ Test :func:`~lib.utils.GetModel._download_model` executes its logic correctly + + Parameters + --------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + error_type: connection error type or ``None`` + Connection error type to mock, or ``None`` for succesful download + error_args: tuple + The arguments to be passed to the exception to be raised + """ + mock_urlopen = mocker.patch("urllib.request.urlopen") + if not error_type: # Model download is successful + get_model_instance._write_zipfile = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + get_model_instance._download_model() + assert mock_urlopen.called + assert get_model_instance._write_zipfile.called + else: # Test that the process exits on download errors + mock_urlopen.side_effect = error_type(*error_args) + with pytest.raises(SystemExit): + get_model_instance._download_model() + mock_urlopen.reset_mock() + + +@pytest.mark.parametrize("dl_type", ["complete", "new", "continue"]) +def test_get_model__write_zipfile(mocker: pytest_mock.MockerFixture, + get_model_instance: GetModel, + dl_type: str) -> None: + """ Test :func:`~lib.utils.GetModel._write_zipfile` executes its logic correctly + + Parameters + --------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + dl_type: str + The type of read to attemp + """ + response = mocker.MagicMock() + assert not os.path.isfile(get_model_instance._model_zip_path) + + downloaded = 10 if dl_type == "complete" else 0 + response.getheader.return_value = 0 + + if dl_type in ("new", "continue"): + chunks = [32, 64, 128, 256, 512, 1024] + data = [b"\x00" * size for size in chunks] + [b""] + response.getheader.return_value = sum(chunks) + response.read.side_effect = data + + if dl_type == "continue": # Write a partial download of the correct size + with open(get_model_instance._model_zip_path, "wb") as partial: + partial.write(b"\x00" * sum(chunks)) + downloaded = os.path.getsize(get_model_instance._model_zip_path) + + get_model_instance._write_zipfile(response, downloaded) + + if dl_type == "complete": # Already downloaded. No more tests + assert not response.read.called + return + + assert response.read.call_count == len(data) # all data read + assert os.path.isfile(get_model_instance._model_zip_path) + downloaded_size = os.path.getsize(get_model_instance._model_zip_path) + downloaded_size = downloaded_size if dl_type == "new" else downloaded_size // 2 + assert downloaded_size == sum(chunks) + + +def test_get_model__unzip_model(mocker: pytest_mock.MockerFixture, + get_model_instance: GetModel) -> None: + """ Test :func:`~lib.utils.GetModel._unzip_model` executes its logic correctly + + Parameters + --------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + """ + mock_zipfile = mocker.patch("zipfile.ZipFile") + # Successful + get_model_instance._unzip_model() + assert mock_zipfile.called + mock_zipfile.reset_mock() + # Error + mock_zipfile.side_effect = zipfile.BadZipFile() + with pytest.raises(SystemExit): + get_model_instance._unzip_model() + mock_zipfile.reset_mock() + + +def test_get_model__write_model(mocker: pytest_mock.MockerFixture, + get_model_instance: GetModel) -> None: + """ Test :func:`~lib.utils.GetModel._write_model` executes its logic correctly + + Parameters + --------- + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in function calls + get_model_instance: `~lib.utils.GetModel` + The patched instance of the class + """ + out_file = os.path.join(get_model_instance._cache_dir, get_model_instance._model_filename[0]) + chunks = [8, 16, 32, 64, 128, 256, 512, 1024] + data = [b"\x00" * size for size in chunks] + [b""] + assert not os.path.isfile(out_file) + mock_zipfile = mocker.patch("zipfile.ZipFile") + mock_zipfile.namelist.return_value = get_model_instance._model_filename + mock_zipfile.open.return_value = mock_zipfile + mock_zipfile.read.side_effect = data + get_model_instance._write_model(mock_zipfile) + assert mock_zipfile.read.call_count == len(data) + assert os.path.isfile(out_file) + assert os.path.getsize(out_file) == sum(chunks) + + +# DebugTimes class +def test_debug_times(): + """ Test :class:`~lib.utils.DebugTimes` executes its logic correctly """ + debug_times = DebugTimes() + + debug_times.step_start("Test1") + time.sleep(0.1) + debug_times.step_end("Test1") + + debug_times.step_start("Test2") + time.sleep(0.2) + debug_times.step_end("Test2") + + debug_times.step_start("Test1") + time.sleep(0.1) + debug_times.step_end("Test1") + + debug_times.summary() + + # Ensure that the summary method prints the min, mean, and max times for each step + assert debug_times._display["min"] is True + assert debug_times._display["mean"] is True + assert debug_times._display["max"] is True + + # Ensure that the summary method includes the correct number of items for each step + assert len(debug_times._times["Test1"]) == 2 + assert len(debug_times._times["Test2"]) == 1 + + # Ensure that the summary method includes the correct min, mean, and max times for each step + # Github workflow for macos-latest can swing out a fair way + threshold = 2e-1 if platform.system() == "Darwin" else 1e-1 + assert min(debug_times._times["Test1"]) == pytest.approx(0.1, abs=threshold) + assert min(debug_times._times["Test2"]) == pytest.approx(0.2, abs=threshold) + assert max(debug_times._times["Test1"]) == pytest.approx(0.1, abs=threshold) + assert max(debug_times._times["Test2"]) == pytest.approx(0.2, abs=threshold) + assert (sum(debug_times._times["Test1"]) / + len(debug_times._times["Test1"])) == pytest.approx(0.1, abs=threshold) + assert (sum(debug_times._times["Test2"]) / + len(debug_times._times["Test2"]) == pytest.approx(0.2, abs=threshold)) diff --git a/_travis/simple_tests.py b/tests/simple_tests.py similarity index 58% rename from _travis/simple_tests.py rename to tests/simple_tests.py index f908675e03..0e6ea127d5 100644 --- a/_travis/simple_tests.py +++ b/tests/simple_tests.py @@ -31,9 +31,8 @@ def print_colored(text, color="OK", bold=False): although travis runs windows stuff in git bash, so it might ? """ color = _COLORS.get(color, color) - print("%s%s%s%s" % ( - color, "" if not bold else _COLORS["BOLD"], text, _COLORS["ENDC"] - )) + fmt = '' if not bold else _COLORS['BOLD'] + print(f"{color}{fmt}{text}{_COLORS['ENDC']}") def print_ok(text): @@ -54,15 +53,15 @@ def print_status(text): def run_test(name, cmd): """ run a test """ global FAIL_COUNT, TEST_COUNT # pylint:disable=global-statement - print_status("[?] running %s" % name) - print("Cmd: %s" % " ".join(cmd)) + print_status(f"[?] running {name}") + print(f"Cmd: {' '.join(cmd)}") TEST_COUNT += 1 try: check_call(cmd) print_ok("[+] Test success") return True except CalledProcessError as err: - print_fail("[-] Test failed with %s" % err) + print_fail(f"[-] Test failed with {err}") FAIL_COUNT += 1 return False @@ -70,57 +69,77 @@ def run_test(name, cmd): def download_file(url, filename): # TODO: retry """ Download a file from given url """ if os.path.isfile(filename): - print_status("[?] '%s' already cached as '%s'" % (url, filename)) + print_status(f"[?] '{url}' already cached as '{filename}'") return filename try: - print_status("[?] Downloading '%s' to '%s'" % (url, filename)) + print_status(f"[?] Downloading '{url}' to '{filename}'") video, _ = urlretrieve(url, filename) return video except urllib.error.URLError as err: - print_fail("[-] Failed downloading: %s" % err) + print_fail(f"[-] Failed downloading: {err}") return None def extract_args(detector, aligner, in_path, out_path, args=None): """ Extraction command """ py_exe = sys.executable - _extract_args = "%s faceswap.py extract -i %s -o %s -D %s -A %s" % ( - py_exe, in_path, out_path, detector, aligner - ) + _extract_args = (f"{py_exe} faceswap.py extract -i {in_path} -o {out_path} -D {detector} " + f"-A {aligner}") if args: - _extract_args += " %s" % args + _extract_args += f" {args}" return _extract_args.split() -def train_args(model, model_path, faces, alignments, iterations=5, batchsize=8, extra_args=""): +def train_args(model, model_path, faces, iterations=1, batchsize=2, extra_args=""): """ Train command """ py_exe = sys.executable - args = "%s faceswap.py train -A %s -B %s -m %s -t %s -bs %i -it %s %s" % ( - py_exe, faces, faces, model_path, model, batchsize, iterations, extra_args - ) + args = (f"{py_exe} faceswap.py train -A {faces} -B {faces} -m {model_path} -t {model} " + f"-b {batchsize} -i {iterations} {extra_args}") return args.split() def convert_args(in_path, out_path, model_path, writer, args=None): """ Convert command """ py_exe = sys.executable - conv_args = "%s faceswap.py convert -i %s -o %s -m %s -w %s" % ( - py_exe, in_path, out_path, model_path, writer - ) + conv_args = (f"{py_exe} faceswap.py convert -i {in_path} -o {out_path} -m {model_path} " + f"-w {writer}") if args: - conv_args += " %s" % args + conv_args += f" {args}" return conv_args.split() # Don't use pathes with spaces ;) -def sort_args(in_path, out_path, sortby="face", groupby="hist", method="rename"): +def sort_args(in_path, out_path, sortby="face", groupby="hist"): """ Sort command """ py_exe = sys.executable - _sort_args = "%s tools.py sort -i %s -o %s -s %s -fp %s -g %s -k" % ( - py_exe, in_path, out_path, sortby, method, groupby - ) + _sort_args = (f"{py_exe} tools.py sort -i {in_path} -o {out_path} -s {sortby} -g {groupby} -k") return _sort_args.split() +def set_train_config(value): + """ Update the mixed_precision and autoclip values to given value + + Parameters + ---------- + value: bool + The value to set the config parameters to. + """ + old_val, new_val = ("False", "True") if value else ("True", "False") + base_path = os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + train_ini = os.path.join(base_path, "config", "train.ini") + try: + cmd = ["sed", "-i", f"s/autoclip = {old_val}/autoclip = {new_val}/", train_ini] + check_call(cmd) + cmd = ["sed", + "-i", + f"s/mixed_precision = {old_val}/mixed_precision = {new_val}/", + train_ini] + check_call(cmd) + print_ok(f"Set autoclip and mixed_precision to `{new_val}`") + except CalledProcessError as err: + print_fail(f"[-] Test failed with {err}") + return False + + def main(): """ Main testing script """ vid_src = "https://faceswap.dev/data/test.mp4" @@ -153,11 +172,17 @@ def main(): ) if vid_extract: + run_test( + "Generate configs and test help output", + ( + py_exe, "faceswap.py", "-h" + ) + ) run_test( "Sort faces.", sort_args( pathjoin(vid_base, "faces"), pathjoin(vid_base, "faces_sorted"), - sortby="face", method="rename" + sortby="face" ) ) @@ -166,26 +191,27 @@ def main(): ( py_exe, "tools.py", "alignments", "-j", "rename", "-a", pathjoin(vid_base, "test_alignments.fsa"), - "-fc", pathjoin(vid_base, "faces_sorted"), + "-c", pathjoin(vid_base, "faces_sorted"), ) ) - + set_train_config(True) run_test( - "Train lightweight model for 1 iteration with WTL.", - train_args( - "lightweight", pathjoin(vid_base, "model"), - pathjoin(vid_base, "faces"), pathjoin(vid_base, "test_alignments.fsa"), - iterations=1, extra_args="-wl" - ) - ) - + "Train lightweight model for 1 iteration with WTL, AutoClip, MixedPrecion", + train_args("lightweight", + pathjoin(vid_base, "model"), + pathjoin(vid_base, "faces"), + iterations=1, + batchsize=1, + extra_args="-M")) + + set_train_config(False) was_trained = run_test( - "Train lightweight model for 5 iterations WITHOUT WTL.", - train_args( - "lightweight", pathjoin(vid_base, "model"), - pathjoin(vid_base, "faces"), pathjoin(vid_base, "test_alignments.fsa") - ) - ) + "Train lightweight model for 1 iterations WITHOUT WTL, AutoClip, MixedPrecion", + train_args("lightweight", + pathjoin(vid_base, "model"), + pathjoin(vid_base, "faces"), + iterations=1, + batchsize=1)) if was_trained: run_test( @@ -205,10 +231,10 @@ def main(): ) if FAIL_COUNT == 0: - print_ok("[+] Failed %i/%i tests." % (FAIL_COUNT, TEST_COUNT)) + print_ok(f"[+] Failed {FAIL_COUNT}/{TEST_COUNT} tests.") sys.exit(0) else: - print_fail("[-] Failed %i/%i tests." % (FAIL_COUNT, TEST_COUNT)) + print_fail(f"[-] Failed {FAIL_COUNT}/{TEST_COUNT} tests.") sys.exit(1) diff --git a/tests/startup_test.py b/tests/startup_test.py index d126358fae..704a96edcf 100644 --- a/tests/startup_test.py +++ b/tests/startup_test.py @@ -2,19 +2,13 @@ """ Sanity checks for Faceswap. """ import inspect - import pytest -from lib.utils import get_backend - -if get_backend() == "amd": - import keras - from keras import backend as K -else: - # Ignore linting errors from Tensorflow's thoroughly broken import system - from tensorflow import keras - from tensorflow.keras import backend as K # pylint:disable=import-error +# Ignore linting errors from Tensorflow's thoroughly broken import system +from tensorflow import keras +from tensorflow.keras import backend as K # pylint:disable=import-error +from lib.utils import get_backend _BACKEND = get_backend() @@ -24,12 +18,11 @@ def test_backend(dummy): # pylint:disable=unused-argument """ Sanity check to ensure that Keras backend is returning the correct object type. """ test_var = K.variable((1, 1, 4, 4)) lib = inspect.getmodule(test_var).__name__.split(".")[0] - assert (_BACKEND == "cpu" and lib == "tensorflow") or (_BACKEND == "amd" and lib == "plaidml") + assert _BACKEND in ("cpu", "directml") and lib == "tensorflow" @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()]) def test_keras(dummy): # pylint:disable=unused-argument - """ Sanity check to ensure that tensorflow keras is being used for CPU and standard - keras for AMD. """ - assert ((_BACKEND == "cpu" and keras.__version__ in ("2.4.0", "2.6.0", "2.7.0", "2.8.0")) or - (_BACKEND == "amd" and keras.__version__ == "2.2.4")) + """ Sanity check to ensure that tensorflow keras is being used for CPU """ + assert (_BACKEND in ("cpu", "directml") + and keras.__version__ in ("2.7.0", "2.8.0", "2.9.0", "2.10.0")) diff --git a/plugins/extract/recognition/.cache/.keep b/tests/tools/__init__.py similarity index 100% rename from plugins/extract/recognition/.cache/.keep rename to tests/tools/__init__.py diff --git a/tests/tools/alignments/media_test.py b/tests/tools/alignments/media_test.py new file mode 100644 index 0000000000..17e45a517e --- /dev/null +++ b/tests/tools/alignments/media_test.py @@ -0,0 +1,863 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`tools.alignments.media` """ +from __future__ import annotations +import os +import typing as T + +from operator import itemgetter +from unittest.mock import MagicMock + +import cv2 +import numpy as np +import pytest +import pytest_mock + +from lib.logger import log_setup +# Need to setup logging to avoid trace/verbose errors +log_setup("DEBUG", f"{__name__}.log", "PyTest, False") + +# pylint:disable=wrong-import-position,protected-access +from lib.utils import FaceswapError # noqa:E402 +from tools.alignments.media import (AlignmentData, Faces, ExtractedFaces, # noqa:E402 + Frames, MediaLoader) + +if T.TYPE_CHECKING: + from collections.abc import Generator + + +class TestAlignmentData: + """ Test for :class:`~tools.alignments.media.AlignmentData` """ + + @pytest.fixture + def alignments_file(self, tmp_path: str) -> Generator[str, None, None]: + """ Fixture for creating dummy alignments files + + Parameters + ---------- + tmp_path: str + pytest temporary path to generate folders + + Yields + ------ + str + Path to a dummy alignments file + """ + alignments_file = os.path.join(tmp_path, "alignments.fsa") + with open(alignments_file, "w", encoding="utf8") as afile: + afile.write("test") + yield alignments_file + os.remove(alignments_file) + + def test_init(self, + alignments_file: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.AlignmentData` __init__ method + + Parameters + ---------- + alignments_file: str + The temporarily generated alignments file + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking the superclass __init__ + """ + alignments_parent_init = mocker.patch("tools.alignments.media.Alignments.__init__") + mocker.patch("tools.alignments.media.Alignments.frames_count", + new_callable=mocker.PropertyMock(return_value=20)) + + AlignmentData(alignments_file) + folder, filename = os.path.split(alignments_file) + alignments_parent_init.assert_called_once_with(folder, filename=filename) + + def test_check_file_exists(self, alignments_file: str) -> None: + """ Test for :class:`~tools.alignments.media.AlignmentData` _check_file_exists method + + Parameters + ---------- + alignments_file: str + The temporarily generated alignments file + """ + assert AlignmentData.check_file_exists(alignments_file) == os.path.split(alignments_file) + fake_file = "/not/possibly/a/real/path/alignments.fsa" + with pytest.raises(SystemExit): + AlignmentData.check_file_exists(fake_file) + + def test_save(self, + alignments_file: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.AlignmentData`save method + + Parameters + ---------- + alignments_file: str + The temporarily generated alignments file + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking the superclass calls + """ + mocker.patch("tools.alignments.media.Alignments.__init__") + mocker.patch("tools.alignments.media.Alignments.frames_count", + new_callable=mocker.PropertyMock(return_value=20)) + alignments_parent_backup = mocker.patch("tools.alignments.media.Alignments.backup") + alignments_parent_save = mocker.patch("tools.alignments.media.Alignments.save") + align_data = AlignmentData(alignments_file) + align_data.save() + alignments_parent_backup.assert_called_once() + alignments_parent_save.assert_called_once() + + +@pytest.fixture(name="folder") +def folder_fixture(tmp_path: str) -> Generator[str, None, None]: + """ Fixture for creating dummy folders + + Parameters + ---------- + tmp_path: str + pytest temporary path to generate folders + + Yields + ------ + str + Path to a dummy folder + """ + folder = os.path.join(tmp_path, "images") + os.mkdir(folder) + for fname in (["a.png", "b.png"]): + with open(os.path.join(folder, fname), "wb"): + pass + yield folder + for fname in (["a.png", "b.png"]): + os.remove(os.path.join(folder, fname)) + os.rmdir(folder) + + +class TestMediaLoader: + """ Test for :class:`~tools.alignments.media.MediaLoader` """ + + @pytest.fixture(name="media_loader_instance") + def media_loader_fixture(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> MediaLoader: + """ An instance of :class:`~tools.alignments.media.MediaLoader` with unimplemented + child methods patched out of __init__ and initialized with a dummy folder containing + 2 images + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking subclass calls + + Returns + ------- + :class:`~tools.alignments.media.MediaLoader` + Initialized instance for testing + """ + mocker.patch("tools.alignments.media.MediaLoader.sorted_items", + return_value=os.listdir(folder)) + mocker.patch("tools.alignments.media.MediaLoader.load_items") + loader = MediaLoader(folder) + return loader + + def test_init(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader`__init__ method + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking subclass calls + """ + sort_patch = mocker.patch("tools.alignments.media.MediaLoader.sorted_items", + return_value=os.listdir(folder)) + load_patch = mocker.patch("tools.alignments.media.MediaLoader.load_items") + loader = MediaLoader(folder) + sort_patch.assert_called_once() + load_patch.assert_called_once() + assert loader.folder == folder + assert loader._count == 2 + assert loader.count == 2 + assert not loader.is_video + + def test_check_input_folder(self, media_loader_instance: MediaLoader) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` check_input_folder method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + """ + media_loader = media_loader_instance + assert media_loader.check_input_folder() is None + media_loader.folder = "" + with pytest.raises(SystemExit): + media_loader.check_input_folder() + media_loader.folder = "/this/path/does/not/exist" + with pytest.raises(SystemExit): + media_loader.check_input_folder() + + def test_valid_extension(self, media_loader_instance: MediaLoader) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` valid_extension method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + """ + media_loader = media_loader_instance + assert media_loader.valid_extension("test.png") + assert media_loader.valid_extension("test.PNG") + assert media_loader.valid_extension("test.jpg") + assert media_loader.valid_extension("test.JPG") + assert not media_loader.valid_extension("test.doc") + assert not media_loader.valid_extension("test.txt") + assert not media_loader.valid_extension("test.mp4") + + def test_load_image(self, + media_loader_instance: MediaLoader, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` load_image method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking loader specific calls + """ + media_loader = media_loader_instance + expected = np.random.rand(256, 256, 3) + media_loader.load_video_frame = T.cast(MagicMock, # type:ignore + mocker.MagicMock(return_value=expected)) + read_image_patch = mocker.patch("tools.alignments.media.read_image", return_value=expected) + filename = "test.png" + output = media_loader.load_image(filename) + np.testing.assert_equal(expected, output) + read_image_patch.assert_called_once_with(os.path.join(media_loader.folder, filename), + raise_error=True) + + mocker.patch("tools.alignments.media.MediaLoader.is_video", + new_callable=mocker.PropertyMock(return_value=True)) + filename = "test.mp4" + output = media_loader.load_image(filename) + np.testing.assert_equal(expected, output) + media_loader.load_video_frame.assert_called_once_with(filename) + + def test_load_video_frame(self, + media_loader_instance: MediaLoader, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` load_video_frame method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking cv2 calls + """ + media_loader = media_loader_instance + filename = "test_0001.png" + with pytest.raises(AssertionError): + media_loader.load_video_frame(filename) + + mocker.patch("tools.alignments.media.MediaLoader.is_video", + new_callable=mocker.PropertyMock(return_value=True)) + expected = np.random.rand(256, 256, 3) + vid_cap = mocker.MagicMock(cv2.VideoCapture) + vid_cap.read.side_effect = ((1, expected), ) + + media_loader._vid_reader = T.cast(MagicMock, vid_cap) # type:ignore + output = media_loader.load_video_frame(filename) + vid_cap.set.assert_called_once() + np.testing.assert_equal(output, expected) + + def test_stream(self, + media_loader_instance: MediaLoader, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` stream method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking loader specific calls + """ + media_loader = media_loader_instance + + loader = mocker.patch("tools.alignments.media.ImagesLoader.load") + expected = [(fname, np.random.rand(256, 256, 3)) + for fname in os.listdir(media_loader.folder)] + loader.side_effect = [expected] + output = list(media_loader.stream()) + assert output == expected + + loader.reset_mock() + + skip_list = [0] + expected = [expected[1]] + loader.side_effect = [expected] + output = list(media_loader.stream(skip_list)) + assert output == expected + assert loader.add_skip_list.called_once_with(skip_list) + + def test_save_image(self, + media_loader_instance: MediaLoader, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.MediaLoader` save_image method + + Parameters + ---------- + media_loader_instance: :class:`~tools.alignments.media.MediaLoader` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking saver specific calls + """ + media_loader = media_loader_instance + out_folder = media_loader.folder + filename = "test_out.jpg" + expected_filename = os.path.join(media_loader.folder, "test_out.png") + img = np.random.rand(256, 256, 3) + metadata = {"test": "data"} + + cv2_write_mock = mocker.patch("cv2.imwrite") + cv2_encode_mock = mocker.patch("cv2.imencode") + png_write_meta_mock = mocker.patch("tools.alignments.media.png_write_meta") + open_mock = mocker.patch("builtins.open") + + media_loader.save_image(out_folder, filename, img, metadata=None) + cv2_write_mock.assert_called_once_with(expected_filename, img) + cv2_encode_mock.assert_not_called() + png_write_meta_mock.assert_not_called() + + cv2_write_mock.reset_mock() + + media_loader.save_image(out_folder, filename, img, metadata=metadata) # type:ignore + cv2_write_mock.assert_not_called() + cv2_encode_mock.assert_called_once_with(".png", img) + png_write_meta_mock.assert_called_once() + open_mock.assert_called_once() + + +class TestFaces: + """ Test for :class:`~tools.alignments.media.Faces` """ + + @pytest.fixture(name="faces_instance") + def faces_fixture(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> Faces: + """ An instance of :class:`~tools.alignments.media.Faces` patching out + read_image_meta_batch so nothing is loaded + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking read_image_meta_batch calls + + Returns + ------- + :class:`~tools.alignments.media.Faces` + Initialized instance for testing + """ + mocker.patch("tools.alignments.media.read_image_meta_batch") + loader = Faces(folder, None) + return loader + + def test_init(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.Faces`__init__ method + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking superclass calls + """ + parent_mock = mocker.patch("tools.alignments.media.super") + alignments_mock = mocker.patch("tools.alignments.media.AlignmentData") + Faces(folder, alignments_mock) + parent_mock.assert_called_once() + + def test__handle_legacy(self, + faces_instance: Faces, + mocker: pytest_mock.MockerFixture, + caplog: pytest.LogCaptureFixture) -> None: + """ Test for :class:`~tools.alignments.media.Faces` _handle_legacy method + + Parameters + ---------- + faces_instance: :class:`~tools.alignments.media.Faces` + Test class instance + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking various objects + caplog: :class:`pytest.LogCaptureFixture + For capturing logging messages + """ + faces = faces_instance + folder = faces.folder + legacy_file = os.path.join(folder, "a.png") + + # No alignments file + with pytest.raises(FaceswapError): + faces._handle_legacy(legacy_file) + + # No returned metadata + alignments_mock = mocker.patch("tools.alignments.media.AlignmentData") + alignments_mock.version = 2.1 + update_mock = mocker.patch("tools.alignments.media.update_legacy_png_header", + return_value={}) + faces = Faces(folder, alignments_mock) + faces.folder = folder + with pytest.raises(FaceswapError): + faces._handle_legacy(legacy_file) + update_mock.assert_called_once_with(legacy_file, alignments_mock) + + # Correct data with logging + caplog.clear() + update_mock.reset_mock() + update_mock.return_value = {"test": "data"} + faces._handle_legacy(legacy_file, log=True) + assert "Legacy faces discovered" in caplog.text + + # Correct data without logging + caplog.clear() + update_mock.reset_mock() + update_mock.return_value = {"test": "data"} + faces._handle_legacy(legacy_file, log=False) + assert "Legacy faces discovered" not in caplog.text + + def test__handle_duplicate(self, faces_instance: Faces) -> None: + """ Test for :class:`~tools.alignments.media.Faces` _handle_duplicate method + + Parameters + ---------- + faces_instance: :class:`~tools.alignments.media.Faces` + The class instance for testing + """ + faces = faces_instance + dupe_dir = os.path.join(faces.folder, "_duplicates") + src_filename = "test_0001.png" + src_face_idx = 0 + paths = [os.path.join(faces.folder, fname) for fname in os.listdir(faces.folder)] + data = {"source": {"source_filename": src_filename, + "face_index": src_face_idx}} + seen: dict[str, list[int]] = {} + + # New item + is_dupe = faces._handle_duplicate(paths[0], data, seen) # type:ignore + assert src_filename in seen and seen[src_filename] == [src_face_idx] + assert not os.path.exists(dupe_dir) + assert not is_dupe + + # Dupe item + is_dupe = faces._handle_duplicate(paths[1], data, seen) # type:ignore + assert src_filename in seen and seen[src_filename] == [src_face_idx] + assert len(seen) == 1 + assert os.path.exists(dupe_dir) + assert not os.path.exists(paths[1]) + assert is_dupe + + # Move everything back for fixture cleanup + os.rename(os.path.join(dupe_dir, os.path.basename(paths[1])), paths[1]) + os.rmdir(dupe_dir) + + def test_process_folder(self, + faces_instance: Faces, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.Faces` process_folder method + + Parameters + ---------- + faces_instance: :class:`~tools.alignments.media.Faces` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking various logic calls + """ + faces = faces_instance + read_image_meta_mock = mocker.patch("tools.alignments.media.read_image_meta_batch") + img_sources = [os.path.join(faces.folder, fname) for fname in os.listdir(faces.folder)] + meta_data = {"itxt": {"source": ({"source_filename": "data.png"})}} + expected = [(fname, meta_data["itxt"]) for fname in os.listdir(faces.folder)] + read_image_meta_mock.side_effect = [[(src, meta_data) for src in img_sources]] + + legacy_mock = mocker.patch("tools.alignments.media.Faces._handle_legacy", + return_value=meta_data["itxt"]) + dupe_mock = mocker.patch("tools.alignments.media.Faces._handle_duplicate", + return_value=False) + + # valid itxt + output = list(faces.process_folder()) + assert read_image_meta_mock.call_count == 1 + assert dupe_mock.call_count == 2 + assert not legacy_mock.called + assert output == expected + + dupe_mock.reset_mock() + read_image_meta_mock.reset_mock() + + # valid itxt with alignemnts data + read_image_meta_mock.side_effect = [[(src, meta_data) for src in img_sources]] + faces._alignments = mocker.MagicMock(AlignmentData) + faces._alignments.version = 2.1 # type:ignore + output = list(faces.process_folder()) + assert faces._alignments.frame_exists.call_count == 2 # type:ignore + assert read_image_meta_mock.call_count == 1 + assert dupe_mock.call_count == 2 + + dupe_mock.reset_mock() + read_image_meta_mock.reset_mock() + faces._alignments = None + + # invalid itxt + read_image_meta_mock.side_effect = [[(src, {}) for src in img_sources]] + output = list(faces.process_folder()) + assert read_image_meta_mock.call_count == 1 + assert legacy_mock.call_count == 2 + assert dupe_mock.call_count == 2 + assert output == expected + + def test_load_items(self, + faces_instance: Faces) -> None: + """ Test for :class:`~tools.alignments.media.Faces` load_items method + + Parameters + ---------- + faces_instance: :class:`~tools.alignments.media.Faces` + The class instance for testing + """ + faces = faces_instance + data = [(f"file{idx}.png", {"source": {"source_filename": f"src{idx}.png", + "face_index": 0}}) + for idx in range(4)] + faces.file_list_sorted = data # type: ignore + expected = {"src0.png": [0], "src1.png": [0], "src2.png": [0], "src3.png": [0]} + result = faces.load_items() + assert result == expected + + data = [(f"file{idx}.png", {"source": {"source_filename": f"src{idx // 2}.png", + "face_index": 0 if idx % 2 == 0 else 1}}) + for idx in range(4)] + faces.file_list_sorted = data # type: ignore + expected = {"src0.png": [0, 1], "src1.png": [0, 1]} + result = faces.load_items() + assert result == expected + + def test_sorted_items(self, + faces_instance: Faces, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.Faces` sorted_items method + + Parameters + ---------- + faces_instance: :class:`~tools.alignments.media.Faces` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking various logic calls + """ + faces = faces_instance + data: list[tuple[str, dict]] = [("file4.png", {}), ("file3.png", {}), + ("file1.png", {}), ("file2.png", {})] + expected = sorted(data) + process_folder_mock = mocker.patch("tools.alignments.media.Faces.process_folder", + side_effect=[data]) + result = faces.sorted_items() + assert process_folder_mock.called + assert result == expected + + +class TestFrames: + """ Test for :class:`~tools.alignments.media.Frames` """ + + def test_process_folder(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.Frames` process_folder method + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking superclass calls + """ + process_video_mock = mocker.patch("tools.alignments.media.Frames.process_video") + process_frames_mock = mocker.patch("tools.alignments.media.Frames.process_frames") + + frames = Frames(folder, None) + frames.process_folder() + process_frames_mock.assert_called_once() + process_video_mock.assert_not_called() + + process_frames_mock.reset_mock() + mocker.patch("tools.alignments.media.Frames.is_video", + new_callable=mocker.PropertyMock(return_value=True)) + frames = Frames(folder, None) + frames.process_folder() + process_frames_mock.assert_not_called() + process_video_mock.assert_called_once() + + def test_process_frames(self, folder: str) -> None: + """ Test for :class:`~tools.alignments.media.Frames` process_frames method + + Parameters + ---------- + folder : str + Dummy media folder + """ + expected = [{"frame_fullname": "a.png", "frame_name": "a", "frame_extension": ".png"}, + {"frame_fullname": "b.png", "frame_name": "b", "frame_extension": ".png"}] + + frames = Frames(folder, None) + returned = sorted(list(frames.process_frames()), key=itemgetter("frame_fullname")) + assert returned == sorted(expected, key=itemgetter("frame_fullname")) + + def test_process_video(self, folder: str) -> None: + """ Test for :class:`~tools.alignments.media.Frames` process_video method + + Parameters + ---------- + folder : str + Dummy media folder + """ + ext = os.path.splitext(folder)[-1] + expected = [{"frame_fullname": f"images_000001{ext}", + "frame_name": "images_000001", + "frame_extension": ext}, + {"frame_fullname": f"images_000002{ext}", + "frame_name": "images_000002", + "frame_extension": ext}] + + frames = Frames(folder, None) + returned = list(frames.process_video()) + assert returned == expected + + def test_load_items(self, folder: str) -> None: + """ Test for :class:`~tools.alignments.media.Frames` load_items method + + Parameters + ---------- + folder : str + Dummy media folder + """ + expected = {"a.png": ("a", ".png"), "b.png": ("b", ".png")} + frames = Frames(folder, None) + result = frames.load_items() + assert result == expected + + def test_sorted_items(self, + folder: str, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.Frames` sorted_items method + + Parameters + ---------- + folder : str + Dummy media folder + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking process_folder call + """ + frames = Frames(folder, None) + data = [{"frame_fullname": "c.png", "frame_name": "c", "frame_extension": ".png"}, + {"frame_fullname": "d.png", "frame_name": "d", "frame_extension": ".png"}, + {"frame_fullname": "b.jpg", "frame_name": "b", "frame_extension": ".jpg"}, + {"frame_fullname": "a.png", "frame_name": "a", "frame_extension": ".png"}] + expected = [{"frame_fullname": "a.png", "frame_name": "a", "frame_extension": ".png"}, + {"frame_fullname": "b.jpg", "frame_name": "b", "frame_extension": ".jpg"}, + {"frame_fullname": "c.png", "frame_name": "c", "frame_extension": ".png"}, + {"frame_fullname": "d.png", "frame_name": "d", "frame_extension": ".png"}] + process_folder_mock = mocker.patch("tools.alignments.media.Frames.process_folder", + side_effect=[data]) + result = frames.sorted_items() + + assert process_folder_mock.called + assert result == expected + + +class TestExtractedFaces: + """ Test for :class:`~tools.alignments.media.ExtractedFaces` """ + + @pytest.fixture(name="extracted_faces_instance") + def extracted_faces_fixture(self, mocker: pytest_mock.MockerFixture) -> ExtractedFaces: + """ An instance of :class:`~tools.alignments.media.ExtractedFaces` patching out Frames and + AlignmentData parameters + + Parameters + ---------- + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking read_image_meta_batch calls + + Returns + ------- + :class:`~tools.alignments.media.ExtractedFaces` + Initialized instance for testing + """ + frames_mock = mocker.MagicMock(Frames) + alignments_mock = mocker.MagicMock(AlignmentData) + return ExtractedFaces(frames_mock, alignments_mock, size=512) + + def test_init(self, extracted_faces_instance: ExtractedFaces) -> None: + """ Test for :class:`~tools.alignments.media.ExtractedFace` __init__ method + + Parameters + ---------- + extracted_faces_instance: :class:`~tools.alignments.media.ExtractedFace` + The class instance for testing + """ + faces = extracted_faces_instance + assert faces.size == 512 + assert faces.padding == int(512 * 0.1875) + assert faces.current_frame is None + assert faces.faces == [] + + def test_get_faces(self, + extracted_faces_instance: ExtractedFaces, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.ExtractedFace` get_faces method + + Parameters + ---------- + extracted_faces_instance: :class:`~tools.alignments.media.ExtractedFace` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking Frames and AlignmentData classes + """ + extract_face_mock = mocker.patch("tools.alignments.media.ExtractedFaces.extract_one_face") + faces = extracted_faces_instance + + frame = "test_frame" + img = np.random.rand(256, 256, 3) + + # No alignment data + faces.alignments.get_faces_in_frame.return_value = [] # type:ignore + faces.get_faces(frame, img) + faces.alignments.get_faces_in_frame.assert_called_once_with(frame) # type:ignore + faces.frames.load_image.assert_not_called() # type:ignore + extract_face_mock.assert_not_called() + assert faces.current_frame is None + + faces.alignments.reset_mock() # type:ignore + + # Alignment data + image + faces.alignments.get_faces_in_frame.return_value = [1, 2, 3] # type:ignore + faces.get_faces(frame, img) + faces.alignments.get_faces_in_frame.assert_called_once_with(frame) # type:ignore + faces.frames.load_image.assert_not_called() # type:ignore + assert extract_face_mock.call_count == 3 + assert faces.current_frame == frame + + faces.alignments.reset_mock() # type:ignore + extract_face_mock.reset_mock() + + # Alignment data + no image + faces.alignments.get_faces_in_frame.return_value = ["data1"] # type:ignore + faces.get_faces(frame, None) + faces.alignments.get_faces_in_frame.assert_called_once_with(frame) # type:ignore + faces.frames.load_image.assert_called_once_with(frame) # type:ignore + assert extract_face_mock.call_count == 1 + assert faces.current_frame == frame + + def test_extract_one_face(self, + extracted_faces_instance: ExtractedFaces, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.ExtractedFace` extract_one_face method + + Parameters + ---------- + extracted_faces_instance: :class:`~tools.alignments.media.ExtractedFace` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking DetectedFace object + """ + detected_face = mocker.patch("tools.alignments.media.DetectedFace") + thumbnail_mock = mocker.patch("tools.alignments.media.generate_thumbnail") + faces = extracted_faces_instance + alignment = {"test"} + img = np.random.rand(256, 256, 3) + returned = faces.extract_one_face(alignment, img) # type:ignore + detected_face.assert_called_once() + detected_face.return_value.from_alignment.assert_called_once_with(alignment, + image=img) + detected_face.return_value.load_aligned.assert_called_once_with(img, + size=512, + centering="head") + thumbnail_mock.assert_called_once() + assert isinstance(returned, MagicMock) + + def test_get_faces_in_frame(self, + extracted_faces_instance: ExtractedFaces, + mocker: pytest_mock.MockerFixture) -> None: + """ Test for :class:`~tools.alignments.media.ExtractedFace` get_faces_in_frame method + + Parameters + ---------- + extracted_faces_instance: :class:`~tools.alignments.media.ExtractedFace` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking get_faces method + """ + faces = extracted_faces_instance + faces.get_faces = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + + frame = "test_frame" + img = None + + faces.get_faces_in_frame(frame, update=False, image=img) + faces.get_faces.assert_called_once_with(frame, image=img) + + faces.get_faces.reset_mock() + + faces.current_frame = frame + faces.get_faces_in_frame(frame, update=False, image=img) + faces.get_faces.assert_not_called() + + faces.get_faces_in_frame(frame, update=True, image=img) + faces.get_faces.assert_called_once_with(frame, image=img) + + _params = [(np.array(([[25, 47], [32, 232], [244, 237], [240, 21]])), 216), + (np.array(([[127, 392], [403, 510], [32, 237], [19, 210]])), 211), + (np.array(([[26, 1927], [112, 1234], [1683, 1433], [78, 1155]])), 773)] + + @pytest.mark.parametrize("roi,expected", _params) + def test_get_roi_size_for_frame(self, + extracted_faces_instance: ExtractedFaces, + mocker: pytest_mock.MockerFixture, + roi: np.ndarray, + expected: int) -> None: + """ Test for :class:`~tools.alignments.media.ExtractedFace` get_roi_size_for_frame method + + Parameters + ---------- + extracted_faces_instance: :class:`~tools.alignments.media.ExtractedFace` + The class instance for testing + mocker: :class:`pytest_mock.MockerFixture` + Fixture for mocking get_faces method and DetectedFace object + roi: :class:`numpy.ndarray` + Test ROI box to feed into the function + expected: int + The expected output for the given ROI box + """ + faces = extracted_faces_instance + faces.get_faces = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + + frame = "test_frame" + faces.get_roi_size_for_frame(frame) + faces.get_faces.assert_called_once_with(frame) + + faces.get_faces.reset_mock() + + faces.current_frame = frame + faces.get_roi_size_for_frame(frame) + faces.get_faces.assert_not_called() + + detected_face = mocker.MagicMock("tools.alignments.media.DetectedFace") + detected_face.aligned = detected_face + detected_face.original_roi = roi + faces.faces = [detected_face] + result = faces.get_roi_size_for_frame(frame) + assert result == [expected] diff --git a/tests/tools/preview/__init__.py b/tests/tools/preview/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/tools/preview/viewer_test.py b/tests/tools/preview/viewer_test.py new file mode 100644 index 0000000000..18fd84d404 --- /dev/null +++ b/tests/tools/preview/viewer_test.py @@ -0,0 +1,483 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`tools.preview.viewer` """ +from __future__ import annotations +import tkinter as tk +import typing as T + +from tkinter import ttk + +from unittest.mock import MagicMock + +import pytest +import pytest_mock +import numpy as np +from PIL import ImageTk + +from lib.logger import log_setup +# Need to setup logging to avoid trace/verbose errors +log_setup("DEBUG", "pytest_viewer.log", "PyTest, False") + +from lib.utils import get_backend # pylint:disable=wrong-import-position # noqa +from tools.preview.viewer import _Faces, FacesDisplay, ImagesCanvas # pylint:disable=wrong-import-position # noqa + +if T.TYPE_CHECKING: + from lib.align.aligned_face import CenteringType + + +# pylint:disable=protected-access + + +def test__faces(): + """ Test the :class:`~tools.preview.viewer._Faces dataclass initializes correctly """ + faces = _Faces() + assert faces.filenames == [] + assert faces.matrix == [] + assert faces.src == [] + assert faces.dst == [] + + +_PARAMS = [(3, 448), (4, 333), (5, 254), (6, 128)] # columns/face_size +_IDS = [f"cols:{c},size:{s}[{get_backend().upper()}]" for c, s in _PARAMS] + + +class TestFacesDisplay(): + """ Test :class:`~tools.preview.viewer.FacesDisplay """ + _padding = 64 + + def get_faces_display_instance(self, columns: int = 5, face_size: int = 256) -> FacesDisplay: + """ Obtain an instance of :class:`~tools.preview.viewer.FacesDisplay` with the given column + and face size layout. + + Parameters + ---------- + columns: int, optional + The number of columns to display in the viewer, default: 5 + face_size: int, optional + The size of each face image to be displayed in the viewer, default: 256 + + Returns + ------- + :class:`~tools.preview.viewer.FacesDisplay` + An instance of the FacesDisplay class at the given settings + """ + app = MagicMock() + retval = FacesDisplay(app, face_size, self._padding) + retval._faces = _Faces( + matrix=[np.random.rand(2, 3) for _ in range(columns)], + src=[np.random.rand(face_size, face_size, 3) for _ in range(columns)], + dst=[np.random.rand(face_size, face_size, 3) for _ in range(columns)]) + return retval + + def test_init(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` __init__ method """ + f_display = self.get_faces_display_instance(face_size=256) + assert f_display._size == 256 + assert f_display._padding == self._padding + assert isinstance(f_display._app, MagicMock) + + assert f_display._display_dims == (1, 1) + assert isinstance(f_display._faces, _Faces) + + assert f_display._centering is None + assert f_display._faces_source.size == 0 + assert f_display._faces_dest.size == 0 + assert f_display._tk_image is None + assert f_display.update_source is False + assert not f_display.source and isinstance(f_display.source, list) + assert not f_display.destination and isinstance(f_display.destination, list) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__total_columns(self, columns: int, face_size: int) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _total_columns property is correctly + calculated + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [None for _ in range(columns)] # type:ignore + assert f_display._total_columns == columns + + def test_set_centering(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` set_centering method """ + f_display = self.get_faces_display_instance() + assert f_display._centering is None + centering: CenteringType = "legacy" + f_display.set_centering(centering) + assert f_display._centering == centering + + def test_set_display_dimensions(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` set_display_dimensions method """ + f_display = self.get_faces_display_instance() + assert f_display._display_dims == (1, 1) + dimensions = (800, 600) + f_display.set_display_dimensions(dimensions) + assert f_display._display_dims == dimensions + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_update_tk_image(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` update_tk_image method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _build_faces_image method called + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._build_faces_image = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._get_scale_size = T.cast(MagicMock, # type:ignore + mocker.MagicMock(return_value=(128, 128))) + f_display._faces_source = np.zeros((face_size, face_size, 3), dtype=np.uint8) + f_display._faces_dest = np.zeros((face_size, face_size, 3), dtype=np.uint8) + + tk.Tk() # tkinter instance needed for image creation + f_display.update_tk_image() + + f_display._build_faces_image.assert_called_once() + f_display._get_scale_size.assert_called_once() + assert isinstance(f_display._tk_image, ImageTk.PhotoImage) + assert f_display._tk_image.width() == 128 + assert f_display._tk_image.height() == 128 + assert f_display.tk_image == f_display._tk_image # public property test + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_get_scale_size(self, columns: int, face_size: int) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` get_scale_size method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.set_display_dimensions((800, 600)) + + img = np.zeros((face_size, face_size, 3), dtype=np.uint8) + size = f_display._get_scale_size(img) + assert size == (600, 600) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__build_faces_image(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _build_faces_image method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking internal methods called + """ + header_size = 32 + + f_display = self.get_faces_display_instance(columns, face_size) + f_display._faces_from_frames = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._header_text = T.cast( # type:ignore + MagicMock, + mocker.MagicMock(return_value=np.random.rand(header_size, face_size * columns, 3))) + f_display._draw_rect = T.cast(MagicMock, # type:ignore + mocker.MagicMock(side_effect=lambda x: x)) + + # Test full update + f_display.update_source = True + f_display._build_faces_image() + + f_display._faces_from_frames.assert_called_once() + f_display._header_text.assert_called_once() + assert f_display._draw_rect.call_count == columns * 2 # src + dst + assert f_display._faces_source.shape == (face_size + header_size, face_size * columns, 3) + assert f_display._faces_dest.shape == (face_size, face_size * columns, 3) + + f_display._faces_from_frames.reset_mock() + f_display._header_text.reset_mock() + f_display._draw_rect.reset_mock() + + # Test dst update only + f_display.update_source = False + f_display._build_faces_image() + + f_display._faces_from_frames.assert_called_once() + assert not f_display._header_text.called + assert f_display._draw_rect.call_count == columns # dst only + assert f_display._faces_dest.shape == (face_size, face_size * columns, 3) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_faces__from_frames(self, + columns, + face_size, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _from_frames method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _build_faces_image method called + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [mocker.MagicMock() for _ in range(3)] + f_display.destination = [np.random.rand(face_size, face_size, 3) for _ in range(3)] + f_display._crop_source_faces = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._crop_destination_faces = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + + # Both src + dst + f_display.update_source = True + f_display._faces_from_frames() + f_display._crop_source_faces.assert_called_once() + f_display._crop_destination_faces.assert_called_once() + + f_display._crop_source_faces.reset_mock() + f_display._crop_destination_faces.reset_mock() + + # Just dst + f_display.update_source = False + f_display._faces_from_frames() + assert not f_display._crop_source_faces.called + f_display._crop_destination_faces.assert_called_once() + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__crop_source_faces(self, + columns: int, + face_size: int, + monkeypatch: pytest.MonkeyPatch, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _crop_source_faces method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + monkeypatch: :class:`pytest.MonkeyPatch` + For patching the transform_image function + mocker: :class:`pytest_mock.MockerFixture` + Mocker for mocking various internal methods + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._centering = "face" + f_display.update_source = True + f_display._faces.src = [] + + transform_image_mock = mocker.MagicMock() + monkeypatch.setattr("tools.preview.viewer.transform_image", transform_image_mock) + + f_display.source = [mocker.MagicMock() for _ in range(columns)] + for idx, mock in enumerate(f_display.source): + assert isinstance(mock, MagicMock) + mock.inbound.detected_faces.__getitem__ = lambda self, x, y=mock: y + mock.aligned.matrix = f"test_matrix_{idx}" + mock.inbound.filename = f"test_filename_{idx}.txt" + + f_display._crop_source_faces() + + assert len(f_display._faces.filenames) == columns + assert len(f_display._faces.matrix) == columns + assert len(f_display._faces.src) == columns + assert not f_display.update_source + assert transform_image_mock.call_count == columns + + for idx in range(columns): + assert f_display._faces.filenames[idx] == f"test_filename_{idx}" + assert f_display._faces.matrix[idx] == f"test_matrix_{idx}" + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__crop_destination_faces(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _crop_destination_faces method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in full frames + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._centering = "face" + f_display._faces.dst = [] # empty object and test populated correctly + + f_display.source = [mocker.MagicMock() for _ in range(columns)] + for item in f_display.source: # type ignore + item.inbound.image = np.random.rand(1280, 720, 3) # type:ignore + + f_display._crop_destination_faces() + assert len(f_display._faces.dst) == columns + assert all(f.shape == (face_size, face_size, 3) for f in f_display._faces.dst) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__header_text(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _header_text method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in cv2 calls + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [None for _ in range(columns)] # type:ignore + f_display._faces.filenames = [f"filename_{idx}.png" for idx in range(columns)] + + cv2_mock = mocker.patch("tools.preview.viewer.cv2") + text_width, text_height = (100, 32) + cv2_mock.getTextSize.return_value = [(text_width, text_height), ] + + header_box = f_display._header_text() + assert cv2_mock.getTextSize.call_count == columns + assert cv2_mock.putText.call_count == columns + assert header_box.shape == (face_size // 8, face_size * columns, 3) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__draw_rect_text(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _draw_rect method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in cv2 calls + """ + f_display = self.get_faces_display_instance(columns, face_size) + cv2_mock = mocker.patch("tools.preview.viewer.cv2") + + image = (np.random.rand(face_size, face_size, 3) * 255.0) + 50 + assert image.max() > 255.0 + output = f_display._draw_rect(image) + cv2_mock.rectangle.assert_called_once() + assert output.max() == 255.0 # np.clip + + +class TestImagesCanvas: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` """ + + @pytest.fixture + def parent(self) -> MagicMock: + """ Mock object to act as the parent widget to the ImagesCanvas + + Returns + -------- + :class:`unittest.mock.MagicMock` + The mocked ttk.PanedWindow widget + """ + retval = MagicMock(spec=ttk.PanedWindow) + retval.tk = retval + retval._w = "mock_ttkPanedWindow" + retval.children = {} + retval.call = retval + retval.createcommand = retval + retval.preview_display = MagicMock(spec=FacesDisplay) + return retval + + @pytest.fixture(name="images_canvas_instance") + def images_canvas_fixture(self, parent) -> ImagesCanvas: + """ Fixture for creating a testing :class:`~tools.preview.viewer.ImagesCanvas` instance + + Parameters + ---------- + parent: :class:`unittest.mock.MagicMock` + The mocked ttk.PanedWindow parent + + Returns + ------- + :class:`~tools.preview.viewer.ImagesCanvas` + The class instance for testing + """ + app = MagicMock() + return ImagesCanvas(app, parent) + + def test_init(self, images_canvas_instance: ImagesCanvas, parent: MagicMock) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` __init__ method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + """ + assert images_canvas_instance._display == parent.preview_display + assert isinstance(images_canvas_instance._canvas, tk.Canvas) + assert images_canvas_instance._canvas.master == images_canvas_instance + assert images_canvas_instance._canvas.winfo_ismapped() + + def test_resize(self, + images_canvas_instance: ImagesCanvas, + parent: MagicMock, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` resize method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in tk calls + """ + event_mock = mocker.MagicMock(spec=tk.Event, width=100, height=200) + images_canvas_instance.reload = T.cast(MagicMock, mocker.MagicMock()) # type:ignore + + images_canvas_instance._resize(event_mock) + + parent.preview_display.set_display_dimensions.assert_called_once_with((100, 200)) + images_canvas_instance.reload.assert_called_once() + + def test_reload(self, + images_canvas_instance: ImagesCanvas, + parent: MagicMock, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` reload method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in tk calls + """ + itemconfig_mock = mocker.patch.object(tk.Canvas, "itemconfig") + + images_canvas_instance.reload() + + parent.preview_display.update_tk_image.assert_called_once() + itemconfig_mock.assert_called_once() diff --git a/tests/utils.py b/tests/utils.py index 248ec0a25b..b357dc13c4 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,7 +2,6 @@ """ Utils imported from Keras as their location changes between Tensorflow Keras and standard Keras. Also ensures testing consistency """ import inspect -import sys import numpy as np @@ -101,25 +100,13 @@ def has_arg(func, name, accept_all=False): bool Whether `func` accepts a `name` keyword argument. """ - if sys.version_info < (3,): - arg_spec = inspect.getargspec(func) - if accept_all and arg_spec.keywords is not None: - return True - return (name in arg_spec.args) - elif sys.version_info < (3, 3): - arg_spec = inspect.getfullargspec(func) - if accept_all and arg_spec.varkw is not None: - return True - return (name in arg_spec.args or - name in arg_spec.kwonlyargs) - else: - signature = inspect.signature(func) - parameter = signature.parameters.get(name) - if parameter is None: - if accept_all: - for param in signature.parameters.values(): - if param.kind == inspect.Parameter.VAR_KEYWORD: - return True - return False - return (parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY)) + signature = inspect.signature(func) + parameter = signature.parameters.get(name) + if parameter is None: + if accept_all: + for param in signature.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + return True + return False + return (parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY)) diff --git a/tools.py b/tools.py index e17dce0d61..c47798c7d0 100755 --- a/tools.py +++ b/tools.py @@ -9,17 +9,13 @@ # Importing the various tools from lib.cli.args import FullHelpArgumentParser - # LOCALES _LANG = gettext.translation("tools", localedir="locales", fallback=True) _ = _LANG.gettext - # Python version check -if sys.version_info[0] < 3: - raise Exception("This program requires at least python3.7") -if sys.version_info[0] == 3 and sys.version_info[1] < 7: - raise Exception("This program requires at least python3.7") +if sys.version_info < (3, 10): + raise ValueError("This program requires at least python 3.10") def bad_args(*args): # pylint:disable=unused-argument @@ -37,15 +33,12 @@ def _get_cli_opts(): if os.path.exists(cli_file): mod = ".".join(("tools", tool_name, "cli")) module = import_module(mod) - cliarg_class = getattr(module, "{}Args".format(tool_name.title())) + cliarg_class = getattr(module, f"{tool_name.title()}Args") help_text = getattr(module, "_HELPTEXT") yield tool_name, help_text, cliarg_class if __name__ == "__main__": - print(_("Please backup your data and/or test the tool you want to use with a smaller data set " - "to make sure you understand how it works.")) - PARSER = FullHelpArgumentParser() SUBPARSER = PARSER.add_subparsers() for tool, helptext, cli_args in _get_cli_opts(): diff --git a/tools/alignments/alignments.py b/tools/alignments/alignments.py index 010b20260e..4c6251f1fe 100644 --- a/tools/alignments/alignments.py +++ b/tools/alignments/alignments.py @@ -1,20 +1,245 @@ #!/usr/bin/env python3 """ Tools for manipulating the alignments serialized file """ import logging +import os +import sys +import typing as T -from typing import TYPE_CHECKING +from argparse import Namespace +from multiprocessing import Process +from lib.utils import FaceswapError, handle_deprecated_cliopts, VIDEO_EXTENSIONS from .media import AlignmentData -from .jobs import (Check, Draw, Extract, FromFaces, Rename, # noqa pylint: disable=unused-import - RemoveFaces, Sort, Spatial) +from .jobs import Check, Export, Sort, Spatial # noqa pylint:disable=unused-import +from .jobs_faces import FromFaces, RemoveFaces, Rename # noqa pylint:disable=unused-import +from .jobs_frames import Draw, Extract # noqa pylint:disable=unused-import -if TYPE_CHECKING: - from argparse import Namespace -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -class Alignments(): # pylint:disable=too-few-public-methods +class Alignments(): + """ The main entry point for Faceswap's Alignments Tool. This tool is part of the Faceswap + Tools suite and should be called from the ``python tools.py alignments`` command. + + The tool allows for manipulation, and working with Faceswap alignments files. + + This parent class handles creating the individual job arguments when running in batch-mode or + triggers the job when not running in batch mode + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + """ + def __init__(self, arguments: Namespace) -> None: + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + self._requires_alignments = ["export", "sort", "spatial"] + self._requires_faces = ["extract", "from-faces"] + self._requires_frames = ["draw", + "extract", + "missing-alignments", + "missing-frames", + "no-faces"] + + self._args = handle_deprecated_cliopts(arguments) + self._batch_mode = self._validate_batch_mode() + self._locations = self._get_locations() + + def _validate_batch_mode(self) -> bool: + """ Validate that the selected job supports batch processing + + Returns + ------- + bool + ``True`` if batch mode has been selected otherwise ``False`` + """ + batch_mode: bool = self._args.batch_mode + if not batch_mode: + logger.debug("Running in standard mode") + return batch_mode + valid = self._requires_alignments + self._requires_faces + self._requires_frames + if self._args.job not in valid: + logger.error("Job '%s' does not support batch mode. Please select a job from %s or " + "disable batch mode", self._args.job, valid) + sys.exit(1) + logger.debug("Running in batch mode") + return batch_mode + + def _get_alignments_locations(self) -> dict[str, list[str | None]]: + """ Obtain the full path to alignments files in a parent (batch) location + + These are jobs that only require an alignments file as input, so frames and face locations + are returned as a list of ``None`` values corresponding to the number of alignments files + detected + + Returns + ------- + dict[str, list[Optional[str]]]: + The list of alignments location paths and None lists for frames and faces locations + """ + if not self._args.alignments_file: + logger.error("Please provide an 'alignments_file' location for '%s' job", + self._args.job) + sys.exit(1) + + alignments = [os.path.join(self._args.alignments_file, fname) + for fname in os.listdir(self._args.alignments_file) + if os.path.splitext(fname)[-1].lower() == ".fsa" + and os.path.splitext(fname)[0].endswith("alignments")] + if not alignments: + logger.error("No alignment files found in '%s'", self._args.alignments_file) + sys.exit(1) + + logger.info("Batch mode selected. Processing alignments: %s", alignments) + retval = {"alignments_file": alignments, + "faces_dir": [None for _ in range(len(alignments))], + "frames_dir": [None for _ in range(len(alignments))]} + return retval + + def _get_frames_locations(self) -> dict[str, list[str | None]]: + """ Obtain the full path to frame locations along with corresponding alignments file + locations contained within the parent (batch) location + + Returns + ------- + dict[str, list[Optional[str]]]: + list of frames and alignments location paths. If the job requires an output faces + location then the faces folders are also returned, otherwise the faces will be a list + of ``Nones`` corresponding to the number of jobs to run + """ + if not self._args.frames_dir: + logger.error("Please provide a 'frames_dir' location for '%s' job", self._args.job) + sys.exit(1) + + frames: list[str] = [] + alignments: list[str] = [] + candidates = [os.path.join(self._args.frames_dir, fname) + for fname in os.listdir(self._args.frames_dir) + if os.path.isdir(os.path.join(self._args.frames_dir, fname)) + or os.path.splitext(fname)[-1].lower() in VIDEO_EXTENSIONS] + logger.debug("Frame candidates: %s", candidates) + + for candidate in candidates: + fname = os.path.join(candidate, "alignments.fsa") + if os.path.isdir(candidate) and os.path.exists(fname): + frames.append(candidate) + alignments.append(fname) + continue + fname = f"{os.path.splitext(candidate)[0]}_alignments.fsa" + if os.path.isfile(candidate) and os.path.exists(fname): + frames.append(candidate) + alignments.append(fname) + continue + logger.warning("Can't locate alignments file for '%s'. Skipping.", candidate) + + if not frames: + logger.error("No valid videos or frames folders found in '%s'", self._args.frames_dir) + sys.exit(1) + + if self._args.job not in self._requires_faces: # faces not required for frames input + faces: list[str | None] = [None for _ in range(len(frames))] + else: + if not self._args.faces_dir: + logger.error("Please provide a 'faces_dir' location for '%s' job", self._args.job) + sys.exit(1) + faces = [os.path.join(self._args.faces_dir, os.path.basename(os.path.splitext(frm)[0])) + for frm in frames] + + logger.info("Batch mode selected. Processing frames: %s", + [os.path.basename(frame) for frame in frames]) + + return {"alignments_file": T.cast(list[str | None], alignments), + "frames_dir": T.cast(list[str | None], frames), + "faces_dir": faces} + + def _get_locations(self) -> dict[str, list[str | None]]: + """ Obtain the full path to any frame, face and alignments input locations for the + selected job when running in batch mode. If not running in batch mode, then the original + passed in values are returned in lists + + Returns + ------- + dict[str, list[Optional[str]]] + A dictionary corresponding to the alignments, frames_dir and faces_dir arguments + with a list of full paths for each job + """ + job: str = self._args.job + if not self._batch_mode: # handle with given arguments + retval = {"alignments_file": [self._args.alignments_file], + "faces_dir": [self._args.faces_dir], + "frames_dir": [self._args.frames_dir]} + + elif job in self._requires_alignments: # Jobs only requiring an alignments file location + retval = self._get_alignments_locations() + + elif job in self._requires_frames: # Jobs that require a frames folder + retval = self._get_frames_locations() + + elif job in self._requires_faces and job not in self._requires_frames: + # Jobs that require faces as input + faces = [os.path.join(self._args.faces_dir, folder) + for folder in os.listdir(self._args.faces_dir) + if os.path.isdir(os.path.join(self._args.faces_dir, folder))] + if not faces: + logger.error("No folders found in '%s'", self._args.faces_dir) + sys.exit(1) + + retval = {"faces_dir": faces, + "frames_dir": [None for _ in range(len(faces))], + "alignments_file": [None for _ in range(len(faces))]} + logger.info("Batch mode selected. Processing faces: %s", + [os.path.basename(folder) for folder in faces]) + else: + raise FaceswapError(f"Unhandled job: {self._args.job}. This is a bug. Please report " + "to the developers") + + logger.debug("File locations: %s", retval) + return retval + + @staticmethod + def _run_process(arguments) -> None: + """ The alignements tool process to be run in a spawned process. + + In some instances, batch-mode memory leaks. Launching each job in a separate process + prevents this leak. + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments to be used for the given job + """ + logger.debug("Starting process: (arguments: %s)", arguments) + tool = _Alignments(arguments) + tool.process() + logger.debug("Finished process: (arguments: %s)", arguments) + + def process(self): + """ The entry point for the Alignments tool from :mod:`lib.tools.alignments.cli`. + + Launches the selected alignments job. + """ + num_jobs = len(self._locations["frames_dir"]) + for idx, (frames, faces, alignments) in enumerate(zip(self._locations["frames_dir"], + self._locations["faces_dir"], + self._locations["alignments_file"])): + if num_jobs > 1: + logger.info("Processing job %s of %s", idx + 1, num_jobs) + + args = Namespace(**self._args.__dict__) + args.frames_dir = frames + args.faces_dir = faces + args.alignments_file = alignments + + if num_jobs > 1: + proc = Process(target=self._run_process, args=(args, )) + proc.start() + proc.join() + else: + self._run_process(args) + + +class _Alignments(): """ The main entry point for Faceswap's Alignments Tool. This tool is part of the Faceswap Tools suite and should be called from the ``python tools.py alignments`` command. @@ -25,22 +250,70 @@ class Alignments(): # pylint:disable=too-few-public-methods arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` """ - def __init__(self, arguments: "Namespace") -> None: + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) - self.args = arguments - job = self.args.job - self.alignments = None if job == "from-faces" else AlignmentData(self.args.alignments_file) + self._args = arguments + job = self._args.job + + if job == "from-faces": + self.alignments = None + else: + self.alignments = AlignmentData(self._find_alignments()) + + if (self.alignments is not None and + arguments.frames_dir and + os.path.isfile(arguments.frames_dir)): + self.alignments.update_legacy_has_source(os.path.basename(arguments.frames_dir)) + logger.debug("Initialized %s", self.__class__.__name__) + def _find_alignments(self) -> str: + """ If an alignments folder is required and hasn't been provided, scan for a file based on + the video folder. + + Exits if an alignments file cannot be located + + Returns + ------- + str + The full path to an alignments file + """ + fname = self._args.alignments_file + frames = self._args.frames_dir + if fname and os.path.isfile(fname) and os.path.splitext(fname)[-1].lower() == ".fsa": + return fname + if fname: + logger.error("Not a valid alignments file: '%s'", fname) + sys.exit(1) + + if not frames or not os.path.exists(frames): + logger.error("Not a valid frames folder: '%s'. Can't scan for alignments.", frames) + sys.exit(1) + + fname = "alignments.fsa" + if os.path.isdir(frames) and os.path.exists(os.path.join(frames, fname)): + return fname + + if os.path.isdir(frames) or os.path.splitext(frames)[-1] not in VIDEO_EXTENSIONS: + logger.error("Can't find a valid alignments file in location: %s", frames) + sys.exit(1) + + fname = f"{os.path.splitext(frames)[0]}_{fname}" + if not os.path.exists(fname): + logger.error("Can't find a valid alignments file for video: %s", frames) + sys.exit(1) + + return fname + def process(self) -> None: """ The entry point for the Alignments tool from :mod:`lib.tools.alignments.cli`. Launches the selected alignments job. """ - if self.args.job in ("missing-alignments", "missing-frames", "multi-faces", "no-faces"): - job = Check + if self._args.job in ("missing-alignments", "missing-frames", "multi-faces", "no-faces"): + job: T.Any = Check else: - job = globals()[self.args.job.title().replace("-", "")] - job = job(self.alignments, self.args) + job = globals()[self._args.job.title().replace("-", "")] + job = job(self.alignments, self._args) logger.debug(job) job.process() diff --git a/tools/alignments/cli.py b/tools/alignments/cli.py index 4ff3a5dfdf..aaf7e7308b 100644 --- a/tools/alignments/cli.py +++ b/tools/alignments/cli.py @@ -1,12 +1,13 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ +import argparse import sys import gettext +import typing as T from lib.cli.args import FaceSwapArgs from lib.cli.actions import DirOrFileFullPaths, DirFullPaths, FileFullPaths, Radio, Slider - # LOCALES _LANG = gettext.translation("tools.alignments.cli", localedir="locales", fallback=True) _ = _LANG.gettext @@ -30,7 +31,8 @@ def get_info() -> str: return _("Alignments tool\nThis tool allows you to perform numerous actions on or using " "an alignments file against its corresponding faceset/frame source.") - def get_argument_list(self) -> dict: + @staticmethod + def get_argument_list() -> list[dict[str, T.Any]]: """ Collect the argparse argument options. Returns @@ -38,127 +40,189 @@ def get_argument_list(self) -> dict: dict The argparse command line options for processing by argparse """ - frames_dir = _(" Must Pass in a frames folder/source video file (-fr).") - faces_dir = _(" Must Pass in a faces folder (-fc).") - frames_or_faces_dir = _(" Must Pass in either a frames folder/source video file OR a" - "faces folder (-fr or -fc).") + frames_dir = _(" Must Pass in a frames folder/source video file (-r).") + faces_dir = _(" Must Pass in a faces folder (-c).") + frames_or_faces_dir = _(" Must Pass in either a frames folder/source video file OR a " + "faces folder (-r or -c).") frames_and_faces_dir = _(" Must Pass in a frames folder/source video file AND a faces " - "folder (-fr and -fc).") + "folder (-r and -c).") output_opts = _(" Use the output option (-o) to process results.") argument_list = [] - argument_list.append(dict( - opts=("-j", "--job"), - action=Radio, - type=str, - choices=("draw", "extract", "from-faces", "missing-alignments", "missing-frames", - "multi-faces", "no-faces", "remove-faces", "rename", "sort", "spatial"), - group=_("processing"), - required=True, - help=_("R|Choose which action you want to perform. NB: All actions require an " - "alignments file (-a) to be passed in." - "\nL|'draw': Draw landmarks on frames in the selected folder/video. A " - "subfolder will be created within the frames folder to hold the output.{0}" - "\nL|'extract': Re-extract faces from the source frames/video based on " - "alignment data. This is a lot quicker than re-detecting faces. Can pass in " - "the '-een' (--extract-every-n) parameter to only extract every nth frame.{1}" - "\nL|'from-faces': Generate alignment file(s) from a folder of extracted " - "faces. if the folder of faces comes from multiple sources, then multiple " - "alignments files will be created. NB: for faces which have been extracted " - "folders of source images, rather than a video, a single alignments file will " - "be created as there is no way for the process to know how many folders of " - "images were originally used. You do not need to provide an alignments file " - "path to run this job. {3}" - "\nL|'missing-alignments': Identify frames that do not exist in the alignments " - "file.{2}{0}" - "\nL|'missing-frames': Identify frames in the alignments file that do not " - "appear within the frames folder/video.{2}{0}" - "\nL|'multi-faces': Identify where multiple faces exist within the alignments " - "file.{2}{4}" - "\nL|'no-faces': Identify frames that exist within the alignment file but no " - "faces were detected.{2}{0}" - "\nL|'remove-faces': Remove deleted faces from an alignments file. The " - "original alignments file will be backed up.{3}" - "\nL|'rename' - Rename faces to correspond with their parent frame and " - "position index in the alignments file (i.e. how they are named after running " - "extract).{3}" - "\nL|'sort': Re-index the alignments from left to right. For alignments with " - "multiple faces this will ensure that the left-most face is at index 0." - "\nL|'spatial': Perform spatial and temporal filtering to smooth alignments " - "(EXPERIMENTAL!)").format(frames_dir, frames_and_faces_dir, output_opts, - faces_dir, frames_or_faces_dir))) - argument_list.append(dict( - opts=("-o", "--output"), - action=Radio, - type=str, - choices=("console", "file", "move"), - group=_("processing"), - default="console", - help=_("R|How to output discovered items ('faces' and 'frames' only):" - "\nL|'console': Print the list of frames to the screen. (DEFAULT)" - "\nL|'file': Output the list of frames to a text file (stored within the " - "source directory)." - "\nL|'move': Move the discovered items to a sub-folder within the source " - "directory."))) - argument_list.append(dict( - opts=("-a", "--alignments_file"), - action=FileFullPaths, - dest="alignments_file", - type=str, - group=_("data"), + argument_list.append({ + "opts": ("-j", "--job"), + "action": Radio, + "type": str, + "choices": ("draw", "extract", "export", "from-faces", "missing-alignments", + "missing-frames", "multi-faces", "no-faces", "remove-faces", "rename", + "sort", "spatial"), + "group": _("processing"), + "required": True, + "help": _( + "R|Choose which action you want to perform. NB: All actions require an " + "alignments file (-a) to be passed in." + "\nL|'draw': Draw landmarks on frames in the selected folder/video. A " + "subfolder will be created within the frames folder to hold the output.{0}" + "\nL|'export': Export the contents of an alignments file to a json file. Can be " + "used for editing alignment information in external tools and then re-importing " + "by using Faceswap's Extract 'Import' plugins. Note: masks and identity vectors " + "will not be included in the exported file, so will be re-generated when the json " + "file is imported back into Faceswap. All data is exported with the origin (0, 0) " + "at the top left of the canvas." + "\nL|'extract': Re-extract faces from the source frames/video based on " + "alignment data. This is a lot quicker than re-detecting faces. Can pass in " + "the '-een' (--extract-every-n) parameter to only extract every nth frame.{1}" + "\nL|'from-faces': Generate alignment file(s) from a folder of extracted " + "faces. if the folder of faces comes from multiple sources, then multiple " + "alignments files will be created. NB: for faces which have been extracted " + "from folders of source images, rather than a video, a single alignments file " + "will be created as there is no way for the process to know how many folders " + "of images were originally used. You do not need to provide an alignments file " + "path to run this job. {3}" + "\nL|'missing-alignments': Identify frames that do not exist in the alignments " + "file.{2}{0}" + "\nL|'missing-frames': Identify frames in the alignments file that do not " + "appear within the frames folder/video.{2}{0}" + "\nL|'multi-faces': Identify where multiple faces exist within the alignments " + "file.{2}{4}" + "\nL|'no-faces': Identify frames that exist within the alignment file but no " + "faces were detected.{2}{0}" + "\nL|'remove-faces': Remove deleted faces from an alignments file. The " + "original alignments file will be backed up.{3}" + "\nL|'rename' - Rename faces to correspond with their parent frame and " + "position index in the alignments file (i.e. how they are named after running " + "extract).{3}" + "\nL|'sort': Re-index the alignments from left to right. For alignments with " + "multiple faces this will ensure that the left-most face is at index 0." + "\nL|'spatial': Perform spatial and temporal filtering to smooth alignments " + "(EXPERIMENTAL!)").format(frames_dir, frames_and_faces_dir, output_opts, + faces_dir, frames_or_faces_dir)}) + argument_list.append({ + "opts": ("-o", "--output"), + "action": Radio, + "type": str, + "choices": ("console", "file", "move"), + "group": _("processing"), + "default": "console", + "help": _( + "R|How to output discovered items ('faces' and 'frames' only):" + "\nL|'console': Print the list of frames to the screen. (DEFAULT)" + "\nL|'file': Output the list of frames to a text file (stored within the " + "source directory)." + "\nL|'move': Move the discovered items to a sub-folder within the source " + "directory.")}) + argument_list.append({ + "opts": ("-a", "--alignments_file"), + "action": FileFullPaths, + "dest": "alignments_file", + "type": str, + "group": _("data"), # hacky solution to not require alignments file if creating alignments from faces: - required="from-faces" not in sys.argv, - filetypes="alignments", - help=_("Full path to the alignments file to be processed. This is required for all " - "jobs except for 'from-faces' when the alignments file will be generated in " - "the specified faces folder."))) - argument_list.append(dict( - opts=("-fc", "-faces_folder"), - action=DirFullPaths, - dest="faces_dir", - group=_("data"), - help=_("Directory containing extracted faces."))) - argument_list.append(dict( - opts=("-fr", "-frames_folder"), - action=DirOrFileFullPaths, - dest="frames_dir", - filetypes="video", - group=_("data"), - help=_("Directory containing source frames that faces were extracted from."))) - argument_list.append(dict( - opts=("-een", "--extract-every-n"), - type=int, - action=Slider, - dest="extract_every_n", - min_max=(1, 100), - default=1, - rounding=1, - group=_("extract"), - help=_("[Extract only] Extract every 'nth' frame. This option will skip frames when " - "extracting faces. For example a value of 1 will extract faces from every " - "frame, a value of 10 will extract faces from every 10th frame."))) - argument_list.append(dict( - opts=("-sz", "--size"), - type=int, - action=Slider, - min_max=(256, 1024), - rounding=64, - default=512, - group=_("extract"), - help=_("[Extract only] The output size of extracted faces."))) - argument_list.append(dict( - opts=("-m", "--min-size"), - type=int, - action=Slider, - min_max=(0, 200), - rounding=1, - default=0, - dest="min_size", - group=_("extract"), - help=_("[Extract only] Only extract faces that have been resized by this percent or " - "more to meet the specified extract size (`-sz`, `--size`). Useful for " - "excluding low-res images from a training set. Set to 0 to extract all faces. " - "Eg: For an extract size of 512px, A setting of 50 will only include faces " - "that have been resized from 256px or above. Setting to 100 will only extract " - "faces that have been resized from 512px or above. A setting of 200 will only " - "extract faces that have been downscaled from 1024px or above."))) + "required": not any(val in sys.argv for val in ["from-faces", + "-r", + "-frames_folder"]), + "filetypes": "alignments", + "help": _( + "Full path to the alignments file to be processed. If you have input a " + "'frames_dir' and don't provide this option, the process will try to find the " + "alignments file at the default location. All jobs require an alignments file " + "with the exception of 'from-faces' when the alignments file will be generated " + "in the specified faces folder.")}) + argument_list.append({ + "opts": ("-c", "-faces_folder"), + "action": DirFullPaths, + "dest": "faces_dir", + "group": ("data"), + "help": ("Directory containing extracted faces.")}) + argument_list.append({ + "opts": ("-r", "-frames_folder"), + "action": DirOrFileFullPaths, + "dest": "frames_dir", + "filetypes": "video", + "group": _("data"), + "help": _("Directory containing source frames that faces were extracted from.")}) + argument_list.append({ + "opts": ("-B", "--batch-mode"), + "action": "store_true", + "dest": "batch_mode", + "default": False, + "group": _("data"), + "help": _( + "R|Run the aligmnents tool on multiple sources. The following jobs support " + "batch mode:" + "\nL|draw, extract, from-faces, missing-alignments, missing-frames, no-faces, " + "sort, spatial." + "\nIf batch mode is selected then the other options should be set as follows:" + "\nL|alignments_file: For 'sort' and 'spatial' this should point to the parent " + "folder containing the alignments files to be processed. For all other jobs " + "this option is ignored, and the alignments files must exist at their default " + "location relative to the original frames folder/video." + "\nL|faces_dir: For 'from-faces' this should be a parent folder, containing " + "sub-folders of extracted faces from which to generate alignments files. For " + "'extract' this should be a parent folder where sub-folders will be created " + "for each extraction to be run. For all other jobs this option is ignored." + "\nL|frames_dir: For 'draw', 'extract', 'missing-alignments', 'missing-frames' " + "and 'no-faces' this should be a parent folder containing video files or sub-" + "folders of images to perform the alignments job on. The alignments file " + "should exist at the default location. For all other jobs this option is " + "ignored.")}) + argument_list.append({ + "opts": ("-N", "--extract-every-n"), + "type": int, + "action": Slider, + "dest": "extract_every_n", + "min_max": (1, 100), + "default": 1, + "rounding": 1, + "group": _("extract"), + "help": _( + "[Extract only] Extract every 'nth' frame. This option will skip frames when " + "extracting faces. For example a value of 1 will extract faces from every frame, " + "a value of 10 will extract faces from every 10th frame.")}) + argument_list.append({ + "opts": ("-z", "--size"), + "type": int, + "action": Slider, + "min_max": (256, 1024), + "rounding": 64, + "default": 512, + "group": _("extract"), + "help": _("[Extract only] The output size of extracted faces.")}) + argument_list.append({ + "opts": ("-m", "--min-size"), + "type": int, + "action": Slider, + "min_max": (0, 200), + "rounding": 1, + "default": 0, + "dest": "min_size", + "group": _("extract"), + "help": _( + "[Extract only] Only extract faces that have been resized by this percent or " + "more to meet the specified extract size (`-sz`, `--size`). Useful for " + "excluding low-res images from a training set. Set to 0 to extract all faces. " + "Eg: For an extract size of 512px, A setting of 50 will only include faces " + "that have been resized from 256px or above. Setting to 100 will only extract " + "faces that have been resized from 512px or above. A setting of 200 will only " + "extract faces that have been downscaled from 1024px or above.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-fc", ), + "type": str, + "dest": "depr_faces_folder_fc_c", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-fr", ), + "type": str, + "dest": "depr_extract-every-n_een_N", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-een", ), + "type": int, + "dest": "depr_faces_folder_fr_r", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-sz", ), + "type": int, + "dest": "depr_size_sz_z", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/alignments/jobs.py b/tools/alignments/jobs.py index f536e3a04a..578ade1ea3 100644 --- a/tools/alignments/jobs.py +++ b/tools/alignments/jobs.py @@ -1,36 +1,35 @@ #!/usr/bin/env python3 """ Tools for manipulating the alignments serialized file """ - +from __future__ import annotations import logging import os import sys -from datetime import datetime -from typing import List, Tuple, TYPE_CHECKING, Optional +import typing as T -from argparse import Namespace +from datetime import datetime -import cv2 import numpy as np from scipy import signal from sklearn import decomposition from tqdm import tqdm -from lib.align import DetectedFace, _EXTRACT_RATIOS -from lib.align.alignments import _VERSION -from lib.image import (encode_image, generate_thumbnail, ImagesSaver, - read_image_meta_batch, update_existing_metadata) -from plugins.extract.pipeline import Extractor, ExtractMedia -from scripts.fsmedia import Alignments +from lib.logger import parse_class_init +from lib.serializer import get_serializer +from lib.utils import FaceswapError -from .media import ExtractedFaces, Faces, Frames +from .media import Faces, Frames +from .jobs_faces import FaceToFile -if TYPE_CHECKING: +if T.TYPE_CHECKING: + from collections.abc import Generator + from argparse import Namespace + from lib.align.alignments import AlignmentFileDict, PNGHeaderDict from .media import AlignmentData -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) -class Check(): +class Check: """ Frames and faces checking tasks. Parameters @@ -40,11 +39,11 @@ class Check(): arguments: :class:`argparse.Namespace` The command line arguments that have called this job """ - def __init__(self, alignments, arguments): - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None: + logger.debug(parse_class_init(locals())) self._alignments = alignments self._job = arguments.job - self._type = None + self._type: T.Literal["faces", "frames"] | None = None self._is_video = False # Set when getting items self._output = arguments.output self._source_dir = self._get_source_dir(arguments) @@ -54,8 +53,19 @@ def __init__(self, alignments, arguments): self.output_message = "" logger.debug("Initialized %s", self.__class__.__name__) - def _get_source_dir(self, arguments): - """ Set the correct source folder """ + def _get_source_dir(self, arguments: Namespace) -> str: + """ Set the correct source folder + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments for the Alignments tool + + Returns + ------- + str + Full path to the source folder + """ if (hasattr(arguments, "faces_dir") and arguments.faces_dir and hasattr(arguments, "frames_dir") and arguments.frames_dir): logger.error("Only select a source frames (-fr) or source faces (-fc) folder") @@ -72,21 +82,38 @@ def _get_source_dir(self, arguments): logger.debug("type: '%s', source_dir: '%s'", self._type, source_dir) return source_dir - def _get_items(self): - """ Set the correct items to process """ - items = globals()[self._type.title()](self._source_dir) + def _get_items(self) -> list[dict[str, str]] | list[tuple[str, PNGHeaderDict]]: + """ Set the correct items to process + + Returns + ------- + list + Sorted list of dictionaries for either faces or frames. If faces the dictionaries + have the current filename as key, with the header source data as value. If frames + the dictionaries will contain the keys 'frame_fullname', 'frame_name', 'extension'. + """ + assert self._type is not None + items: Frames | Faces = globals()[self._type.title()](self._source_dir) self._is_video = items.is_video - return items.file_list_sorted + return T.cast(list[dict[str, str]] | list[tuple[str, "PNGHeaderDict"]], + items.file_list_sorted) - def process(self): + def process(self) -> None: """ Process the frames check against the alignments file """ + assert self._type is not None logger.info("[CHECK %s]", self._type.upper()) items_output = self._compile_output() + + if self._type == "faces": + filelist = T.cast(list[tuple[str, "PNGHeaderDict"]], self._items) + check_update = FaceToFile(self._alignments, [val[1] for val in filelist]) + if check_update(): + self._alignments.save() + self._output_results(items_output) - def _validate(self): - """ Check that the selected type is valid for - selected task and job """ + def _validate(self) -> None: + """ Check that the selected type is valid for selected task and job """ if self._job == "missing-frames" and self._output == "move": logger.warning("Missing_frames was selected with move output, but there will " "be nothing to move. Defaulting to output: console") @@ -96,74 +123,131 @@ def _validate(self): "supported for 'multi-faces'") sys.exit(1) - def _compile_output(self): - """ Compile list of frames that meet criteria """ + def _compile_output(self) -> list[str] | list[tuple[str, int]]: + """ Compile list of frames that meet criteria + + Returns + ------- + list + List of filenames or filenames and face indices for the selected criteria + """ action = self._job.replace("-", "_") processor = getattr(self, f"_get_{action}") logger.debug("Processor: %s", processor) return [item for item in processor()] # pylint:disable=unnecessary-comprehension - def _get_no_faces(self): - """ yield each frame that has no face match in alignments file """ + def _get_no_faces(self) -> Generator[str, None, None]: + """ yield each frame that has no face match in alignments file + + Yields + ------ + str + The frame name of any frames which have no faces + """ self.output_message = "Frames with no faces" - for frame in tqdm(self._items, desc=self.output_message): - logger.trace(frame) + for frame in tqdm(T.cast(list[dict[str, str]], self._items), + desc=self.output_message, + leave=False): + logger.trace(frame) # type:ignore frame_name = frame["frame_fullname"] if not self._alignments.frame_has_faces(frame_name): logger.debug("Returning: '%s'", frame_name) yield frame_name - def _get_multi_faces(self): - """ yield each frame or face that has multiple faces - matched in alignments file """ + def _get_multi_faces(self) -> (Generator[str, None, None] | + Generator[tuple[str, int], None, None]): + """ yield each frame or face that has multiple faces matched in alignments file + + Yields + ------ + str or tuple + The frame name of any frames which have multiple faces and potentially the face id + """ process_type = getattr(self, f"_get_multi_faces_{self._type}") for item in process_type(): yield item - def _get_multi_faces_frames(self): - """ Return Frames that contain multiple faces """ + def _get_multi_faces_frames(self) -> Generator[str, None, None]: + """ Return Frames that contain multiple faces + + Yields + ------ + str + The frame name of any frames which have multiple faces + """ self.output_message = "Frames with multiple faces" - for item in tqdm(self._items, desc=self.output_message): + for item in tqdm(T.cast(list[dict[str, str]], self._items), + desc=self.output_message, + leave=False): filename = item["frame_fullname"] if not self._alignments.frame_has_multiple_faces(filename): continue - logger.trace("Returning: '%s'", filename) + logger.trace("Returning: '%s'", filename) # type:ignore yield filename - def _get_multi_faces_faces(self): - """ Return Faces when there are multiple faces in a frame """ + def _get_multi_faces_faces(self) -> Generator[tuple[str, int], None, None]: + """ Return Faces when there are multiple faces in a frame + + Yields + ------ + tuple + The frame name and the face id of any frames which have multiple faces + """ self.output_message = "Multiple faces in frame" - for item in tqdm(self._items, desc=self.output_message): - if not self._alignments.frame_has_multiple_faces(item["source_filename"]): + for item in tqdm(T.cast(list[tuple[str, "PNGHeaderDict"]], self._items), + desc=self.output_message, + leave=False): + src = item[1]["source"] + if not self._alignments.frame_has_multiple_faces(src["source_filename"]): continue - retval = (item["current_filename"], item["face_index"]) - logger.trace("Returning: '%s'", retval) + retval = (item[0], src["face_index"]) + logger.trace("Returning: '%s'", retval) # type:ignore yield retval - def _get_missing_alignments(self): - """ yield each frame that does not exist in alignments file """ + def _get_missing_alignments(self) -> Generator[str, None, None]: + """ yield each frame that does not exist in alignments file + + Yields + ------ + str + The frame name of any frames missing alignments + """ self.output_message = "Frames missing from alignments file" exclude_filetypes = set(["yaml", "yml", "p", "json", "txt"]) - for frame in tqdm(self._items, desc=self.output_message): + for frame in tqdm(T.cast(dict[str, str], self._items), + desc=self.output_message, + leave=False): frame_name = frame["frame_fullname"] if (frame["frame_extension"] not in exclude_filetypes and not self._alignments.frame_exists(frame_name)): logger.debug("Returning: '%s'", frame_name) yield frame_name - def _get_missing_frames(self): - """ yield each frame in alignments that does - not have a matching file """ + def _get_missing_frames(self) -> Generator[str, None, None]: + """ yield each frame in alignments that does not have a matching file + + Yields + ------ + str + The frame name of any frames in alignments with no matching file + """ self.output_message = "Missing frames that are in alignments file" - frames = set(item["frame_fullname"] for item in self._items) - for frame in tqdm(self._alignments.data.keys(), desc=self.output_message): + frames = set(item["frame_fullname"] for item in T.cast(list[dict[str, str]], self._items)) + for frame in tqdm(self._alignments.data.keys(), desc=self.output_message, leave=False): if frame not in frames: logger.debug("Returning: '%s'", frame) yield frame - def _output_results(self, items_output): - """ Output the results in the requested format """ - logger.trace("items_output: %s", items_output) + def _output_results(self, items_output: list[str] | list[tuple[str, int]]) -> None: + """ Output the results in the requested format + + Parameters + ---------- + items_output + The list of frame names, and potentially face ids, of any items which met the + selection criteria + """ + logger.trace("items_output: %s", items_output) # type:ignore if self._output == "move" and self._is_video and self._type == "frames": logger.warning("Move was selected with an input video. This is not possible so " "falling back to console output") @@ -176,33 +260,54 @@ def _output_results(self, items_output): return if self._job == "multi-faces" and self._type == "faces": # Strip the index for printed/file output - items_output = [item[0] for item in items_output] + final_output = [item[0] for item in items_output] + else: + final_output = T.cast(list[str], items_output) output_message = "-----------------------------------------------\r\n" - output_message += f" {self.output_message} ({len(items_output)})\r\n" + output_message += f" {self.output_message} ({len(final_output)})\r\n" output_message += "-----------------------------------------------\r\n" - output_message += "\r\n".join(items_output) + output_message += "\r\n".join(final_output) if self._output == "console": for line in output_message.splitlines(): logger.info(line) if self._output == "file": - self.output_file(output_message, len(items_output)) + self.output_file(output_message, len(final_output)) + + def _get_output_folder(self) -> str: + """ Return output folder. Needs to be in the root if input is a video and processing + frames - def _get_output_folder(self): - """ Return output folder. Needs to be in the root if input is a - video and processing frames """ + Returns + ------- + str + Full path to the output folder + """ if self._is_video and self._type == "frames": return os.path.dirname(self._source_dir) return self._source_dir - def _get_filename_prefix(self): - """ Video name needs to be prefixed to filename if input is a - video and processing frames """ + def _get_filename_prefix(self) -> str: + """ Video name needs to be prefixed to filename if input is a video and processing frames + + Returns + ------- + str + The common filename prefix to use + """ if self._is_video and self._type == "frames": return f"{os.path.basename(self._source_dir)}_" return "" - def output_file(self, output_message, items_discovered): - """ Save the output to a text file in the frames directory """ + def output_file(self, output_message: str, items_discovered: int) -> None: + """ Save the output to a text file in the frames directory + + Parameters + ---------- + output_message: str + The message to write out to file + items_discovered: int + The number of items which matched the criteria + """ now = datetime.now().strftime("%Y%m%d_%H%M%S") dst_dir = self._get_output_folder() filename = (f"{self._get_filename_prefix()}{self.output_message.replace(' ', '_').lower()}" @@ -212,8 +317,14 @@ def output_file(self, output_message, items_discovered): with open(output_file, "w", encoding="utf8") as f_output: f_output.write(output_message) - def _move_file(self, items_output): - """ Move the identified frames to a new sub folder """ + def _move_file(self, items_output: list[str] | list[tuple[str, int]]) -> None: + """ Move the identified frames to a new sub folder + + Parameters + ---------- + items_output: list + List of items to move + """ now = datetime.now().strftime("%Y%m%d_%H%M%S") folder_name = (f"{self._get_filename_prefix()}" f"{self.output_message.replace(' ','_').lower()}_{now}") @@ -225,8 +336,16 @@ def _move_file(self, items_output): logger.debug("Move function: %s", move) move(output_folder, items_output) - def _move_frames(self, output_folder, items_output): - """ Move frames into single sub folder """ + def _move_frames(self, output_folder: str, items_output: list[str]) -> None: + """ Move frames into single sub folder + + Parameters + ---------- + output_folder: str + The folder to move the output to + items_output: list + List of items to move + """ logger.info("Moving %s frame(s) to '%s'", len(items_output), output_folder) for frame in items_output: src = os.path.join(self._source_dir, frame) @@ -234,9 +353,16 @@ def _move_frames(self, output_folder, items_output): logger.debug("Moving: '%s' to '%s'", src, dst) os.rename(src, dst) - def _move_faces(self, output_folder, items_output): - """ Make additional sub folders for each face that appears - Enables easier manual sorting """ + def _move_faces(self, output_folder: str, items_output: list[tuple[str, int]]) -> None: + """ Make additional sub folders for each face that appears Enables easier manual sorting + + Parameters + ---------- + output_folder: str + The folder to move the output to + items_output: list + List of items and face indices to move + """ logger.info("Moving %s faces(s) to '%s'", len(items_output), output_folder) for frame, idx in items_output: src = os.path.join(self._source_dir, frame) @@ -249,865 +375,81 @@ def _move_faces(self, output_folder, items_output): os.rename(src, dst) -class Draw(): # pylint:disable=too-few-public-methods - """ Draws annotations onto original frames and saves into a sub-folder next to the original - frames. - - Parameters - --------- - alignments: :class:`tools.alignments.media.AlignmentsData` - The loaded alignments corresponding to the frames to be annotated - arguments: :class:`argparse.Namespace` - The command line arguments that have called this job - """ - def __init__(self, alignments, arguments): - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) - self._alignments = alignments - self._frames = Frames(arguments.frames_dir) - self._output_folder = self._set_output() - self._mesh_areas = dict(mouth=(48, 68), - right_eyebrow=(17, 22), - left_eyebrow=(22, 27), - right_eye=(36, 42), - left_eye=(42, 48), - nose=(27, 36), - jaw=(0, 17), - chin=(8, 11)) - logger.debug("Initialized %s", self.__class__.__name__) - - def _set_output(self): - """ Set the output folder path. - - If annotating a folder of frames, output will be placed in a sub folder within the frames - folder. If annotating a video, output will be a folder next to the original video. - - Returns - ------- - str - Full path to the output folder - - """ - now = datetime.now().strftime("%Y%m%d_%H%M%S") - folder_name = f"drawn_landmarks_{now}" - if self._frames.is_video: - dest_folder = os.path.dirname(self._frames.folder) - else: - dest_folder = self._frames.folder - output_folder = os.path.join(dest_folder, folder_name) - logger.debug("Creating folder: '%s'", output_folder) - os.makedirs(output_folder) - return output_folder - - def process(self): - """ Runs the process to draw face annotations onto original source frames. """ - logger.info("[DRAW LANDMARKS]") # Tidy up cli output - frames_drawn = 0 - for frame in tqdm(self._frames.file_list_sorted, desc="Drawing landmarks"): - frame_name = frame["frame_fullname"] - - if not self._alignments.frame_exists(frame_name): - logger.verbose("Skipping '%s' - Alignments not found", frame_name) - continue - - self._annotate_image(frame_name) - frames_drawn += 1 - logger.info("%s Frame(s) output", frames_drawn) - - def _annotate_image(self, frame_name): - """ Annotate the frame with each face that appears in the alignments file. - - Parameters - ---------- - frame_name: str - The full path to the original frame - """ - logger.trace("Annotating frame: '%s'", frame_name) - image = self._frames.load_image(frame_name) - - for idx, alignment in enumerate(self._alignments.get_faces_in_frame(frame_name)): - face = DetectedFace() - face.from_alignment(alignment, image=image) - # Bounding Box - cv2.rectangle(image, (face.left, face.top), (face.right, face.bottom), (255, 0, 0), 1) - self._annotate_landmarks(image, np.rint(face.landmarks_xy).astype("int32")) - self._annotate_extract_boxes(image, face, idx) - self._annotate_pose(image, face) # Pose (head is still loaded) - - self._frames.save_image(self._output_folder, frame_name, image) - - def _annotate_landmarks(self, image, landmarks): - """ Annotate the extract boxes onto the frame. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The frame that extract boxes are to be annotated on to - landmarks: :class:`numpy.ndarray` - The 68 point landmarks that are to be annotated onto the frame - index: int - The face index for the given face - """ - # Mesh - for area, indices in self._mesh_areas.items(): - fill = area in ("right_eye", "left_eye", "mouth") - cv2.polylines(image, [landmarks[indices[0]:indices[1]]], fill, (255, 255, 0), 1) - # Landmarks - for (pos_x, pos_y) in landmarks: - cv2.circle(image, (pos_x, pos_y), 1, (0, 255, 255), -1) - - @classmethod - def _annotate_extract_boxes(cls, image, face, index): - """ Annotate the mesh and landmarks boxes onto the frame. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The frame that mesh and landmarks are to be annotated on to - face: :class:`lib.align.AlignedFace` - The aligned face - """ - for area in ("face", "head"): - face.load_aligned(image, centering=area, force=True) - color = (0, 255, 0) if area == "face" else (0, 0, 255) - top_left = face.aligned.original_roi[0] # pylint:disable=unsubscriptable-object - top_left = (top_left[0], top_left[1] - 10) - cv2.putText(image, str(index), top_left, cv2.FONT_HERSHEY_DUPLEX, 1.0, color, 1) - cv2.polylines(image, [face.aligned.original_roi], True, color, 1) - - @classmethod - def _annotate_pose(cls, image, face): - """ Annotate the pose onto the frame. +class Export: + """ Export alignments from a Faceswap .fsa file to a json formatted file. Parameters - ---------- - image: :class:`numpy.ndarray` - The frame that pose is to be annotated on to - face: :class:`lib.align.AlignedFace` - The aligned face loaded for head centering - """ - center = np.int32((face.aligned.size / 2, face.aligned.size / 2)).reshape(1, 2) - center = np.rint(face.aligned.transform_points(center, invert=True)).astype("int32") - points = face.aligned.pose.xyz_2d * face.aligned.size - points = np.rint(face.aligned.transform_points(points, invert=True)).astype("int32") - cv2.line(image, tuple(center), tuple(points[1]), (0, 255, 0), 2) - cv2.line(image, tuple(center), tuple(points[0]), (255, 0, 0), 2) - cv2.line(image, tuple(center), tuple(points[2]), (0, 0, 255), 2) - - -class Extract(): # pylint:disable=too-few-public-methods - """ Re-extract faces from source frames based on Alignment data - - Parameters ---------- alignments: :class:`tools.lib_alignments.media.AlignmentData` The alignments data loaded from an alignments file for this rename job arguments: :class:`argparse.Namespace` - The :mod:`argparse` arguments as passed in from :mod:`tools.py` + The :mod:`argparse` arguments as passed in from :mod:`tools.py`. Unused """ - def __init__(self, alignments: "AlignmentData", arguments: Namespace) -> None: - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) - self._arguments = arguments + def __init__(self, + alignments: AlignmentData, + arguments: Namespace) -> None: # pylint:disable=unused-argument + logger.debug(parse_class_init(locals())) self._alignments = alignments - self._is_legacy = self._alignments.version == 1.0 # pylint:disable=protected-access - self._mask_pipeline = None - self._faces_dir = arguments.faces_dir - self._min_size = self._get_min_size(arguments.size, arguments.min_size) - - self._frames = Frames(arguments.frames_dir, self._get_count()) - self._extracted_faces = ExtractedFaces(self._frames, - self._alignments, - size=arguments.size) - self._saver = None + self._serializer = get_serializer("json") + self._output_file = self._get_output_file() logger.debug("Initialized %s", self.__class__.__name__) - @classmethod - def _get_min_size(cls, extract_size: int, min_size: int) -> int: - """ Obtain the minimum size that a face has been resized from to be included as a valid - extract. - - Parameters - ---------- - extract_size: int - The requested size of the extracted images - min_size: int - The percentage amount that has been supplied for valid faces (as a percentage of - extract size) - - Returns - ------- - int - The minimum size, in pixels, that a face is resized from to be considered valid - """ - retval = 0 if min_size == 0 else max(4, int(extract_size * (min_size / 100.))) - logger.debug("Extract size: %s, min percentage size: %s, min_size: %s", - extract_size, min_size, retval) - return retval - - def _get_count(self) -> Optional[int]: - """ If the alignments file has been run through the manual tool, then it will hold video - meta information, meaning that the count of frames in the alignment file can be relied - on to be accurate. - - Returns - ------- - int or ``None`` - For video input which contain video meta-data in the alignments file then the count of - frames is returned. In all other cases ``None`` is returned - """ - has_meta = all(val is not None for val in self._alignments.video_meta_data.values()) - retval = len(self._alignments.video_meta_data["pts_time"]) if has_meta else None - logger.debug("Frame count from alignments file: (has_meta: %s, %s", has_meta, retval) - return retval - - def process(self) -> None: - """ Run the re-extraction from Alignments file process""" - logger.info("[EXTRACT FACES]") # Tidy up cli output - self._check_folder() - if self._is_legacy: - self._legacy_check() - self._saver = ImagesSaver(self._faces_dir, as_bytes=True) - - if self._min_size > 0: - logger.info("Only selecting faces that have been resized from a minimum resolution " - "of %spx", self._min_size) - - self._export_faces() - - def _check_folder(self) -> None: - """ Check that the faces folder doesn't pre-exist and create. """ - err = None - if not self._faces_dir: - err = "ERROR: Output faces folder not provided." - elif not os.path.isdir(self._faces_dir): - logger.debug("Creating folder: '%s'", self._faces_dir) - os.makedirs(self._faces_dir) - elif os.listdir(self._faces_dir): - err = f"ERROR: Output faces folder should be empty: '{self._faces_dir}'" - if err: - logger.error(err) - sys.exit(0) - logger.verbose("Creating output folder at '%s'", self._faces_dir) - - def _legacy_check(self) -> None: - """ Check whether the alignments file was created with the legacy extraction method. - - If so, force user to re-extract all faces if any options have been specified, otherwise - raise the appropriate warnings and set the legacy options. - """ - if self._min_size > 0 or self._arguments.extract_every_n != 1: - logger.warning("This alignments file was generated with the legacy extraction method.") - logger.warning("You should run this extraction job, but with 'min_size' set to 0 and " - "'extract-every-n' set to 1 to update the alignments file.") - logger.warning("You can then re-run this extraction job with your chosen options.") - sys.exit(0) - - maskers = ["components", "extended"] - nn_masks = [mask for mask in list(self._alignments.mask_summary) if mask not in maskers] - logtype = logger.warning if nn_masks else logger.info - logtype("This alignments file was created with the legacy extraction method and will be " - "updated.") - logtype("Faces will be extracted using the new method and landmarks based masks will be " - "regenerated.") - if nn_masks: - logtype("However, the NN based masks '%s' will be cropped to the legacy extraction " - "method, so you may want to run the mask tool to regenerate these " - "masks.", "', '".join(nn_masks)) - self._mask_pipeline = Extractor(None, None, maskers, multiprocess=True) - self._mask_pipeline.launch() - # Update alignments versioning - self._alignments._version = _VERSION # pylint:disable=protected-access - - def _export_faces(self) -> None: - """ Export the faces to the output folder. """ - extracted_faces = 0 - skip_list = self._set_skip_list() - count = self._frames.count if skip_list is None else self._frames.count - len(skip_list) - - for filename, image in tqdm(self._frames.stream(skip_list=skip_list), - total=count, desc="Saving extracted faces"): - frame_name = os.path.basename(filename) - if not self._alignments.frame_exists(frame_name): - logger.verbose("Skipping '%s' - Alignments not found", frame_name) - continue - extracted_faces += self._output_faces(frame_name, image) - if self._is_legacy and extracted_faces != 0 and self._min_size == 0: - self._alignments.save() - logger.info("%s face(s) extracted", extracted_faces) - - def _set_skip_list(self) -> Optional[List[int]]: - """ Set the indices for frames that should be skipped based on the `extract_every_n` - command line option. - - Returns - ------- - list or ``None`` - A list of indices to be skipped if extract_every_n is not `1` otherwise - returns ``None`` - """ - skip_num = self._arguments.extract_every_n - if skip_num == 1: - logger.debug("Not skipping any frames") - return None - skip_list = [] - for idx, item in enumerate(self._frames.file_list_sorted): - if idx % skip_num != 0: - logger.trace("Adding image '%s' to skip list due to extract_every_n = %s", - item["frame_fullname"], skip_num) - skip_list.append(idx) - logger.debug("Adding skip list: %s", skip_list) - return skip_list - - def _output_faces(self, filename: str, image: np.ndarray) -> int: - """ For each frame save out the faces - - Parameters - ---------- - filename: str - The filename (without the full path) of the current frame - image: :class:`numpy.ndarray` - The full frame that faces are to be extracted from - - Returns - ------- - int - The total number of faces that have been extracted - """ - logger.trace("Outputting frame: %s", filename) - face_count = 0 - frame_name = os.path.splitext(filename)[0] - faces = self._select_valid_faces(filename, image) - if not faces: - return face_count - if self._is_legacy: - faces = self._process_legacy(filename, image, faces) - - for idx, face in enumerate(faces): - output = f"{frame_name}_{idx}.png" - meta = dict(alignments=face.to_png_meta(), - source=dict(alignments_version=self._alignments.version, - original_filename=output, - face_index=idx, - source_filename=filename, - source_is_video=self._frames.is_video, - source_frame_dims=image.shape[:2])) - self._saver.save(output, encode_image(face.aligned.face, ".png", metadata=meta)) - if self._min_size == 0 and self._is_legacy: - face.thumbnail = generate_thumbnail(face.aligned.face, size=96, quality=60) - self._alignments.data[filename]["faces"][idx] = face.to_alignment() - face_count += 1 - self._saver.close() - return face_count - - def _select_valid_faces(self, frame: str, image: np.ndarray) -> List[DetectedFace]: - """ Return the aligned faces from a frame that meet the selection criteria, - - Parameters - ---------- - frame: str - The filename (without the full path) of the current frame - image: :class:`numpy.ndarray` - The full frame that faces are to be extracted from - - Returns - ------- - list: - List of valid :class:`lib,align.DetectedFace` objects - """ - faces = self._extracted_faces.get_faces_in_frame(frame, image=image) - if self._min_size == 0: - valid_faces = faces - else: - sizes = self._extracted_faces.get_roi_size_for_frame(frame) - valid_faces = [faces[idx] for idx, size in enumerate(sizes) - if size >= self._min_size] - logger.trace("frame: '%s', total_faces: %s, valid_faces: %s", - frame, len(faces), len(valid_faces)) - return valid_faces - - def _process_legacy(self, - filename: str, - image: np.ndarray, - detected_faces: List[DetectedFace]) -> List[DetectedFace]: - """ Process legacy face extractions to new extraction method. - - Updates stored masks to new extract size - - Parameters - ---------- - filename: str - The current frame filename - image: :class:`numpy.ndarray` - The current image the contains the faces - detected_faces: list - list of :class:`lib.align.DetectedFace` objects for the current frame + def _get_output_file(self) -> str: + """ Obtain the name of an output file. If a file of the request name exists, then append a + digit to the end until a unique filename is found Returns ------- - list - The updated list of :class:`lib.align.DetectedFace` objects for the current frame - """ - # Update landmarks based masks for face centering - mask_item = ExtractMedia(filename, image, detected_faces=detected_faces) - self._mask_pipeline.input_queue.put(mask_item) - faces = next(self._mask_pipeline.detected_faces()).detected_faces - - # Pad and shift Neural Network based masks to face centering - for face in faces: - self._pad_legacy_masks(face) - return faces - - @classmethod - def _pad_legacy_masks(cls, detected_face: DetectedFace) -> None: - """ Recenter legacy Neural Network based masks from legacy centering to face centering - and pad accordingly. - - Update the masks back into the detected face objects. - - Parameters - ---------- - detected_face: :class:`lib.align.DetectedFace` - The detected face to update the masks for - """ - offset = detected_face.aligned.pose.offset["face"] - for name, mask in detected_face.mask.items(): # Re-center mask and pad to face size - if name in ("components", "extended"): - continue - old_mask = mask.mask.astype("float32") / 255.0 - size = old_mask.shape[0] - new_size = int(size + (size * _EXTRACT_RATIOS["face"]) / 2) - - shift = np.rint(offset * (size - (size * _EXTRACT_RATIOS["face"]))).astype("int32") - pos = np.array([(new_size // 2 - size // 2) - shift[1], - (new_size // 2) + (size // 2) - shift[1], - (new_size // 2 - size // 2) - shift[0], - (new_size // 2) + (size // 2) - shift[0]]) - bounds = np.array([max(0, pos[0]), min(new_size, pos[1]), - max(0, pos[2]), min(new_size, pos[3])]) - - slice_in = [slice(0 - (pos[0] - bounds[0]), size - (pos[1] - bounds[1])), - slice(0 - (pos[2] - bounds[2]), size - (pos[3] - bounds[3]))] - slice_out = [slice(bounds[0], bounds[1]), slice(bounds[2], bounds[3])] - - new_mask = np.zeros((new_size, new_size, 1), dtype="float32") - new_mask[slice_out[0], slice_out[1], :] = old_mask[slice_in[0], slice_in[1], :] - - mask.replace_mask(new_mask) - # Get the affine matrix from recently generated components mask - # pylint:disable=protected-access - mask._affine_matrix = detected_face.mask["components"].affine_matrix - - -class FromFaces(): # pylint:disable=too-few-public-methods - """ Scan a folder of Faceswap Extracted Faces and re-create the associated alignments file(s) - - Parameters - ---------- - alignments: NoneType - Parameter included for standard job naming convention, but not used for this process. - arguments: :class:`argparse.Namespace` - The :mod:`argparse` arguments as passed in from :mod:`tools.py` - """ - def __init__(self, alignments: None, arguments: Namespace) -> None: - logger.debug("Initializing %s: (alignments: %s, arguments: %s)", - self.__class__.__name__, alignments, arguments) - self._faces_dir = arguments.faces_dir - self._filelist = self._get_filenames() - logger.debug("Initialized %s", self.__class__.__name__) - - def _get_filenames(self) -> List[str]: - """ Obtain the full path to all filenames in the specified faces folder. - - Only png files will be returned, any other files will be ignored. An error is output if - the returned filelist is not valid - - Returns - ------- - list - Full path list to face png files + str + Full path to an output json file """ - err = None - if not self._faces_dir: - err = "A faces folder must be provided." - elif not os.path.isdir(self._faces_dir): - err = f"The Faces location '{self._faces_dir}' does not exit" - else: - filelist = [os.path.join(self._faces_dir, fname) - for fname in os.listdir(self._faces_dir) - if os.path.splitext(fname.lower())[1] == ".png"] - if not err and not filelist: - err = "Faces folder should contain Faceswap extracted PNG files" - if err: - logger.error(err) - sys.exit(0) - logger.debug("Collected %s png images from folder '%s'", len(filelist), self._faces_dir) - return filelist - - def process(self) -> None: - """ Run the job to read faces from a folder to create alignments file(s). """ - logger.info("[CREATE ALIGNMENTS FROM FACES]") # Tidy up cli output - skip_count = 0 - d_align = {} - for filename, meta in tqdm(read_image_meta_batch(self._filelist), - desc="Generating Alignments", - total=len(self._filelist), - leave=False): - - if "itxt" not in meta or "alignments" not in meta["itxt"]: - logger.verbose("skipping invalid file: '%s'", filename) - skip_count += 1 - continue - - align_fname = self._get_alignments_filename(meta["itxt"]["source"]) - source_name, f_idx, alignment = self._extract_alignment(meta) - full_info = (f_idx, alignment, filename, meta["itxt"]["source"]) - - d_align.setdefault(align_fname, {}).setdefault(source_name, []).append(full_info) - - alignments = self._sort_alignments(d_align) - self._save_alignments(alignments) - if skip_count > 1: - logger.warning("%s of %s files skipped that do not contain valid alignment data", - skip_count, len(self._filelist)) - logger.warning("Run the process in verbose mode to see which files were skipped") + in_file = self._alignments.file + base_filename = f"{os.path.splitext(in_file)[0]}_export" + out_file = f"{base_filename}.json" + idx = 1 + while True: + if not os.path.exists(out_file): + break + logger.debug("Output file exists: '%s'", out_file) + out_file = f"{base_filename}_{idx}.json" + idx += 1 + logger.debug("Setting output file to '%s'", out_file) + return out_file @classmethod - def _get_alignments_filename(cls, source_data: dict) -> str: - """ Obtain the name of the alignments file from the source information contained within the - PNG metadata. + def _format_face(cls, face: AlignmentFileDict) -> dict[str, list[int] | list[list[float]]]: + """ Format the relevant keys from an alignment file's face into the correct format for + export/import Parameters ---------- - source_data: dict - The source information contained within a Faceswap extracted PNG + face: :class:`~lib.align.alignments.AlignmentFileDict` + The alignment dictionary for a face to process Returns ------- - str: - If the face was generated from a video file, the filename will be - `'_alignments.fsa'`. If it was extracted from an image file it will be - `'alignments.fsa'` + dict[str, list[int] | list[list[float]]] + The face formatted for exporting to a json file """ - is_video = source_data["source_is_video"] - src_name = source_data["source_filename"] - prefix = f"{src_name.rpartition('_')[0]}_" if is_video else "" - retval = f"{prefix}alignments.fsa" - logger.trace("Extracted alignments file filename: '%s'", retval) + lms = face["landmarks_xy"] + assert isinstance(lms, np.ndarray) + retval = {"detected": [int(round(face["x"], 0)), + int(round(face["y"], 0)), + int(round(face["x"] + face["w"], 0)), + int(round(face["y"] + face["h"], 0))], + "landmarks_2d": lms.tolist()} return retval - def _extract_alignment(self, metadata: dict) -> Tuple[str, int, dict]: - """ Extract alignment data from a PNG image's itxt header. - - Formats the landmarks into a numpy array and adds in mask centering information if it is - from an older extract. - - Parameters - ---------- - metadata: dict - An extracted faces PNG Header data - - Returns - ------- - tuple - The alignment's source frame name in position 0. The index of the face within the - alignment file in position 1. The alignment data correctly formatted for writing to an - alignments file in positin 2 - """ - alignment = metadata["itxt"]["alignments"] - alignment["landmarks_xy"] = np.array(alignment["landmarks_xy"], dtype="float32") - - src = metadata["itxt"]["source"] - frame_name = src["source_filename"] - face_index = int(src["face_index"]) - version = src["alignments_version"] - - if version < 2.2: - logger.trace("Updating mask centering for frame '%s', face index: %s, version: %s", - frame_name, face_index, version) - self._update_mask_centering(alignment) - - logger.trace("Extracted alignment for frame: '%s', face index: %s", frame_name, face_index) - return frame_name, face_index, alignment - - @classmethod - def _update_mask_centering(cls, alignment: dict) -> None: - """ Prior to alignment version 2.2 all masks were stored with face centering. - - Update the existing masks with correct centering parameter. - - Parameters - ---------- - alignment: dict - The alignment for the face to have the mask centering parameter updated - """ - if "mask" not in alignment: - alignment["mask"] = {} - for mask in alignment["mask"].values(): - mask["stored_centering"] = "face" - - def _sort_alignments(self, alignments: dict) -> dict: - """ Sort the faces into face index order as they appeared in the original alignments file. - - If the face index stored in the png header does not match it's position in the alignments - file (i.e. A face has been removed from a frame) then update the header of the - corresponding png to the correct index as exists in the newly created alignments file. - - Parameters - ---------- - alignments: dict - The unsorted alignments file(s) as generated from the face PNG headers, including the - face index of the face within it's respective frame, the original face filename and - the orignal face header source information - - Returns - ------- - dict - The alignments file dictionaries sorted into the correct face order, ready for savind - """ - logger.info("Sorting and checking faces...") - aln_sorted = {} - for fname, frames in alignments.items(): - this_file = {} - for frame in tqdm(sorted(frames), desc=f"Sorting {fname}", leave=False): - this_file[frame] = [] - for real_idx, (f_id, alignment, f_path, f_src) in enumerate(sorted(frames[frame])): - if real_idx != f_id: - self._update_png_header(f_path, real_idx, alignment, f_src) - this_file[frame].append(alignment) - aln_sorted[fname] = this_file - return aln_sorted - - @classmethod - def _update_png_header(cls, - face_path: str, - new_index: int, - alignment: dict, - source_info: dict) -> None: - """ Update the PNG header for faces where the stored index does not correspond with the - alignments file. This can occur when frames with multiple faces have had some faces deleted - from the faces folder. - - Updates the original filename and index in the png header. - - Parameters - ---------- - face_path: str - Full path to the saved face image that requires updating - new_index: int - The new index as it appears in the newly generated alignments file - alignment: dict - The alignment information to store in the png header - source_info: dict - The face source information as extracted from the original face png file - """ - face = DetectedFace() - face.from_alignment(alignment) - new_filename = f"{os.path.splitext(source_info['source_filename'])[0]}_{new_index}.png" - - logger.trace("Updating png header for '%s': (face index from %s to %s, original filename " - "from '%s' to '%s'", face_path, source_info["face_index"], new_index, - source_info["original_filename"], new_filename) - - source_info["face_index"] = new_index - source_info["original_filename"] = new_filename - meta = dict(alignments=face.to_png_meta(), source=source_info) - update_existing_metadata(face_path, meta) - - def _save_alignments(self, all_alignments: dict) -> None: - """ Save the newely generated alignments file(s). - - If an alignments file already exists in the source faces folder, back it up rather than - overwriting - - Parameters - ---------- - all_alignments: dict - The alignment(s) dictionaries found in the faces folder. Alignment filename as key, - corresponding alignments as value. - """ - for fname, alignments in all_alignments.items(): - alignments_path = os.path.join(self._faces_dir, fname) - dummy_args = Namespace(alignments_path=alignments_path) - aln = Alignments(dummy_args, is_extract=True) - aln._data = alignments # pylint:disable=protected-access - aln.backup() - aln.save() - - -class RemoveFaces(): # pylint:disable=too-few-public-methods - """ Remove items from alignments file. - - Parameters - --------- - alignments: :class:`tools.alignments.media.AlignmentsData` - The loaded alignments containing faces to be removed - arguments: :class:`argparse.Namespace` - The command line arguments that have called this job - """ - def __init__(self, alignments: "AlignmentData", arguments: Namespace) -> None: - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) - self._alignments = alignments - - kwargs = {} - if alignments.version < 2.1: - # Update headers of faces generated with hash based alignments - kwargs["alignments"] = alignments - self._items = Faces(arguments.faces_dir, **kwargs) - logger.debug("Initialized %s", self.__class__.__name__) - - def process(self) -> None: - """ Run the job to remove faces from an alignments file that do not exist within a faces - folder. """ - logger.info("[REMOVE FACES FROM ALIGNMENTS]") # Tidy up cli output - - if not self._items.items: - logger.error("No matching faces found in your faces folder. This would remove all " - "faces from your alignments file. Process aborted.") - return - - pre_face_count = self._alignments.faces_count - self._alignments.filter_faces(self._items.items, filter_out=False) - del_count = pre_face_count - self._alignments.faces_count - if del_count == 0: - logger.info("No changes made to alignments file. Exiting") - return - - logger.info("%s alignment(s) were removed from alignments file", del_count) - - self._update_png_headers() - self._alignments.save() - - rename = Rename(self._alignments, None, self._items) - rename.process() - - def _update_png_headers(self) -> None: - """ Update the EXIF iTXt field of any face PNGs that have had their face index changed. - - Notes - ----- - This could be quicker if parellizing in threads, however, Windows (at least) does not seem - to like this and has a tendency to throw permission errors, so this remains single threaded - for now. - """ - to_update = [ # Items whose face index has changed - x for x in self._items.file_list_sorted - if x["face_index"] != self._items.items[x["source_filename"]].index(x["face_index"])] - - for file_info in tqdm(to_update, desc="Updating PNG Headers", leave=False): - frame = file_info["source_filename"] - face_index = file_info["face_index"] - new_index = self._items.items[frame].index(face_index) - - fullpath = os.path.join(self._items.folder, file_info["current_filename"]) - logger.debug("Updating png header for '%s': face index from %s to %s", - fullpath, face_index, new_index) - - # Update file_list_sorted for rename task - orig_filename = f"{os.path.splitext(frame)[0]}_{new_index}.png" - file_info["face_index"] = new_index - file_info["original_filename"] = orig_filename - - face = DetectedFace() - face.from_alignment(self._alignments.get_faces_in_frame(frame)[new_index]) - meta = dict(alignments=face.to_png_meta(), - source=dict(alignments_version=file_info["alignments_version"], - original_filename=orig_filename, - face_index=new_index, - source_filename=frame, - source_is_video=file_info["source_is_video"], - source_frame_dims=file_info.get("source_frame_dims"))) - update_existing_metadata(fullpath, meta) - - logger.info("%s Extracted face(s) had their header information updated", len(to_update)) - - -class Rename(): # pylint:disable=too-few-public-methods - """ Rename faces in a folder to match their filename as stored in an alignments file. - - Parameters - ---------- - alignments: :class:`tools.lib_alignments.media.AlignmentData` - The alignments data loaded from an alignments file for this rename job - arguments: :class:`argparse.Namespace` - The :mod:`argparse` arguments as passed in from :mod:`tools.py` - faces: :class:`tools.lib_alignments.media.Faces`, Optional - An optional faces object, if the rename task is being called by another job. - Default: ``None`` - """ - def __init__(self, - alignments: "AlignmentData", - arguments: Namespace, - faces: Optional[Faces] = None) -> None: - logger.debug("Initializing %s: (arguments: %s, faces: %s)", - self.__class__.__name__, arguments, faces) - self._alignments = alignments - - kwargs = {} - if alignments.version < 2.1: - # Update headers of faces generated with hash based alignments - kwargs["alignments"] = alignments - self._faces = faces if faces else Faces(arguments.faces_dir, **kwargs) - logger.debug("Initialized %s", self.__class__.__name__) - def process(self) -> None: - """ Process the face renaming """ - logger.info("[RENAME FACES]") # Tidy up cli output - rename_mappings = sorted([(face["current_filename"], face["original_filename"]) - for face in self._faces.file_list_sorted - if face["current_filename"] != face["original_filename"]], - key=lambda x: x[1]) - rename_count = self._rename_faces(rename_mappings) - logger.info("%s faces renamed", rename_count) + """ Parse the imported alignments file and output relevant information to a json file """ + logger.info("[EXPORTING ALIGNMENTS]") # Tidy up cli output + formatted = {key: [self._format_face(face) for face in val["faces"]] + for key, val in self._alignments.data.items()} + logger.info("Saving export alignments to '%s'...", self._output_file) + self._serializer.save(self._output_file, formatted) - def _rename_faces(self, filename_mappings: List[Tuple[str, str]]) -> int: - """ Rename faces back to their original name as exists in the alignments file. - If the source and destination filename are the same then skip that file. - - Parameters - ---------- - filename_mappings: list - List of tuples of (`source filename`, `destination filename`) ordered by destination - filename - - Returns - ------- - int - The number of faces that have been renamed - """ - if not filename_mappings: - return 0 - - rename_count = 0 - conflicts = [] - for src, dst in tqdm(filename_mappings, desc="Renaming Faces"): - old = os.path.join(self._faces.folder, src) - new = os.path.join(self._faces.folder, dst) - - if os.path.exists(new): - # Interim add .tmp extension to files that will cause a rename conflict, to - # process afterwards - logger.debug("interim renaming file to avoid conflict: (src: '%s', dst: '%s')", - src, dst) - new = new + ".tmp" - conflicts.append(new) - - logger.verbose("Renaming '%s' to '%s'", old, new) - os.rename(old, new) - rename_count += 1 - if conflicts: - for old in tqdm(conflicts, desc="Renaming Faces"): - new = old[:-4] # Remove .tmp extension - if os.path.exists(new): - # This should only be running on faces. If there is still a conflict - # then the user has done something stupid, so we will delete the file and - # replace. They can always re-extract :/ - os.remove(new) - logger.verbose("Renaming '%s' to '%s'", old, new) - os.rename(old, new) - return rename_count - - -class Sort(): +class Sort: """ Sort alignments' index by the order they appear in an image in left to right order. Parameters @@ -1115,10 +457,12 @@ class Sort(): alignments: :class:`tools.lib_alignments.media.AlignmentData` The alignments data loaded from an alignments file for this rename job arguments: :class:`argparse.Namespace` - The :mod:`argparse` arguments as passed in from :mod:`tools.py` + The :mod:`argparse` arguments as passed in from :mod:`tools.py`. Unused """ - def __init__(self, alignments: "AlignmentData", arguments: Namespace) -> None: - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + def __init__(self, + alignments: AlignmentData, + arguments: Namespace) -> None: # pylint:disable=unused-argument + logger.debug(parse_class_init(locals())) self._alignments = alignments logger.debug("Initialized %s", self.__class__.__name__) @@ -1131,41 +475,53 @@ def process(self) -> None: logger.warning("If you have a face-set corresponding to the alignment file you " "processed then you should run the 'Extract' job to regenerate it.") - def reindex_faces(self) -> None: + def reindex_faces(self) -> int: """ Re-Index the faces """ reindexed = 0 for alignment in tqdm(self._alignments.yield_faces(), - desc="Sort alignment indexes", total=self._alignments.frames_count): + desc="Sort alignment indexes", + total=self._alignments.frames_count, + leave=False): frame, alignments, count, key = alignment if count <= 1: - logger.trace("0 or 1 face in frame. Not sorting: '%s'", frame) + logger.trace("0 or 1 face in frame. Not sorting: '%s'", frame) # type:ignore continue sorted_alignments = sorted(alignments, key=lambda x: (x["x"])) if sorted_alignments == alignments: - logger.trace("Alignments already in correct order. Not sorting: '%s'", frame) + logger.trace("Alignments already in correct order. Not " # type:ignore + "sorting: '%s'", frame) continue - logger.trace("Sorting alignments for frame: '%s'", frame) + logger.trace("Sorting alignments for frame: '%s'", frame) # type:ignore self._alignments.data[key]["faces"] = sorted_alignments reindexed += 1 logger.info("%s Frames had their faces reindexed", reindexed) return reindexed -class Spatial(): +class Spatial: """ Apply spatial temporal filtering to landmarks - Adapted from: - https://www.kaggle.com/selfishgene/animating-and-smoothing-3d-facial-keypoints/notebook """ - def __init__(self, alignments, arguments): - logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + Parameters + ---------- + alignments: :class:`tools.lib_alignments.media.AlignmentData` + The alignments data loaded from an alignments file for this rename job + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + + Reference + --------- + https://www.kaggle.com/selfishgene/animating-and-smoothing-3d-facial-keypoints/notebook + """ + def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None: + logger.debug(parse_class_init(locals())) self.arguments = arguments self._alignments = alignments - self.mappings = {} - self.normalized = {} - self.shapes_model = None + self._mappings: dict[int, str] = {} + self._normalized: dict[str, np.ndarray] = {} + self._shapes_model: decomposition.PCA | None = None logger.debug("Initialized %s", self.__class__.__name__) - def process(self): + def process(self) -> None: """ Perform spatial filtering """ logger.info("[SPATIO-TEMPORAL FILTERING]") # Tidy up cli output logger.info("NB: The process only processes the alignments for the first " @@ -1173,19 +529,35 @@ def process(self): "there is only a single face in the alignments file and all false positives " "have been removed") - self.normalize() - self.shape_model() - landmarks = self.spatially_filter() - landmarks = self.temporally_smooth(landmarks) - self.update_alignments(landmarks) + self._normalize() + self._shape_model() + landmarks = self._spatially_filter() + landmarks = self._temporally_smooth(landmarks) + self._update_alignments(landmarks) self._alignments.save() logger.warning("If you have a face-set corresponding to the alignment file you " "processed then you should run the 'Extract' job to regenerate it.") # Define shape normalization utility functions @staticmethod - def normalize_shapes(shapes_im_coords): - """ Normalize a 2D or 3D shape """ + def _normalize_shapes(shapes_im_coords: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ Normalize a 2D or 3D shape + + Parameters + ---------- + shaped_im_coords: :class:`numpy.ndarray` + The facial landmarks + + Returns + ------- + shapes_normalized: :class:`numpy.ndarray` + The normalized shapes + scale_factors: :class:`numpy.ndarray` + The scale factors + mean_coords: :class:`numpy.ndarray` + The mean coordinates + """ logger.debug("Normalize shapes") (num_pts, num_dims, _) = shapes_im_coords.shape @@ -1204,8 +576,25 @@ def normalize_shapes(shapes_im_coords): return shapes_normalized, scale_factors, mean_coords @staticmethod - def normalized_to_original(shapes_normalized, scale_factors, mean_coords): - """ Transform a normalized shape back to original image coordinates """ + def _normalized_to_original(shapes_normalized: np.ndarray, + scale_factors: np.ndarray, + mean_coords: np.ndarray) -> np.ndarray: + """ Transform a normalized shape back to original image coordinates + + Parameters + ---------- + shapes_normalized: :class:`numpy.ndarray` + The normalized shapes + scale_factors: :class:`numpy.ndarray` + The scale factors + mean_coords: :class:`numpy.ndarray` + The mean coordinates + + Returns + ------- + :class:`numpy.ndarray` + The normalized shape transformed back to original coordinates + """ logger.debug("Normalize to original") (num_pts, num_dims, _) = shapes_normalized.shape @@ -1217,72 +606,97 @@ def normalized_to_original(shapes_normalized, scale_factors, mean_coords): logger.debug("Normalized to original: %s", shapes_im_coords) return shapes_im_coords - def normalize(self): + def _normalize(self) -> None: """ Compile all original and normalized alignments """ logger.debug("Normalize") count = sum(1 for val in self._alignments.data.values() if val["faces"]) - landmarks_all = np.zeros((68, 2, int(count))) + + sample_lm = next((val["faces"][0]["landmarks_xy"] + for val in self._alignments.data.values() if val["faces"]), 68) + assert isinstance(sample_lm, np.ndarray) + lm_count = sample_lm.shape[0] + if lm_count != 68: + raise FaceswapError("Spatial smoothing only supports 68 point facial landmarks") + + landmarks_all = np.zeros((lm_count, 2, int(count))) end = 0 - for key in tqdm(sorted(self._alignments.data.keys()), desc="Compiling"): + for key in tqdm(sorted(self._alignments.data.keys()), desc="Compiling", leave=False): val = self._alignments.data[key]["faces"] if not val: continue # We should only be normalizing a single face, so just take # the first landmarks found - landmarks = np.array(val[0]["landmarks_xy"]).reshape((68, 2, 1)) + landmarks = np.array(val[0]["landmarks_xy"]).reshape((lm_count, 2, 1)) start = end end = start + landmarks.shape[2] # Store in one big array landmarks_all[:, :, start:end] = landmarks # Make sure we keep track of the mapping to the original frame - self.mappings[start] = key + self._mappings[start] = key # Normalize shapes - normalized_shape = self.normalize_shapes(landmarks_all) - self.normalized["landmarks"] = normalized_shape[0] - self.normalized["scale_factors"] = normalized_shape[1] - self.normalized["mean_coords"] = normalized_shape[2] - logger.debug("Normalized: %s", self.normalized) + normalized_shape = self._normalize_shapes(landmarks_all) + self._normalized["landmarks"] = normalized_shape[0] + self._normalized["scale_factors"] = normalized_shape[1] + self._normalized["mean_coords"] = normalized_shape[2] + logger.debug("Normalized: %s", self._normalized) - def shape_model(self): + def _shape_model(self) -> None: """ build 2D shape model """ logger.debug("Shape model") - landmarks_norm = self.normalized["landmarks"] + landmarks_norm = self._normalized["landmarks"] num_components = 20 normalized_shapes_tbl = np.reshape(landmarks_norm, [68*2, landmarks_norm.shape[2]]).T - self.shapes_model = decomposition.PCA(n_components=num_components, - whiten=True, - random_state=1).fit(normalized_shapes_tbl) - explained = self.shapes_model.explained_variance_ratio_.sum() + self._shapes_model = decomposition.PCA(n_components=num_components, + whiten=True, + random_state=1).fit(normalized_shapes_tbl) + explained = self._shapes_model.explained_variance_ratio_.sum() logger.info("Total explained percent by PCA model with %s components is %s%%", num_components, round(100 * explained, 1)) logger.debug("Shaped model") - def spatially_filter(self): - """ interpret the shapes using our shape model - (project and reconstruct) """ + def _spatially_filter(self) -> np.ndarray: + """ interpret the shapes using our shape model (project and reconstruct) + + Returns + ------- + :class:`numpy.ndarray` + The filtered landmarks in original coordinate space + """ logger.debug("Spatially Filter") - landmarks_norm = self.normalized["landmarks"] + assert self._shapes_model is not None + landmarks_norm = self._normalized["landmarks"] # Convert to matrix form landmarks_norm_table = np.reshape(landmarks_norm, [68 * 2, landmarks_norm.shape[2]]).T # Project onto shapes model and reconstruct - landmarks_norm_table_rec = self.shapes_model.inverse_transform( - self.shapes_model.transform(landmarks_norm_table)) + landmarks_norm_table_rec = self._shapes_model.inverse_transform( + self._shapes_model.transform(landmarks_norm_table)) # Convert back to shapes (numKeypoint, num_dims, numFrames) landmarks_norm_rec = np.reshape(landmarks_norm_table_rec.T, [68, 2, landmarks_norm.shape[2]]) # Transform back to image co-ordinates - retval = self.normalized_to_original(landmarks_norm_rec, - self.normalized["scale_factors"], - self.normalized["mean_coords"]) + retval = self._normalized_to_original(landmarks_norm_rec, + self._normalized["scale_factors"], + self._normalized["mean_coords"]) logger.debug("Spatially Filtered: %s", retval) return retval @staticmethod - def temporally_smooth(landmarks): - """ apply temporal filtering on the 2D points """ + def _temporally_smooth(landmarks: np.ndarray) -> np.ndarray: + """ apply temporal filtering on the 2D points + + Parameters + ---------- + landmarks: :class:`numpy.ndarray` + 68 point landmarks to be temporally smoothed + + Returns + ------- + :class: `numpy.ndarray` + The temporally smoothed landmarks + """ logger.debug("Temporally Smooth") filter_half_length = 2 temporal_filter = np.ones((1, 1, 2 * filter_half_length + 1)) @@ -1296,13 +710,20 @@ def temporally_smooth(landmarks): logger.debug("Temporally Smoothed: %s", retval) return retval - def update_alignments(self, landmarks): - """ Update smoothed landmarks back to alignments """ + def _update_alignments(self, landmarks: np.ndarray) -> None: + """ Update smoothed landmarks back to alignments + + Parameters + ---------- + landmarks: :class:`numpy.ndarray` + The smoothed landmarks + """ logger.debug("Update alignments") - for idx, frame in tqdm(self.mappings.items(), desc="Updating"): - logger.trace("Updating: (frame: %s)", frame) + for idx, frame in tqdm(self._mappings.items(), desc="Updating", leave=False): + logger.trace("Updating: (frame: %s)", frame) # type:ignore landmarks_update = landmarks[:, :, idx] landmarks_xy = landmarks_update.reshape(68, 2).tolist() self._alignments.data[frame]["faces"][0]["landmarks_xy"] = landmarks_xy - logger.trace("Updated: (frame: '%s', landmarks: %s)", frame, landmarks_xy) + logger.trace("Updated: (frame: '%s', landmarks: %s)", # type:ignore + frame, landmarks_xy) logger.debug("Updated alignments") diff --git a/tools/alignments/jobs_faces.py b/tools/alignments/jobs_faces.py new file mode 100644 index 0000000000..ac2205f89c --- /dev/null +++ b/tools/alignments/jobs_faces.py @@ -0,0 +1,484 @@ +#!/usr/bin/env python3 +""" Tools for manipulating the alignments using extracted Faces as a source """ +from __future__ import annotations +import logging +import os +import typing as T + +from argparse import Namespace +from operator import itemgetter + +import numpy as np +from tqdm import tqdm + +from lib.align import DetectedFace +from lib.image import update_existing_metadata # TODO remove +from scripts.fsmedia import Alignments + +from .media import Faces + +if T.TYPE_CHECKING: + from .media import AlignmentData + from lib.align.alignments import (AlignmentDict, AlignmentFileDict, + PNGHeaderDict, PNGHeaderAlignmentsDict) + +logger = logging.getLogger(__name__) + + +class FromFaces(): + """ Scan a folder of Faceswap Extracted Faces and re-create the associated alignments file(s) + + Parameters + ---------- + alignments: NoneType + Parameter included for standard job naming convention, but not used for this process. + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + """ + def __init__(self, alignments: None, arguments: Namespace) -> None: + logger.debug("Initializing %s: (alignments: %s, arguments: %s)", + self.__class__.__name__, alignments, arguments) + self._faces_dir = arguments.faces_dir + self._faces = Faces(arguments.faces_dir) + logger.debug("Initialized %s", self.__class__.__name__) + + def process(self) -> None: + """ Run the job to read faces from a folder to create alignments file(s). """ + logger.info("[CREATE ALIGNMENTS FROM FACES]") # Tidy up cli output + + all_versions: dict[str, list[float]] = {} + d_align: dict[str, dict[str, list[tuple[int, AlignmentFileDict, str, dict]]]] = {} + filelist = T.cast(list[tuple[str, "PNGHeaderDict"]], self._faces.file_list_sorted) + for filename, meta in tqdm(filelist, + desc="Generating Alignments", + total=len(filelist), + leave=False): + + align_fname = self._get_alignments_filename(meta["source"]) + source_name, f_idx, alignment = self._extract_alignment(meta) + full_info = (f_idx, alignment, filename, meta["source"]) + + d_align.setdefault(align_fname, {}).setdefault(source_name, []).append(full_info) + all_versions.setdefault(align_fname, []).append(meta["source"]["alignments_version"]) + + versions = {k: min(v) for k, v in all_versions.items()} + alignments = self._sort_alignments(d_align) + self._save_alignments(alignments, versions) + + @classmethod + def _get_alignments_filename(cls, source_data: dict) -> str: + """ Obtain the name of the alignments file from the source information contained within the + PNG metadata. + + Parameters + ---------- + source_data: dict + The source information contained within a Faceswap extracted PNG + + Returns + ------- + str: + If the face was generated from a video file, the filename will be + `'_alignments.fsa'`. If it was extracted from an image file it will be + `'alignments.fsa'` + """ + is_video = source_data["source_is_video"] + src_name = source_data["source_filename"] + prefix = f"{src_name.rpartition('_')[0]}_" if is_video else "" + retval = f"{prefix}alignments.fsa" + logger.trace("Extracted alignments file filename: '%s'", retval) # type:ignore + return retval + + def _extract_alignment(self, metadata: dict) -> tuple[str, int, AlignmentFileDict]: + """ Extract alignment data from a PNG image's itxt header. + + Formats the landmarks into a numpy array and adds in mask centering information if it is + from an older extract. + + Parameters + ---------- + metadata: dict + An extracted faces PNG Header data + + Returns + ------- + tuple + The alignment's source frame name in position 0. The index of the face within the + alignment file in position 1. The alignment data correctly formatted for writing to an + alignments file in positin 2 + """ + alignment = metadata["alignments"] + alignment["landmarks_xy"] = np.array(alignment["landmarks_xy"], dtype="float32") + + src = metadata["source"] + frame_name = src["source_filename"] + face_index = int(src["face_index"]) + + logger.trace("Extracted alignment for frame: '%s', face index: %s", # type:ignore + frame_name, face_index) + return frame_name, face_index, alignment + + def _sort_alignments(self, + alignments: dict[str, dict[str, list[tuple[int, + AlignmentFileDict, + str, + dict]]]] + ) -> dict[str, dict[str, AlignmentDict]]: + """ Sort the faces into face index order as they appeared in the original alignments file. + + If the face index stored in the png header does not match it's position in the alignments + file (i.e. A face has been removed from a frame) then update the header of the + corresponding png to the correct index as exists in the newly created alignments file. + + Parameters + ---------- + alignments: dict + The unsorted alignments file(s) as generated from the face PNG headers, including the + face index of the face within it's respective frame, the original face filename and + the orignal face header source information + + Returns + ------- + dict + The alignments file dictionaries sorted into the correct face order, ready for saving + """ + logger.info("Sorting and checking faces...") + aln_sorted: dict[str, dict[str, AlignmentDict]] = {} + for fname, frames in alignments.items(): + this_file: dict[str, AlignmentDict] = {} + for frame in tqdm(sorted(frames), desc=f"Sorting {fname}", leave=False): + this_file[frame] = {"video_meta": {}, "faces": []} + for real_idx, (f_id, almt, f_path, f_src) in enumerate(sorted(frames[frame], + key=itemgetter(0))): + if real_idx != f_id: + full_path = os.path.join(self._faces_dir, f_path) + self._update_png_header(full_path, real_idx, almt, f_src) + this_file[frame]["faces"].append(almt) + aln_sorted[fname] = this_file + return aln_sorted + + @classmethod + def _update_png_header(cls, + face_path: str, + new_index: int, + alignment: AlignmentFileDict, + source_info: dict) -> None: + """ Update the PNG header for faces where the stored index does not correspond with the + alignments file. This can occur when frames with multiple faces have had some faces deleted + from the faces folder. + + Updates the original filename and index in the png header. + + Parameters + ---------- + face_path: str + Full path to the saved face image that requires updating + new_index: int + The new index as it appears in the newly generated alignments file + alignment: dict + The alignment information to store in the png header + source_info: dict + The face source information as extracted from the original face png file + """ + face = DetectedFace() + face.from_alignment(alignment) + new_filename = f"{os.path.splitext(source_info['source_filename'])[0]}_{new_index}.png" + + logger.trace("Updating png header for '%s': (face index from %s to %s, " # type:ignore + "original filename from '%s' to '%s'", face_path, source_info["face_index"], + new_index, source_info["original_filename"], new_filename) + + source_info["face_index"] = new_index + source_info["original_filename"] = new_filename + meta = {"alignments": face.to_png_meta(), "source": source_info} + update_existing_metadata(face_path, meta) + + def _save_alignments(self, + all_alignments: dict[str, dict[str, AlignmentDict]], + versions: dict[str, float]) -> None: + """ Save the newely generated alignments file(s). + + If an alignments file already exists in the source faces folder, back it up rather than + overwriting + + Parameters + ---------- + all_alignments: dict + The alignment(s) dictionaries found in the faces folder. Alignment filename as key, + corresponding alignments as value. + versions: dict + The minimum version number that exists in a face set for each alignments file to be + generated + """ + for fname, alignments in all_alignments.items(): + version = versions[fname] + alignments_path = os.path.join(self._faces_dir, fname) + dummy_args = Namespace(alignments_path=alignments_path) + aln = Alignments(dummy_args, is_extract=True) + aln.update_from_dict(alignments) + aln._io._version = version # pylint:disable=protected-access + aln._io.update_legacy() # pylint:disable=protected-access + aln.backup() + aln.save() + + +class Rename(): + """ Rename faces in a folder to match their filename as stored in an alignments file. + + Parameters + ---------- + alignments: :class:`tools.lib_alignments.media.AlignmentData` + The alignments data loaded from an alignments file for this rename job + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + faces: :class:`tools.lib_alignments.media.Faces`, Optional + An optional faces object, if the rename task is being called by another job. + Default: ``None`` + """ + def __init__(self, + alignments: AlignmentData, + arguments: Namespace | None, + faces: Faces | None = None) -> None: + logger.debug("Initializing %s: (arguments: %s, faces: %s)", + self.__class__.__name__, arguments, faces) + self._alignments = alignments + + kwargs = {} + if alignments.version < 2.1: + # Update headers of faces generated with hash based alignments + kwargs["alignments"] = alignments + if faces: + self._faces = faces + else: + assert arguments is not None + self._faces = Faces(arguments.faces_dir, **kwargs) # type:ignore # needs TypedDict :/ + logger.debug("Initialized %s", self.__class__.__name__) + + def process(self) -> None: + """ Process the face renaming """ + logger.info("[RENAME FACES]") # Tidy up cli output + filelist = T.cast(list[tuple[str, "PNGHeaderDict"]], self._faces.file_list_sorted) + rename_mappings = sorted([(face[0], face[1]["source"]["original_filename"]) + for face in filelist + if face[0] != face[1]["source"]["original_filename"]], + key=lambda x: x[1]) + rename_count = self._rename_faces(rename_mappings) + logger.info("%s faces renamed", rename_count) + + filelist = T.cast(list[tuple[str, "PNGHeaderDict"]], self._faces.file_list_sorted) + copyback = FaceToFile(self._alignments, [val[1] for val in filelist]) + if copyback(): + self._alignments.save() + + def _rename_faces(self, filename_mappings: list[tuple[str, str]]) -> int: + """ Rename faces back to their original name as exists in the alignments file. + + If the source and destination filename are the same then skip that file. + + Parameters + ---------- + filename_mappings: list + List of tuples of (`source filename`, `destination filename`) ordered by destination + filename + + Returns + ------- + int + The number of faces that have been renamed + """ + if not filename_mappings: + return 0 + + rename_count = 0 + conflicts = [] + for src, dst in tqdm(filename_mappings, desc="Renaming Faces", leave=False): + old = os.path.join(self._faces.folder, src) + new = os.path.join(self._faces.folder, dst) + + if os.path.exists(new): + # Interim add .tmp extension to files that will cause a rename conflict, to + # process afterwards + logger.debug("interim renaming file to avoid conflict: (src: '%s', dst: '%s')", + src, dst) + new = new + ".tmp" + conflicts.append(new) + + logger.verbose("Renaming '%s' to '%s'", old, new) # type:ignore + os.rename(old, new) + rename_count += 1 + if conflicts: + for old in tqdm(conflicts, desc="Renaming Faces", leave=False): + new = old[:-4] # Remove .tmp extension + if os.path.exists(new): + # This should only be running on faces. If there is still a conflict + # then the user has done something stupid, so we will delete the file and + # replace. They can always re-extract :/ + os.remove(new) + logger.verbose("Renaming '%s' to '%s'", old, new) # type:ignore + os.rename(old, new) + return rename_count + + +class RemoveFaces(): + """ Remove items from alignments file. + + Parameters + --------- + alignments: :class:`tools.alignments.media.AlignmentsData` + The loaded alignments containing faces to be removed + arguments: :class:`argparse.Namespace` + The command line arguments that have called this job + """ + def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None: + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + self._alignments = alignments + + self._items = Faces(arguments.faces_dir, alignments=alignments) + logger.debug("Initialized %s", self.__class__.__name__) + + def process(self) -> None: + """ Run the job to remove faces from an alignments file that do not exist within a faces + folder. """ + logger.info("[REMOVE FACES FROM ALIGNMENTS]") # Tidy up cli output + + if not self._items.items: + logger.error("No matching faces found in your faces folder. This would remove all " + "faces from your alignments file. Process aborted.") + return + + items = T.cast(dict[str, list[int]], self._items.items) + pre_face_count = self._alignments.faces_count + self._alignments.filter_faces(items, filter_out=False) + del_count = pre_face_count - self._alignments.faces_count + if del_count == 0: + logger.info("No changes made to alignments file. Exiting") + return + + logger.info("%s alignment(s) were removed from alignments file", del_count) + + self._update_png_headers() + self._alignments.save() + + rename = Rename(self._alignments, None, self._items) + rename.process() + + def _update_png_headers(self) -> None: + """ Update the EXIF iTXt field of any face PNGs that have had their face index changed. + + Notes + ----- + This could be quicker if parellizing in threads, however, Windows (at least) does not seem + to like this and has a tendency to throw permission errors, so this remains single threaded + for now. + """ + items = T.cast(dict[str, list[int]], self._items.items) + srcs = [(x[0], x[1]["source"]) + for x in T.cast(list[tuple[str, "PNGHeaderDict"]], self._items.file_list_sorted)] + to_update = [ # Items whose face index has changed + x for x in srcs + if x[1]["face_index"] != items[x[1]["source_filename"]].index(x[1]["face_index"])] + + for item in tqdm(to_update, desc="Updating PNG Headers", leave=False): + filename, file_info = item + frame = file_info["source_filename"] + face_index = file_info["face_index"] + new_index = items[frame].index(face_index) + + fullpath = os.path.join(self._items.folder, filename) + logger.debug("Updating png header for '%s': face index from %s to %s", + fullpath, face_index, new_index) + + # Update file_list_sorted for rename task + orig_filename = f"{os.path.splitext(frame)[0]}_{new_index}.png" + file_info["face_index"] = new_index + file_info["original_filename"] = orig_filename + + face = DetectedFace() + face.from_alignment(self._alignments.get_faces_in_frame(frame)[new_index]) + meta = {"alignments": face.to_png_meta(), + "source": {"alignments_version": file_info["alignments_version"], + "original_filename": orig_filename, + "face_index": new_index, + "source_filename": frame, + "source_is_video": file_info["source_is_video"], + "source_frame_dims": file_info.get("source_frame_dims")}} + update_existing_metadata(fullpath, meta) + + logger.info("%s Extracted face(s) had their header information updated", len(to_update)) + + +class FaceToFile(): + """ Updates any optional/missing keys in the alignments file with any data that has been + populated in a PNGHeader. Includes masks and identity fields. + + Parameters + --------- + alignments: :class:`tools.alignments.media.AlignmentsData` + The loaded alignments containing faces to be removed + face_data: list + List of :class:`PNGHeaderDict` objects + """ + def __init__(self, alignments: AlignmentData, face_data: list[PNGHeaderDict]) -> None: + logger.debug("Initializing %s: alignments: %s, face_data: %s", + self.__class__.__name__, alignments, len(face_data)) + self._alignments = alignments + self._face_alignments = face_data + self._updatable_keys: list[T.Literal["identity", "mask"]] = ["identity", "mask"] + self._counts: dict[str, int] = {} + logger.debug("Initialized %s", self.__class__.__name__) + + def _check_and_update(self, + alignment: PNGHeaderAlignmentsDict, + face: AlignmentFileDict) -> None: + """ Check whether the key requires updating and update it. + + alignment: dict + The alignment dictionary from the PNG Header + face: dict + The alignment dictionary for the face from the alignments file + """ + for key in self._updatable_keys: + if key == "mask": + exist_masks = face["mask"] + for mask_name, mask_data in alignment["mask"].items(): + if mask_name in exist_masks: + continue + exist_masks[mask_name] = mask_data + count_key = f"mask_{mask_name}" + self._counts[count_key] = self._counts.get(count_key, 0) + 1 + continue + + if not face.get(key, {}) and alignment.get(key): + face[key] = alignment[key] + self._counts[key] = self._counts.get(key, 0) + 1 + + def __call__(self) -> bool: + """ Parse through the face data updating any entries in the alignments file. + + Returns + ------- + bool + ``True`` if any alignment information was updated otherwise ``False`` + """ + for meta in tqdm(self._face_alignments, + desc="Updating Alignments File from PNG Header", + leave=False): + src = meta["source"] + alignment = meta["alignments"] + if not any(alignment.get(key, {}) for key in self._updatable_keys): + continue + + faces = self._alignments.get_faces_in_frame(src["source_filename"]) + if len(faces) < src["face_index"] + 1: # list index out of range + logger.debug("Skipped face '%s'. Index does not exist in alignments file", + src["original_filename"]) + continue + + face = faces[src["face_index"]] + self._check_and_update(alignment, face) + + retval = False + if self._counts: + retval = True + logger.info("Updated alignments file from PNG Data: %s", self._counts) + return retval diff --git a/tools/alignments/jobs_frames.py b/tools/alignments/jobs_frames.py new file mode 100644 index 0000000000..3c25b48121 --- /dev/null +++ b/tools/alignments/jobs_frames.py @@ -0,0 +1,476 @@ +#!/usr/bin/env python3 +""" Tools for manipulating the alignments using Frames as a source """ +from __future__ import annotations +import logging +import os +import sys +import typing as T + +from datetime import datetime + +import cv2 +import numpy as np +from tqdm import tqdm + +from lib.align import DetectedFace, EXTRACT_RATIOS, LANDMARK_PARTS, LandmarkType +from lib.align.alignments import _VERSION, PNGHeaderDict +from lib.image import encode_image, generate_thumbnail, ImagesSaver +from plugins.extract import ExtractMedia, Extractor +from .media import ExtractedFaces, Frames + +if T.TYPE_CHECKING: + from argparse import Namespace + from .media import AlignmentData + +logger = logging.getLogger(__name__) + + +class Draw(): + """ Draws annotations onto original frames and saves into a sub-folder next to the original + frames. + + Parameters + --------- + alignments: :class:`tools.alignments.media.AlignmentsData` + The loaded alignments corresponding to the frames to be annotated + arguments: :class:`argparse.Namespace` + The command line arguments that have called this job + """ + def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None: + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + self._alignments = alignments + self._frames = Frames(arguments.frames_dir) + self._output_folder = self._set_output() + logger.debug("Initialized %s", self.__class__.__name__) + + def _set_output(self) -> str: + """ Set the output folder path. + + If annotating a folder of frames, output will be placed in a sub folder within the frames + folder. If annotating a video, output will be a folder next to the original video. + + Returns + ------- + str + Full path to the output folder + + """ + now = datetime.now().strftime("%Y%m%d_%H%M%S") + folder_name = f"drawn_landmarks_{now}" + if self._frames.is_video: + dest_folder = os.path.dirname(self._frames.folder) + else: + dest_folder = self._frames.folder + output_folder = os.path.join(dest_folder, folder_name) + logger.debug("Creating folder: '%s'", output_folder) + os.makedirs(output_folder) + return output_folder + + def process(self) -> None: + """ Runs the process to draw face annotations onto original source frames. """ + logger.info("[DRAW LANDMARKS]") # Tidy up cli output + frames_drawn = 0 + for frame in tqdm(self._frames.file_list_sorted, desc="Drawing landmarks", leave=False): + frame_name = frame["frame_fullname"] + + if not self._alignments.frame_exists(frame_name): + logger.verbose("Skipping '%s' - Alignments not found", frame_name) # type:ignore + continue + + self._annotate_image(frame_name) + frames_drawn += 1 + logger.info("%s Frame(s) output", frames_drawn) + + def _annotate_image(self, frame_name: str) -> None: + """ Annotate the frame with each face that appears in the alignments file. + + Parameters + ---------- + frame_name: str + The full path to the original frame + """ + logger.trace("Annotating frame: '%s'", frame_name) # type:ignore + image = self._frames.load_image(frame_name) + + for idx, alignment in enumerate(self._alignments.get_faces_in_frame(frame_name)): + face = DetectedFace() + face.from_alignment(alignment, image=image) + # Bounding Box + assert face.left is not None + assert face.top is not None + cv2.rectangle(image, (face.left, face.top), (face.right, face.bottom), (255, 0, 0), 1) + self._annotate_landmarks(image, np.rint(face.landmarks_xy).astype("int32")) + self._annotate_extract_boxes(image, face, idx) + self._annotate_pose(image, face) # Pose (head is still loaded) + + self._frames.save_image(self._output_folder, frame_name, image) + + def _annotate_landmarks(self, image: np.ndarray, landmarks: np.ndarray) -> None: + """ Annotate the extract boxes onto the frame. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The frame that extract boxes are to be annotated on to + landmarks: :class:`numpy.ndarray` + The facial landmarks that are to be annotated onto the frame + """ + # Mesh + for start, end, fill in LANDMARK_PARTS[LandmarkType.from_shape(landmarks.shape)].values(): + cv2.polylines(image, [landmarks[start:end]], fill, (255, 255, 0), 1) + # Landmarks + for (pos_x, pos_y) in landmarks: + cv2.circle(image, (pos_x, pos_y), 1, (0, 255, 255), -1) + + @classmethod + def _annotate_extract_boxes(cls, image: np.ndarray, face: DetectedFace, index: int) -> None: + """ Annotate the mesh and landmarks boxes onto the frame. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The frame that mesh and landmarks are to be annotated on to + face: :class:`lib.align.DetectedFace` + The aligned face + index: int + The face index for the given face + """ + for area in T.get_args(T.Literal["face", "head"]): + face.load_aligned(image, centering=area, force=True) + color = (0, 255, 0) if area == "face" else (0, 0, 255) + top_left = face.aligned.original_roi[0] + top_left = (top_left[0], top_left[1] - 10) + cv2.putText(image, str(index), top_left, cv2.FONT_HERSHEY_DUPLEX, 1.0, color, 1) + cv2.polylines(image, [face.aligned.original_roi], True, color, 1) + + @classmethod + def _annotate_pose(cls, image: np.ndarray, face: DetectedFace) -> None: + """ Annotate the pose onto the frame. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The frame that pose is to be annotated on to + face: :class:`lib.align.DetectedFace` + The aligned face loaded for head centering + """ + center = np.array((face.aligned.size / 2, + face.aligned.size / 2)).astype("int32").reshape(1, 2) + center = np.rint(face.aligned.transform_points(center, invert=True)).astype("int32") + points = face.aligned.pose.xyz_2d * face.aligned.size + points = np.rint(face.aligned.transform_points(points, invert=True)).astype("int32") + cv2.line(image, tuple(center), tuple(points[1]), (0, 255, 0), 2) + cv2.line(image, tuple(center), tuple(points[0]), (255, 0, 0), 2) + cv2.line(image, tuple(center), tuple(points[2]), (0, 0, 255), 2) + + +class Extract(): + """ Re-extract faces from source frames based on Alignment data + + Parameters + ---------- + alignments: :class:`tools.lib_alignments.media.AlignmentData` + The alignments data loaded from an alignments file for this rename job + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + """ + def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None: + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + self._arguments = arguments + self._alignments = alignments + self._is_legacy = self._alignments.version == 1.0 # pylint:disable=protected-access + self._mask_pipeline: Extractor | None = None + self._faces_dir = arguments.faces_dir + self._min_size = self._get_min_size(arguments.size, arguments.min_size) + + self._frames = Frames(arguments.frames_dir, self._get_count()) + self._extracted_faces = ExtractedFaces(self._frames, + self._alignments, + size=arguments.size) + self._saver: ImagesSaver | None = None + logger.debug("Initialized %s", self.__class__.__name__) + + @classmethod + def _get_min_size(cls, extract_size: int, min_size: int) -> int: + """ Obtain the minimum size that a face has been resized from to be included as a valid + extract. + + Parameters + ---------- + extract_size: int + The requested size of the extracted images + min_size: int + The percentage amount that has been supplied for valid faces (as a percentage of + extract size) + + Returns + ------- + int + The minimum size, in pixels, that a face is resized from to be considered valid + """ + retval = 0 if min_size == 0 else max(4, int(extract_size * (min_size / 100.))) + logger.debug("Extract size: %s, min percentage size: %s, min_size: %s", + extract_size, min_size, retval) + return retval + + def _get_count(self) -> int | None: + """ If the alignments file has been run through the manual tool, then it will hold video + meta information, meaning that the count of frames in the alignment file can be relied + on to be accurate. + + Returns + ------- + int or ``None`` + For video input which contain video meta-data in the alignments file then the count of + frames is returned. In all other cases ``None`` is returned + """ + meta = self._alignments.video_meta_data + has_meta = all(val is not None for val in meta.values()) + if has_meta: + retval: int | None = len(T.cast(dict[str, list[int] | list[float]], meta["pts_time"])) + else: + retval = None + logger.debug("Frame count from alignments file: (has_meta: %s, %s", has_meta, retval) + return retval + + def process(self) -> None: + """ Run the re-extraction from Alignments file process""" + logger.info("[EXTRACT FACES]") # Tidy up cli output + self._check_folder() + if self._is_legacy: + self._legacy_check() + self._saver = ImagesSaver(self._faces_dir, as_bytes=True) + + if self._min_size > 0: + logger.info("Only selecting faces that have been resized from a minimum resolution " + "of %spx", self._min_size) + + self._export_faces() + + def _check_folder(self) -> None: + """ Check that the faces folder doesn't pre-exist and create. """ + err = None + if not self._faces_dir: + err = "ERROR: Output faces folder not provided." + elif not os.path.isdir(self._faces_dir): + logger.debug("Creating folder: '%s'", self._faces_dir) + os.makedirs(self._faces_dir) + elif os.listdir(self._faces_dir): + err = f"ERROR: Output faces folder should be empty: '{self._faces_dir}'" + if err: + logger.error(err) + sys.exit(0) + logger.verbose("Creating output folder at '%s'", self._faces_dir) # type:ignore + + def _legacy_check(self) -> None: + """ Check whether the alignments file was created with the legacy extraction method. + + If so, force user to re-extract all faces if any options have been specified, otherwise + raise the appropriate warnings and set the legacy options. + """ + if self._min_size > 0 or self._arguments.extract_every_n != 1: + logger.warning("This alignments file was generated with the legacy extraction method.") + logger.warning("You should run this extraction job, but with 'min_size' set to 0 and " + "'extract-every-n' set to 1 to update the alignments file.") + logger.warning("You can then re-run this extraction job with your chosen options.") + sys.exit(0) + + maskers = ["components", "extended"] + nn_masks = [mask for mask in list(self._alignments.mask_summary) if mask not in maskers] + logtype = logger.warning if nn_masks else logger.info + logtype("This alignments file was created with the legacy extraction method and will be " + "updated.") + logtype("Faces will be extracted using the new method and landmarks based masks will be " + "regenerated.") + if nn_masks: + logtype("However, the NN based masks '%s' will be cropped to the legacy extraction " + "method, so you may want to run the mask tool to regenerate these " + "masks.", "', '".join(nn_masks)) + self._mask_pipeline = Extractor(None, None, maskers, multiprocess=True) + self._mask_pipeline.launch() + # Update alignments versioning + self._alignments._io._version = _VERSION # pylint:disable=protected-access + + def _export_faces(self) -> None: + """ Export the faces to the output folder. """ + extracted_faces = 0 + skip_list = self._set_skip_list() + count = self._frames.count if skip_list is None else self._frames.count - len(skip_list) + + for filename, image in tqdm(self._frames.stream(skip_list=skip_list), + total=count, desc="Saving extracted faces", + leave=False): + frame_name = os.path.basename(filename) + if not self._alignments.frame_exists(frame_name): + logger.verbose("Skipping '%s' - Alignments not found", frame_name) # type:ignore + continue + extracted_faces += self._output_faces(frame_name, image) + if self._is_legacy and extracted_faces != 0 and self._min_size == 0: + self._alignments.save() + logger.info("%s face(s) extracted", extracted_faces) + + def _set_skip_list(self) -> list[int] | None: + """ Set the indices for frames that should be skipped based on the `extract_every_n` + command line option. + + Returns + ------- + list or ``None`` + A list of indices to be skipped if extract_every_n is not `1` otherwise + returns ``None`` + """ + skip_num = self._arguments.extract_every_n + if skip_num == 1: + logger.debug("Not skipping any frames") + return None + skip_list = [] + for idx, item in enumerate(T.cast(list[dict[str, str]], self._frames.file_list_sorted)): + if idx % skip_num != 0: + logger.trace("Adding image '%s' to skip list due to " # type:ignore + "extract_every_n = %s", item["frame_fullname"], skip_num) + skip_list.append(idx) + logger.debug("Adding skip list: %s", skip_list) + return skip_list + + def _output_faces(self, filename: str, image: np.ndarray) -> int: + """ For each frame save out the faces + + Parameters + ---------- + filename: str + The filename (without the full path) of the current frame + image: :class:`numpy.ndarray` + The full frame that faces are to be extracted from + + Returns + ------- + int + The total number of faces that have been extracted + """ + logger.trace("Outputting frame: %s", filename) # type:ignore + face_count = 0 + frame_name = os.path.splitext(filename)[0] + faces = self._select_valid_faces(filename, image) + assert self._saver is not None + if not faces: + return face_count + if self._is_legacy: + faces = self._process_legacy(filename, image, faces) + + for idx, face in enumerate(faces): + output = f"{frame_name}_{idx}.png" + meta: PNGHeaderDict = { + "alignments": face.to_png_meta(), + "source": {"alignments_version": self._alignments.version, + "original_filename": output, + "face_index": idx, + "source_filename": filename, + "source_is_video": self._frames.is_video, + "source_frame_dims": T.cast(tuple[int, int], image.shape[:2])}} + assert face.aligned.face is not None + self._saver.save(output, encode_image(face.aligned.face, ".png", metadata=meta)) + if self._min_size == 0 and self._is_legacy: + face.thumbnail = generate_thumbnail(face.aligned.face, size=96, quality=60) + self._alignments.data[filename]["faces"][idx] = face.to_alignment() + face_count += 1 + self._saver.close() + return face_count + + def _select_valid_faces(self, frame: str, image: np.ndarray) -> list[DetectedFace]: + """ Return the aligned faces from a frame that meet the selection criteria, + + Parameters + ---------- + frame: str + The filename (without the full path) of the current frame + image: :class:`numpy.ndarray` + The full frame that faces are to be extracted from + + Returns + ------- + list: + List of valid :class:`lib,align.DetectedFace` objects + """ + faces = self._extracted_faces.get_faces_in_frame(frame, image=image) + if self._min_size == 0: + valid_faces = faces + else: + sizes = self._extracted_faces.get_roi_size_for_frame(frame) + valid_faces = [faces[idx] for idx, size in enumerate(sizes) + if size >= self._min_size] + logger.trace("frame: '%s', total_faces: %s, valid_faces: %s", # type:ignore + frame, len(faces), len(valid_faces)) + return valid_faces + + def _process_legacy(self, + filename: str, + image: np.ndarray, + detected_faces: list[DetectedFace]) -> list[DetectedFace]: + """ Process legacy face extractions to new extraction method. + + Updates stored masks to new extract size + + Parameters + ---------- + filename: str + The current frame filename + image: :class:`numpy.ndarray` + The current image the contains the faces + detected_faces: list + list of :class:`lib.align.DetectedFace` objects for the current frame + + Returns + ------- + list + The updated list of :class:`lib.align.DetectedFace` objects for the current frame + """ + # Update landmarks based masks for face centering + assert self._mask_pipeline is not None + mask_item = ExtractMedia(filename, image, detected_faces=detected_faces) + self._mask_pipeline.input_queue.put(mask_item) + faces = next(self._mask_pipeline.detected_faces()).detected_faces + + # Pad and shift Neural Network based masks to face centering + for face in faces: + self._pad_legacy_masks(face) + return faces + + @classmethod + def _pad_legacy_masks(cls, detected_face: DetectedFace) -> None: + """ Recenter legacy Neural Network based masks from legacy centering to face centering + and pad accordingly. + + Update the masks back into the detected face objects. + + Parameters + ---------- + detected_face: :class:`lib.align.DetectedFace` + The detected face to update the masks for + """ + offset = detected_face.aligned.pose.offset["face"] + for name, mask in detected_face.mask.items(): # Re-center mask and pad to face size + if name in ("components", "extended"): + continue + old_mask = mask.mask.astype("float32") / 255.0 + size = old_mask.shape[0] + new_size = int(size + (size * EXTRACT_RATIOS["face"]) / 2) + + shift = np.rint(offset * (size - (size * EXTRACT_RATIOS["face"]))).astype("int32") + pos = np.array([(new_size // 2 - size // 2) - shift[1], + (new_size // 2) + (size // 2) - shift[1], + (new_size // 2 - size // 2) - shift[0], + (new_size // 2) + (size // 2) - shift[0]]) + bounds = np.array([max(0, pos[0]), min(new_size, pos[1]), + max(0, pos[2]), min(new_size, pos[3])]) + + slice_in = [slice(0 - (pos[0] - bounds[0]), size - (pos[1] - bounds[1])), + slice(0 - (pos[2] - bounds[2]), size - (pos[3] - bounds[3]))] + slice_out = [slice(bounds[0], bounds[1]), slice(bounds[2], bounds[3])] + + new_mask = np.zeros((new_size, new_size, 1), dtype="float32") + new_mask[slice_out[0], slice_out[1], :] = old_mask[slice_in[0], slice_in[1], :] + + mask.replace_mask(new_mask) + # Get the affine matrix from recently generated components mask + # pylint:disable=protected-access + mask._affine_matrix = detected_face.mask["components"].affine_matrix diff --git a/tools/alignments/media.py b/tools/alignments/media.py index 196601a551..a0d6a94365 100644 --- a/tools/alignments/media.py +++ b/tools/alignments/media.py @@ -1,10 +1,12 @@ #!/usr/bin/env python3 """ Media items (Alignments, Faces, Frames) for alignments tool """ - +from __future__ import annotations import logging +from operator import itemgetter import os import sys +import typing as T import cv2 from tqdm import tqdm @@ -15,55 +17,62 @@ from lib.align import Alignments, DetectedFace, update_legacy_png_header from lib.image import (count_frames, generate_thumbnail, ImagesLoader, png_write_meta, read_image, read_image_meta_batch) -from lib.utils import _image_extensions, _video_extensions, FaceswapError +from lib.utils import IMAGE_EXTENSIONS, VIDEO_EXTENSIONS, FaceswapError + +if T.TYPE_CHECKING: + from collections.abc import Generator + import numpy as np + from lib.align.alignments import AlignmentFileDict, PNGHeaderDict -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class AlignmentData(Alignments): - """ Class to hold the alignment data """ + """ Class to hold the alignment data - def __init__(self, alignments_file): + Parameters + ---------- + alignments_file: str + Full path to an alignments file + """ + def __init__(self, alignments_file: str) -> None: logger.debug("Initializing %s: (alignments file: '%s')", self.__class__.__name__, alignments_file) logger.info("[ALIGNMENT DATA]") # Tidy up cli output folder, filename = self.check_file_exists(alignments_file) super().__init__(folder, filename=filename) - logger.verbose("%s items loaded", self.frames_count) + logger.verbose("%s items loaded", self.frames_count) # type: ignore logger.debug("Initialized %s", self.__class__.__name__) @staticmethod - def check_file_exists(alignments_file): - """ Check the alignments file exists""" + def check_file_exists(alignments_file: str) -> tuple[str, str]: + """ Check if the alignments file exists, and returns a tuple of the folder and filename. + + Parameters + ---------- + alignments_file: str + Full path to an alignments file + + Returns + ------- + folder: str + The full path to the folder containing the alignments file + filename: str + The filename of the alignments file + """ folder, filename = os.path.split(alignments_file) if not os.path.isfile(alignments_file): logger.error("ERROR: alignments file not found at: '%s'", alignments_file) sys.exit(0) if folder: - logger.verbose("Alignments file exists at '%s'", alignments_file) + logger.verbose("Alignments file exists at '%s'", alignments_file) # type: ignore return folder, filename - def save(self): + def save(self) -> None: """ Backup copy of old alignments and save new alignments """ self.backup() super().save() - def reload(self): - """ Read the alignments data from the correct format """ - logger.debug("Re-loading alignments") - self._data = self._load() - logger.debug("Re-loaded alignments") - - def set_filename(self, filename): - """ Set the :attr:`_file` to the given filename. - - Parameters - ---------- - filename: str - The full path and filename to set the alignments file name to - """ - self._file = filename - class MediaLoader(): """ Class to load images. @@ -75,26 +84,27 @@ class MediaLoader(): count: int or ``None``, optional If the total frame count is known it can be passed in here which will skip analyzing a video file. If the count is not passed in, it will be calculated. + Default: ``None`` """ - def __init__(self, folder, count=None): + def __init__(self, folder: str, count: int | None = None): logger.debug("Initializing %s: (folder: '%s')", self.__class__.__name__, folder) logger.info("[%s DATA]", self.__class__.__name__.upper()) self._count = count self.folder = folder - self.vid_reader = self.check_input_folder() + self._vid_reader = self.check_input_folder() self.file_list_sorted = self.sorted_items() self.items = self.load_items() - logger.verbose("%s items loaded", self.count) + logger.verbose("%s items loaded", self.count) # type: ignore logger.debug("Initialized %s", self.__class__.__name__) @property - def is_video(self): - """ Return whether source is a video or not """ - return self.vid_reader is not None + def is_video(self) -> bool: + """ bool: Return whether source is a video or not """ + return self._vid_reader is not None @property - def count(self): - """ Number of faces or frames """ + def count(self) -> int: + """ int: Number of faces or frames """ if self._count is not None: return self._count if self.is_video: @@ -103,78 +113,106 @@ def count(self): self._count = len(self.file_list_sorted) return self._count - def check_input_folder(self): - """ makes sure that the frames or faces folder exists - If frames folder contains a video file return imageio reader object """ + def check_input_folder(self) -> cv2.VideoCapture | None: + """ Ensure that the frames or faces folder exists and is valid. + If frames folder contains a video file return imageio reader object + + Returns + ------- + :class:`cv2.VideoCapture` + Object for reading a video stream + """ err = None loadtype = self.__class__.__name__ if not self.folder: - err = "ERROR: A {} folder must be specified".format(loadtype) + err = f"ERROR: A {loadtype} folder must be specified" elif not os.path.exists(self.folder): - err = ("ERROR: The {} location {} could not be " - "found".format(loadtype, self.folder)) + err = f"ERROR: The {loadtype} location {self.folder} could not be found" if err: logger.error(err) sys.exit(0) if (loadtype == "Frames" and os.path.isfile(self.folder) and - os.path.splitext(self.folder)[1].lower() in _video_extensions): - logger.verbose("Video exists at: '%s'", self.folder) - retval = cv2.VideoCapture(self.folder) # pylint: disable=no-member + os.path.splitext(self.folder)[1].lower() in VIDEO_EXTENSIONS): + logger.verbose("Video exists at: '%s'", self.folder) # type: ignore + retval = cv2.VideoCapture(self.folder) # pylint:disable=no-member # TODO ImageIO single frame seek seems slow. Look into this # retval = imageio.get_reader(self.folder, "ffmpeg") else: - logger.verbose("Folder exists at '%s'", self.folder) + logger.verbose("Folder exists at '%s'", self.folder) # type: ignore retval = None return retval @staticmethod - def valid_extension(filename): - """ Check whether passed in file has a valid extension """ + def valid_extension(filename) -> bool: + """ bool: Check whether passed in file has a valid extension """ extension = os.path.splitext(filename)[1] - retval = extension.lower() in _image_extensions - logger.trace("Filename has valid extension: '%s': %s", filename, retval) + retval = extension.lower() in IMAGE_EXTENSIONS + logger.trace("Filename has valid extension: '%s': %s", filename, retval) # type: ignore return retval - @staticmethod - def sorted_items(): + def sorted_items(self) -> list[dict[str, str]] | list[tuple[str, PNGHeaderDict]]: """ Override for specific folder processing """ - return list() + raise NotImplementedError() - @staticmethod - def process_folder(): + def process_folder(self) -> (Generator[dict[str, str], None, None] | + Generator[tuple[str, PNGHeaderDict], None, None]): """ Override for specific folder processing """ - return list() + raise NotImplementedError() - @staticmethod - def load_items(): + def load_items(self) -> dict[str, list[int]] | dict[str, tuple[str, str]]: """ Override for specific item loading """ - return dict() + raise NotImplementedError() + + def load_image(self, filename: str) -> np.ndarray: + """ Load an image + + Parameters + ---------- + filename: str + The filename of the image to load - def load_image(self, filename): - """ Load an image """ + Returns + ------- + :class:`numpy.ndarray` + The loaded image + """ if self.is_video: image = self.load_video_frame(filename) else: src = os.path.join(self.folder, filename) - logger.trace("Loading image: '%s'", src) + logger.trace("Loading image: '%s'", src) # type: ignore image = read_image(src, raise_error=True) return image - def load_video_frame(self, filename): - """ Load a requested frame from video """ + def load_video_frame(self, filename: str) -> np.ndarray: + """ Load a requested frame from video + + Parameters + ---------- + filename: str + The frame name to load + + Returns + ------- + :class:`numpy.ndarray` + The loaded image + """ + assert self._vid_reader is not None frame = os.path.splitext(filename)[0] - logger.trace("Loading video frame: '%s'", frame) + logger.trace("Loading video frame: '%s'", frame) # type: ignore frame_no = int(frame[frame.rfind("_") + 1:]) - 1 - self.vid_reader.set(cv2.CAP_PROP_POS_FRAMES, frame_no) # pylint: disable=no-member - _, image = self.vid_reader.read() + self._vid_reader.set(cv2.CAP_PROP_POS_FRAMES, frame_no) # pylint:disable=no-member + + _, image = self._vid_reader.read() # TODO imageio single frame seek seems slow. Look into this - # self.vid_reader.set_image_index(frame_no) - # image = self.vid_reader.get_next_data()[:, :, ::-1] + # self._vid_reader.set_image_index(frame_no) + # image = self._vid_reader.get_next_data()[:, :, ::-1] return image - def stream(self, skip_list=None): + def stream(self, skip_list: list[int] | None = None + ) -> Generator[tuple[str, np.ndarray], None, None]: """ Load the images in :attr:`folder` in the order they are received from :class:`lib.image.ImagesLoader` in a background thread. @@ -198,18 +236,21 @@ def stream(self, skip_list=None): yield filename, image @staticmethod - def save_image(output_folder, filename, image, metadata=None): + def save_image(output_folder: str, + filename: str, + image: np.ndarray, + metadata: PNGHeaderDict | None = None) -> None: """ Save an image """ output_file = os.path.join(output_folder, filename) output_file = os.path.splitext(output_file)[0] + ".png" - logger.trace("Saving image: '%s'", output_file) + logger.trace("Saving image: '%s'", output_file) # type: ignore if metadata: - encoded_image = cv2.imencode(".png", image)[1] - encoded_image = png_write_meta(encoded_image.tobytes(), metadata) + encoded = cv2.imencode(".png", image)[1] + encoded_image = png_write_meta(encoded.tobytes(), metadata) with open(output_file, "wb") as out_file: out_file.write(encoded_image) else: - cv2.imwrite(output_file, image) # pylint: disable=no-member + cv2.imwrite(output_file, image) # pylint:disable=no-member class Faces(MediaLoader): @@ -220,15 +261,91 @@ class Faces(MediaLoader): folder: str The folder to load faces from alignments: :class:`lib.align.Alignments`, optional - The alignments object that contains the faces. Used to update legacy hash based faces - for None: self._alignments = alignments super().__init__(folder) - def process_folder(self): + def _handle_legacy(self, fullpath: str, log: bool = False) -> PNGHeaderDict: + """Handle facesets that are legacy (i.e. do not contain alignment information in the + header data) + + Parameters + ---------- + fullpath : str + The full path to the extracted face image + log : bool, optional + Whether to log a message that legacy updating is occurring + + Returns + ------- + :class:`~lib.align.alignments.PNGHeaderDict` + The Alignments information from the face in PNG Header dict format + + Raises + ------ + FaceswapError + If legacy faces can't be updated because the alignments file does not exist or some of + the faces do not appear in the provided alignments file + """ + if self._alignments is None: # Can't update legacy + raise FaceswapError(f"The folder '{self.folder}' contains images that do not include " + "Faceswap metadata.\nAll images in the provided folder should " + "contain faces generated from Faceswap's extraction process.\n" + "Please double check the source and try again.") + if log: + logger.warning("Legacy faces discovered. These faces will be updated") + + data = update_legacy_png_header(fullpath, self._alignments) + if not data: + raise FaceswapError( + f"Some of the faces being passed in from '{self.folder}' could not be " + f"matched to the alignments file '{self._alignments.file}'\nPlease double " + "check your sources and try again.") + return data + + def _handle_duplicate(self, + fullpath: str, + header_dict: PNGHeaderDict, + seen: dict[str, list[int]]) -> bool: + """ Check whether the given face has already been seen for the source frame and face index + from an existing face. Can happen when filenames have changed due to sorting etc. and users + have done multiple extractions/copies and placed all of the faces in the same folder + + Parameters + ---------- + fullpath : str + The full path to the face image that is being checked + header_dict : class:`~lib.align.alignments.PNGHeaderDict` + The PNG header dictionary for the given face + seen : dict[str, list[int]] + Dictionary of original source filename and face indices that have already been seen and + will be updated with the face processing now + + Returns + ------- + bool + ``True`` if the face was a duplicate and has been removed, otherwise ``False`` + """ + src_filename = header_dict["source"]["source_filename"] + face_index = header_dict["source"]["face_index"] + + if src_filename in seen and face_index in seen[src_filename]: + dupe_dir = os.path.join(self.folder, "_duplicates") + os.makedirs(dupe_dir, exist_ok=True) + filename = os.path.basename(fullpath) + logger.trace("Moving duplicate: %s", filename) # type:ignore + os.rename(fullpath, os.path.join(dupe_dir, filename)) + return True + + seen.setdefault(src_filename, []).append(face_index) + return False + + def process_folder(self) -> Generator[tuple[str, PNGHeaderDict], None, None]: """ Iterate through the faces folder pulling out various information for each face. Yields @@ -238,8 +355,11 @@ def process_folder(self): :class:`lib.image.read_image_meta_batch` """ logger.info("Loading file list from %s", self.folder) + filter_count = 0 + dupe_count = 0 + seen: dict[str, list[int]] = {} - if self._alignments is not None: # Legacy updating + if self._alignments is not None and self._alignments.version < 2.1: # Legacy updating filelist = [os.path.join(self.folder, face) for face in os.listdir(self.folder) if self.valid_extension(face)] @@ -254,30 +374,33 @@ def process_folder(self): desc="Reading Face Data"): if "itxt" not in metadata or "source" not in metadata["itxt"]: - if self._alignments is None: # Can't update legacy - raise FaceswapError( - f"The folder '{self.folder}' contains images that do not include Faceswap " - "metadata.\nAll images in the provided folder should contain faces " - "generated from Faceswap's extraction process.\nPlease double check the " - "source and try again.") - - if not log_once: - logger.warning("Legacy faces discovered. These faces will be updated") - log_once = True - data = update_legacy_png_header(fullpath, self._alignments) - if not data: - raise FaceswapError( - "Some of the faces being passed in from '{}' could not be matched to the " - "alignments file '{}'\nPlease double check your sources and try " - "again.".format(self.folder, self._alignments.file)) - retval = data["source"] + sub_dict = self._handle_legacy(fullpath, not log_once) + log_once = True else: - retval = metadata["itxt"]["source"] + sub_dict = T.cast("PNGHeaderDict", metadata["itxt"]) - retval["current_filename"] = os.path.basename(fullpath) + if self._handle_duplicate(fullpath, sub_dict, seen): + dupe_count += 1 + continue + + if (self._alignments is not None and # filter existing + not self._alignments.frame_exists(sub_dict["source"]["source_filename"])): + filter_count += 1 + continue + + retval = (os.path.basename(fullpath), sub_dict) yield retval - def load_items(self): + if self._alignments is not None: + logger.debug("Faces filtered out that did not exist in alignments file: %s", + filter_count) + + if dupe_count > 0: + logger.warning("%s Duplicate face images were found. These files have been moved to " + "'%s' from where they can be safely deleted", + dupe_count, os.path.join(self.folder, "_duplicates")) + + def load_items(self) -> dict[str, list[int]]: """ Load the face names into dictionary. Returns @@ -285,13 +408,14 @@ def load_items(self): dict The source filename as key with list of face indices for the frame as value """ - faces = dict() - for face in self.file_list_sorted: - faces.setdefault(face["source_filename"], list()).append(face["face_index"]) - logger.trace(faces) + faces: dict[str, list[int]] = {} + for face in T.cast(list[tuple[str, "PNGHeaderDict"]], self.file_list_sorted): + src = face[1]["source"] + faces.setdefault(src["source_filename"], []).append(src["face_index"]) + logger.trace(faces) # type: ignore return faces - def sorted_items(self): + def sorted_items(self) -> list[tuple[str, PNGHeaderDict]]: """ Return the items sorted by the saved file name. Returns @@ -299,22 +423,34 @@ def sorted_items(self): list List of `dict` objects for each face found, sorted by the face's current filename """ - items = sorted(self.process_folder(), key=lambda x: (x["current_filename"])) - logger.trace(items) + items = sorted(self.process_folder(), key=itemgetter(0)) + logger.trace(items) # type: ignore return items class Frames(MediaLoader): """ Object to hold the frames that are to be checked against """ - def process_folder(self): - """ Iterate through the frames folder pulling the base filename """ + def process_folder(self) -> Generator[dict[str, str], None, None]: + """ Iterate through the frames folder pulling the base filename + + Yields + ------ + dict + The full framename, the filename and the file extension of the frame + """ iterator = self.process_video if self.is_video else self.process_frames for item in iterator(): yield item - def process_frames(self): - """ Process exported Frames """ + def process_frames(self) -> Generator[dict[str, str], None, None]: + """ Process exported Frames + + Yields + ------ + dict + The full framename, the filename and the file extension of the frame + """ logger.info("Loading file list from %s", self.folder) for frame in os.listdir(self.folder): if not self.valid_extension(frame): @@ -325,69 +461,122 @@ def process_frames(self): retval = {"frame_fullname": frame, "frame_name": filename, "frame_extension": file_extension} - logger.trace(retval) + logger.trace(retval) # type: ignore yield retval - def process_video(self): - """Dummy in frames for video """ + def process_video(self) -> Generator[dict[str, str], None, None]: + """Dummy in frames for video + + Yields + ------ + dict + The full framename, the filename and the file extension of the frame + """ logger.info("Loading video frames from %s", self.folder) - vidname = os.path.splitext(os.path.basename(self.folder))[0] + vidname, ext = os.path.splitext(os.path.basename(self.folder)) for i in range(self.count): idx = i + 1 # Keep filename format for outputted face - filename = "{}_{:06d}".format(vidname, idx) - retval = {"frame_fullname": "{}.png".format(filename), + filename = f"{vidname}_{idx:06d}" + retval = {"frame_fullname": f"{filename}{ext}", "frame_name": filename, - "frame_extension": ".png"} - logger.trace(retval) + "frame_extension": ext} + logger.trace(retval) # type: ignore yield retval - def load_items(self): - """ Load the frame info into dictionary """ - frames = dict() - for frame in self.file_list_sorted: + def load_items(self) -> dict[str, tuple[str, str]]: + """ Load the frame info into dictionary + + Returns + ------- + dict + Fullname as key, tuple of frame name and extension as value + """ + frames: dict[str, tuple[str, str]] = {} + for frame in T.cast(list[dict[str, str]], self.file_list_sorted): frames[frame["frame_fullname"]] = (frame["frame_name"], frame["frame_extension"]) - logger.trace(frames) + logger.trace(frames) # type: ignore return frames - def sorted_items(self): - """ Return the items sorted by filename """ + def sorted_items(self) -> list[dict[str, str]]: + """ Return the items sorted by filename + + Returns + ------- + list + The sorted list of frame information + """ items = sorted(self.process_folder(), key=lambda x: (x["frame_name"])) - logger.trace(items) + logger.trace(items) # type: ignore return items class ExtractedFaces(): - """ Holds the extracted faces and matrix for - alignments """ - def __init__(self, frames, alignments, size=512): - logger.trace("Initializing %s: size: %s", self.__class__.__name__, size) + """ Holds the extracted faces and matrix for alignments + + Parameters + ---------- + frames: :class:`Frames` + The frames object to extract faces from + alignments: :class:`AlignmentData` + The alignment data corresponding to the frames + size: int, optional + The extract face size. Default: 512 + """ + def __init__(self, frames: Frames, alignments: AlignmentData, size: int = 512) -> None: + logger.trace("Initializing %s: size: %s", # type: ignore + self.__class__.__name__, size) self.size = size self.padding = int(size * 0.1875) self.alignments = alignments self.frames = frames - self.current_frame = None - self.faces = list() - logger.trace("Initialized %s", self.__class__.__name__) + self.current_frame: str | None = None + self.faces: list[DetectedFace] = [] + logger.trace("Initialized %s", self.__class__.__name__) # type: ignore - def get_faces(self, frame, image=None): - """ Return faces and transformed landmarks - for each face in a given frame with it's alignments""" - logger.trace("Getting faces for frame: '%s'", frame) + def get_faces(self, frame: str, image: np.ndarray | None = None) -> None: + """ Obtain faces and transformed landmarks for each face in a given frame with its + alignments + + Parameters + ---------- + frame: str + The frame name to obtain faces for + image: :class:`numpy.ndarray`, optional + The image to extract the face from, if we already have it, otherwise ``None`` to + load the image. Default: ``None`` + """ + logger.trace("Getting faces for frame: '%s'", frame) # type: ignore self.current_frame = None alignments = self.alignments.get_faces_in_frame(frame) - logger.trace("Alignments for frame: (frame: '%s', alignments: %s)", frame, alignments) + logger.trace("Alignments for frame: (frame: '%s', alignments: %s)", # type: ignore + frame, alignments) if not alignments: - self.faces = list() + self.faces = [] return image = self.frames.load_image(frame) if image is None else image self.faces = [self.extract_one_face(alignment, image) for alignment in alignments] self.current_frame = frame - def extract_one_face(self, alignment, image): - """ Extract one face from image """ - logger.trace("Extracting one face: (frame: '%s', alignment: %s)", + def extract_one_face(self, + alignment: AlignmentFileDict, + image: np.ndarray) -> DetectedFace: + """ Extract one face from image + + Parameters + ---------- + alignment: dict + The alignment for a single face + image: :class:`numpy.ndarray` + The image to extract the face from + + Returns + ------- + :class:`~lib.align.DetectedFace` + The detected face object for the given alignment with the aligned face loaded + """ + logger.trace("Extracting one face: (frame: '%s', alignment: %s)", # type: ignore self.current_frame, alignment) face = DetectedFace() face.from_alignment(alignment, image=image) @@ -395,20 +584,51 @@ def extract_one_face(self, alignment, image): face.thumbnail = generate_thumbnail(face.aligned.face, size=80, quality=60) return face - def get_faces_in_frame(self, frame, update=False, image=None): - """ Return the faces for the selected frame """ - logger.trace("frame: '%s', update: %s", frame, update) + def get_faces_in_frame(self, + frame: str, + update: bool = False, + image: np.ndarray | None = None) -> list[DetectedFace]: + """ Return the faces for the selected frame + + Parameters + ---------- + frame: str + The frame name to get the faces for + update: bool, optional + ``True`` if the faces should be refreshed regardless of current frame. ``False`` to not + force a refresh. Default ``False`` + image: :class:`numpy.ndarray`, optional + Image to load faces from if it exists, otherwise ``None`` to load the image. + Default: ``None`` + + Returns + ------- + list + List of :class:`~lib.align.DetectedFace` objects for the frame, with the aligned face + loaded + """ + logger.trace("frame: '%s', update: %s", frame, update) # type: ignore if self.current_frame != frame or update: self.get_faces(frame, image=image) return self.faces - def get_roi_size_for_frame(self, frame): - """ Return the size of the original extract box for - the selected frame """ - logger.trace("frame: '%s'", frame) + def get_roi_size_for_frame(self, frame: str) -> list[int]: + """ Return the size of the original extract box for the selected frame. + + Parameters + ---------- + frame: str + The frame to obtain the original sized bounding boxes for + + Returns + ------- + list + List of original pixel sizes of faces held within the frame + """ + logger.trace("frame: '%s'", frame) # type: ignore if self.current_frame != frame: self.get_faces(frame) - sizes = list() + sizes = [] for face in self.faces: roi = face.aligned.original_roi.squeeze() top_left, top_right = roi[0], roi[3] @@ -419,5 +639,5 @@ def get_roi_size_for_frame(self, frame): else: length = int(((len_x ** 2) + (len_y ** 2)) ** 0.5) sizes.append(length) - logger.trace("sizes: '%s'", sizes) + logger.trace("sizes: '%s'", sizes) # type: ignore return sizes diff --git a/tools/effmpeg/cli.py b/tools/effmpeg/cli.py index ececeeaa39..ac7647f8e4 100644 --- a/tools/effmpeg/cli.py +++ b/tools/effmpeg/cli.py @@ -1,20 +1,48 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ +import argparse import gettext from lib.cli.args import FaceSwapArgs from lib.cli.actions import ContextFullPaths, FileFullPaths, Radio -from lib.utils import _image_extensions +from lib.utils import IMAGE_EXTENSIONS # LOCALES _LANG = gettext.translation("tools.effmpeg.cli", localedir="locales", fallback=True) _ = _LANG.gettext - _HELPTEXT = _("This command allows you to easily execute common ffmpeg tasks.") +def __parse_transpose(value: str) -> str: + """ Parse transpose option + + Parameters + ---------- + value: str + The value to parse + + Returns + ------- + str + The option item for the given value + """ + index = 0 + opts = ["(0, 90CounterClockwise&VerticalFlip)", + "(1, 90Clockwise)", + "(2, 90CounterClockwise)", + "(3, 90Clockwise&VerticalFlip)"] + if len(value) == 1: + index = int(value) + else: + for i in range(5): + if value in opts[i]: + index = i + break + return opts[index] + + class EffmpegArgs(FaceSwapArgs): """ Class to parse the command line arguments for EFFMPEG tool """ @@ -24,167 +52,184 @@ def get_info(): return _("A wrapper for ffmpeg for performing image <> video converting.") @staticmethod - def __parse_transpose(value): - index = 0 - opts = ["(0, 90CounterClockwise&VerticalFlip)", - "(1, 90Clockwise)", - "(2, 90CounterClockwise)", - "(3, 90Clockwise&VerticalFlip)"] - if len(value) == 1: - index = int(value) - else: - for i in range(5): - if value in opts[i]: - index = i - break - return opts[index] - - def get_argument_list(self): - argument_list = list() - argument_list.append(dict( - opts=('-a', '--action'), - action=Radio, - dest="action", - choices=("extract", "gen-vid", "get-fps", "get-info", "mux-audio", "rescale", "rotate", - "slice"), - default="extract", - help=_("R|Choose which action you want ffmpeg ffmpeg to do." - "\nL|'extract': turns videos into images " - "\nL|'gen-vid': turns images into videos " - "\nL|'get-fps' returns the chosen video's fps." - "\nL|'get-info' returns information about a video." - "\nL|'mux-audio' add audio from one video to another." - "\nL|'rescale' resize video." - "\nL|'rotate' rotate video." - "\nL|'slice' cuts a portion of the video into a separate video file."))) - argument_list.append(dict( - opts=('-i', '--input'), - action=ContextFullPaths, - dest="input", - default="input", - help=_("Input file."), - group=_("data"), - required=True, - action_option="-a", - filetypes="video")) - argument_list.append(dict( - opts=('-o', '--output'), - action=ContextFullPaths, - group=_("data"), - default="", - dest="output", - help=_("Output file. If no output is specified then: if the output is meant to be a " - "video then a video called 'out.mkv' will be created in the input directory; " - "if the output is meant to be a directory then a directory called 'out' will " - "be created inside the input directory. Note: the chosen output file extension " - "will determine the file encoding."), - action_option="-a", - filetypes="video")) - argument_list.append(dict( - opts=('-r', '--reference-video'), - action=FileFullPaths, - dest="ref_vid", - group=_("data"), - default=None, - help=_("Path to reference video if 'input' was not a video."), - filetypes="video")) - argument_list.append(dict( - opts=('-fps', '--fps'), - type=str, - dest="fps", - group=_("output"), - default="-1.0", - help=_("Provide video fps. Can be an integer, float or fraction. Negative values will " - "will make the program try to get the fps from the input or reference " - "videos."))) - argument_list.append(dict( - opts=("-ef", "--extract-filetype"), - action=Radio, - choices=_image_extensions, - dest="extract_ext", - group=_("output"), - default=".png", - help=_("Image format that extracted images should be saved as. '.bmp' will offer the " - "fastest extraction speed, but will take the most storage space. '.png' will " - "be slower but will take less storage."))) - argument_list.append(dict( - opts=('-s', '--start'), - type=str, - dest="start", - group=_("clip"), - default="00:00:00", - help=_("Enter the start time from which an action is to be applied. Default: " - "00:00:00, in HH:MM:SS format. You can also enter the time with or without the " - "colons, e.g. 00:0000 or 026010."))) - argument_list.append(dict( - opts=('-e', '--end'), - type=str, - dest="end", - group=_("clip"), - default="00:00:00", - help=_("Enter the end time to which an action is to be applied. If both an end time " - "and duration are set, then the end time will be used and the duration will be " - "ignored. Default: 00:00:00, in HH:MM:SS."))) - argument_list.append(dict( - opts=('-d', '--duration'), - type=str, - dest="duration", - group=_("clip"), - default="00:00:00", - help=_("Enter the duration of the chosen action, for example if you enter 00:00:10 " - "for slice, then the first 10 seconds after and including the start time will " - "be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. You can " - "also enter the time with or without the colons, e.g. 00:0000 or 026010."))) - argument_list.append(dict( - opts=('-m', '--mux-audio'), - action="store_true", - dest="mux_audio", - group=_("output"), - default=False, - help=_("Mux the audio from the reference video into the input video. This option is " - "only used for the 'gen-vid' action. 'mux-audio' action has this turned on " - "implicitly."))) - argument_list.append(dict( - opts=('-tr', '--transpose'), - choices=("(0, 90CounterClockwise&VerticalFlip)", - "(1, 90Clockwise)", - "(2, 90CounterClockwise)", - "(3, 90Clockwise&VerticalFlip)"), - type=lambda v: self.__parse_transpose(v), # pylint:disable=unnecessary-lambda - dest="transpose", - group=_("rotate"), - default=None, - help=_("Transpose the video. If transpose is set, then degrees will be ignored. For " - "cli you can enter either the number or the long command name, e.g. to use (1, " - "90Clockwise) -tr 1 or -tr 90Clockwise"))) - argument_list.append(dict( - opts=('-de', '--degrees'), - type=str, - dest="degrees", - default=None, - group=_("rotate"), - help=_("Rotate the video clockwise by the given number of degrees."))) - argument_list.append(dict( - opts=('-sc', '--scale'), - type=str, - dest="scale", - group=_("output"), - default="1920x1080", - help=_("Set the new resolution scale if the chosen action is 'rescale'."))) - argument_list.append(dict( - opts=('-q', '--quiet'), - action="store_true", - dest="quiet", - group=_("settings"), - default=False, - help=_("Reduces output verbosity so that only serious errors are printed. If both " - "quiet and verbose are set, verbose will override quiet."))) - argument_list.append(dict( - opts=('-v', '--verbose'), - action="store_true", - dest="verbose", - group=_("settings"), - default=False, - help=_("Increases output verbosity. If both quiet and verbose are set, verbose will " - "override quiet."))) + def get_argument_list(): + argument_list = [] + argument_list.append({ + "opts": ('-a', '--action'), + "action": Radio, + "dest": "action", + "choices": ("extract", "gen-vid", "get-fps", "get-info", "mux-audio", "rescale", + "rotate", "slice"), + "default": "extract", + "help": _("R|Choose which action you want ffmpeg ffmpeg to do." + "\nL|'extract': turns videos into images " + "\nL|'gen-vid': turns images into videos " + "\nL|'get-fps' returns the chosen video's fps." + "\nL|'get-info' returns information about a video." + "\nL|'mux-audio' add audio from one video to another." + "\nL|'rescale' resize video." + "\nL|'rotate' rotate video." + "\nL|'slice' cuts a portion of the video into a separate video file.")}) + argument_list.append({ + "opts": ('-i', '--input'), + "action": ContextFullPaths, + "dest": "input", + "default": "input", + "help": _("Input file."), + "group": _("data"), + "required": True, + "action_option": "-a", + "filetypes": "video"}) + argument_list.append({ + "opts": ('-o', '--output'), + "action": ContextFullPaths, + "group": _("data"), + "default": "", + "dest": "output", + "help": _("Output file. If no output is specified then: if the output is meant to be " + "a video then a video called 'out.mkv' will be created in the input " + "directory; if the output is meant to be a directory then a directory " + "called 'out' will be created inside the input directory. Note: the chosen " + "output file extension will determine the file encoding."), + "action_option": "-a", + "filetypes": "video"}) + argument_list.append({ + "opts": ('-r', '--reference-video'), + "action": FileFullPaths, + "dest": "ref_vid", + "group": _("data"), + "default": None, + "help": _("Path to reference video if 'input' was not a video."), + "filetypes": "video"}) + argument_list.append({ + "opts": ('-R', '--fps'), + "type": str, + "dest": "fps", + "group": _("output"), + "default": "-1.0", + "help": _("Provide video fps. Can be an integer, float or fraction. Negative values " + "will will make the program try to get the fps from the input or reference " + "videos.")}) + argument_list.append({ + "opts": ("-E", "--extract-filetype"), + "action": Radio, + "choices": IMAGE_EXTENSIONS, + "dest": "extract_ext", + "group": _("output"), + "default": ".png", + "help": _("Image format that extracted images should be saved as. '.bmp' will offer " + "the fastest extraction speed, but will take the most storage space. '.png' " + "will be slower but will take less storage.")}) + argument_list.append({ + "opts": ('-s', '--start'), + "type": str, + "dest": "start", + "group": _("clip"), + "default": "00:00:00", + "help": _("Enter the start time from which an action is to be applied. Default: " + "00:00:00, in HH:MM:SS format. You can also enter the time with or without " + "the colons, e.g. 00:0000 or 026010.")}) + argument_list.append({ + "opts": ('-e', '--end'), + "type": str, + "dest": "end", + "group": _("clip"), + "default": "00:00:00", + "help": _("Enter the end time to which an action is to be applied. If both an end " + "time and duration are set, then the end time will be used and the duration " + "will be ignored. Default: 00:00:00, in HH:MM:SS.")}) + argument_list.append({ + "opts": ('-d', '--duration'), + "type": str, + "dest": "duration", + "group": _("clip"), + "default": "00:00:00", + "help": _("Enter the duration of the chosen action, for example if you enter 00:00:10 " + "for slice, then the first 10 seconds after and including the start time " + "will be cut out into a new video. Default: 00:00:00, in HH:MM:SS format. " + "You can also enter the time with or without the colons, e.g. 00:0000 or " + "026010.")}) + argument_list.append({ + "opts": ('-m', '--mux-audio'), + "action": "store_true", + "dest": "mux_audio", + "group": _("output"), + "default": False, + "help": _("Mux the audio from the reference video into the input video. This option " + "is only used for the 'gen-vid' action. 'mux-audio' action has this turned " + "on implicitly.")}) + argument_list.append({ + "opts": ('-T', '--transpose'), + "choices": ("(0, 90CounterClockwise&VerticalFlip)", + "(1, 90Clockwise)", + "(2, 90CounterClockwise)", + "(3, 90Clockwise&VerticalFlip)"), + "type": lambda v: __parse_transpose(v), # pylint:disable=unnecessary-lambda + "dest": "transpose", + "group": _("rotate"), + "default": None, + "help": _("Transpose the video. If transpose is set, then degrees will be ignored. " + "For cli you can enter either the number or the long command name, e.g. to " + "use (1, 90Clockwise) -tr 1 or -tr 90Clockwise")}) + argument_list.append({ + "opts": ('-D', '--degrees'), + "type": str, + "dest": "degrees", + "default": None, + "group": _("rotate"), + "help": _("Rotate the video clockwise by the given number of degrees.")}) + argument_list.append({ + "opts": ('-S', '--scale'), + "type": str, + "dest": "scale", + "group": _("output"), + "default": "1920x1080", + "help": _("Set the new resolution scale if the chosen action is 'rescale'.")}) + argument_list.append({ + "opts": ('-q', '--quiet'), + "action": "store_true", + "dest": "quiet", + "group": _("settings"), + "default": False, + "help": _("Reduces output verbosity so that only serious errors are printed. If both " + "quiet and verbose are set, verbose will override quiet.")}) + argument_list.append({ + "opts": ('-v', '--verbose'), + "action": "store_true", + "dest": "verbose", + "group": _("settings"), + "default": False, + "help": _("Increases output verbosity. If both quiet and verbose are set, verbose " + "will override quiet.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ('-fps', ), + "type": str, + "dest": "depr_fps_fps_R", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-ef", ), + "type": str, + "choices": IMAGE_EXTENSIONS, + "dest": "depr_extract-filetype_et_E", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ('-tr', ), + "choices": ("(0, 90CounterClockwise&VerticalFlip)", + "(1, 90Clockwise)", + "(2, 90CounterClockwise)", + "(3, 90Clockwise&VerticalFlip)"), + "type": lambda v: __parse_transpose(v), # pylint:disable=unnecessary-lambda + "dest": "depr_transpose_tr_T", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ('-de', ), + "type": str, + "dest": "depr_degrees_de_D", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ('-sc', ), + "type": str, + "dest": "depr_scale_sc_S", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/effmpeg/effmpeg.py b/tools/effmpeg/effmpeg.py index 025b7d4a31..187f0a08aa 100644 --- a/tools/effmpeg/effmpeg.py +++ b/tools/effmpeg/effmpeg.py @@ -17,9 +17,9 @@ from ffmpy import FFmpeg, FFRuntimeError # faceswap imports -from lib.utils import _image_extensions, _video_extensions +from lib.utils import handle_deprecated_cliopts, IMAGE_EXTENSIONS, VIDEO_EXTENSIONS -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class DataItem(): @@ -27,10 +27,10 @@ class DataItem(): A simple class used for storing the media data items and directories that Effmpeg uses for 'input', 'output' and 'ref_vid'. """ - vid_ext = _video_extensions + vid_ext = VIDEO_EXTENSIONS # future option in effmpeg to use audio file for muxing - audio_ext = ['.aiff', '.flac', '.mp3', '.wav'] - img_ext = _image_extensions + audio_ext = [".aiff", ".flac", ".mp3", ".wav"] + img_ext = IMAGE_EXTENSIONS def __init__(self, path=None, name=None, item_type=None, ext=None, fps=None): @@ -68,11 +68,11 @@ def set_type_ext(self, path=None): if self.path is not None: item_ext = os.path.splitext(self.path)[1].lower() if item_ext in DataItem.vid_ext: - item_type = 'vid' + item_type = "vid" elif item_ext in DataItem.audio_ext: - item_type = 'audio' + item_type = "audio" else: - item_type = 'dir' + item_type = "dir" self.type = item_type self.ext = item_ext logger.debug("path: '%s', type: '%s', ext: '%s'", self.path, self.type, self.ext) @@ -140,16 +140,16 @@ class Effmpeg(): # Class variable that stores the common ffmpeg arguments based on verbosity __common_ffmpeg_args_dict = {"normal": "-hide_banner ", "quiet": "-loglevel panic -hide_banner ", - "verbose": ''} + "verbose": ""} # _common_ffmpeg_args is the class variable that will get used by various # actions and it will be set by the process_arguments() method based on # passed verbosity - _common_ffmpeg_args = '' + _common_ffmpeg_args = "" def __init__(self, arguments): logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) - self.args = arguments + self.args = handle_deprecated_cliopts(arguments) self.exe = im_ffm.get_ffmpeg_exe() self.input = DataItem() self.output = DataItem() @@ -160,17 +160,8 @@ def __init__(self, arguments): self.print_ = False logger.debug("Initialized %s", self.__class__.__name__) - def process(self): - """ EFFMPEG Process """ - logger.debug("Running Effmpeg") - # Format action to match the method name - self.args.action = self.args.action.replace('-', '_') - logger.debug("action: '%s", self.args.action) - - # Instantiate input DataItem object - self.input = DataItem(path=self.args.input) - - # Instantiate output DataItem object + def _set_output(self) -> None: + """ Set :attr:`output` based on input arguments """ if self.args.action in self._actions_have_dir_output: self.output = DataItem(path=self.__get_default_output()) elif self.args.action in self._actions_have_vid_output: @@ -180,70 +171,101 @@ def process(self): else: self.output = DataItem(path=self.__get_default_output()) - if self.args.ref_vid is None \ - or self.args.ref_vid == '': + def _set_ref_video(self) -> None: + """ Set :attr:`ref_vid` based on input arguments """ + if self.args.ref_vid is None or self.args.ref_vid == "": self.args.ref_vid = None - # Instantiate ref_vid DataItem object self.ref_vid = DataItem(path=self.args.ref_vid) - # Check that correct input and output arguments were provided + def _check_inputs(self) -> None: + """ Validate provided arguments are valid + + Raises + ------ + ValueError + If provided arguments are not valid + """ + if self.args.action in self._actions_have_dir_input and not self.input.is_type("dir"): - raise ValueError("The chosen action requires a directory as its " - "input, but you entered: " - "{}".format(self.input.path)) + raise ValueError("The chosen action requires a directory as its input, but you " + f"entered: {self.input.path}") if self.args.action in self._actions_have_vid_input and not self.input.is_type("vid"): - raise ValueError("The chosen action requires a video as its " - "input, but you entered: " - "{}".format(self.input.path)) + raise ValueError("The chosen action requires a video as its input, but you entered: " + f"{self.input.path}") if self.args.action in self._actions_have_dir_output and not self.output.is_type("dir"): - raise ValueError("The chosen action requires a directory as its " - "output, but you entered: " - "{}".format(self.output.path)) + raise ValueError("The chosen action requires a directory as its output, but you " + f"entered: {self.output.path}") if self.args.action in self._actions_have_vid_output and not self.output.is_type("vid"): - raise ValueError("The chosen action requires a video as its " - "output, but you entered: " - "{}".format(self.output.path)) + raise ValueError("The chosen action requires a video as its output, but you entered: " + f"{self.output.path}") # Check that ref_vid is a video when it needs to be if self.args.action in self._actions_req_ref_video: if self.ref_vid.is_type("none"): - raise ValueError("The file chosen as the reference video is " - "not a video, either leave the field blank " - "or type 'None': " - "{}".format(self.ref_vid.path)) + raise ValueError("The file chosen as the reference video is not a video, either " + f"leave the field blank or type 'None': {self.ref_vid.path}") elif self.args.action in self._actions_can_use_ref_video: if self.ref_vid.is_type("none"): logger.warning("Warning: no reference video was supplied, even though " "one may be used with the chosen action. If this is " "intentional then ignore this warning.") - # Process start and duration arguments + def _set_times(self) -> None: + """Set start, end and duration attributes """ self.start = self.parse_time(self.args.start) self.end = self.parse_time(self.args.end) if not self.__check_equals_time(self.args.end, "00:00:00"): self.duration = self.__get_duration(self.start, self.end) else: self.duration = self.parse_time(str(self.args.duration)) + + def _set_fps(self) -> None: + """ Set :attr:`arguments.fps` based on input arguments""" # If fps was left blank in gui, set it to default -1.0 value - if self.args.fps == '': + if self.args.fps == "": self.args.fps = str(-1.0) # Try to set fps automatically if needed and not supplied by user if self.args.action in self._actions_req_fps \ and self.__convert_fps(self.args.fps) <= 0: - if self.__check_have_fps(['r', 'i']): + if self.__check_have_fps(["r", "i"]): _error_str = "No fps, input or reference video was supplied, " _error_str += "hence it's not possible to " - _error_str += "'{}'.".format(self.args.action) + _error_str += f"'{self.args.action}'." raise ValueError(_error_str) - if self.output.fps is not None and self.__check_have_fps(['r', 'i']): + if self.output.fps is not None and self.__check_have_fps(["r", "i"]): self.args.fps = self.output.fps - elif self.ref_vid.fps is not None and self.__check_have_fps(['i']): + elif self.ref_vid.fps is not None and self.__check_have_fps(["i"]): self.args.fps = self.ref_vid.fps - elif self.input.fps is not None and self.__check_have_fps(['r']): + elif self.input.fps is not None and self.__check_have_fps(["r"]): self.args.fps = self.input.fps + def process(self): + """ EFFMPEG Process """ + logger.debug("Running Effmpeg") + # Format action to match the method name + self.args.action = self.args.action.replace("-", "_") + logger.debug("action: '%s'", self.args.action) + + # Instantiate input DataItem object + self.input = DataItem(path=self.args.input) + + # Instantiate output DataItem object + self._set_output() + + # Instantiate ref_vid DataItem object + self._set_ref_video() + + # Check that correct input and output arguments were provided + self._check_inputs() + + # Process start and duration arguments + self._set_times() + + # Set fps + self._set_fps() + # Processing transpose if self.args.transpose is None or \ self.args.transpose.lower() == "none": @@ -254,7 +276,7 @@ def process(self): # Processing degrees if self.args.degrees is None \ or self.args.degrees.lower() == "none" \ - or self.args.degrees == '': + or self.args.degrees == "": self.args.degrees = None elif self.args.transpose is None: try: @@ -300,7 +322,7 @@ def extract(input_=None, output=None, fps=None, # pylint:disable=unused-argumen input_, output, fps, extract_ext, start, duration) _input_opts = Effmpeg._common_ffmpeg_args[:] if start is not None and duration is not None: - _input_opts += '-ss {} -t {}'.format(start, duration) + _input_opts += f"-ss {start} -t {duration}" _input = {input_.path: _input_opts} _output_opts = '-y -vf fps="' + str(fps) + '" -q:v 1' _output_path = output.path + "/" + input_.name + "_%05d" + extract_ext @@ -318,12 +340,12 @@ def gen_vid(input_=None, output=None, fps=None, # pylint:disable=unused-argumen filename = Effmpeg.__get_extracted_filename(input_.path) _input_opts = Effmpeg._common_ffmpeg_args[:] _input_path = os.path.join(input_.path, filename) - _fps_arg = '-r ' + str(fps) + ' ' + _fps_arg = "-r " + str(fps) + " " _input_opts += _fps_arg + "-f image2 " - _output_opts = '-y ' + _fps_arg + ' -c:v libx264' + _output_opts = "-y " + _fps_arg + " -c:v libx264" if mux_audio: - _ref_vid_opts = '-c copy -map 0:0 -map 1:1' - _output_opts = _ref_vid_opts + ' ' + _output_opts + _ref_vid_opts = "-c copy -map 0:0 -map 1:1" + _output_opts = _ref_vid_opts + " " + _output_opts _inputs = OrderedDict([(_input_path, _input_opts), (ref_vid.path, None)]) else: _inputs = {_input_path: _input_opts} @@ -377,13 +399,12 @@ def rotate(input_=None, output=None, degrees=None, # pylint:disable=unused-argu transpose=None, exe=None, **kwargs): """ Rotate Video """ if transpose is None and degrees is None: - raise ValueError("You have not supplied a valid transpose or " - "degrees value:\ntranspose: {}\ndegrees: " - "{}".format(transpose, degrees)) + raise ValueError("You have not supplied a valid transpose or degrees value:\n" + f"transpose: {transpose}\ndegrees: {degrees}") _input_opts = Effmpeg._common_ffmpeg_args[:] - _output_opts = '-y -c:a copy -vf ' - _bilinear = '' + _output_opts = "-y -c:a copy -vf " + _bilinear = "" if transpose is not None: _output_opts += 'transpose="' + str(transpose) + '"' elif int(degrees) != 0: @@ -402,7 +423,7 @@ def mux_audio(input_=None, output=None, ref_vid=None, # pylint:disable=unused-a """ Mux Audio """ _input_opts = Effmpeg._common_ffmpeg_args[:] _ref_vid_opts = None - _output_opts = '-y -c copy -map 0:0 -map 1:1 -shortest' + _output_opts = "-y -c copy -map 0:0 -map 1:1 -shortest" _inputs = OrderedDict([(input_.path, _input_opts), (ref_vid.path, _ref_vid_opts)]) _outputs = {output.path: _output_opts} Effmpeg.__run_ffmpeg(exe=exe, inputs=_inputs, outputs=_outputs) @@ -433,28 +454,28 @@ def __get_default_output(self): if the user didn't specify it. """ if self.args.output == "": if self.args.action in self._actions_have_dir_output: - retval = os.path.join(self.input.dirname, 'out') + retval = os.path.join(self.input.dirname, "out") elif self.args.action in self._actions_have_vid_output: if self.input.is_type("media"): # Using the same extension as input leads to very poor # output quality, hence the default is mkv for now retval = os.path.join(self.input.dirname, "out.mkv") # + self.input.ext) else: # case if input was a directory - retval = os.path.join(self.input.dirname, 'out.mkv') + retval = os.path.join(self.input.dirname, "out.mkv") else: retval = self.args.output logger.debug(retval) return retval def __check_have_fps(self, items): - items_to_check = list() + items_to_check = [] for i in items: - if i == 'r': - items_to_check.append('ref_vid') - elif i == 'i': - items_to_check.append('input') - elif i == 'o': - items_to_check.append('output') + if i == "r": + items_to_check.append("ref_vid") + elif i == "i": + items_to_check.append("input") + elif i == "o": + items_to_check.append("output") return all(getattr(self, i).fps is None for i in items_to_check) @@ -470,8 +491,7 @@ def __run_ffmpeg(exe=im_ffm.get_ffmpeg_exe(), inputs=None, outputs=None): if ffe.exit_code == 255: pass else: - raise ValueError("An unexpected FFRuntimeError occurred: " - "{}".format(ffe)) + raise ValueError(f"An unexpected FFRuntimeError occurred: {ffe}") from ffe except KeyboardInterrupt: pass # Do nothing if voluntary interruption logger.debug("ffmpeg finished") @@ -479,8 +499,8 @@ def __run_ffmpeg(exe=im_ffm.get_ffmpeg_exe(), inputs=None, outputs=None): @staticmethod def __convert_fps(fps): """ Convert to Frames per Second """ - if '/' in fps: - _fps = fps.split('/') + if "/" in fps: + _fps = fps.split("/") retval = float(_fps[0]) / float(_fps[1]) else: retval = float(fps) @@ -490,15 +510,13 @@ def __convert_fps(fps): @staticmethod def __get_duration(start_time, end_time): """ Get the duration """ - start = [int(i) for i in start_time.split(':')] - end = [int(i) for i in end_time.split(':')] + start = [int(i) for i in start_time.split(":")] + end = [int(i) for i in end_time.split(":")] start = datetime.timedelta(hours=start[0], minutes=start[1], seconds=start[2]) end = datetime.timedelta(hours=end[0], minutes=end[1], seconds=end[2]) delta = end - start secs = delta.total_seconds() - retval = '{:02}:{:02}:{:02}'.format(int(secs // 3600), - int(secs % 3600 // 60), - int(secs % 60)) + retval = f"{int(secs // 3600):02}:{int(secs % 3600 // 60):02}:{int(secs % 60):02}" logger.debug(retval) return retval @@ -506,7 +524,7 @@ def __get_duration(start_time, end_time): def __get_extracted_filename(path): """ Get the extracted filename """ logger.debug("path: '%s'", path) - filename = '' + filename = "" for file in os.listdir(path): if any(i in file for i in DataItem.img_ext): filename = file @@ -515,7 +533,7 @@ def __get_extracted_filename(path): filename, img_ext = os.path.splitext(filename) zero_pad = Effmpeg.__get_zero_pad(filename) name = filename[:-zero_pad] - retval = "{}%{}d{}".format(name, zero_pad, img_ext) + retval = f"{name}%{zero_pad}d{img_ext}" logger.debug("filename: %s, img_ext: '%s', zero_pad: %s, name: '%s'", filename, img_ext, zero_pad, name) logger.debug(retval) @@ -527,25 +545,17 @@ def __get_zero_pad(filename): chkstring = filename[::-1] logger.trace("filename: %s, chkstring: %s", filename, chkstring) pos = 0 - for pos in range(len(chkstring)): - if not chkstring[pos].isdigit(): + for char in chkstring: + if not char.isdigit(): break logger.debug("filename: '%s', pos: %s", filename, pos) return pos - @staticmethod - def __check_is_valid_time(value): - """ Check valid time """ - val = value.replace(':', '') - retval = val.isdigit() - logger.debug("value: '%s', retval: %s", value, retval) - return retval - @staticmethod def __check_equals_time(value, time): """ Check equals time """ - val = value.replace(':', '') - tme = time.replace(':', '') + val = value.replace(":", "") + tme = time.replace(":", "") retval = val.zfill(6) == tme.zfill(6) logger.debug("value: '%s', time: %s, retval: %s", value, time, retval) return retval @@ -553,10 +563,10 @@ def __check_equals_time(value, time): @staticmethod def parse_time(txt): """ Parse Time """ - clean_txt = txt.replace(':', '') + clean_txt = txt.replace(":", "") hours = clean_txt[0:2] minutes = clean_txt[2:4] seconds = clean_txt[4:6] - retval = hours + ':' + minutes + ':' + seconds + retval = hours + ":" + minutes + ":" + seconds logger.debug("txt: '%s', retval: %s", txt, retval) return retval diff --git a/tools/manual/cli.py b/tools/manual/cli.py index 3cb571efc1..db27d785b6 100644 --- a/tools/manual/cli.py +++ b/tools/manual/cli.py @@ -1,15 +1,15 @@ #!/usr/bin/env python3 """ The Command Line Arguments for the Manual Editor tool. """ +import argparse import gettext -from lib.cli.args import FaceSwapArgs, DirOrFileFullPaths, FileFullPaths - +from lib.cli.args import FaceSwapArgs +from lib.cli.actions import DirOrFileFullPaths, FileFullPaths # LOCALES _LANG = gettext.translation("tools.manual", localedir="locales", fallback=True) _ = _LANG.gettext - _HELPTEXT = _("This command lets you perform various actions on frames, " "faces and alignments files using visual tools.") @@ -26,39 +26,54 @@ def get_info(): @staticmethod def get_argument_list(): """ Generate the command line argument list for the Manual Tool. """ - argument_list = list() - argument_list.append(dict( - opts=("-al", "--alignments"), - action=FileFullPaths, - filetypes="alignments", - type=str, - group=_("data"), - dest="alignments_path", - help=_("Path to the alignments file for the input, if not at the default location"))) - argument_list.append(dict( - opts=("-fr", "--frames"), - action=DirOrFileFullPaths, - filetypes="video", - required=True, - group=_("data"), - help=_("Video file or directory containing source frames that faces were extracted " - "from."))) - argument_list.append(dict( - opts=("-t", "--thumb-regen"), - action="store_true", - dest="thumb_regen", - default=False, - group=_("options"), - help=_("Force regeneration of the low resolution jpg thumbnails in the alignments " - "file."))) - argument_list.append(dict( - opts=("-s", "--single-process"), - action="store_true", - dest="single_process", - default=False, - group=_("options"), - help=_("The process attempts to speed up generation of thumbnails by extracting from " - "the video in parallel threads. For some videos, this causes the caching " - "process to hang. If this happens, then set this option to generate the " - "thumbnails in a slower, but more stable single thread."))) + argument_list = [] + argument_list.append({ + "opts": ("-a", "--alignments"), + "action": FileFullPaths, + "filetypes": "alignments", + "type": str, + "group": _("data"), + "dest": "alignments_path", + "help": _( + "Path to the alignments file for the input, if not at the default location")}) + argument_list.append({ + "opts": ("-f", "--frames"), + "action": DirOrFileFullPaths, + "filetypes": "video", + "required": True, + "group": _("data"), + "help": _( + "Video file or directory containing source frames that faces were extracted " + "from.")}) + argument_list.append({ + "opts": ("-t", "--thumb-regen"), + "action": "store_true", + "dest": "thumb_regen", + "default": False, + "group": _("options"), + "help": _( + "Force regeneration of the low resolution jpg thumbnails in the alignments " + "file.")}) + argument_list.append({ + "opts": ("-s", "--single-process"), + "action": "store_true", + "dest": "single_process", + "default": False, + "group": _("options"), + "help": _( + "The process attempts to speed up generation of thumbnails by extracting from the " + "video in parallel threads. For some videos, this causes the caching process to " + "hang. If this happens, then set this option to generate the thumbnails in a " + "slower, but more stable single thread.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-al", ), + "type": str, + "dest": "depr_alignments_al_a", + "help": argparse.SUPPRESS}) + argument_list.append({ + "opts": ("-fr", ), + "type": str, + "dest": "depr_frames_fr_f", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/manual/detected_faces.py b/tools/manual/detected_faces.py index 43138f51da..7dcd90fc83 100644 --- a/tools/manual/detected_faces.py +++ b/tools/manual/detected_faces.py @@ -2,31 +2,30 @@ """ Alignments handling for Faceswap's Manual Adjustments tool. Handles the conversion of alignments data to :class:`~lib.align.DetectedFace` objects, and the update of these faces when edits are made in the GUI. """ - +from __future__ import annotations import logging import os import sys import tkinter as tk +import typing as T from copy import deepcopy from queue import Queue, Empty -from time import sleep -from threading import Lock import cv2 -import imageio import numpy as np -from tqdm import tqdm from lib.align import Alignments, AlignedFace, DetectedFace from lib.gui.custom_widgets import PopupProgress from lib.gui.utils import FileHandler -from lib.image import (SingleFrameLoader, ImagesLoader, ImagesSaver, encode_image, - generate_thumbnail) +from lib.image import ImagesLoader, ImagesSaver, encode_image, generate_thumbnail from lib.multithreading import MultiThread from lib.utils import get_folder +if T.TYPE_CHECKING: + from . import manual + from lib.align.alignments import AlignmentFileDict, PNGHeaderDict -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class DetectedFaces(): @@ -46,102 +45,124 @@ class DetectedFaces(): extractor: :class:`~tools.manual.manual.Aligner` The pipeline for passing faces through the aligner and retrieving results """ - def __init__(self, tk_globals, alignments_path, input_location, extractor): + def __init__(self, + tk_globals: manual.TkGlobals, + alignments_path: str, + input_location: str, + extractor: manual.Aligner) -> None: logger.debug("Initializing %s: (tk_globals: %s. alignments_path: %s, input_location: %s " "extractor: %s)", self.__class__.__name__, tk_globals, alignments_path, input_location, extractor) self._globals = tk_globals - self._frame_faces = [] - self._updated_frame_indices = set() + self._frame_faces: list[list[DetectedFace]] = [] + self._updated_frame_indices: set[int] = set() + + self._alignments: Alignments = self._get_alignments(alignments_path, input_location) + self._alignments.update_legacy_has_source(os.path.basename(input_location)) - self._alignments = self._get_alignments(alignments_path, input_location) self._extractor = extractor self._tk_vars = self._set_tk_vars() - self._children = dict(io=_DiskIO(self, input_location), - update=FaceUpdate(self), - filter=Filter(self)) + + self._io = _DiskIO(self, input_location) + self._update = FaceUpdate(self) + self._filter = Filter(self) logger.debug("Initialized %s", self.__class__.__name__) # <<<< PUBLIC PROPERTIES >>>> # - # << SUBCLASSES >> # @property - def extractor(self): + def extractor(self) -> manual.Aligner: """ :class:`~tools.manual.manual.Aligner`: The pipeline for passing faces through the aligner and retrieving results. """ return self._extractor @property - def filter(self): + def filter(self) -> Filter: """ :class:`Filter`: Handles returning of faces and stats based on the current user set navigation mode filter. """ - return self._children["filter"] + return self._filter @property - def update(self): + def update(self) -> FaceUpdate: """ :class:`FaceUpdate`: Handles the adding, removing and updating of :class:`~lib.align.DetectedFace` stored within the alignments file. """ - return self._children["update"] + return self._update # << TKINTER VARIABLES >> # @property - def tk_unsaved(self): + def tk_unsaved(self) -> tk.BooleanVar: """ :class:`tkinter.BooleanVar`: The variable indicating whether the alignments have been updated since the last save. """ return self._tk_vars["unsaved"] @property - def tk_edited(self): + def tk_edited(self) -> tk.BooleanVar: """ :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred meaning a GUI redraw needs to be triggered. """ return self._tk_vars["edited"] @property - def tk_face_count_changed(self): + def tk_face_count_changed(self) -> tk.BooleanVar: """ :class:`tkinter.BooleanVar`: The variable indicating whether a face has been added or removed meaning the :class:`FaceViewer` grid redraw needs to be triggered. """ return self._tk_vars["face_count_changed"] # << STATISTICS >> # @property - def available_masks(self): - """ dict: The mask type names stored in the alignments; type as key with the number - of faces which possess the mask type as value. """ + def frame_list(self) -> list[str]: + """ list[str]: The list of all frame names that appear in the alignments file """ + return list(self._alignments.data) + + @property + def available_masks(self) -> dict[str, int]: + """ dict[str, int]: The mask type names stored in the alignments; type as key with the + number of faces which possess the mask type as value. """ return self._alignments.mask_summary @property - def current_faces(self): - """ list: The most up to date full list of :class:`~lib.align.DetectedFace` - objects. """ + def current_faces(self) -> list[list[DetectedFace]]: + """ list[list[:class:`~lib.align.DetectedFace`]]: The most up to date full list of detected + face objects. """ return self._frame_faces @property - def video_meta_data(self): - """ dict: The frame meta data stored in the alignments file. If data does not exist in the - alignments file then ``None`` is returned for each Key """ + def video_meta_data(self) -> dict[str, list[int] | list[float] | None]: + """ dict[str, list[int] | list[float] | None]: The frame meta data stored in the alignments + file. If data does not exist in the alignments file then ``None`` is returned for each + Key """ return self._alignments.video_meta_data @property - def face_count_per_index(self): - """ list: Count of faces for each frame. List is in frame index order. + def face_count_per_index(self) -> list[int]: + """ list[int]: Count of faces for each frame. List is in frame index order. The list needs to be calculated on the fly as the number of faces in a frame can change based on user actions. """ return [len(faces) for faces in self._frame_faces] # <<<< PUBLIC METHODS >>>> # - def is_frame_updated(self, frame_index): - """ bool: ``True`` if the given frame index has updated faces within it otherwise - ``False`` """ + def is_frame_updated(self, frame_index: int) -> bool: + """ Check whether the given frame index has been updated + + Parameters + ---------- + frame_index: int + The frame index to check + + Returns + ------- + bool: + ``True`` if the given frame index has updated faces within it otherwise ``False`` + """ return frame_index in self._updated_frame_indices - def load_faces(self): + def load_faces(self) -> None: """ Load the faces as :class:`~lib.align.DetectedFace` objects from the alignments file. """ - self._children["io"].load() + self._io.load() - def save(self): + def save(self) -> None: """ Save the alignments file with the latest edits. """ - self._children["io"].save() + self._io.save() def revert_to_saved(self, frame_index): """ Revert the frame's alignments to their saved version for the given frame index. @@ -151,23 +172,23 @@ def revert_to_saved(self, frame_index): frame_index: int The frame that should have their faces reverted to their saved version """ - self._children["io"].revert_to_saved(frame_index) + self._io.revert_to_saved(frame_index) - def extract(self): + def extract(self) -> None: """ Extract the faces in the current video to a user supplied folder. """ - self._children["io"].extract() + self._io.extract() - def save_video_meta_data(self, pts_time, keyframes): + def save_video_meta_data(self, pts_time: list[float], keyframes: list[int]) -> None: """ Save video meta data to the alignments file. This is executed if the video meta data does not already exist in the alignments file, so the video does not need to be scanned on every use of the Manual Tool. Parameters ---------- - pts_time: list - A list of presentation timestamps (`float`) in frame index order for every frame in - the input video - keyframes: list + pts_time: list[float] + A list of presentation timestamps in frame index order for every frame in the input + video + keyframes: list[int] A list of frame indices corresponding to the key frames in the input video. """ if self._globals.is_video: @@ -176,7 +197,8 @@ def save_video_meta_data(self, pts_time, keyframes): # <<<< PRIVATE METHODS >>> # # << INIT >> # @staticmethod - def _set_tk_vars(): + def _set_tk_vars() -> dict[T.Literal["unsaved", "edited", "face_count_changed"], + tk.BooleanVar]: """ Set the required tkinter variables. The alignments specific `unsaved` and `edited` are set here. @@ -188,15 +210,15 @@ def _set_tk_vars(): dict The internal variable name as key with the tkinter variable as value """ - retval = dict() - for name in ("unsaved", "edited", "face_count_changed"): + retval = {} + for name in T.get_args(T.Literal["unsaved", "edited", "face_count_changed"]): var = tk.BooleanVar() var.set(False) retval[name] = var logger.debug(retval) return retval - def _get_alignments(self, alignments_path, input_location): + def _get_alignments(self, alignments_path: str, input_location: str) -> Alignments: """ Get the :class:`~lib.align.Alignments` object for the given location. Parameters @@ -219,7 +241,7 @@ def _get_alignments(self, alignments_path, input_location): filename = "alignments.fsa" if self._globals.is_video: folder, vid = os.path.split(os.path.splitext(input_location)[0]) - filename = "{}_{}".format(vid, filename) + filename = f"{vid}_{filename}" else: folder = input_location retval = Alignments(folder, filename) @@ -232,7 +254,7 @@ def _get_alignments(self, alignments_path, input_location): return retval -class _DiskIO(): # pylint:disable=too-few-public-methods +class _DiskIO(): """ Handles the loading of :class:`~lib.align.DetectedFaces` from the alignments file into :class:`DetectedFaces` and the saving of this data (in the opposite direction) to an alignments file. @@ -244,7 +266,7 @@ class _DiskIO(): # pylint:disable=too-few-public-methods input_location: str The location of the input folder of frames or video file """ - def __init__(self, detected_faces, input_location): + def __init__(self, detected_faces: DetectedFaces, input_location: str) -> None: logger.debug("Initializing %s: (detected_faces: %s, input_location: %s)", self.__class__.__name__, detected_faces, input_location) self._input_location = input_location @@ -257,14 +279,14 @@ def __init__(self, detected_faces, input_location): self._globals = detected_faces._globals # Must be populated after loading faces as video_meta_data may have increased frame count - self._sorted_frame_names = None + self._sorted_frame_names: list[str] = [] logger.debug("Initialized %s", self.__class__.__name__) - def load(self): + def load(self) -> None: """ Load the faces from the alignments file, convert to :class:`~lib.align.DetectedFace`. objects and add to :attr:`_frame_faces`. """ for key in sorted(self._alignments.data): - this_frame_faces = [] + this_frame_faces: list[DetectedFace] = [] for item in self._alignments.data[key]["faces"]: face = DetectedFace() face.from_alignment(item, with_thumb=True) @@ -274,16 +296,18 @@ def load(self): self._frame_faces.append(this_frame_faces) self._sorted_frame_names = sorted(self._alignments.data) - def save(self): + def save(self) -> None: """ Convert updated :class:`~lib.align.DetectedFace` objects to alignments format and save the alignments file. """ if not self._tk_unsaved.get(): logger.debug("Alignments not updated. Returning") return frames = list(self._updated_frame_indices) - logger.verbose("Saving alignments for %s updated frames", len(frames)) + logger.verbose("Saving alignments for %s updated frames", # type:ignore[attr-defined] + len(frames)) - for idx, faces in zip(frames, np.array(self._frame_faces)[np.array(frames)]): + for idx, faces in zip(frames, + np.array(self._frame_faces, dtype="object")[np.array(frames)]): frame = self._sorted_frame_names[idx] self._alignments.data[frame]["faces"] = [face.to_alignment() for face in faces] @@ -292,7 +316,7 @@ def save(self): self._updated_frame_indices.clear() self._tk_unsaved.set(False) - def revert_to_saved(self, frame_index): + def revert_to_saved(self, frame_index: int) -> None: """ Revert the frame's alignments to their saved version for the given frame index. Parameters @@ -303,7 +327,8 @@ def revert_to_saved(self, frame_index): if frame_index not in self._updated_frame_indices: logger.debug("Alignments not amended. Returning") return - logger.verbose("Reverting alignments for frame_index %s", frame_index) + logger.verbose("Reverting alignments for frame_index %s", # type:ignore[attr-defined] + frame_index) alignments = self._alignments.data[self._sorted_frame_names[frame_index]]["faces"] faces = self._frame_faces[frame_index] @@ -322,12 +347,28 @@ def revert_to_saved(self, frame_index): self._tk_face_count_changed.set(True) else: self._tk_edited.set(True) - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) @classmethod - def _add_remove_faces(cls, alignments, faces): + def _add_remove_faces(cls, + alignments: list[AlignmentFileDict], + faces: list[DetectedFace]) -> bool: """ On a revert, ensure that the alignments and detected face object counts for each frame - are in sync. """ + are in sync. + + Parameters + ---------- + alignments: list[:class:`~lib.align.alignments.AlignmentFileDict`] + Alignments stored for a frame + + faces: list[:class:`~lib.align.DetectedFace`] + List of detected faces for a frame + + Returns + ------- + bool + ``True`` if a face was added or removed otherwise ``False`` + """ num_alignments = len(alignments) num_faces = len(faces) if num_alignments == num_faces: @@ -340,7 +381,7 @@ def _add_remove_faces(cls, alignments, faces): retval = True return retval - def extract(self): + def extract(self) -> None: """ Extract the current faces to a folder. To stop the GUI becoming completely unresponsive (particularly in Windows) the extract is @@ -354,24 +395,27 @@ def extract(self): return logger.debug(dirname) - queue = Queue() + queue: Queue = Queue() pbar = PopupProgress("Extracting Faces...", self._alignments.frames_count + 1) thread = MultiThread(self._background_extract, dirname, queue) thread.start() self._monitor_extract(thread, queue, pbar) - def _monitor_extract(self, thread, queue, progress_bar): + def _monitor_extract(self, + thread: MultiThread, + queue: Queue, + progress_bar: PopupProgress) -> None: """ Monitor the extraction thread, and update the progress bar. On completion, save alignments and clear progress bar. Parameters ---------- - thread: :class:`lib.multithreading.MultiThread` + thread: :class:`~lib.multithreading.MultiThread` The thread that is performing the extraction task queue: :class:`queue.Queue` The queue that the worker thread is putting it's incremental counts to - progress_bar: :class:`lib.gui.custom_widget.PopupProgress` + progress_bar: :class:`~lib.gui.custom_widget.PopupProgress` The popped up progress bar """ thread.check_and_raise_error() @@ -387,7 +431,7 @@ def _monitor_extract(self, thread, queue, progress_bar): break progress_bar.after(100, self._monitor_extract, thread, queue, progress_bar) - def _background_extract(self, output_folder, progress_queue): + def _background_extract(self, output_folder: str, progress_queue: Queue) -> None: """ Perform the background extraction in a thread so GUI doesn't become unresponsive. Parameters @@ -397,32 +441,32 @@ def _background_extract(self, output_folder, progress_queue): progress_queue: :class:`queue.Queue` The queue to place incremental counts to for updating the GUI's progress bar """ - _io = dict(saver=ImagesSaver(get_folder(output_folder), as_bytes=True), - loader=ImagesLoader(self._input_location, count=self._alignments.frames_count)) - - for frame_idx, (filename, image) in enumerate(_io["loader"].load()): - logger.trace("Outputting frame: %s: %s", frame_idx, filename) + saver = ImagesSaver(get_folder(output_folder), as_bytes=True) + loader = ImagesLoader(self._input_location, count=self._alignments.frames_count) + for frame_idx, (filename, image) in enumerate(loader.load()): + logger.trace("Outputting frame: %s: %s", # type:ignore[attr-defined] + frame_idx, filename) src_filename = os.path.basename(filename) - frame_name = os.path.splitext(src_filename)[0] progress_queue.put(1) for face_idx, face in enumerate(self._frame_faces[frame_idx]): - output = "{}_{}{}".format(frame_name, str(face_idx), ".png") + output = f"{os.path.splitext(src_filename)[0]}_{face_idx}.png" aligned = AlignedFace(face.landmarks_xy, image=image, centering="head", size=512) # TODO user selectable size - meta = dict(alignments=face.to_png_meta(), - source=dict(alignments_version=self._alignments.version, - original_filename=output, - face_index=face_idx, - source_filename=src_filename, - source_is_video=self._globals.is_video, - source_frame_dims=image.shape[:2])) - + meta: PNGHeaderDict = {"alignments": face.to_png_meta(), + "source": {"alignments_version": self._alignments.version, + "original_filename": output, + "face_index": face_idx, + "source_filename": src_filename, + "source_is_video": self._globals.is_video, + "source_frame_dims": image.shape[:2]}} + + assert aligned.face is not None b_image = encode_image(aligned.face, ".png", metadata=meta) - _io["saver"].save(output, b_image) - _io["saver"].close() + saver.save(output, b_image) + saver.close() class Filter(): @@ -434,7 +478,7 @@ class Filter(): detected_faces: :class:`DetectedFaces` The parent :class:`DetectedFaces` object """ - def __init__(self, detected_faces): + def __init__(self, detected_faces: DetectedFaces) -> None: logger.debug("Initializing %s: (detected_faces: %s)", self.__class__.__name__, detected_faces) self._globals = detected_faces._globals @@ -442,58 +486,62 @@ def __init__(self, detected_faces): logger.debug("Initialized %s", self.__class__.__name__) @property - def frame_meets_criteria(self): + def frame_meets_criteria(self) -> bool: """ bool: ``True`` if the current frame meets the selected filter criteria otherwise ``False`` """ - filter_mode = self._globals.filter_mode + filter_mode = self._globals.var_filter_mode.get() frame_faces = self._detected_faces.current_faces[self._globals.frame_index] distance = self._filter_distance + retval = ( filter_mode == "All Frames" or (filter_mode == "No Faces" and not frame_faces) or - (filter_mode == "Has Face(s)" and frame_faces) or + (filter_mode == "Has Face(s)" and len(frame_faces) > 0) or (filter_mode == "Multiple Faces" and len(frame_faces) > 1) or (filter_mode == "Misaligned Faces" and any(face.aligned.average_distance > distance - for face in frame_faces))) - logger.trace("filter_mode: %s, frame meets criteria: %s", filter_mode, retval) + for face in frame_faces))) + assert isinstance(retval, bool) + logger.trace("filter_mode: %s, frame meets criteria: %s", # type:ignore[attr-defined] + filter_mode, retval) return retval @property - def _filter_distance(self): + def _filter_distance(self) -> float: """ float: The currently selected distance when Misaligned Faces filter is selected. """ try: - retval = self._globals.tk_filter_distance.get() + retval = self._globals.var_filter_distance.get() except tk.TclError: # Suppress error when distance box is empty retval = 0 return retval / 100. @property - def count(self): + def count(self) -> int: """ int: The number of frames that meet the filter criteria returned by - :attr:`~tools.manual.manual.TkGlobals.filter_mode`. """ + :attr:`~tools.manual.manual.TkGlobals.var_filter_mode.get()`. """ face_count_per_index = self._detected_faces.face_count_per_index - if self._globals.filter_mode == "No Faces": + if self._globals.var_filter_mode.get() == "No Faces": retval = sum(1 for fcount in face_count_per_index if fcount == 0) - elif self._globals.filter_mode == "Has Face(s)": + elif self._globals.var_filter_mode.get() == "Has Face(s)": retval = sum(1 for fcount in face_count_per_index if fcount != 0) - elif self._globals.filter_mode == "Multiple Faces": + elif self._globals.var_filter_mode.get() == "Multiple Faces": retval = sum(1 for fcount in face_count_per_index if fcount > 1) - elif self._globals.filter_mode == "Misaligned Faces": + elif self._globals.var_filter_mode.get() == "Misaligned Faces": distance = self._filter_distance retval = sum(1 for frame in self._detected_faces.current_faces if any(face.aligned.average_distance > distance for face in frame)) else: retval = len(face_count_per_index) - logger.trace("filter mode: %s, frame count: %s", self._globals.filter_mode, retval) + logger.trace("filter mode: %s, frame count: %s", # type:ignore[attr-defined] + self._globals.var_filter_mode.get(), retval) return retval @property - def raw_indices(self): - """ dict: The frame and face indices that meet the current filter criteria for each - displayed face. """ - frame_indices = [] - face_indices = [] + def raw_indices(self) -> dict[T.Literal["frame", "face"], list[int]]: + """ dict[str, int]: The frame and face indices that meet the current filter criteria for + each displayed face. """ + frame_indices: list[int] = [] + face_indices: list[int] = [] face_counts = self._detected_faces.face_count_per_index # Copy to avoid recalculations for frame_idx in self.frames_list: @@ -501,28 +549,31 @@ def raw_indices(self): frame_indices.append(frame_idx) face_indices.append(face_idx) - retval = dict(frame=frame_indices, face=face_indices) - logger.trace("frame_indices: %s, face_indices: %s", frame_indices, face_indices) + retval: dict[T.Literal["frame", "face"], list[int]] = {"frame": frame_indices, + "face": face_indices} + logger.trace("frame_indices: %s, face_indices: %s", # type:ignore[attr-defined] + frame_indices, face_indices) return retval @property - def frames_list(self): - """ list: The list of frame indices that meet the filter criteria returned by - :attr:`~tools.manual.manual.TkGlobals.filter_mode`. """ + def frames_list(self) -> list[int]: + """ list[int]: The list of frame indices that meet the filter criteria returned by + :attr:`~tools.manual.manual.TkGlobals.var_filter_mode.get()`. """ face_count_per_index = self._detected_faces.face_count_per_index - if self._globals.filter_mode == "No Faces": + if self._globals.var_filter_mode.get() == "No Faces": retval = [idx for idx, count in enumerate(face_count_per_index) if count == 0] - elif self._globals.filter_mode == "Multiple Faces": + elif self._globals.var_filter_mode.get() == "Multiple Faces": retval = [idx for idx, count in enumerate(face_count_per_index) if count > 1] - elif self._globals.filter_mode == "Has Face(s)": + elif self._globals.var_filter_mode.get() == "Has Face(s)": retval = [idx for idx, count in enumerate(face_count_per_index) if count != 0] - elif self._globals.filter_mode == "Misaligned Faces": + elif self._globals.var_filter_mode.get() == "Misaligned Faces": distance = self._filter_distance retval = [idx for idx, frame in enumerate(self._detected_faces.current_faces) if any(face.aligned.average_distance > distance for face in frame)] else: - retval = range(len(face_count_per_index)) - logger.trace("filter mode: %s, number_frames: %s", self._globals.filter_mode, len(retval)) + retval = list(range(len(face_count_per_index))) + logger.trace("filter mode: %s, number_frames: %s", # type:ignore[attr-defined] + self._globals.var_filter_mode.get(), len(retval)) return retval @@ -535,7 +586,7 @@ class FaceUpdate(): detected_faces: :class:`DetectedFaces` The parent :class:`DetectedFaces` object """ - def __init__(self, detected_faces): + def __init__(self, detected_faces: DetectedFaces) -> None: logger.debug("Initializing %s: (detected_faces: %s)", self.__class__.__name__, detected_faces) self._detected_faces = detected_faces @@ -547,7 +598,7 @@ def __init__(self, detected_faces): logger.debug("Initialized %s", self.__class__.__name__) @property - def _tk_edited(self): + def _tk_edited(self) -> tk.BooleanVar: """ :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred meaning a GUI redraw needs to be triggered. @@ -558,7 +609,7 @@ def _tk_edited(self): return self._detected_faces.tk_edited @property - def _tk_face_count_changed(self): + def _tk_face_count_changed(self) -> tk.BooleanVar: """ :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred meaning a GUI redraw needs to be triggered. @@ -568,7 +619,7 @@ def _tk_face_count_changed(self): """ return self._detected_faces.tk_face_count_changed - def _faces_at_frame_index(self, frame_index): + def _faces_at_frame_index(self, frame_index: int) -> list[DetectedFace]: """ Checks whether the frame has already been added to :attr:`_updated_frame_indices` and adds it. Triggers the unsaved variable if this is the first edited frame. Returns the detected face objects for the given frame. @@ -589,7 +640,7 @@ def _faces_at_frame_index(self, frame_index): retval = self._frame_faces[frame_index] return retval - def add(self, frame_index, pnt_x, width, pnt_y, height): + def add(self, frame_index: int, pnt_x: int, width: int, pnt_y: int, height: int) -> None: """ Add a :class:`~lib.align.DetectedFace` object to the current frame with the given dimensions. @@ -615,7 +666,7 @@ def add(self, frame_index, pnt_x, width, pnt_y, height): face.load_aligned(None) self._tk_face_count_changed.set(True) - def delete(self, frame_index, face_index): + def delete(self, frame_index: int, face_index: int) -> None: """ Delete the :class:`~lib.align.DetectedFace` object for the given frame and face indices. @@ -630,9 +681,16 @@ def delete(self, frame_index, face_index): faces = self._faces_at_frame_index(frame_index) del faces[face_index] self._tk_face_count_changed.set(True) - self._globals.tk_update.set(True) - - def bounding_box(self, frame_index, face_index, pnt_x, width, pnt_y, height, aligner="FAN"): + self._globals.var_full_update.set(True) + + def bounding_box(self, + frame_index: int, + face_index: int, + pnt_x: int, + width: int, + pnt_y: int, + height: int, + aligner: manual.TypeManualExtractor = "FAN") -> None: """ Update the bounding box for the :class:`~lib.align.DetectedFace` object at the given frame and face indices, with the given dimensions and update the 68 point landmarks from the :class:`~tools.manual.manual.Aligner` for the updated bounding box. @@ -654,17 +712,23 @@ def bounding_box(self, frame_index, face_index, pnt_x, width, pnt_y, height, ali aligner: ["cv2-dnn", "FAN"], optional The aligner to use to generate the landmarks. Default: "FAN" """ - logger.trace("frame_index: %s, face_index %s, pnt_x %s, width %s, pnt_y %s, height %s, " - "aligner: %s", frame_index, face_index, pnt_x, width, pnt_y, height, aligner) + logger.trace("frame_index: %s, face_index %s, pnt_x %s, " # type:ignore[attr-defined] + "width %s, pnt_y %s, height %s, aligner: %s", + frame_index, face_index, pnt_x, width, pnt_y, height, aligner) face = self._faces_at_frame_index(frame_index)[face_index] - face.x = pnt_x - face.w = width - face.y = pnt_y - face.h = height - face.landmarks_xy = self._extractor.get_landmarks(frame_index, face_index, aligner) - self._globals.tk_update.set(True) - - def landmark(self, frame_index, face_index, landmark_index, shift_x, shift_y, is_zoomed): + face.left = pnt_x + face.width = width + face.top = pnt_y + face.height = height + face.add_landmarks_xy(self._extractor.get_landmarks(frame_index, face_index, aligner)) + self._globals.var_full_update.set(True) + + def landmark(self, + frame_index: int, face_index: int, + landmark_index: int, + shift_x: int, + shift_y: int, + is_zoomed: bool) -> None: """ Shift a single landmark point for the :class:`~lib.align.DetectedFace` object at the given frame and face indices by the given x and y values. @@ -689,7 +753,7 @@ def landmark(self, frame_index, face_index, landmark_index, shift_x, shift_y, is aligned = AlignedFace(face.landmarks_xy, centering="face", size=min(self._globals.frame_display_dims)) - landmark = aligned.landmarks[landmark_index] # pylint:disable=unsubscriptable-object + landmark = aligned.landmarks[landmark_index] landmark += (shift_x, shift_y) matrix = aligned.adjusted_matrix matrix = cv2.invertAffineTransform(matrix) @@ -698,15 +762,15 @@ def landmark(self, frame_index, face_index, landmark_index, shift_x, shift_y, is landmark = cv2.transform(landmark, matrix, landmark.shape).squeeze() face.landmarks_xy[landmark_index] = landmark else: - for lmk, idx in zip(landmark, landmark_index): + for lmk, idx in zip(landmark, landmark_index): # type:ignore[call-overload] lmk = np.reshape(lmk, (1, 1, 2)) lmk = cv2.transform(lmk, matrix, lmk.shape).squeeze() face.landmarks_xy[idx] = lmk else: face.landmarks_xy[landmark_index] += (shift_x, shift_y) - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) - def landmarks(self, frame_index, face_index, shift_x, shift_y): + def landmarks(self, frame_index: int, face_index: int, shift_x: int, shift_y: int) -> None: """ Shift all of the landmarks and bounding box for the :class:`~lib.align.DetectedFace` object at the given frame and face indices by the given x and y values and update the masks. @@ -728,12 +792,17 @@ def landmarks(self, frame_index, face_index, shift_x, shift_y): aligned with the newly adjusted landmarks. """ face = self._faces_at_frame_index(frame_index)[face_index] - face.x += shift_x - face.y += shift_y - face.landmarks_xy += (shift_x, shift_y) - self._globals.tk_update.set(True) - - def landmarks_rotate(self, frame_index, face_index, angle, center): + assert face.left is not None and face.top is not None + face.left += shift_x + face.top += shift_y + face.add_landmarks_xy(face.landmarks_xy + (shift_x, shift_y)) + self._globals.var_full_update.set(True) + + def landmarks_rotate(self, + frame_index: int, + face_index: int, + angle: float, + center: np.ndarray) -> None: """ Rotate the landmarks on an Extract Box rotate for the :class:`~lib.align.DetectedFace` object at the given frame and face indices for the given angle from the given center point. @@ -744,18 +813,22 @@ def landmarks_rotate(self, frame_index, face_index, angle, center): The frame that the face is being set for face_index: int The face index within the frame - angle: :class:`numpy.ndarray` + angle: float The angle, in radians to rotate the points by center: :class:`numpy.ndarray` The center point of the Landmark's Extract Box """ face = self._faces_at_frame_index(frame_index)[face_index] rot_mat = cv2.getRotationMatrix2D(tuple(center.astype("float32")), angle, 1.) - face.landmarks_xy = cv2.transform(np.expand_dims(face.landmarks_xy, axis=0), - rot_mat).squeeze() - self._globals.tk_update.set(True) - - def landmarks_scale(self, frame_index, face_index, scale, center): + face.add_landmarks_xy(cv2.transform(np.expand_dims(face.landmarks_xy, axis=0), + rot_mat).squeeze()) + self._globals.var_full_update.set(True) + + def landmarks_scale(self, + frame_index: int, + face_index: int, + scale: np.ndarray, + center: np.ndarray) -> None: """ Scale the landmarks on an Extract Box resize for the :class:`~lib.align.DetectedFace` object at the given frame and face indices from the given center point. @@ -772,10 +845,10 @@ def landmarks_scale(self, frame_index, face_index, scale, center): The center point of the Landmark's Extract Box """ face = self._faces_at_frame_index(frame_index)[face_index] - face.landmarks_xy = ((face.landmarks_xy - center) * scale) + center - self._globals.tk_update.set(True) + face.add_landmarks_xy(((face.landmarks_xy - center) * scale) + center) + self._globals.var_full_update.set(True) - def mask(self, frame_index, face_index, mask, mask_type): + def mask(self, frame_index: int, face_index: int, mask: np.ndarray, mask_type: str) -> None: """ Update the mask on an edit for the :class:`~lib.align.DetectedFace` object at the given frame and face indices, for the given mask and mask type. @@ -793,9 +866,9 @@ def mask(self, frame_index, face_index, mask, mask_type): face = self._faces_at_frame_index(frame_index)[face_index] face.mask[mask_type].replace_mask(mask) self._tk_edited.set(True) - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) - def copy(self, frame_index, direction): + def copy(self, frame_index: int, direction: T.Literal["prev", "next"]) -> None: """ Copy the alignments from the previous or next frame that has alignments to the current frame. @@ -825,7 +898,7 @@ def copy(self, frame_index, direction): # aligned_face cannot be deep copied, so remove and recreate to_copy = self._faces_at_frame_index(idx) for face in to_copy: - face.aligned = None + face._aligned = None # pylint:disable=protected-access copied = deepcopy(to_copy) for old_face, new_face in zip(to_copy, copied): @@ -834,9 +907,9 @@ def copy(self, frame_index, direction): faces.extend(copied) self._tk_face_count_changed.set(True) - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) - def post_edit_trigger(self, frame_index, face_index): + def post_edit_trigger(self, frame_index: int, face_index: int) -> None: """ Update the jpg thumbnail, the viewport thumbnail, the landmark masks and the aligned face on a face edit. @@ -850,240 +923,14 @@ def post_edit_trigger(self, frame_index, face_index): face = self._frame_faces[frame_index][face_index] face.load_aligned(None, force=True) # Update average distance face.mask = self._extractor.get_masks(frame_index, face_index) + face.clear_all_identities() aligned = AlignedFace(face.landmarks_xy, - image=self._globals.current_frame["image"], + image=self._globals.current_frame.image, centering="head", size=96) + assert aligned.face is not None face.thumbnail = generate_thumbnail(aligned.face, size=96) - if self._globals.filter_mode == "Misaligned Faces": + if self._globals.var_filter_mode.get() == "Misaligned Faces": self._detected_faces.tk_face_count_changed.set(True) self._tk_edited.set(True) - - -class ThumbsCreator(): - """ Background loader to generate thumbnails for the alignments file. Generates low resolution - thumbnails in parallel threads for faster processing. - - Parameters - ---------- - detected_faces: :class:`~tool.manual.faces.DetectedFaces` - The :class:`~lib.align.DetectedFace` objects for this video - input_location: str - The location of the input folder of frames or video file - """ - def __init__(self, detected_faces, input_location, single_process): - logger.debug("Initializing %s: (detected_faces: %s, input_location: %s, " - "single_process: %s)", self.__class__.__name__, detected_faces, - input_location, single_process) - self._size = 80 - self._pbar = dict(pbar=None, lock=Lock()) - self._meta = dict(key_frames=detected_faces.video_meta_data.get("keyframes", None), - pts_times=detected_faces.video_meta_data.get("pts_time", None)) - self._location = input_location - self._alignments = detected_faces._alignments - self._frame_faces = detected_faces._frame_faces - - self._is_video = all(val is not None for val in self._meta.values()) - self._num_threads = os.cpu_count() - 2 - if self._is_video and single_process: - self._num_threads = 1 - elif self._is_video and not single_process: - self._num_threads = min(self._num_threads, len(self._meta["key_frames"])) - else: - self._num_threads = max(self._num_threads, 32) - self._threads = [] - logger.debug("Initialized %s", self.__class__.__name__) - - @property - def has_thumbs(self): - """ bool: ``True`` if the underlying alignments file holds thumbnail images - otherwise ``False``. """ - return self._alignments.thumbnails.has_thumbnails - - def generate_cache(self): - """ Extract the face thumbnails from a video or folder of images into the - alignments file. """ - self._pbar["pbar"] = tqdm(desc="Caching Thumbnails", - leave=False, - total=len(self._frame_faces)) - if self._is_video: - self._launch_video() - else: - self._launch_folder() - while True: - self._check_and_raise_error() - if all(not thread.is_alive() for thread in self._threads): - break - sleep(1) - self._join_threads() - self._pbar["pbar"].close() - self._alignments.save() - - # << PRIVATE METHODS >> # - def _check_and_raise_error(self): - """ Monitor the loading threads for errors and raise if any occur. """ - for thread in self._threads: - thread.check_and_raise_error() - - def _join_threads(self): - """ Join the loading threads """ - logger.debug("Joining face viewer loading threads") - for thread in self._threads: - thread.join() - - def _launch_video(self): - """ Launch multiple :class:`lib.multithreading.MultiThread` objects to load faces from - a video file. - - Splits the video into segments and passes each of these segments to separate background - threads for some speed up. - """ - key_frame_split = len(self._meta["key_frames"]) // self._num_threads - key_frames = self._meta["key_frames"] - pts_times = self._meta["pts_times"] - for idx in range(self._num_threads): - is_final = idx == self._num_threads - 1 - start_idx = idx * key_frame_split - keyframe_idx = len(key_frames) - 1 if is_final else start_idx + key_frame_split - end_idx = key_frames[keyframe_idx] - start_pts = pts_times[key_frames[start_idx]] - end_pts = False if idx + 1 == self._num_threads else pts_times[end_idx] - starting_index = pts_times.index(start_pts) - if end_pts: - segment_count = len(pts_times[key_frames[start_idx]:end_idx]) - else: - segment_count = len(pts_times[key_frames[start_idx]:]) - logger.debug("thread index: %s, start_idx: %s, end_idx: %s, start_pts: %s, " - "end_pts: %s, starting_index: %s, segment_count: %s", idx, start_idx, - end_idx, start_pts, end_pts, starting_index, segment_count) - thread = MultiThread(self._load_from_video, - start_pts, - end_pts, - starting_index, - segment_count) - thread.start() - self._threads.append(thread) - - def _launch_folder(self): - """ Launch :class:`lib.multithreading.MultiThread` to retrieve faces from a - folder of images. - - Goes through the file list one at a time, passing each file to a separate background - thread for some speed up. - """ - reader = SingleFrameLoader(self._location) - num_threads = min(reader.count, self._num_threads) - frame_split = reader.count // self._num_threads - logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s", - reader.count, num_threads, frame_split) - for idx in range(num_threads): - is_final = idx == num_threads - 1 - start_idx = idx * frame_split - end_idx = reader.count if is_final else start_idx + frame_split - thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx) - thread.start() - self._threads.append(thread) - - def _load_from_video(self, pts_start, pts_end, start_index, segment_count): - """ Loads faces from video for the given segment of the source video. - - Each segment of the video is extracted from in a different background thread. - - Parameters - ---------- - pts_start: float - The start time to cut the segment out of the video - pts_end: float - The end time to cut the segment out of the video - start_index: int - The frame index that this segment starts from. Used for calculating the actual frame - index of each frame extracted - segment_count: int - The number of frames that appear in this segment. Used for ending early in case more - frames come out of the segment than should appear (sometimes more frames are picked up - at the end of the segment, so these are discarded) - """ - logger.debug("pts_start: %s, pts_end: %s, start_index: %s, segment_count: %s", - pts_start, pts_end, start_index, segment_count) - reader = self._get_reader(pts_start, pts_end) - idx = 0 - sample_filename = next(fname for fname in self._alignments.data) - vidname = sample_filename[:sample_filename.rfind("_")] - for idx, frame in enumerate(reader): - frame_idx = idx + start_index - filename = "{}_{:06d}.png".format(vidname, frame_idx + 1) - self._set_thumbail(filename, frame[..., ::-1], frame_idx) - if idx == segment_count - 1: - # Sometimes extra frames are picked up at the end of a segment, so stop - # processing when segment frame count has been hit. - break - reader.close() - logger.debug("Segment complete: (starting_frame_index: %s, processed_count: %s)", - start_index, idx) - - def _get_reader(self, pts_start, pts_end): - """ Get an imageio iterator for this thread's segment. - - Parameters - ---------- - pts_start: float - The start time to cut the segment out of the video - pts_end: float - The end time to cut the segment out of the video - - Returns - ------- - :class:`imageio.Reader` - A reader iterator for the requested segment of video - """ - input_params = ["-ss", str(pts_start)] - if pts_end: - input_params.extend(["-to", str(pts_end)]) - logger.debug("pts_start: %s, pts_end: %s, input_params: %s", - pts_start, pts_end, input_params) - return imageio.get_reader(self._location, "ffmpeg", input_params=input_params) - - def _load_from_folder(self, reader, start_index, end_index): - """ Loads faces from the given range of frame indices from a folder of images. - - Each frame range is extracted in a different background thread. - - Parameters - ---------- - reader: :class:`lib.image.SingleFrameLoader` - The reader that is used to retrieve the requested frame - start_index: int - The starting frame index for the images to extract faces from - end_index: int - The end frame index for the images to extract faces from - """ - logger.debug("reader: %s, start_index: %s, end_index: %s", - reader, start_index, end_index) - for frame_index in range(start_index, end_index): - filename, frame = reader.image_from_index(frame_index) - self._set_thumbail(filename, frame, frame_index) - logger.debug("Segment complete: (start_index: %s, processed_count: %s)", - start_index, end_index - start_index) - - def _set_thumbail(self, filename, frame, frame_index): - """ Extracts the faces from the frame and adds to alignments file - - Parameters - ---------- - filename: str - The filename of the frame within the alignments file - frame: :class:`numpy.ndarray` - The frame that contains the faces - frame_index: int - The frame index of this frame in the :attr:`_frame_faces` - """ - for face_idx, face in enumerate(self._frame_faces[frame_index]): - aligned = AlignedFace(face.landmarks_xy, - image=frame, - centering="head", - size=96) - face.thumbnail = generate_thumbnail(aligned.face, size=96) - self._alignments.thumbnails.add_thumbnail(filename, face_idx, face.thumbnail) - with self._pbar["lock"]: - self._pbar["pbar"].update(1) diff --git a/tools/manual/faceviewer/frame.py b/tools/manual/faceviewer/frame.py index 3c07b364b0..5c6c8f024b 100644 --- a/tools/manual/faceviewer/frame.py +++ b/tools/manual/faceviewer/frame.py @@ -1,11 +1,13 @@ #!/usr/bin/env python3 """ The Faces Viewer Frame and Canvas for Faceswap's Manual Tool. """ +from __future__ import annotations import colorsys import gettext import logging import platform import tkinter as tk from tkinter import ttk +import typing as T from math import floor, ceil from threading import Thread, Event @@ -14,10 +16,16 @@ from lib.gui.custom_widgets import RightClickMenu, Tooltip from lib.gui.utils import get_config, get_images from lib.image import hex_to_rgb, rgb_to_hex +from lib.logger import parse_class_init from .viewport import Viewport -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from tools.manual.detected_faces import DetectedFaces + from tools.manual.frameviewer.frame import DisplayFrame + from tools.manual.manual import TkGlobals + +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("tools.manual", localedir="locales", fallback=True) @@ -30,7 +38,7 @@ class FacesFrame(ttk.Frame): # pylint:disable=too-many-ancestors Parameters ---------- - parent: :class:`ttk.PanedWindow` + parent: :class:`ttk.Frame` The paned window that the faces frame resides in tk_globals: :class:`~tools.manual.manual.TkGlobals` The tkinter variables that apply to the whole of the GUI @@ -39,16 +47,18 @@ class FacesFrame(ttk.Frame): # pylint:disable=too-many-ancestors display_frame: :class:`~tools.manual.frameviewer.frame.DisplayFrame` The section of the Manual Tool that holds the frames viewer """ - def __init__(self, parent, tk_globals, detected_faces, display_frame): - logger.debug("Initializing %s: (parent: %s, tk_globals: %s, detected_faces: %s, " - "display_frame: %s)", self.__class__.__name__, parent, tk_globals, - detected_faces, display_frame) + def __init__(self, + parent: ttk.Frame, + tk_globals: TkGlobals, + detected_faces: DetectedFaces, + display_frame: DisplayFrame) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent) self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) self._actions_frame = FacesActionsFrame(self) self._faces_frame = ttk.Frame(self) - self._faces_frame.pack_propagate(0) + self._faces_frame.pack_propagate(False) self._faces_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True) self._event = Event() self._canvas = FacesViewer(self._faces_frame, @@ -60,7 +70,7 @@ def __init__(self, parent, tk_globals, detected_faces, display_frame): self._add_scrollbar() logger.debug("Initialized %s", self.__class__.__name__) - def _add_scrollbar(self): + def _add_scrollbar(self) -> None: """ Add a scrollbar to the faces frame """ logger.debug("Add Faces Viewer Scrollbar") scrollbar = ttk.Scrollbar(self._faces_frame, command=self._on_scroll) @@ -69,9 +79,8 @@ def _add_scrollbar(self): self.bind("", self._update_viewport) logger.debug("Added Faces Viewer Scrollbar") self.update_idletasks() # Update so scrollbar width is correct - return scrollbar.winfo_width() - def _on_scroll(self, *event): + def _on_scroll(self, *event: tk.Event) -> None: """ Callback on scrollbar scroll. Updates the canvas location and displays/hides thumbnail images. @@ -83,7 +92,7 @@ def _on_scroll(self, *event): self._canvas.yview(*event) self._canvas.viewport.update() - def _update_viewport(self, event): # pylint: disable=unused-argument + def _update_viewport(self, event: tk.Event) -> None: # pylint:disable=unused-argument """ Update the faces viewport and scrollbar. Parameters @@ -94,7 +103,7 @@ def _update_viewport(self, event): # pylint: disable=unused-argument self._canvas.viewport.update() self._canvas.configure(scrollregion=self._canvas.bbox("backdrop")) - def canvas_scroll(self, direction): + def canvas_scroll(self, direction: T.Literal["up", "down", "page-up", "page-down"]) -> None: """ Scroll the canvas on an up/down or page-up/page-down key press. Notes @@ -110,9 +119,11 @@ def canvas_scroll(self, direction): """ if self._event.is_set(): - logger.trace("Update already running. Aborting repeated keypress") + logger.trace("Update already running. " # type:ignore[attr-defined] + "Aborting repeated keypress") return - logger.trace("Running update on received key press: %s", direction) + logger.trace("Running update on received key press: %s", # type:ignore[attr-defined] + direction) amount = 1 if direction.endswith("down") else -1 units = "pages" if direction.startswith("page") else "units" @@ -121,7 +132,7 @@ def canvas_scroll(self, direction): args=(amount, units, self._event)) thread.start() - def set_annotation_display(self, key): + def set_annotation_display(self, key: str) -> None: """ Set the optional annotation overlay based on keyboard shortcut. Parameters @@ -140,33 +151,33 @@ class FacesActionsFrame(ttk.Frame): # pylint:disable=too-many-ancestors parent: :class:`FacesFrame` The Faces frame that this actions frame reside in """ - def __init__(self, parent): - logger.debug("Initializing %s: (parent: %s)", - self.__class__.__name__, parent) + def __init__(self, parent: FacesFrame) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent) self.pack(side=tk.LEFT, fill=tk.Y, padx=(2, 4), pady=2) - self._tk_vars = dict() + self._tk_vars: dict[T.Literal["mesh", "mask"], tk.BooleanVar] = {} self._configure_styles() self._buttons = self._add_buttons() logger.debug("Initialized %s", self.__class__.__name__) @property - def key_bindings(self): + def key_bindings(self) -> dict[str, T.Literal["mask", "mesh"]]: """ dict: The mapping of key presses to optional annotations to display. Keyboard shortcuts utilize the function keys. """ - return {"F{}".format(idx + 9): display for idx, display in enumerate(("mesh", "mask"))} + return {f"F{idx + 9}": display + for idx, display in enumerate(T.get_args(T.Literal["mesh", "mask"]))} @property - def _helptext(self): + def _helptext(self) -> dict[T.Literal["mask", "mesh"], str]: """ dict: `button key`: `button helptext`. The help text to display for each button. """ inverse_keybindings = {val: key for key, val in self.key_bindings.items()} - retval = dict(mesh=_("Display the landmarks mesh"), - mask=_("Display the mask")) + retval: dict[T.Literal["mask", "mesh"], str] = {"mesh": _('Display the landmarks mesh'), + "mask": _('Display the mask')} for item in retval: - retval[item] += " ({})".format(inverse_keybindings[item]) + retval[item] += f" ({inverse_keybindings[item]})" return retval - def _configure_styles(self): + def _configure_styles(self) -> None: """ Configure the background color for button frame and the button styles. """ style = ttk.Style() style.configure("display.TFrame", background='#d3d3d3') @@ -174,17 +185,17 @@ def _configure_styles(self): style.configure("display_deselected.TButton", relief="flat") self.config(style="display.TFrame") - def _add_buttons(self): + def _add_buttons(self) -> dict[T.Literal["mesh", "mask"], ttk.Button]: """ Add the display buttons to the Faces window. Returns ------- - dict + dict[Literal["mesh", "mask"], tk.Button]] The display name and its associated button. """ frame = ttk.Frame(self) frame.pack(side=tk.TOP, fill=tk.Y) - buttons = dict() + buttons = {} for display in self.key_bindings.values(): var = tk.BooleanVar() var.set(False) @@ -193,7 +204,7 @@ def _add_buttons(self): lookup = "landmarks" if display == "mesh" else display button = ttk.Button(frame, image=get_images().icons[lookup], - command=lambda t=display: self.on_click(t), + command=T.cast(T.Callable, lambda t=display: self.on_click(t)), style="display_deselected.TButton") button.state(["!pressed", "!focus"]) button.pack() @@ -201,13 +212,13 @@ def _add_buttons(self): buttons[display] = button return buttons - def on_click(self, display): + def on_click(self, display: T.Literal["mesh", "mask"]) -> None: """ Click event for the optional annotation buttons. Loads and unloads the annotations from the faces viewer. Parameters ---------- - display: str + display: Literal["mesh", "mask"] The display name for the button that has called this event as exists in :attr:`_buttons` """ @@ -239,16 +250,19 @@ class FacesViewer(tk.Canvas): # pylint:disable=too-many-ancestors event: :class:`threading.Event` The threading event object for repeated key press protection """ - def __init__(self, parent, tk_globals, tk_action_vars, detected_faces, display_frame, event): - logger.debug("Initializing %s: (parent: %s, tk_globals: %s, tk_action_vars: %s, " - "detected_faces: %s, display_frame: %s, event: %s)", self.__class__.__name__, - parent, tk_globals, tk_action_vars, detected_faces, display_frame, event) + def __init__(self, parent: ttk.Frame, + tk_globals: TkGlobals, + tk_action_vars: dict[T.Literal["mesh", "mask"], tk.BooleanVar], + detected_faces: DetectedFaces, + display_frame: DisplayFrame, + event: Event) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent, bd=0, highlightthickness=0, bg=get_config().user_theme["group_panel"]["panel_background"]) self.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, anchor=tk.E) - self._sizes = dict(tiny=32, small=64, medium=96, large=128, extralarge=192) + self._sizes = {"tiny": 32, "small": 64, "medium": 96, "large": 128, "extralarge": 192} self._globals = tk_globals self._tk_optional_annotations = tk_action_vars @@ -256,8 +270,8 @@ def __init__(self, parent, tk_globals, tk_action_vars, detected_faces, display_f self._display_frame = display_frame self._grid = Grid(self, detected_faces) self._view = Viewport(self, detected_faces.tk_edited) - self._annotation_colors = dict(mesh=self.get_muted_color("Mesh"), - box=self.control_colors["ExtractBox"]) + self._annotation_colors = {"mesh": self.get_muted_color("Mesh"), + "box": self.control_colors["ExtractBox"]} ContextMenu(self, detected_faces) self._bind_mouse_wheel_scrolling() @@ -265,66 +279,74 @@ def __init__(self, parent, tk_globals, tk_action_vars, detected_faces, display_f logger.debug("Initialized %s", self.__class__.__name__) @property - def face_size(self): + def face_size(self) -> int: """ int: The currently selected thumbnail size in pixels """ scaling = get_config().scaling_factor - size = self._sizes[self._globals.tk_faces_size.get().lower().replace(" ", "")] + size = self._sizes[self._globals.var_faces_size.get().lower().replace(" ", "")] scaled = size * scaling return int(round(scaled / 2) * 2) @property - def viewport(self): + def viewport(self) -> Viewport: """ :class:`~tools.manual.faceviewer.viewport.Viewport`: The viewport area of the faces viewer. """ return self._view @property - def grid(self): + def layout(self) -> Grid: """ :class:`Grid`: The grid for the current :class:`FacesViewer`. """ return self._grid @property - def optional_annotations(self): - """ dict: The values currently set for the selectable optional annotations. """ + def optional_annotations(self) -> dict[T.Literal["mesh", "mask"], bool]: + """ dict[Literal["mesh", "mask"], bool]: The values currently set for the + selectable optional annotations. """ return {opt: val.get() for opt, val in self._tk_optional_annotations.items()} @property - def selected_mask(self): + def selected_mask(self) -> str: """ str: The currently selected mask from the display frame control panel. """ return self._display_frame.tk_selected_mask.get().lower() @property - def control_colors(self): - """ :dict: The frame Editor name as key with the current user selected hex code as + def control_colors(self) -> dict[str, str]: + """dict[str, str]: The frame Editor name as key with the current user selected hex code as value. """ return ({key: val.get() for key, val in self._display_frame.tk_control_colors.items()}) # << CALLBACK FUNCTIONS >> # - def _set_tk_callbacks(self, detected_faces): + def _set_tk_callbacks(self, detected_faces: DetectedFaces): """ Set the tkinter variable call backs. + Parameters + ---------- + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The Manual Tool's Detected Faces object + Redraw the grid on a face size change, a filter change or on add/remove faces. Updates the annotation colors when user amends a color drop down. Updates the mask type when the user changes the selected mask types Toggles the face viewer annotations on an optional annotation button press. """ - for var in (self._globals.tk_faces_size, self._globals.tk_filter_mode): - var.trace("w", lambda *e, v=var: self.refresh_grid(v)) - var = detected_faces.tk_face_count_changed - var.trace("w", lambda *e, v=var: self.refresh_grid(v, retain_position=True)) - - self._display_frame.tk_control_colors["Mesh"].trace( - "w", lambda *e: self._update_mesh_color()) - self._display_frame.tk_control_colors["ExtractBox"].trace( - "w", lambda *e: self._update_box_color()) - self._display_frame.tk_selected_mask.trace("w", lambda *e: self._update_mask_type()) + for strvar in (self._globals.var_faces_size, self._globals.var_filter_mode): + strvar.trace_add("write", lambda *e, v=strvar: self.refresh_grid(v)) + boolvar = detected_faces.tk_face_count_changed + boolvar.trace_add("write", + lambda *e, v=boolvar: self.refresh_grid(v, retain_position=True)) + + self._display_frame.tk_control_colors["Mesh"].trace_add( + "write", lambda *e: self._update_mesh_color()) + self._display_frame.tk_control_colors["ExtractBox"].trace_add( + "write", lambda *e: self._update_box_color()) + self._display_frame.tk_selected_mask.trace_add( + "write", lambda *e: self._update_mask_type()) for opt, var in self._tk_optional_annotations.items(): - var.trace("w", lambda *e, o=opt: self._toggle_annotations(o)) + var.trace_add("write", lambda *e, o=opt: self._toggle_annotations(o)) self.bind("", lambda *e: self._view.update()) - def refresh_grid(self, trigger_var, retain_position=False): + def refresh_grid(self, trigger_var: tk.BooleanVar, retain_position: bool = False) -> None: """ Recalculate the full grid and redraw. Used when the active filter pull down is used, a face has been added or removed, or the face thumbnail size has changed. @@ -351,15 +373,16 @@ def refresh_grid(self, trigger_var, retain_position=False): if not size_change: trigger_var.set(False) - def _update_mask_type(self): + def _update_mask_type(self) -> None: """ Update the displayed mask in the :class:`FacesViewer` canvas when the user changes the mask type. """ + state: T.Literal["normal", "hidden"] state = "normal" if self.optional_annotations["mask"] else "hidden" logger.debug("Updating mask type: (mask_type: %s. state: %s)", self.selected_mask, state) self._view.toggle_mask(state, self.selected_mask) # << MOUSE HANDLING >> - def _bind_mouse_wheel_scrolling(self): + def _bind_mouse_wheel_scrolling(self) -> None: """ Bind mouse wheel to scroll the :class:`FacesViewer` canvas. """ if platform.system() == "Linux": self.bind("", self._scroll) @@ -367,7 +390,7 @@ def _bind_mouse_wheel_scrolling(self): else: self.bind("", self._scroll) - def _scroll(self, event): + def _scroll(self, event: tk.Event) -> None: """ Handle mouse wheel scrolling over the :class:`FacesViewer` canvas. Update is run in a thread to avoid repeated scroll actions stacking and locking up the GUI. @@ -378,12 +401,13 @@ def _scroll(self, event): The event fired by the mouse scrolling """ if self._event.is_set(): - logger.trace("Update already running. Aborting repeated mousewheel") + logger.trace("Update already running. " # type:ignore[attr-defined] + "Aborting repeated mousewheel") return if platform.system() == "Darwin": adjust = event.delta elif platform.system() == "Windows": - adjust = event.delta / 120 + adjust = int(event.delta / 120) elif event.num == 5: adjust = -1 else: @@ -392,14 +416,14 @@ def _scroll(self, event): thread = Thread(target=self.canvas_scroll, args=(-1 * adjust, "units", self._event)) thread.start() - def canvas_scroll(self, amount, units, event): + def canvas_scroll(self, amount: int, units: T.Literal["pages", "units"], event: Event) -> None: """ Scroll the canvas on an up/down or page-up/page-down key press. Parameters ---------- amount: int The number of units to scroll the canvas - units: ["page", "units"] + units: Literal["pages", "units"] The unit type to scroll by event: :class:`threading.Event` event to indicate to the calling process whether the scroll is still updating @@ -410,7 +434,7 @@ def canvas_scroll(self, amount, units, event): event.clear() # << OPTIONAL ANNOTATION METHODS >> # - def _update_mesh_color(self): + def _update_mesh_color(self) -> None: """ Update the mesh color when user updates the control panel. """ color = self.get_muted_color("Mesh") if self._annotation_colors["mesh"] == color: @@ -423,7 +447,7 @@ def _update_mesh_color(self): self.itemconfig("active_mesh_line", fill=highlight_color) self._annotation_colors["mesh"] = color - def _update_box_color(self): + def _update_box_color(self) -> None: """ Update the active box color when user updates the control panel. """ color = self.control_colors["ExtractBox"] @@ -432,13 +456,18 @@ def _update_box_color(self): self.itemconfig("active_highlighter", outline=color) self._annotation_colors["box"] = color - def get_muted_color(self, color_key): + def get_muted_color(self, color_key: str) -> str: """ Creates a muted version of the given annotation color for non-active faces. Parameters ---------- color_key: str The annotation key to obtain the color for from :attr:`control_colors` + + Returns + ------- + str + The hex color code of the muted color """ scale = 0.65 hls = np.array(colorsys.rgb_to_hls(*hex_to_rgb(self.control_colors[color_key]))) @@ -448,7 +477,7 @@ def get_muted_color(self, color_key): retval = rgb_to_hex(rgb) return retval - def _toggle_annotations(self, annotation): + def _toggle_annotations(self, annotation: T.Literal["mesh", "mask"]) -> None: """ Toggle optional annotations on or off after the user depresses an optional button. Parameters @@ -456,6 +485,7 @@ def _toggle_annotations(self, annotation): annotation: ["mesh", "mask"] The optional annotation to toggle on or off """ + state: T.Literal["hidden", "normal"] state = "normal" if self.optional_annotations[annotation] else "hidden" logger.debug("Toggle annotation: (annotation: %s, state: %s)", annotation, state) if annotation == "mesh": @@ -473,23 +503,22 @@ class Grid(): Parameters ---------- - canvas: :class:`tkinter.Canvas` + canvas: :class:`~FacesViewer` The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` The :class:`~lib.align.DetectedFace` objects for this video """ - def __init__(self, canvas, detected_faces): - logger.debug("Initializing %s: (detected_faces: %s)", - self.__class__.__name__, detected_faces) + def __init__(self, canvas: FacesViewer, detected_faces: DetectedFaces): + logger.debug(parse_class_init(locals())) self._canvas = canvas self._detected_faces = detected_faces self._raw_indices = detected_faces.filter.raw_indices self._frames_list = detected_faces.filter.frames_list - self._is_valid = False - self._face_size = None - self._grid = None - self._display_faces = None + self._is_valid: bool = False + self._face_size: int = 0 + self._grid: np.ndarray | None = None + self._display_faces: np.ndarray | None = None self._canvas.update_idletasks() self._canvas.create_rectangle(0, 0, 0, 0, tags=["backdrop"]) @@ -497,64 +526,76 @@ def __init__(self, canvas, detected_faces): logger.debug("Initialized %s", self.__class__.__name__) @property - def face_size(self): + def face_size(self) -> int: """ int: The pixel size of each thumbnail within the face viewer. """ return self._face_size @property - def is_valid(self): + def is_valid(self) -> bool: """ bool: ``True`` if the current filter means that the grid holds faces. ``False`` if there are no faces displayed in the grid. """ return self._is_valid @property - def columns_rows(self): + def columns_rows(self) -> tuple[int, int]: """ tuple: the (`columns`, `rows`) required to hold all display images. """ - retval = tuple(reversed(self._grid.shape[1:])) if self._is_valid else (0, 0) - return retval + if not self._is_valid: + return (0, 0) + assert self._grid is not None + retval = tuple(reversed(self._grid.shape[1:])) + return T.cast(tuple[int, int], retval) @property - def dimensions(self): + def dimensions(self) -> tuple[int, int]: """ tuple: The (`width`, `height`) required to hold all display images. """ if self._is_valid: + assert self._grid is not None retval = tuple(dim * self._face_size for dim in reversed(self._grid.shape[1:])) + assert len(retval) == 2 else: retval = (0, 0) - return retval + return T.cast(tuple[int, int], retval) @property - def _visible_row_indices(self): + def _visible_row_indices(self) -> tuple[int, int]: """tuple: A 1 dimensional array of the (`top_row_index`, `bottom_row_index`) of the grid currently in the viewable area. """ height = self.dimensions[1] visible = (max(0, floor(height * self._canvas.yview()[0]) - self._face_size), ceil(height * self._canvas.yview()[1])) - logger.trace("height: %s, yview: %s, face_size: %s, visible: %s", - height, self._canvas.yview(), self._face_size, visible) + logger.trace("height: %s, yview: %s, face_size: %s, " # type:ignore[attr-defined] + "visible: %s", height, self._canvas.yview(), self._face_size, visible) + assert self._grid is not None y_points = self._grid[3, :, 1] top = np.searchsorted(y_points, visible[0], side="left") bottom = np.searchsorted(y_points, visible[1], side="right") - return top, bottom + return int(top), int(bottom) @property - def visible_area(self): - """:class:`numpy.ndarray`: A numpy array of shape (`4`, `rows`, `columns`) corresponding + def visible_area(self) -> tuple[np.ndarray, np.ndarray]: + """tuple[:class:`numpy.ndarray`, :class:`numpy.ndarray`]: Tuple containing 2 arrays. + + 1st array contains an array of shape (`4`, `rows`, `columns`) corresponding to the viewable area of the display grid. 1st dimension contains frame indices, 2nd dimension face indices. The 3rd and 4th dimension contain the x and y position of the top left corner of the face respectively. + 2nd array contains :class:`~lib.align.DetectedFace` objects laid out in (rows, columns) + Any locations that are not populated by a face will have a frame and face index of -1 """ if not self._is_valid: - retval = None, None + retval = np.zeros((4, 0, 0)), np.zeros((0, 0)) else: + assert self._grid is not None + assert self._display_faces is not None top, bottom = self._visible_row_indices retval = self._grid[:, top:bottom, :], self._display_faces[top:bottom, :] - logger.trace([r if r is None else r.shape for r in retval]) + logger.trace([r if r is None else r.shape for r in retval]) # type:ignore[attr-defined] return retval - def y_coord_from_frame(self, frame_index): + def y_coord_from_frame(self, frame_index: int) -> int: """ Return the y coordinate for the first face that appears in the given frame. Parameters @@ -567,9 +608,10 @@ def y_coord_from_frame(self, frame_index): int The y coordinate of the first face for the given frame """ + assert self._grid is not None return min(self._grid[3][np.where(self._grid[0] == frame_index)]) - def frame_has_faces(self, frame_index): + def frame_has_faces(self, frame_index: int) -> bool | np.bool_: """ Check whether the given frame index contains any faces. Parameters @@ -582,9 +624,12 @@ def frame_has_faces(self, frame_index): bool ``True`` if there are faces in the given frame otherwise ``False`` """ - return self._is_valid and np.any(self._grid[0] == frame_index) + if not self._is_valid: + return False + assert self._grid is not None + return np.any(self._grid[0] == frame_index) - def update(self): + def update(self) -> None: """ Update the underlying grid. Called on initialization, on a filter change or on add/remove faces. Recalculates the @@ -597,25 +642,23 @@ def update(self): self._get_grid() self._get_display_faces() self._canvas.coords("backdrop", 0, 0, *self.dimensions) - self._canvas.configure(scrollregion=(self._canvas.bbox("backdrop"))) + self._canvas.configure(scrollregion=self._canvas.bbox("backdrop")) self._canvas.yview_moveto(0.0) - def _get_grid(self): + def _get_grid(self) -> None: """ Get the grid information for faces currently displayed in the :class:`FacesViewer`. + and set to :attr:`_grid`. Creates a numpy array of shape (`4`, `rows`, `columns`) + corresponding to the display grid. 1st dimension contains frame indices, 2nd dimension face + indices. The 3rd and 4th dimension contain the x and y position of the top left corner of + the face respectively. - Returns - :class:`numpy.ndarray` - A numpy array of shape (`4`, `rows`, `columns`) corresponding to the display grid. - 1st dimension contains frame indices, 2nd dimension face indices. The 3rd and 4th - dimension contain the x and y position of the top left corner of the face respectively. - - Any locations that are not populated by a face will have a frame and face index of -1 - """ + Any locations that are not populated by a face will have a frame and face index of -1""" labels = self._get_labels() if not self._is_valid: logger.debug("Setting grid to None for no faces.") self._grid = None return + assert labels is not None x_coords = np.linspace(0, labels.shape[2] * self._face_size, num=labels.shape[2], @@ -629,12 +672,12 @@ def _get_grid(self): self._grid = np.array((*labels, *np.meshgrid(x_coords, y_coords)), dtype="int") logger.debug(self._grid.shape) - def _get_labels(self): + def _get_labels(self) -> np.ndarray | None: """ Get the frame and face index for each grid position for the current filter. Returns ------- - :class:`numpy.ndarray` + :class:`numpy.ndarray` | None Array of dimensions (2, rows, columns) corresponding to the display grid, with frame index as the first dimension and face index within the frame as the 2nd dimension. @@ -657,17 +700,12 @@ def _get_labels(self): return labels def _get_display_faces(self): - """ Get the detected faces for the current filter and arrange to grid. + """ Get the detected faces for the current filter, arrange to grid and set to + :attr:`_display_faces`. This is an array of dimensions (rows, columns) corresponding to the + display grid, containing the corresponding :class:`lib.align.DetectFace` object - Returns - ------- - :class:`numpy.ndarray` - Array of dimensions (rows, columns) corresponding to the display grid, containing the - corresponding :class:`lib.align.DetectFace` object - - Any remaining placeholders at the end of the grid which are not populated with a face - are replaced with ``None`` - """ + Any remaining placeholders at the end of the grid which are not populated with a face are + replaced with ``None``""" if not self._is_valid: logger.debug("Setting display_faces to None for no faces.") self._display_faces = None @@ -684,7 +722,7 @@ def _get_display_faces(self): logger.debug("faces: (shape: %s, dtype: %s)", self._display_faces.shape, self._display_faces.dtype) - def transport_index_from_frame(self, frame_index): + def transport_index_from_frame(self, frame_index: int) -> int | None: """ Return the main frame's transport index for the given frame index based on the current filter criteria. @@ -695,11 +733,13 @@ def transport_index_from_frame(self, frame_index): Returns ------- - int - The index of the requested frame within the filtered frames view. + int | None + The index of the requested frame within the filtered frames view. None if no valid + frames """ retval = self._frames_list.index(frame_index) if frame_index in self._frames_list else None - logger.trace("frame_index: %s, transport_index: %s", frame_index, retval) + logger.trace("frame_index: %s, transport_index: %s", # type:ignore[attr-defined] + frame_index, retval) return retval @@ -737,17 +777,17 @@ def _pop_menu(self, event): frame_idx, face_idx = self._canvas.viewport.face_from_point( self._canvas.canvasx(event.x), self._canvas.canvasy(event.y))[:2] if frame_idx == -1: - logger.trace("No valid item under mouse") + logger.trace("No valid item under mouse") # type:ignore[attr-defined] self._frame_index = self._face_index = None return self._frame_index = frame_idx self._face_index = face_idx - logger.trace("Popping right click menu") + logger.trace("Popping right click menu") # type:ignore[attr-defined] self._menu.popup(event) def _delete_face(self): """ Delete the selected face on a right click mouse delete action. """ - logger.trace("Right click delete received. frame_id: %s, face_id: %s", - self._frame_index, self._face_index) + logger.trace("Right click delete received. frame_id: %s, " # type:ignore[attr-defined] + "face_id: %s", self._frame_index, self._face_index) self._detected_faces.update.delete(self._frame_index, self._face_index) self._frame_index = self._face_index = None diff --git a/tools/manual/faceviewer/interact.py b/tools/manual/faceviewer/interact.py new file mode 100644 index 0000000000..124629320c --- /dev/null +++ b/tools/manual/faceviewer/interact.py @@ -0,0 +1,423 @@ +#!/usr/bin/env python3 +""" Handles the viewport area for mouse hover actions and the active frame """ +from __future__ import annotations +import logging +import tkinter as tk +import typing as T +from dataclasses import dataclass + +import numpy as np + +from lib.logger import parse_class_init + +if T.TYPE_CHECKING: + from lib.align import DetectedFace + from .viewport import Viewport + +logger = logging.getLogger(__name__) + + +class HoverBox(): + """ Handle the current mouse location when over the :class:`Viewport`. + + Highlights the face currently underneath the cursor and handles actions when clicking + on a face. + + Parameters + ---------- + viewport: :class:`Viewport` + The viewport object for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + """ + def __init__(self, viewport: Viewport) -> None: + logger.debug(parse_class_init(locals())) + self._viewport = viewport + self._canvas = viewport._canvas + self._grid = viewport._canvas.layout + self._globals = viewport._canvas._globals + self._navigation = viewport._canvas._display_frame.navigation + self._box = self._canvas.create_rectangle(0., # type:ignore[call-overload] + 0., + float(self._size), + float(self._size), + outline="#0000ff", + width=2, + state="hidden", + fill="#0000ff", + stipple="gray12", + tags="hover_box") + self._current_frame_index = None + self._current_face_index = None + self._canvas.bind("", lambda e: self._clear()) + self._canvas.bind("", self.on_hover) + self._canvas.bind("", lambda e: self._select_frame()) + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def _size(self) -> int: + """ int: the currently set viewport face size in pixels. """ + return self._viewport.face_size + + def on_hover(self, event: tk.Event | None) -> None: + """ Highlight the face and set the mouse cursor for the mouse's current location. + + Parameters + ---------- + event: :class:`tkinter.Event` or ``None`` + The tkinter mouse event. Provides the current location of the mouse cursor. If ``None`` + is passed as the event (for example when this function is being called outside of a + mouse event) then the location of the cursor will be calculated + """ + if event is None: + pnts = np.array((self._canvas.winfo_pointerx(), self._canvas.winfo_pointery())) + pnts -= np.array((self._canvas.winfo_rootx(), self._canvas.winfo_rooty())) + else: + pnts = np.array((event.x, event.y)) + + coords = (int(self._canvas.canvasx(pnts[0])), int(self._canvas.canvasy(pnts[1]))) + face = self._viewport.face_from_point(*coords) + frame_idx, face_idx = face[:2] + + if frame_idx == self._current_frame_index and face_idx == self._current_face_index: + return + + is_zoomed = self._globals.is_zoomed + if (-1 in face or (frame_idx == self._globals.frame_index + and (not is_zoomed or + (is_zoomed and face_idx == self._globals.face_index)))): + self._clear() + self._canvas.config(cursor="") + self._current_frame_index = None + self._current_face_index = None + return + + logger.debug("Viewport hover: frame_idx: %s, face_idx: %s", frame_idx, face_idx) + + self._canvas.config(cursor="hand2") + self._highlight(face[2:]) + self._current_frame_index = frame_idx + self._current_face_index = face_idx + + def _clear(self) -> None: + """ Hide the hover box when the mouse is not over a face. """ + if self._canvas.itemcget(self._box, "state") != "hidden": + self._canvas.itemconfig(self._box, state="hidden") + + def _highlight(self, top_left: np.ndarray) -> None: + """ Display the hover box around the face that the mouse is currently over. + + Parameters + ---------- + top_left: :class:`np.ndarray` + The top left point of the highlight box location + """ + coords = (*top_left, *[x + self._size for x in top_left]) + self._canvas.coords(self._box, *coords) + self._canvas.itemconfig(self._box, state="normal") + self._canvas.tag_raise(self._box) + + def _select_frame(self) -> None: + """ Select the face and the subsequent frame (in the editor view) when a face is clicked + on in the :class:`Viewport`. """ + frame_id = self._current_frame_index + is_zoomed = self._globals.is_zoomed + logger.debug("Face clicked. Global frame index: %s, Current frame_id: %s, is_zoomed: %s", + self._globals.frame_index, frame_id, is_zoomed) + if frame_id is None or (frame_id == self._globals.frame_index and not is_zoomed): + return + face_idx = self._current_face_index if is_zoomed else 0 + self._globals.set_face_index(face_idx) + transport_id = self._grid.transport_index_from_frame(frame_id) + logger.trace("frame_index: %s, transport_id: %s, face_idx: %s", + frame_id, transport_id, face_idx) + if transport_id is None: + return + self._navigation.stop_playback() + self._globals.var_transport_index.set(transport_id) + self._viewport.move_active_to_top() + self.on_hover(None) + + +@dataclass +class Asset: + """ Holds all of the display assets identifiers for the active frame's face viewer objects + + Parameters + ---------- + images: list[int] + Indices for a frame's tk image ids displayed in the active frame + meshes: list[dict[Literal["polygon", "line"], list[int]]] + Indices for a frame's tk line/polygon object ids displayed in the active frame + faces: list[:class:`~lib.align.detected_faces.DetectedFace`] + DetectedFace objects that exist in the current frame + boxes: list[int] + Indices for a frame's bounding box object ids displayed in the active frame + """ + images: list[int] + """list[int]: Indices for a frame's tk image ids displayed in the active frame""" + meshes: list[dict[T.Literal["polygon", "line"], list[int]]] + """list[dict[Literal["polygon", "line"], list[int]]]: Indices for a frame's tk line/polygon + object ids displayed in the active frame""" + faces: list[DetectedFace] + """list[:class:`~lib.align.detected_faces.DetectedFace`]: DetectedFace objects that exist + in the current frame""" + boxes: list[int] + """list[int]: Indices for a frame's bounding box object ids displayed in the active + frame""" + + +class ActiveFrame(): + """ Handles the display of faces and annotations for the currently active frame. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + tk_edited_variable: :class:`tkinter.BooleanVar` + The tkinter callback variable indicating that a face has been edited + """ + def __init__(self, viewport: Viewport, tk_edited_variable: tk.BooleanVar) -> None: + logger.debug(parse_class_init(locals())) + self._objects = viewport._objects + self._viewport = viewport + self._grid = viewport._grid + self._tk_faces = viewport._tk_faces + self._canvas = viewport._canvas + self._globals = viewport._canvas._globals + self._navigation = viewport._canvas._display_frame.navigation + self._last_execution: dict[T.Literal["frame_index", "size"], + int] = {"frame_index": -1, "size": viewport.face_size} + self._tk_vars: dict[T.Literal["selected_editor", "edited"], + tk.StringVar | tk.BooleanVar] = { + "selected_editor": self._canvas._display_frame.tk_selected_action, + "edited": tk_edited_variable} + self._assets: Asset = Asset([], [], [], []) + + self._globals.var_update_active_viewport.trace_add("write", + lambda *e: self._reload_callback()) + tk_edited_variable.trace_add("write", lambda *e: self._update_on_edit()) + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def frame_index(self) -> int: + """ int: The frame index of the currently displayed frame. """ + return self._globals.frame_index + + @property + def current_frame(self) -> np.ndarray: + """ :class:`numpy.ndarray`: A BGR version of the frame currently being displayed. """ + return self._globals.current_frame.image + + @property + def _size(self) -> int: + """ int: The size of the thumbnails displayed in the viewport, in pixels. """ + return self._viewport.face_size + + @property + def _optional_annotations(self) -> dict[T.Literal["mesh", "mask"], bool]: + """ dict[Literal["mesh", "mask"], bool]: The currently selected optional + annotations """ + return self._canvas.optional_annotations + + def _reload_callback(self) -> None: + """ If a frame has changed, triggering the variable, then update the active frame. Return + having done nothing if the variable is resetting. """ + if self._globals.var_update_active_viewport.get(): + self.reload_annotations() + + def reload_annotations(self) -> None: + """ Handles the reloading of annotations for the currently active faces. + + Highlights the faces within the viewport of those faces that exist in the currently + displaying frame. Applies annotations based on the optional annotations and current + editor selections. + """ + logger.trace("Reloading annotations") # type:ignore[attr-defined] + if self._assets.images: + self._clear_previous() + + self._set_active_objects() + self._check_active_in_view() + + if not self._assets.images: + logger.trace("No active faces. Returning") # type:ignore[attr-defined] + self._last_execution["frame_index"] = self.frame_index + return + + if self._last_execution["frame_index"] != self.frame_index: + self.move_to_top() + self._create_new_boxes() + + self._update_face() + self._canvas.tag_raise("active_highlighter") + self._globals.var_update_active_viewport.set(False) + self._last_execution["frame_index"] = self.frame_index + + def _clear_previous(self) -> None: + """ Reverts the previously selected annotations to their default state. """ + logger.trace("Clearing previous active frame") # type:ignore[attr-defined] + self._canvas.itemconfig("active_highlighter", state="hidden") + + for key in T.get_args(T.Literal["polygon", "line"]): + tag = f"active_mesh_{key}" + self._canvas.itemconfig(tag, **self._viewport.mesh_kwargs[key], width=1) + self._canvas.dtag(tag) + + if self._viewport.selected_editor == "mask" and not self._optional_annotations["mask"]: + for name, tk_face in self._tk_faces.items(): + if name.startswith(f"{self._last_execution['frame_index']}_"): + tk_face.update_mask(None) + + def _set_active_objects(self) -> None: + """ Collect the objects that exist in the currently active frame from the main grid. """ + if self._grid.is_valid: + rows, cols = np.where(self._objects.visible_grid[0] == self.frame_index) + logger.trace("Setting active objects: (rows: %s, " # type:ignore[attr-defined] + "columns: %s)", rows, cols) + self._assets.images = self._objects.images[rows, cols].tolist() + self._assets.meshes = self._objects.meshes[rows, cols].tolist() + self._assets.faces = self._objects.visible_faces[rows, cols].tolist() + else: + logger.trace("No valid grid. Clearing active objects") # type:ignore[attr-defined] + self._assets.images = [] + self._assets.meshes = [] + self._assets.faces = [] + + def _check_active_in_view(self) -> None: + """ If the frame has changed, there are faces in the frame, but they don't appear in the + viewport, then bring the active faces to the top of the viewport. """ + if (not self._assets.images and + self._last_execution["frame_index"] != self.frame_index and + self._grid.frame_has_faces(self.frame_index)): + y_coord = self._grid.y_coord_from_frame(self.frame_index) + logger.trace("Active not in view. Moving to: %s", y_coord) # type:ignore[attr-defined] + self._canvas.yview_moveto(y_coord / self._canvas.bbox("backdrop")[3]) + self._viewport.update() + + def move_to_top(self) -> None: + """ Move the currently selected frame's faces to the top of the viewport if they are moving + off the bottom of the viewer. """ + height = self._canvas.bbox("backdrop")[3] + bot = int(self._canvas.coords(self._assets.images[-1])[1] + self._size) + + y_top, y_bot = (int(round(pnt * height)) for pnt in self._canvas.yview()) + + if y_top < bot < y_bot: # bottom face is still in fully visible area + logger.trace("Active faces in frame. Returning") # type:ignore[attr-defined] + return + + top = int(self._canvas.coords(self._assets.images[0])[1]) + if y_top == top: + logger.trace("Top face already on top row. Returning") # type:ignore[attr-defined] + return + + if self._canvas.winfo_height() > self._size: + logger.trace("Viewport taller than single face height. " # type:ignore[attr-defined] + "Moving Active faces to top: %s", top) + self._canvas.yview_moveto(top / height) + self._viewport.update() + elif self._canvas.winfo_height() <= self._size and y_top != top: + logger.trace("Viewport shorter than single face height. " # type:ignore[attr-defined] + "Moving Active faces to top: %s", top) + self._canvas.yview_moveto(top / height) + self._viewport.update() + + def _create_new_boxes(self) -> None: + """ The highlight boxes (border around selected faces) are the only additional annotations + that are required for the highlighter. If more faces are displayed in the current frame + than highlight boxes are available, then new boxes are created to accommodate the + additional faces. """ + new_boxes_count = max(0, len(self._assets.images) - len(self._assets.boxes)) + if new_boxes_count == 0: + return + logger.debug("new_boxes_count: %s", new_boxes_count) + for _ in range(new_boxes_count): + box = self._canvas.create_rectangle(0., # type:ignore[call-overload] + 0., + float(self._viewport.face_size), + float(self._viewport.face_size), + outline="#00FF00", + width=2, + state="hidden", + tags=["active_highlighter"]) + logger.trace("Created new highlight_box: %s", box) # type:ignore[attr-defined] + self._assets.boxes.append(box) + + def _update_on_edit(self) -> None: + """ Update the active faces on a frame edit. """ + if not self._tk_vars["edited"].get(): + return + self._set_active_objects() + self._update_face() + assert isinstance(self._tk_vars["edited"], tk.BooleanVar) + self._tk_vars["edited"].set(False) + + def _update_face(self) -> None: + """ Update the highlighted annotations for faces in the currently selected frame. """ + for face_idx, (image_id, mesh_ids, box_id, det_face), in enumerate( + zip(self._assets.images, + self._assets.meshes, + self._assets.boxes, + self._assets.faces)): + if det_face is None: + continue + top_left = self._canvas.coords(image_id) + coords = [*top_left, *[x + self._size for x in top_left]] + tk_face = self._viewport.get_tk_face(self.frame_index, face_idx, det_face) + self._canvas.itemconfig(image_id, image=tk_face.photo) + self._show_box(box_id, coords) + self._show_mesh(mesh_ids, face_idx, det_face, top_left) + self._last_execution["size"] = self._viewport.face_size + + def _show_box(self, item_id: int, coordinates: list[float]) -> None: + """ Display the highlight box around the given coordinates. + + Parameters + ---------- + item_id: int + The tkinter canvas object identifier for the highlight box + coordinates: list[float] + The (x, y, x1, y1) coordinates of the top left corner of the box + """ + self._canvas.coords(item_id, *coordinates) + self._canvas.itemconfig(item_id, state="normal") + + def _show_mesh(self, + mesh_ids: dict[T.Literal["polygon", "line"], list[int]], + face_index: int, + detected_face: DetectedFace, + top_left: list[float]) -> None: + """ Display the mesh annotation for the given face, at the given location. + + Parameters + ---------- + mesh_ids: dict[Literal["polygon", "line"], list[int]] + Dictionary containing the `polygon` and `line` tkinter canvas identifiers that make up + the mesh for the given face + face_index: int + The face index within the frame for the given face + detected_face: :class:`~lib.align.DetectedFace` + The detected face object that contains the landmarks for generating the mesh + top_left: list[float] + The (x, y) top left co-ordinates of the mesh's bounding box + """ + state = "normal" if (self._tk_vars["selected_editor"].get() != "Mask" or + self._optional_annotations["mesh"]) else "hidden" + kwargs: dict[T.Literal["polygon", "line"], dict[str, T.Any]] = { + "polygon": {"fill": "", "width": 2, "outline": self._canvas.control_colors["Mesh"]}, + "line": {"fill": self._canvas.control_colors["Mesh"], "width": 2}} + + assert isinstance(self._tk_vars["edited"], tk.BooleanVar) + edited = (self._tk_vars["edited"].get() and + self._tk_vars["selected_editor"].get() not in ("Mask", "View")) + landmarks = self._viewport.get_landmarks(self.frame_index, + face_index, + detected_face, + top_left, + edited) + for key, kwarg in kwargs.items(): + if key not in mesh_ids: + continue + for idx, mesh_id in enumerate(mesh_ids[key]): + self._canvas.coords(mesh_id, *landmarks[key][idx].flatten()) + self._canvas.itemconfig(mesh_id, state=state, **kwarg) + self._canvas.addtag_withtag(f"active_mesh_{key}", mesh_id) diff --git a/tools/manual/faceviewer/viewport.py b/tools/manual/faceviewer/viewport.py index 8c566f9a97..94486a82f8 100644 --- a/tools/manual/faceviewer/viewport.py +++ b/tools/manual/faceviewer/viewport.py @@ -1,16 +1,24 @@ #!/usr/bin/env python3 """ Handles the visible area of the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas. """ - +from __future__ import annotations import logging import tkinter as tk +import typing as T import cv2 import numpy as np from PIL import Image, ImageTk -from lib.align import AlignedFace +from lib.align import AlignedFace, LANDMARK_PARTS, LandmarkType +from lib.logger import parse_class_init + +from .interact import ActiveFrame, HoverBox + +if T.TYPE_CHECKING: + from lib.align import CenteringType, DetectedFace + from .frame import FacesViewer -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Viewport(): @@ -23,74 +31,64 @@ class Viewport(): tk_edited_variable: :class:`tkinter.BooleanVar` The variable that indicates that a face has been edited """ - def __init__(self, canvas, tk_edited_variable): - logger.debug("Initializing: %s: (canvas: %s, tk_edited_variable: %s)", - self.__class__.__name__, canvas, tk_edited_variable) + def __init__(self, canvas: FacesViewer, tk_edited_variable: tk.BooleanVar) -> None: + logger.debug(parse_class_init(locals())) self._canvas = canvas - self._grid = canvas.grid - self._centering = "face" + self._grid = canvas.layout + self._centering: CenteringType = "face" self._tk_selected_editor = canvas._display_frame.tk_selected_action - self._landmark_mapping = dict(mouth_inner=(60, 68), - mouth_outer=(48, 60), - right_eyebrow=(17, 22), - left_eyebrow=(22, 27), - right_eye=(36, 42), - left_eye=(42, 48), - nose=(27, 36), - jaw=(0, 17), - chin=(8, 11)) - self._landmarks = dict() - self._tk_faces = dict() + self._landmarks: dict[str, dict[T.Literal["polygon", "line"], list[np.ndarray]]] = {} + self._tk_faces: dict[str, TKFace] = {} self._objects = VisibleObjects(self) self._hoverbox = HoverBox(self) self._active_frame = ActiveFrame(self, tk_edited_variable) self._tk_selected_editor.trace( "w", lambda *e: self._active_frame.reload_annotations()) + logger.debug("Initialized %s", self.__class__.__name__) @property - def face_size(self): + def face_size(self) -> int: """ int: The pixel size of each thumbnail """ return self._grid.face_size @property - def mesh_kwargs(self): - """ dict: The color and state keyword arguments for the objects that make up a single - face's mesh annotation based on the current user selected options. Key is the object - type (`polygon` or `line`), value are the keyword arguments for that type. """ + def mesh_kwargs(self) -> dict[T.Literal["polygon", "line"], dict[str, T.Any]]: + """ dict[Literal["polygon", "line"], str | int]: Dynamic keyword arguments defining the + color and state for the objects that make up a single face's mesh annotation based on the + current user selected options. Values are the keyword arguments for that given type. """ state = "normal" if self._canvas.optional_annotations["mesh"] else "hidden" color = self._canvas.control_colors["Mesh"] - kwargs = dict(polygon=dict(fill="", outline=color, state=state), - line=dict(fill=color, state=state)) - return kwargs + return {"polygon": {"fill": "", "outline": color, "state": state}, + "line": {"fill": color, "state": state}} @property - def hover_box(self): + def hover_box(self) -> HoverBox: """ :class:`HoverBox`: The hover box for the viewport. """ return self._hoverbox @property - def selected_editor(self): + def selected_editor(self) -> str: """ str: The currently selected editor. """ return self._tk_selected_editor.get().lower() - def toggle_mesh(self, state): + def toggle_mesh(self, state: T.Literal["hidden", "normal"]) -> None: """ Toggles the mesh optional annotations on and off. Parameters ---------- - state: ["hidden", "normal"] + state: Literal["hidden", "normal"] The state to set the mesh annotations to """ logger.debug("Toggling mesh annotations to: %s", state) self._canvas.itemconfig("viewport_mesh", state=state) self.update() - def toggle_mask(self, state, mask_type): + def toggle_mask(self, state: T.Literal["hidden", "normal"], mask_type: str) -> None: """ Toggles the mask optional annotation on and off. Parameters ---------- - state: ["hidden", "normal"] + state: Literal["hidden", "normal"] Whether the mask should be displayed or hidden mask_type: str The type of mask to overlay onto the face @@ -108,7 +106,7 @@ def toggle_mask(self, state, mask_type): self.update() @classmethod - def _obtain_mask(cls, detected_face, mask_type): + def _obtain_mask(cls, detected_face: DetectedFace, mask_type: str) -> np.ndarray | None: """ Obtain the mask for the correct "face" centering that is used in the thumbnail display. Parameters @@ -128,16 +126,17 @@ def _obtain_mask(cls, detected_face, mask_type): return None if mask.stored_centering != "face": face = AlignedFace(detected_face.landmarks_xy) - mask.set_sub_crop(face.pose.offset["face"] - face.pose.offset[mask.stored_centering], + mask.set_sub_crop(face.pose.offset[mask.stored_centering], + face.pose.offset["face"], centering="face") return mask.mask.squeeze() - def reset(self): + def reset(self) -> None: """ Reset all the cached objects on a face size change. """ - self._landmarks = dict() - self._tk_faces = dict() + self._landmarks = {} + self._tk_faces = {} - def update(self, refresh_annotations=False): + def update(self, refresh_annotations: bool = False) -> None: """ Update the viewport. Parameters @@ -152,7 +151,7 @@ def update(self, refresh_annotations=False): self._update_viewport(refresh_annotations) self._active_frame.reload_annotations() - def _update_viewport(self, refresh_annotations): + def _update_viewport(self, refresh_annotations: bool) -> None: """ Update the viewport Parameters @@ -173,12 +172,12 @@ def _update_viewport(self, refresh_annotations): self._objects.meshes, self._objects.visible_faces): for (frame_idx, face_idx, pnt_x, pnt_y), image_id, mesh_ids, face in zip(*collection): - top_left = np.array((pnt_x, pnt_y)) if frame_idx == self._active_frame.frame_index and not refresh_annotations: - logger.trace("Skipping active frame: %s", frame_idx) + logger.trace("Skipping active frame: %s", # type:ignore[attr-defined] + frame_idx) continue if frame_idx == -1: - logger.trace("Blanking non-existant face") + logger.trace("Blanking non-existant face") # type:ignore[attr-defined] self._canvas.itemconfig(image_id, image="") for area in mesh_ids.values(): for mesh_id in area: @@ -191,20 +190,21 @@ def _update_viewport(self, refresh_annotations): if (self._canvas.optional_annotations["mesh"] or frame_idx == self._active_frame.frame_index or refresh_annotations): - landmarks = self.get_landmarks(frame_idx, face_idx, face, top_left, + landmarks = self.get_landmarks(frame_idx, face_idx, face, [pnt_x, pnt_y], refresh=True) self._locate_mesh(mesh_ids, landmarks) - def _discard_tk_faces(self): + def _discard_tk_faces(self) -> None: """ Remove any :class:`TKFace` objects from the cache that are not currently displayed. """ - keys = ["{}_{}".format(pnt_x, pnt_y) + keys = [f"{pnt_x}_{pnt_y}" for pnt_x, pnt_y in self._objects.visible_grid[:2].T.reshape(-1, 2)] for key in list(self._tk_faces): if key not in keys: del self._tk_faces[key] - logger.trace("keys: %s allocated_faces: %s", keys, len(self._tk_faces)) + logger.trace("keys: %s allocated_faces: %s", # type:ignore[attr-defined] + keys, len(self._tk_faces)) - def get_tk_face(self, frame_index, face_index, face): + def get_tk_face(self, frame_index: int, face_index: int, face: DetectedFace) -> TKFace: """ Obtain the :class:`TKFace` object for the given face from the cache. If the face does not exist in the cache, then it is generated and added prior to returning. @@ -226,26 +226,33 @@ def get_tk_face(self, frame_index, face_index, face): is_active = frame_index == self._active_frame.frame_index key = "_".join([str(frame_index), str(face_index)]) if key not in self._tk_faces or is_active: - logger.trace("creating new tk_face: (key: %s, is_active: %s)", key, is_active) + logger.trace("creating new tk_face: (key: %s, " # type:ignore[attr-defined] + "is_active: %s)", key, is_active) if is_active: image = AlignedFace(face.landmarks_xy, image=self._active_frame.current_frame, centering=self._centering, size=self.face_size).face else: + thumb = face.thumbnail + assert thumb is not None image = AlignedFace(face.landmarks_xy, - image=cv2.imdecode(face.thumbnail, cv2.IMREAD_UNCHANGED), + image=cv2.imdecode(thumb, cv2.IMREAD_UNCHANGED), centering=self._centering, size=self.face_size, is_aligned=True).face + assert image is not None tk_face = self._get_tk_face_object(face, image, is_active) self._tk_faces[key] = tk_face else: - logger.trace("tk_face exists: %s", key) + logger.trace("tk_face exists: %s", key) # type:ignore[attr-defined] tk_face = self._tk_faces[key] return tk_face - def _get_tk_face_object(self, face, image, is_active): + def _get_tk_face_object(self, + face: DetectedFace, + image: np.ndarray, + is_active: bool) -> TKFace: """ Obtain an existing unallocated, or a newly created :class:`TKFace` and populate it with face information from the requested frame and face index. @@ -271,10 +278,16 @@ def _get_tk_face_object(self, face, image, is_active): (is_active and self.selected_editor == "mask")) mask = self._obtain_mask(face, self._canvas.selected_mask) if get_mask else None tk_face = TKFace(image, size=self.face_size, mask=mask) - logger.trace("face: %s, tk_face: %s", face, tk_face) + logger.trace("face: %s, tk_face: %s", face, tk_face) # type:ignore[attr-defined] return tk_face - def get_landmarks(self, frame_index, face_index, face, top_left, refresh=False): + def get_landmarks(self, + frame_index: int, + face_index: int, + face: DetectedFace, + top_left: list[float], + refresh: bool = False + ) -> dict[T.Literal["polygon", "line"], list[np.ndarray]]: """ Obtain the landmark points for each mesh annotation. First tries to obtain the aligned landmarks from the cache. If the landmarks do not exist @@ -289,7 +302,7 @@ def get_landmarks(self, frame_index, face_index, face, top_left, refresh=False): The face index of the face within the requested frame face: :class:`lib.align.DetectedFace` The detected face object to obtain landmarks for - top_left: tuple + top_left: list[float] The top left (x, y) points of the face's bounding box within the viewport refresh: bool, optional Whether to force a reload of the face's aligned landmarks, even if they already exist @@ -302,17 +315,16 @@ def get_landmarks(self, frame_index, face_index, face, top_left, refresh=False): (`polygon`, `line`). The value is a list containing the (x, y) coordinates of each part of the mesh annotation, from the top left corner location. """ - key = "{}_{}".format(frame_index, face_index) + key = f"{frame_index}_{face_index}" landmarks = self._landmarks.get(key, None) if not landmarks or refresh: aligned = AlignedFace(face.landmarks_xy, centering=self._centering, size=self.face_size) - landmarks = dict(polygon=[], line=[]) - for area, val in self._landmark_mapping.items(): - # pylint:disable=unsubscriptable-object - points = aligned.landmarks[val[0]:val[1]] + top_left - shape = "polygon" if area.endswith("eye") or area.startswith("mouth") else "line" + landmarks = {"polygon": [], "line": []} + for start, end, fill in LANDMARK_PARTS[aligned.landmark_type].values(): + points = aligned.landmarks[start:end] + top_left + shape: T.Literal["polygon", "line"] = "polygon" if fill else "line" landmarks[shape].append(points) self._landmarks[key] = landmarks return landmarks @@ -328,10 +340,12 @@ def _locate_mesh(self, mesh_ids, landmarks): The mesh point groupings and whether each group should be a line or a polygon """ for key, area in landmarks.items(): + if key not in mesh_ids: + continue for coords, mesh_id in zip(area, mesh_ids[key]): self._canvas.coords(mesh_id, *coords.flatten()) - def face_from_point(self, point_x, point_y): + def face_from_point(self, point_x: int, point_y: int) -> np.ndarray: """ Given an (x, y) point on the :class:`Viewport`, obtain the face information at that location. @@ -360,15 +374,117 @@ def face_from_point(self, point_x, point_y): retval = np.array((-1, -1, -1, -1)) else: retval = self._objects.visible_grid[:, y_idx, x_idx] - logger.trace(retval) + logger.trace(retval) # type:ignore[attr-defined] return retval - def move_active_to_top(self): + def move_active_to_top(self) -> None: """ Check whether the active frame is going off the bottom of the viewport, if so: move it to the top of the viewport. """ self._active_frame.move_to_top() +class Recycler: + """ Tkinter can slow down when constantly creating new objects. + + This class delivers recycled objects, if stale objects are available, otherwise creates a new + object + + Parameters + ---------- + :class:`~tools.manual.faceviewe.frame.FacesViewer` + The canvas that holds the faces display + """ + def __init__(self, canvas: FacesViewer) -> None: + self._canvas = canvas + self._assets: dict[T.Literal["image", "line", "polygon"], + list[int]] = {"image": [], "line": [], "polygon": []} + self._mesh_methods: dict[T.Literal["line", "polygon"], + T.Callable] = {"line": canvas.create_line, + "polygon": canvas.create_polygon} + + def recycle_assets(self, asset_ids: list[int]) -> None: + """ Recycle assets that are no longer required + + Parameters + ---------- + asset_ids: list[int] + The IDs of the assets to be recycled + """ + logger.trace("Recycling %s objects", len(asset_ids)) # type:ignore[attr-defined] + for asset_id in asset_ids: + asset_type = self._canvas.type(asset_id) + assert asset_type in self._assets + coords = (0, 0, 0, 0) if asset_type == "line" else (0, 0) + self._canvas.coords(asset_id, *coords) + + if asset_type == "image": + self._canvas.itemconfig(asset_id, image="") + + self._assets[asset_type].append(asset_id) + logger.trace("Recycled objects: %s", self._assets) # type:ignore[attr-defined] + + def get_image(self, coordinates: tuple[float | int, float | int]) -> int: + """ Obtain a recycled or new image object ID + + Parameters + ---------- + coordinates: tuple[float | int, float | int] + The co-ordinates that the image should be displayed at + + Returns + ------- + int + The canvas object id for the created image + """ + if self._assets["image"]: + retval = self._assets["image"].pop() + self._canvas.coords(retval, *coordinates) + logger.trace("Recycled image: %s", retval) # type:ignore[attr-defined] + else: + retval = self._canvas.create_image(*coordinates, + anchor=tk.NW, + tags=["viewport", "viewport_image"]) + logger.trace("Created new image: %s", retval) # type:ignore[attr-defined] + return retval + + def get_mesh(self, face: DetectedFace) -> dict[T.Literal["polygon", "line"], list[int]]: + """ Get the mesh annotation for the landmarks. This is made up of a series of polygons + or lines, depending on which part of the face is being annotated. Creates a new series of + objects, or pulls existing objects from the recycled objects pool if they are available. + + Parameters + ---------- + face: :class:`~lib.align.detected_face.DetectedFace` + The detected face object to obrain the mesh for + + Returns + ------- + dict[Literal["polygon", "line"], list[int]] + The dictionary of line and polygon tkinter canvas object ids for the mesh annotation + """ + mesh_kwargs = self._canvas.viewport.mesh_kwargs + mesh_parts = LANDMARK_PARTS[LandmarkType.from_shape(face.landmarks_xy.shape)] + retval: dict[T.Literal["polygon", "line"], list[int]] = {} + for _, _, fill in mesh_parts.values(): + asset_type: T.Literal["polygon", "line"] = "polygon" if fill else "line" + kwargs = mesh_kwargs[asset_type] + if self._assets[asset_type]: + asset_id = self._assets[asset_type].pop() + self._canvas.itemconfig(asset_id, **kwargs) + logger.trace("Recycled mesh %s: %s", # type:ignore[attr-defined] + asset_type, asset_id) + else: + coords = (0, 0) if asset_type == "polygon" else (0, 0, 0, 0) + tags = ["viewport", "viewport_mesh", f"viewport_{asset_type}"] + asset_id = self._mesh_methods[asset_type](coords, width=1, tags=tags, **kwargs) + logger.trace("Created new mesh %s: %s", # type:ignore[attr-defined] + asset_type, asset_id) + + retval.setdefault(asset_type, []).append(asset_id) + logger.trace("Got mesh: %s", retval) # type:ignore[attr-defined] + return retval + + class VisibleObjects(): """ Holds the objects from the :class:`~tools.manual.faceviewer.frame.Grid` that appear in the viewable area of the :class:`Viewport`. @@ -378,20 +494,22 @@ class VisibleObjects(): viewport: :class:`Viewport` The viewport object for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas """ - def __init__(self, viewport): + def __init__(self, viewport: Viewport) -> None: + logger.debug(parse_class_init(locals())) self._viewport = viewport self._canvas = viewport._canvas self._grid = viewport._grid self._size = viewport.face_size - self._visible_grid = None - self._visible_faces = None - self._images = [] - self._meshes = [] - self._recycled = dict(images=[], meshes=[]) + self._visible_grid = np.zeros((4, 0, 0)) + self._visible_faces = np.zeros((0, 0)) + self._recycler = Recycler(self._canvas) + self._images = np.zeros((0, 0), dtype=np.int64) + self._meshes = np.zeros((0, 0)) + logger.debug("Initialized: %s", self.__class__.__name__) @property - def visible_grid(self): + def visible_grid(self) -> np.ndarray: """ :class:`numpy.ndarray`: The currently visible section of the :class:`~tools.manual.faceviewer.frame.Grid` @@ -403,7 +521,7 @@ def visible_grid(self): return self._visible_grid @property - def visible_faces(self): + def visible_faces(self) -> np.ndarray: """ :class:`numpy.ndarray`: The currently visible :class:`~lib.align.DetectedFace` objects. @@ -414,7 +532,7 @@ def visible_faces(self): return self._visible_faces @property - def images(self): + def images(self) -> np.ndarray: """ :class:`numpy.ndarray`: The viewport's tkinter canvas image objects. A numpy array of shape (`rows`, `columns`) corresponding to the viewable area of the @@ -423,7 +541,7 @@ def images(self): return self._images @property - def meshes(self): + def meshes(self) -> np.ndarray: """ :class:`numpy.ndarray`: The viewport's tkinter canvas mesh annotation objects. A numpy array of shape (`rows`, `columns`) corresponding to the viewable area of the @@ -433,28 +551,29 @@ def meshes(self): return self._meshes @property - def _top_left(self): + def _top_left(self) -> np.ndarray: """ :class:`numpy.ndarray`: The canvas (`x`, `y`) position of the face currently in the viewable area's top left position. """ - if self._images is None or not np.any(self._images): - retval = [0, 0] + if not np.any(self._images): + retval = [0.0, 0.0] else: retval = self._canvas.coords(self._images[0][0]) return np.array(retval, dtype="int") - def update(self): + def update(self) -> None: """ Load and unload thumbnails in the visible area of the faces viewer. """ if self._canvas.optional_annotations["mesh"]: # Display any hidden end of row meshes self._canvas.itemconfig("viewport_mesh", state="normal") self._visible_grid, self._visible_faces = self._grid.visible_area - if (isinstance(self._images, np.ndarray) and isinstance(self._visible_grid, np.ndarray) + if (np.any(self._images) and np.any(self._visible_grid) and self._visible_grid.shape[1:] != self._images.shape): self._reset_viewport() required_rows = self._visible_grid.shape[1] if self._grid.is_valid else 0 existing_rows = len(self._images) - logger.trace("existing_rows: %s. required_rows: %s", existing_rows, required_rows) + logger.trace("existing_rows: %s. required_rows: %s", # type:ignore[attr-defined] + existing_rows, required_rows) if existing_rows > required_rows: self._remove_rows(existing_rows, required_rows) @@ -463,43 +582,20 @@ def update(self): self._shift() - def _reset_viewport(self): + def _reset_viewport(self) -> None: """ Reset all objects in the viewport on a column count change. Reset the viewport size to the newly specified face size. """ logger.debug("Resetting Viewport") self._size = self._viewport.face_size images = self._images.flatten().tolist() - meshes = self._meshes.flatten().tolist() - self._recycle_objects(images, meshes) - self._images = [] - self._meshes = [] - - def _recycle_objects(self, images, meshes): - """ Reset the visible property and position of the given objects and add to the recycle - bin. - - Parameters - --------- - images: list - List of image_ids to be recycled - meshes: list - List of dictionaries containing the mesh annotation ids to be recycled - """ - logger.debug("Recycling objects: (images: %s, meshes: %s)", len(images), len(meshes)) - for image_id in images: - self._canvas.itemconfig(image_id, image="") - self._canvas.coords(image_id, 0, 0) - for mesh in meshes: - for key, mesh_ids in mesh.items(): - coords = (0, 0, 0, 0) if key == "line" else (0, 0) - for mesh_id in mesh_ids: - self._canvas.coords(mesh_id, *coords) - - self._recycled["images"].extend(images) - self._recycled["meshes"].extend(meshes) - logger.trace("Recycled objects: %s", self._recycled) - - def _remove_rows(self, existing_rows, required_rows): + meshes = [parts for mesh in [mesh.values() for mesh in self._meshes.flatten()] + for parts in mesh] + mesh_ids = [asset for mesh in meshes for asset in mesh] + self._recycler.recycle_assets(images + mesh_ids) + self._images = np.zeros((0, 0), np.int64) + self._meshes = np.zeros((0, 0)) + + def _remove_rows(self, existing_rows: int, required_rows: int) -> None: """ Remove and recycle rows from the viewport that are not in the view area. Parameters @@ -511,14 +607,19 @@ def _remove_rows(self, existing_rows, required_rows): """ logger.debug("Removing rows from viewport: (existing_rows: %s, required_rows: %s)", existing_rows, required_rows) - self._recycle_objects(self._images[required_rows: existing_rows].flatten().tolist(), - self._meshes[required_rows: existing_rows].flatten().tolist()) + images = self._images[required_rows: existing_rows].flatten().tolist() + meshes = [parts + for mesh in [mesh.values() + for mesh in self._meshes[required_rows: existing_rows].flatten()] + for parts in mesh] + mesh_ids = [asset for mesh in meshes for asset in mesh] + self._recycler.recycle_assets(images + mesh_ids) self._images = self._images[:required_rows] self._meshes = self._meshes[:required_rows] - logger.trace("self._images: %s, self._meshes: %s", + logger.trace("self._images: %s, self._meshes: %s", # type:ignore[attr-defined] self._images.shape, self._meshes.shape) - def _add_rows(self, existing_rows, required_rows): + def _add_rows(self, existing_rows: int, required_rows: int) -> None: """ Add rows to the viewport. Parameters @@ -531,92 +632,42 @@ def _add_rows(self, existing_rows, required_rows): logger.debug("Adding rows to viewport: (existing_rows: %s, required_rows: %s)", existing_rows, required_rows) columns = self._grid.columns_rows[0] - if not isinstance(self._images, np.ndarray): - base_coords = [(col * self._size, 0) for col in range(columns)] + + base_coords: list[list[float | int]] + + if not np.any(self._images): + base_coords = [[col * self._size, 0] for col in range(columns)] else: base_coords = [self._canvas.coords(item_id) for item_id in self._images[0]] - logger.trace("existing rows: %s, required_rows: %s, base_coords: %s", - existing_rows, required_rows, base_coords) + logger.trace("existing rows: %s, required_rows: %s, " # type:ignore[attr-defined] + "base_coords: %s", existing_rows, required_rows, base_coords) images = [] meshes = [] for row in range(existing_rows, required_rows): y_coord = base_coords[0][1] + (row * self._size) - images.append(np.array([self._get_image((coords[0], y_coord)) - for coords in base_coords])) - meshes.append(np.array([self._get_mesh() for _ in range(columns)])) - images = np.array(images) - meshes = np.array(meshes) + images.append([self._recycler.get_image((coords[0], y_coord)) + for coords in base_coords]) + meshes.append([{} if face is None else self._recycler.get_mesh(face) + for face in self._visible_faces[row]]) + + a_images = np.array(images) + a_meshes = np.array(meshes) - if not isinstance(self._images, np.ndarray): + if not np.any(self._images): logger.debug("Adding initial viewport objects: (image shapes: %s, mesh shapes: %s)", - images.shape, meshes.shape) - self._images = images - self._meshes = meshes + a_images.shape, a_meshes.shape) + self._images = a_images + self._meshes = a_meshes else: logger.debug("Adding new viewport objects: (image shapes: %s, mesh shapes: %s)", - images.shape, meshes.shape) - self._images = np.concatenate((self._images, images)) - self._meshes = np.concatenate((self._meshes, meshes)) - logger.trace("self._images: %s, self._meshes: %s", self._images.shape, self._meshes.shape) + a_images.shape, a_meshes.shape) + self._images = np.concatenate((self._images, a_images)) + self._meshes = np.concatenate((self._meshes, a_meshes)) - def _get_image(self, coordinates): - """ Create or recycle a tkinter canvas image object with the given coordinates. - - Parameters - ---------- - coordinates: tuple - The (`x`, `y`) coordinates for the top left corner of the image - - Returns - ------- - int - The canvas object id for the created image - """ - if self._recycled["images"]: - image_id = self._recycled["images"].pop() - self._canvas.coords(image_id, *coordinates) - logger.trace("Recycled image: %s", image_id) - else: - image_id = self._canvas.create_image(*coordinates, - anchor=tk.NW, - tags=["viewport", "viewport_image"]) - logger.trace("Created new image: %s", image_id) - return image_id - - def _get_mesh(self): - """ Get the mesh annotation for the landmarks. This is made up of a series of polygons - or lines, depending on which part of the face is being annotated. Creates a new series of - objects, or pulls existing objects from the recycled objects pool if they are available. + logger.trace("self._images: %s, self._meshes: %s", # type:ignore[attr-defined] + self._images.shape, self._meshes.shape) - Returns - ------- - dict - The dictionary of line and polygon tkinter canvas object ids for the mesh annotation - """ - kwargs = self._viewport.mesh_kwargs - logger.trace("self.mesh_kwargs: %s", kwargs) - if self._recycled["meshes"]: - mesh = self._recycled["meshes"].pop() - for key, mesh_ids in mesh.items(): - for mesh_id in mesh_ids: - self._canvas.itemconfig(mesh_id, **kwargs[key]) - logger.trace("Recycled mesh: %s", mesh) - else: - tags = ["viewport", "viewport_mesh"] - mesh = dict(polygon=[self._canvas.create_polygon(0, 0, - width=1, - tags=tags + ["viewport_polygon"], - **kwargs["polygon"]) - for _ in range(4)], - line=[self._canvas.create_line(0, 0, 0, 0, - width=1, - tags=tags + ["viewport_line"], - **kwargs["line"]) - for _ in range(5)]) - logger.trace("Created new mesh: %s", mesh) - return mesh - - def _shift(self): + def _shift(self) -> bool: """ Shift the viewport in the y direction if required Returns @@ -625,378 +676,18 @@ def _shift(self): ``True`` if the viewport was shifted otherwise ``False`` """ current_y = self._top_left[1] - required_y = self._visible_grid[3, 0, 0] if self._grid.is_valid else 0 - logger.trace("current_y: %s, required_y: %s", current_y, required_y) + required_y = self.visible_grid[3, 0, 0] if self._grid.is_valid else 0 + logger.trace("current_y: %s, required_y: %s", # type:ignore[attr-defined] + current_y, required_y) if current_y == required_y: - logger.trace("No move required") + logger.trace("No move required") # type:ignore[attr-defined] return False shift_amount = required_y - current_y - logger.trace("Shifting viewport: %s", shift_amount) + logger.trace("Shifting viewport: %s", shift_amount) # type:ignore[attr-defined] self._canvas.move("viewport", 0, shift_amount) return True -class HoverBox(): # pylint:disable=too-few-public-methods - """ Handle the current mouse location when over the :class:`Viewport`. - - Highlights the face currently underneath the cursor and handles actions when clicking - on a face. - - Parameters - ---------- - viewport: :class:`Viewport` - The viewport object for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas - """ - def __init__(self, viewport): - logger.debug("Initializing: %s (viewport: %s)", self.__class__.__name__, viewport) - self._viewport = viewport - self._canvas = viewport._canvas - self._grid = viewport._canvas.grid - self._globals = viewport._canvas._globals - self._navigation = viewport._canvas._display_frame.navigation - self._box = self._canvas.create_rectangle(0, 0, self._size, self._size, - outline="#0000ff", - width=2, - state="hidden", - fill="#0000ff", - stipple="gray12", - tags="hover_box") - self._current_frame_index = None - self._current_face_index = None - self._canvas.bind("", lambda e: self._clear()) - self._canvas.bind("", self.on_hover) - self._canvas.bind("", lambda e: self._select_frame()) - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def _size(self): - """ int: the currently set viewport face size in pixels. """ - return self._viewport.face_size - - def on_hover(self, event): - """ Highlight the face and set the mouse cursor for the mouse's current location. - - Parameters - ---------- - event: :class:`tkinter.Event` or ``None`` - The tkinter mouse event. Provides the current location of the mouse cursor. If ``None`` - is passed as the event (for example when this function is being called outside of a - mouse event) then the location of the cursor will be calculated - """ - if event is None: - pnts = np.array((self._canvas.winfo_pointerx(), self._canvas.winfo_pointery())) - pnts -= np.array((self._canvas.winfo_rootx(), self._canvas.winfo_rooty())) - else: - pnts = (event.x, event.y) - - coords = (int(self._canvas.canvasx(pnts[0])), int(self._canvas.canvasy(pnts[1]))) - face = self._viewport.face_from_point(*coords) - frame_idx, face_idx = face[:2] - - if frame_idx == self._current_frame_index and face_idx == self._current_face_index: - return - - is_zoomed = self._globals.is_zoomed - if (-1 in face or (frame_idx == self._globals.frame_index - and (not is_zoomed or - (is_zoomed and face_idx == self._globals.tk_face_index.get())))): - self._clear() - self._canvas.config(cursor="") - self._current_frame_index = None - self._current_face_index = None - return - - logger.debug("Viewport hover: frame_idx: %s, face_idx: %s", frame_idx, face_idx) - - self._canvas.config(cursor="hand2") - self._highlight(face[2:]) - self._current_frame_index = frame_idx - self._current_face_index = face_idx - - def _clear(self): - """ Hide the hover box when the mouse is not over a face. """ - if self._canvas.itemcget(self._box, "state") != "hidden": - self._canvas.itemconfig(self._box, state="hidden") - - def _highlight(self, top_left): - """ Display the hover box around the face that the mouse is currently over. - - Parameters - ---------- - top_left: tuple - The top left point of the highlight box location - """ - coords = (*top_left, *top_left + self._size) - self._canvas.coords(self._box, *coords) - self._canvas.itemconfig(self._box, state="normal") - self._canvas.tag_raise(self._box) - - def _select_frame(self): - """ Select the face and the subsequent frame (in the editor view) when a face is clicked - on in the :class:`Viewport`. - """ - frame_id = self._current_frame_index - is_zoomed = self._globals.is_zoomed - logger.debug("Face clicked. Global frame index: %s, Current frame_id: %s, is_zoomed: %s", - self._globals.frame_index, frame_id, is_zoomed) - if frame_id is None or (frame_id == self._globals.frame_index and not is_zoomed): - return - face_idx = self._current_face_index if is_zoomed else 0 - self._globals.tk_face_index.set(face_idx) - transport_id = self._grid.transport_index_from_frame(frame_id) - logger.trace("frame_index: %s, transport_id: %s, face_idx: %s", - frame_id, transport_id, face_idx) - if transport_id is None: - return - self._navigation.stop_playback() - self._globals.tk_transport_index.set(transport_id) - self._viewport.move_active_to_top() - self.on_hover(None) - - -class ActiveFrame(): - """ Handles the display of faces and annotations for the currently active frame. - - Parameters - ---------- - canvas: :class:`tkinter.Canvas` - The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas - tk_edited_variable: :class:`tkinter.BooleanVar` - The tkinter callback variable indicating that a face has been edited - """ - def __init__(self, viewport, tk_edited_variable): - logger.debug("Initializing: %s (viewport: %s, tk_edited_variable: %s)", - self.__class__.__name__, viewport, tk_edited_variable) - self._objects = viewport._objects - self._viewport = viewport - self._grid = viewport._grid - self._tk_faces = viewport._tk_faces - self._canvas = viewport._canvas - self._globals = viewport._canvas._globals - self._navigation = viewport._canvas._display_frame.navigation - self._last_execution = dict(frame_index=-1, size=viewport.face_size) - self._tk_vars = dict(selected_editor=self._canvas._display_frame.tk_selected_action, - edited=tk_edited_variable) - self._assets = dict(images=[], meshes=[], faces=[], boxes=[]) - - self._globals.tk_update_active_viewport.trace("w", lambda *e: self._reload_callback()) - tk_edited_variable.trace("w", lambda *e: self._update_on_edit()) - logger.debug("Initialized: %s", self.__class__.__name__) - - @property - def frame_index(self): - """ int: The frame index of the currently displayed frame. """ - return self._globals.frame_index - - @property - def current_frame(self): - """ :class:`numpy.ndarray`: A BGR version of the frame currently being displayed. """ - return self._globals.current_frame["image"] - - @property - def _size(self): - """ int: The size of the thumbnails displayed in the viewport, in pixels. """ - return self._viewport.face_size - - @property - def _optional_annotations(self): - """ dict: The currently selected optional annotations """ - return self._canvas.optional_annotations - - def _reload_callback(self): - """ If a frame has changed, triggering the variable, then update the active frame. Return - having done nothing if the variable is resetting. """ - if self._globals.tk_update_active_viewport.get(): - self.reload_annotations() - - def reload_annotations(self): - """ Handles the reloading of annotations for the currently active faces. - - Highlights the faces within the viewport of those faces that exist in the currently - displaying frame. Applies annotations based on the optional annotations and current - editor selections. - """ - logger.trace("Reloading annotations") - if np.any(self._assets["images"]): - self._clear_previous() - - self._set_active_objects() - self._check_active_in_view() - - if not np.any(self._assets["images"]): - logger.trace("No active faces. Returning") - self._last_execution["frame_index"] = self.frame_index - return - - if self._last_execution["frame_index"] != self.frame_index: - self.move_to_top() - self._create_new_boxes() - - self._update_face() - self._canvas.tag_raise("active_highlighter") - self._globals.tk_update_active_viewport.set(False) - self._last_execution["frame_index"] = self.frame_index - - def _clear_previous(self): - """ Reverts the previously selected annotations to their default state. """ - logger.trace("Clearing previous active frame") - self._canvas.itemconfig("active_highlighter", state="hidden") - - for key in ("polygon", "line"): - tag = "active_mesh_{}".format(key) - self._canvas.itemconfig(tag, **self._viewport.mesh_kwargs[key], width=1) - self._canvas.dtag(tag) - - if self._viewport.selected_editor == "mask" and not self._optional_annotations["mask"]: - for key, tk_face in self._tk_faces.items(): - if key.startswith("{}_".format(self._last_execution["frame_index"])): - tk_face.update_mask(None) - - def _set_active_objects(self): - """ Collect the objects that exist in the currently active frame from the main grid. """ - if self._grid.is_valid: - rows, cols = np.where(self._objects.visible_grid[0] == self.frame_index) - logger.trace("Setting active objects: (rows: %s, columns: %s)", rows, cols) - self._assets["images"] = self._objects.images[rows, cols] - self._assets["meshes"] = self._objects.meshes[rows, cols] - self._assets["faces"] = self._objects.visible_faces[rows, cols] - else: - logger.trace("No valid grid. Clearing active objects") - self._assets["images"] = [] - self._assets["meshes"] = [] - self._assets["faces"] = [] - - def _check_active_in_view(self): - """ If the frame has changed, there are faces in the frame, but they don't appear in the - viewport, then bring the active faces to the top of the viewport. """ - if (not np.any(self._assets["images"]) and - self._last_execution["frame_index"] != self.frame_index and - self._grid.frame_has_faces(self.frame_index)): - y_coord = self._grid.y_coord_from_frame(self.frame_index) - logger.trace("Active not in view. Moving to: %s", y_coord) - self._canvas.yview_moveto(y_coord / self._canvas.bbox("backdrop")[3]) - self._viewport.update() - - def move_to_top(self): - """ Move the currently selected frame's faces to the top of the viewport if they are moving - off the bottom of the viewer. """ - height = self._canvas.bbox("backdrop")[3] - bot = int(self._canvas.coords(self._assets["images"][-1])[1] + self._size) - - y_top, y_bot = (int(round(pnt * height)) for pnt in self._canvas.yview()) - - if y_top < bot < y_bot: # bottom face is still in fully visible area - logger.trace("Active faces in frame. Returning") - return - - top = int(self._canvas.coords(self._assets["images"][0])[1]) - if y_top == top: - logger.trace("Top face already on top row. Returning") - return - - if self._canvas.winfo_height() > self._size: - logger.trace("Viewport taller than single face height. Moving Active faces to top: %s", - top) - self._canvas.yview_moveto(top / height) - self._viewport.update() - elif self._canvas.winfo_height() <= self._size and y_top != top: - logger.trace("Viewport shorter than single face height. Moving Active faces to " - "top: %s", top) - self._canvas.yview_moveto(top / height) - self._viewport.update() - - def _create_new_boxes(self): - """ The highlight boxes (border around selected faces) are the only additional annotations - that are required for the highlighter. If more faces are displayed in the current frame - than highlight boxes are available, then new boxes are created to accommodate the - additional faces. """ - new_boxes_count = max(0, len(self._assets["images"]) - len(self._assets["boxes"])) - if new_boxes_count == 0: - return - logger.debug("new_boxes_count: %s", new_boxes_count) - for _ in range(new_boxes_count): - box = self._canvas.create_rectangle(0, - 0, - self._viewport.face_size, self._viewport.face_size, - outline="#00FF00", - width=2, - state="hidden", - tags=["active_highlighter"]) - logger.trace("Created new highlight_box: %s", box) - self._assets["boxes"].append(box) - - def _update_on_edit(self): - """ Update the active faces on a frame edit. """ - if not self._tk_vars["edited"].get(): - return - self._set_active_objects() - self._update_face() - self._tk_vars["edited"].set(False) - - def _update_face(self): - """ Update the highlighted annotations for faces in the currently selected frame. """ - for face_idx, (image_id, mesh_ids, box_id, det_face), in enumerate( - zip(self._assets["images"], - self._assets["meshes"], - self._assets["boxes"], - self._assets["faces"])): - if det_face is None: - continue - top_left = np.array(self._canvas.coords(image_id)) - coords = (*top_left, *top_left + self._size) - tk_face = self._viewport.get_tk_face(self.frame_index, face_idx, det_face) - self._canvas.itemconfig(image_id, image=tk_face.photo) - self._show_box(box_id, coords) - self._show_mesh(mesh_ids, face_idx, det_face, top_left) - self._last_execution["size"] = self._viewport.face_size - - def _show_box(self, item_id, coordinates): - """ Display the highlight box around the given coordinates. - - Parameters - ---------- - item_id: int - The tkinter canvas object identifier for the highlight box - coordinates: :class:`numpy.ndarray` - The (x, y, x1, y1) coordinates of the top left corner of the box - """ - self._canvas.coords(item_id, *coordinates) - self._canvas.itemconfig(item_id, state="normal") - - def _show_mesh(self, mesh_ids, face_index, detected_face, top_left): - """ Display the mesh annotation for the given face, at the given location. - - Parameters - ---------- - mesh_ids: dict - Dictionary containing the `polygon` and `line` tkinter canvas identifiers that make up - the mesh for the given face - face_index: int - The face index within the frame for the given face - detected_face: :class:`~lib.align.DetectedFace` - The detected face object that contains the landmarks for generating the mesh - top_left: tuple - The (x, y) top left co-ordinates of the mesh's bounding box - """ - state = "normal" if (self._tk_vars["selected_editor"].get() != "Mask" or - self._optional_annotations["mesh"]) else "hidden" - kwargs = dict(polygon=dict(fill="", width=2, outline=self._canvas.control_colors["Mesh"]), - line=dict(fill=self._canvas.control_colors["Mesh"], width=2)) - - edited = (self._tk_vars["edited"].get() and - self._tk_vars["selected_editor"].get() not in ("Mask", "View")) - landmarks = self._viewport.get_landmarks(self.frame_index, - face_index, - detected_face, - top_left, - edited) - for key, kwarg in kwargs.items(): - for idx, mesh_id in enumerate(mesh_ids[key]): - self._canvas.coords(mesh_id, *landmarks[key][idx].flatten()) - self._canvas.itemconfig(mesh_id, state=state, **kwarg) - self._canvas.addtag_withtag("active_mesh_{}".format(key), mesh_id) - - class TKFace(): """ An object that holds a single :class:`tkinter.PhotoImage` face, ready for placement in the :class:`Viewport`, Handles the placement of and removal of masks for the face as well as @@ -1013,12 +704,8 @@ class TKFace(): The mask to be applied to the face image. Pass ``None`` if no mask is to be used. Default ``None`` """ - def __init__(self, face, size=128, mask=None): - logger.trace("Initializing %s: (face: %s, size: %s, mask: %s)", - self.__class__.__name__, - face if face is None else face.shape, - size, - mask if mask is None else mask.shape) + def __init__(self, face: np.ndarray, size: int = 128, mask: np.ndarray | None = None) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] self._size = size if face.ndim == 2 and face.shape[1] == 1: self._face = self._image_from_jpg(face) @@ -1026,17 +713,17 @@ def __init__(self, face, size=128, mask=None): self._face = face[..., 2::-1] self._photo = ImageTk.PhotoImage(self._generate_tk_face_data(mask)) - logger.trace("Initialized %s", self.__class__.__name__) + logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] # << PUBLIC PROPERTIES >> # @property - def photo(self): + def photo(self) -> tk.PhotoImage: """ :class:`tkinter.PhotoImage`: The face in a format that can be placed on the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas. """ return self._photo # << PUBLIC METHODS >> # - def update(self, face, mask): + def update(self, face: np.ndarray, mask: np.ndarray) -> None: """ Update the :attr:`photo` with the given face and mask. Parameters @@ -1049,7 +736,7 @@ def update(self, face, mask): self._face = face[..., 2::-1] self._photo.paste(self._generate_tk_face_data(mask)) - def update_mask(self, mask): + def update_mask(self, mask: np.ndarray | None) -> None: """ Update the mask in the 4th channel of :attr:`photo` to the given mask. Parameters @@ -1060,7 +747,7 @@ def update_mask(self, mask): self._photo.paste(self._generate_tk_face_data(mask)) # << PRIVATE METHODS >> # - def _image_from_jpg(self, face): + def _image_from_jpg(self, face: np.ndarray) -> np.ndarray: """ Convert an encoded jpg into 3 channel BGR image. Parameters @@ -1079,7 +766,7 @@ def _image_from_jpg(self, face): face = cv2.resize(face, (self._size, self._size), interpolation=interp) return face[..., 2::-1] - def _generate_tk_face_data(self, mask): + def _generate_tk_face_data(self, mask: np.ndarray | None) -> tk.PhotoImage: """ Create the :class:`tkinter.PhotoImage` from the currant :attr:`_face`. Parameters diff --git a/tools/manual/frameviewer/control.py b/tools/manual/frameviewer/control.py index 54737783b3..8315cf6a3a 100644 --- a/tools/manual/frameviewer/control.py +++ b/tools/manual/frameviewer/control.py @@ -11,7 +11,7 @@ from lib.align import AlignedFace -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) class Navigation(): @@ -48,11 +48,11 @@ def nav_scale_callback(self, *args, reset_progress=True): # pylint:disable=unus frame_count = self._det_faces.filter.count if self._current_nav_frame_count == frame_count: logger.trace("Filtered count has not changed. Returning") - if self._globals.tk_filter_mode.get() == "Misaligned Faces": + if self._globals.var_filter_mode.get() == "Misaligned Faces": self._det_faces.tk_face_count_changed.set(True) self._update_total_frame_count() if reset_progress: - self._globals.tk_transport_index.set(0) + self._globals.var_transport_index.set(0) def _update_total_frame_count(self, *args): # pylint:disable=unused-argument """ Update the displayed number of total frames that meet the current filter criteria. @@ -70,7 +70,7 @@ def _update_total_frame_count(self, *args): # pylint:disable=unused-argument logger.debug("Filtered frame count has changed. Updating from %s to %s", self._current_nav_frame_count, frame_count) self._nav["scale"].config(to=max_frame) - self._nav["label"].config(text="/{}".format(max_frame)) + self._nav["label"].config(text=f"/{max_frame}") state = "disabled" if max_frame == 0 else "normal" self._nav["entry"].config(state=state) @@ -106,7 +106,7 @@ def increment_frame(self, frame_count=None, is_playing=False): logger.debug("End of Stream. Not incrementing") self.stop_playback() return - self._globals.tk_transport_index.set(min(position + 1, max(0, frame_count - 1))) + self._globals.var_transport_index.set(min(position + 1, max(0, frame_count - 1))) def decrement_frame(self): """ Update The frame navigation position to the previous frame based on filter. """ @@ -116,11 +116,11 @@ def decrement_frame(self): if not face_count_change and (self._det_faces.filter.count == 0 or position == 0): logger.debug("End of Stream. Not decrementing") return - self._globals.tk_transport_index.set(min(max(0, self._det_faces.filter.count - 1), - max(0, position - 1))) + self._globals.var_transport_index.set(min(max(0, self._det_faces.filter.count - 1), + max(0, position - 1))) def _get_safe_frame_index(self): - """ Obtain the current frame position from the tk_transport_index variable in + """ Obtain the current frame position from the var_transport_index variable in a safe manner (i.e. handle for non-numeric) Returns @@ -129,35 +129,35 @@ def _get_safe_frame_index(self): The current transport frame index """ try: - retval = self._globals.tk_transport_index.get() + retval = self._globals.var_transport_index.get() except tk.TclError as err: if "expected floating-point" not in str(err): raise - val = str(err).split(" ")[-1].replace("\"", "") + val = str(err).rsplit(" ", maxsplit=1)[-1].replace("\"", "") retval = "".join(ch for ch in val if ch.isdigit()) retval = 0 if not retval else int(retval) - self._globals.tk_transport_index.set(retval) + self._globals.var_transport_index.set(retval) return retval def goto_first_frame(self): """ Go to the first frame that meets the filter criteria. """ self.stop_playback() - position = self._globals.tk_transport_index.get() + position = self._globals.var_transport_index.get() if position == 0: return - self._globals.tk_transport_index.set(0) + self._globals.var_transport_index.set(0) def goto_last_frame(self): """ Go to the last frame that meets the filter criteria. """ self.stop_playback() - position = self._globals.tk_transport_index.get() + position = self._globals.var_transport_index.get() frame_count = self._det_faces.filter.count if position == frame_count - 1: return - self._globals.tk_transport_index.set(frame_count - 1) + self._globals.var_transport_index.set(frame_count - 1) -class BackgroundImage(): # pylint:disable=too-few-public-methods +class BackgroundImage(): """ The background image of the canvas """ def __init__(self, canvas): self._canvas = canvas @@ -190,7 +190,7 @@ def refresh(self, view_mode): """ self._switch_image(view_mode) logger.trace("Updating background frame") - getattr(self, "_update_tk_{}".format(self._current_view_mode))() + getattr(self, f"_update_tk_{self._current_view_mode}")() def _switch_image(self, view_mode): """ Switch the image between the full frame image and the zoomed face image. @@ -206,10 +206,10 @@ def _switch_image(self, view_mode): self._zoomed_centering = self._canvas.active_editor.zoomed_centering logger.trace("Switching background image from '%s' to '%s'", self._current_view_mode, view_mode) - img = getattr(self, "_tk_{}".format(view_mode)) + img = getattr(self, f"_tk_{view_mode}") self._canvas.itemconfig(self._image, image=img) - self._globals.tk_is_zoomed.set(view_mode == "face") - self._globals.tk_face_index.set(0) + self._globals.set_zoomed(view_mode == "face") + self._globals.set_face_index(0) def _update_tk_face(self): """ Update the currently zoomed face. """ @@ -239,14 +239,14 @@ def _get_zoomed_face(self): if face_idx + 1 > faces_in_frame: logger.debug("Resetting face index to 0 for more faces in frame than current index: (" "faces_in_frame: %s, zoomed_face_index: %s", faces_in_frame, face_idx) - self._globals.tk_face_index.set(0) + self._globals.set_face_index(0) if faces_in_frame == 0: face = np.ones((size, size, 3), dtype="uint8") else: det_face = self._det_faces.current_faces[frame_idx][face_idx] face = AlignedFace(det_face.landmarks_xy, - image=self._globals.current_frame["image"], + image=self._globals.current_frame.image, centering=self._zoomed_centering, size=size).face logger.trace("face shape: %s", face.shape) @@ -254,9 +254,9 @@ def _get_zoomed_face(self): def _update_tk_frame(self): """ Place the currently held frame into :attr:`_tk_frame`. """ - img = cv2.resize(self._globals.current_frame["image"], - self._globals.current_frame["display_dims"], - interpolation=self._globals.current_frame["interpolation"])[..., 2::-1] + img = cv2.resize(self._globals.current_frame.image, + self._globals.current_frame.display_dims, + interpolation=self._globals.current_frame.interpolation)[..., 2::-1] padding = self._get_padding(img.shape[:2]) if any(padding): img = cv2.copyMakeBorder(img, *padding, cv2.BORDER_CONSTANT) diff --git a/tools/manual/frameviewer/editor/__init__.py b/tools/manual/frameviewer/editor/__init__.py index 8a7244abe4..7902b47227 100644 --- a/tools/manual/frameviewer/editor/__init__.py +++ b/tools/manual/frameviewer/editor/__init__.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 """ The Frame Viewer for Faceswap's Manual Tool. """ -from ._base import View # noqa -from .bounding_box import BoundingBox # noqa -from .extract_box import ExtractBox # noqa -from .landmarks import Landmarks, Mesh # noqa -from .mask import Mask # noqa +from ._base import View +from .bounding_box import BoundingBox +from .extract_box import ExtractBox +from .landmarks import Landmarks, Mesh +from .mask import Mask diff --git a/tools/manual/frameviewer/editor/_base.py b/tools/manual/frameviewer/editor/_base.py index 46c19b5bd2..d295c0d847 100644 --- a/tools/manual/frameviewer/editor/_base.py +++ b/tools/manual/frameviewer/editor/_base.py @@ -11,7 +11,7 @@ from lib.gui.control_helper import ControlPanelOption -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("tools.manual", localedir="locales", fallback=True) @@ -41,9 +41,9 @@ def __init__(self, canvas, detected_faces, control_text="", key_bindings=None): self._globals = canvas._globals self._det_faces = detected_faces - self._current_color = dict() + self._current_color = {} self._actions = OrderedDict() - self._controls = dict(header=control_text, controls=[]) + self._controls = {"header": control_text, "controls": []} self._add_key_bindings(key_bindings) self._add_actions() @@ -51,7 +51,7 @@ def __init__(self, canvas, detected_faces, control_text="", key_bindings=None): self._add_annotation_format_controls() self._mouse_location = None - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None self.bind_mouse_motion() logger.debug("Initialized %s", self.__class__.__name__) @@ -80,7 +80,7 @@ def _is_active(self): def view_mode(self): """ ["frame", "face"]: The view mode for the currently selected editor. If the editor does not have a view mode that can be updated, then `"frame"` will be returned. """ - tk_var = self._actions.get("magnify", dict()).get("tk_var", None) + tk_var = self._actions.get("magnify", {}).get("tk_var", None) retval = "frame" if tk_var is None or not tk_var.get() else "face" return retval @@ -106,7 +106,7 @@ def _zoomed_dims(self): @property def _control_vars(self): """ dict: The tk control panel variables for the currently selected editor. """ - return self._canvas.control_tk_vars.get(self.__class__.__name__, dict()) + return self._canvas.control_tk_vars.get(self.__class__.__name__, {}) @property def controls(self): @@ -155,7 +155,7 @@ def _add_key_bindings(self, key_bindings): for key, method in key_bindings.items(): logger.debug("Binding key '%s' to method %s for editor '%s'", key, method, self.__class__.__name__) - self._canvas.key_bindings.setdefault(key, dict())["bound_to"] = None + self._canvas.key_bindings.setdefault(key, {})["bound_to"] = None self._canvas.key_bindings[key][self.__class__.__name__] = method @staticmethod @@ -187,7 +187,7 @@ def _get_anchor_points(bounding_box): for cnr in bounding_box) return display_anchors, grab_anchors - def update_annotation(self): # pylint:disable=no-self-use + def update_annotation(self): """ Update the display annotations for the current objects. Override for specific editors. @@ -233,7 +233,7 @@ def _object_tracker(self, key, object_type, face_index, """ object_color_keys = self._get_object_color_keys(key, object_type) tracking_id = "_".join((key, str(face_index))) - face_tag = "face_{}".format(face_index) + face_tag = f"face_{face_index}" face_objects = set(self._canvas.find_withtag(face_tag)) annotation_objects = set(self._canvas.find_withtag(key)) existing_object = tuple(face_objects.intersection(annotation_objects)) @@ -311,7 +311,7 @@ def _add_new_object(self, key, object_type, face_index, coordinates, object_kwar coordinates, object_kwargs) object_kwargs["tags"] = self._set_object_tags(face_index, key) item_id = getattr(self._canvas, - "create_{}".format(object_type))(*coordinates, **object_kwargs) + f"create_{object_type}")(*coordinates, **object_kwargs) return item_id def _set_object_tags(self, face_index, key): @@ -329,17 +329,17 @@ def _set_object_tags(self, face_index, key): list The generated tags for the current object """ - tags = ["face_{}".format(face_index), + tags = [f"face_{face_index}", self.__class__.__name__, - "{}_face_{}".format(self.__class__.__name__, face_index), + f"{self.__class__.__name__}_face_{face_index}", key, - "{}_face_{}".format(key, face_index)] + f"{key}_face_{face_index}"] if "_" in key: split_key = key.split("_") if split_key[-1].isdigit(): base_tag = "_".join(split_key[:-1]) tags.append(base_tag) - tags.append("{}_face_{}".format(base_tag, face_index)) + tags.append(f"{base_tag}_face_{face_index}") return tags def _update_existing_object(self, item_id, coordinates, object_kwargs, @@ -366,11 +366,11 @@ def _update_existing_object(self, item_id, coordinates, object_kwargs, """ update_color = (object_color_keys and object_kwargs[object_color_keys[0]] != self._current_color[tracking_id]) - update_kwargs = dict(state=object_kwargs.get("state", "normal")) + update_kwargs = {"state": object_kwargs.get("state", "normal")} if update_color: for key in object_color_keys: update_kwargs[key] = object_kwargs[object_color_keys[0]] - if self._canvas.type(item_id) == "image" and "image" in object_kwargs: + if self._canvas.type(item_id) == "image" and "image" in object_kwargs: # noqa:E721 update_kwargs["image"] = object_kwargs["image"] logger.trace("Updating coordinates: (item_id: '%s', object_kwargs: %s, " "coordinates: %s, update_kwargs: %s", item_id, object_kwargs, @@ -389,7 +389,7 @@ def bind_mouse_motion(self): """ self._canvas.bind("", self._update_cursor) - def _update_cursor(self, event): # pylint: disable=unused-argument + def _update_cursor(self, event): # pylint:disable=unused-argument """ The mouse cursor display as bound to the mouse's event.. The default is to always return a standard cursor, so this method should be overridden for @@ -433,7 +433,7 @@ def _drag_start(self, event): # pylint:disable=unused-argument The tkinter mouse event. Unused but for default action, but available for editor specific actions """ - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None def _drag(self, event): @@ -461,7 +461,7 @@ def _drag_stop(self, event): # pylint:disable=unused-argument event: :class:`tkinter.Event` The tkinter mouse event. Unused but required """ - self._drag_data = dict() + self._drag_data = {} def _scale_to_display(self, points): """ Scale and offset the given points to the current display scale and offset values. @@ -476,7 +476,7 @@ def _scale_to_display(self, points): :class:`numpy.ndarray` The adjusted x, y co-ordinates for display purposes rounded to the nearest integer """ - retval = np.rint((points * self._globals.current_frame["scale"]) + retval = np.rint((points * self._globals.current_frame.scale) + self._canvas.offset).astype("int32") logger.trace("Original points: %s, scaled points: %s", points, retval) return retval @@ -499,7 +499,7 @@ def scale_from_display(self, points, do_offset=True): integer """ offset = self._canvas.offset if do_offset else (0, 0) - retval = np.rint((points - offset) / self._globals.current_frame["scale"]).astype("int32") + retval = np.rint((points - offset) / self._globals.current_frame.scale).astype("int32") logger.trace("Original points: %s, scaled points: %s", points, retval) return retval @@ -532,7 +532,11 @@ def _add_action(self, title, icon, helptext, group=None, hotkey=None): Default: ``None`` """ var = tk.BooleanVar() - action = dict(icon=icon, helptext=helptext, group=group, tk_var=var, hotkey=hotkey) + action = {"icon": icon, + "helptext": helptext, + "group": group, + "tk_var": var, + "hotkey": hotkey} logger.debug("Adding action: %s", action) self._actions[title] = action @@ -567,7 +571,7 @@ def _add_control(self, option, global_control=False): group_key = "none" if group_key == "_master" else group_key annotation_key = option.title.replace(" ", "") self._canvas.control_tk_vars.setdefault( - editor_key, dict()).setdefault(group_key, dict())[annotation_key] = option.tk_var + editor_key, {}).setdefault(group_key, {})[annotation_key] = option.tk_var def _add_annotation_format_controls(self): """ Add the annotation display (color/size) controls to :attr:`_annotation_formats`. @@ -594,7 +598,7 @@ def _add_annotation_format_controls(self): default=self._default_colors[annotation_key], helptext="Set the annotation color") colors.set(self._default_colors[annotation_key]) - self._annotation_formats.setdefault(annotation_key, dict())["color"] = colors + self._annotation_formats.setdefault(annotation_key, {})["color"] = colors self._annotation_formats[annotation_key]["mask_opacity"] = opacity for editor in editors: @@ -627,4 +631,6 @@ def _add_actions(self): """ Add the optional action buttons to the viewer. Current actions are Zoom. """ self._add_action("magnify", "zoom", _("Magnify/Demagnify the View"), group=None, hotkey="M") - self._actions["magnify"]["tk_var"].trace("w", lambda *e: self._globals.tk_update.set(True)) + self._actions["magnify"]["tk_var"].trace_add( + "write", + lambda *e: self._globals.var_full_update.set(True)) diff --git a/tools/manual/frameviewer/editor/bounding_box.py b/tools/manual/frameviewer/editor/bounding_box.py index 9eeef016fe..d546feb172 100644 --- a/tools/manual/frameviewer/editor/bounding_box.py +++ b/tools/manual/frameviewer/editor/bounding_box.py @@ -105,7 +105,7 @@ def update_annotation(self): for idx, face in enumerate(self._face_iterator): box = np.array([(face.left, face.top), (face.right, face.bottom)]) box = self._scale_to_display(box).astype("int32").flatten() - kwargs = dict(outline=color, width=1) + kwargs = {"outline": color, "width": 1} logger.trace("frame_index: %s, face_index: %s, box: %s, kwargs: %s", self._globals.frame_index, idx, box, kwargs) self._object_tracker(key, "rectangle", idx, box, kwargs) @@ -137,10 +137,10 @@ def _update_anchor_annotation(self, face_index, bounding_box, color): (bounding_box[2], bounding_box[3]), (bounding_box[0], bounding_box[3]))) for idx, (anc_dsp, anc_grb) in enumerate(zip(*anchor_points)): - dsp_kwargs = dict(outline=color, fill=fill_color, width=1) - grb_kwargs = dict(outline="", fill="", width=1, activefill=activefill_color) - dsp_key = "bb_anc_dsp_{}".format(idx) - grb_key = "bb_anc_grb_{}".format(idx) + dsp_kwargs = {"outline": color, "fill": fill_color, "width": 1} + grb_kwargs = {"outline": '', "fill": '', "width": 1, "activefill": activefill_color} + dsp_key = f"bb_anc_dsp_{idx}" + grb_key = f"bb_anc_grb_{idx}" self._object_tracker(dsp_key, "oval", face_index, anc_dsp, dsp_kwargs) self._object_tracker(grb_key, "oval", face_index, anc_grb, grb_kwargs) logger.trace("Updated bounding box anchor annotations") @@ -193,8 +193,9 @@ def _check_cursor_anchors(self): corner_idx = int(next(tag for tag in tags if tag.startswith("bb_anc_grb_") and "face_" not in tag).split("_")[-1]) - self._canvas.config(cursor="{}_{}_corner".format(*self._corner_order[corner_idx])) - self._mouse_location = ("anchor", "{}_{}".format(face_idx, corner_idx)) + pos_x, pos_y = self._corner_order[corner_idx] + self._canvas.config(cursor=f"{pos_x}_{pos_y}_corner") + self._mouse_location = ("anchor", f"{face_idx}_{corner_idx}") return True def _check_cursor_bounding_box(self, event): @@ -242,7 +243,7 @@ def _check_cursor_image(self, event): """ if self._globals.frame_index == -1: return False - display_dims = self._globals.current_frame["display_dims"] + display_dims = self._globals.current_frame.display_dims if (self._canvas.offset[0] <= event.x <= display_dims[0] + self._canvas.offset[0] and self._canvas.offset[1] <= event.y <= display_dims[1] + self._canvas.offset[1]): self._canvas.config(cursor="plus") @@ -275,7 +276,7 @@ def _drag_start(self, event): The tkinter mouse event. """ if self._mouse_location is None: - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None return if self._mouse_location[0] == "anchor": @@ -291,7 +292,7 @@ def _drag_start(self, event): self._update_cursor(event) self._drag_start(event) - def _drag_stop(self, event): # pylint: disable=unused-argument + def _drag_stop(self, event): # pylint:disable=unused-argument """ Trigger a viewport thumbnail update on click + drag release Parameters @@ -315,7 +316,7 @@ def _create_new_bounding_box(self, event): event: :class:`tkinter.Event` The tkinter mouse event """ - size = min(self._globals.current_frame["display_dims"]) // 8 + size = min(self._globals.current_frame.display_dims) // 8 box = (event.x - size, event.y - size, event.x + size, event.y + size) logger.debug("Creating new bounding box: %s ", box) self._det_faces.update.add(self._globals.frame_index, *self._coords_to_bounding_box(box)) @@ -329,7 +330,7 @@ def _resize(self, event): The tkinter mouse event. """ face_idx = int(self._mouse_location[1].split("_")[0]) - face_tag = "bb_box_face_{}".format(face_idx) + face_tag = f"bb_box_face_{face_idx}" box = self._canvas.coords(face_tag) logger.trace("Face Index: %s, Corner Index: %s. Original ROI: %s", face_idx, self._drag_data["corner"], box) @@ -361,7 +362,7 @@ def _move(self, event): face_idx = int(self._mouse_location[1]) shift = (event.x - self._drag_data["current_location"][0], event.y - self._drag_data["current_location"][1]) - face_tag = "bb_box_face_{}".format(face_idx) + face_tag = f"bb_box_face_{face_idx}" coords = np.array(self._canvas.coords(face_tag)) + (*shift, *shift) logger.trace("face_tag: %s, shift: %s, new co-ords: %s", face_tag, shift, coords) self._det_faces.update.bounding_box(self._globals.frame_index, diff --git a/tools/manual/frameviewer/editor/extract_box.py b/tools/manual/frameviewer/editor/extract_box.py index eb739545a2..ffe8bf4734 100644 --- a/tools/manual/frameviewer/editor/extract_box.py +++ b/tools/manual/frameviewer/editor/extract_box.py @@ -61,9 +61,9 @@ def update_annotation(self): aligned = AlignedFace(face.landmarks_xy, centering="face") box = self._scale_to_display(aligned.original_roi).flatten() top_left = box[:2] - 10 - kwargs = dict(fill=color, font=("Default", 20, "bold"), text=str(idx)) + kwargs = {"fill": color, "font": ('Default', 20, 'bold'), "text": str(idx)} self._object_tracker("eb_text", "text", idx, top_left, kwargs) - kwargs = dict(fill="", outline=color, width=1) + kwargs = {"fill": '', "outline": color, "width": 1} self._object_tracker("eb_box", "polygon", idx, box, kwargs) self._update_anchor_annotation(idx, box, color) logger.trace("Updated extract box annotations") @@ -93,10 +93,10 @@ def _update_anchor_annotation(self, face_index, extract_box, color): extract_box[4:6], extract_box[6:])) for idx, (anc_dsp, anc_grb) in enumerate(zip(*anchor_points)): - dsp_kwargs = dict(outline=color, fill=fill_color, width=1) - grb_kwargs = dict(outline="", fill="", width=1, activefill=activefill_color) - dsp_key = "eb_anc_dsp_{}".format(idx) - grb_key = "eb_anc_grb_{}".format(idx) + dsp_kwargs = {"outline": color, "fill": fill_color, "width": 1} + grb_kwargs = {"outline": '', "fill": '', "width": 1, "activefill": activefill_color} + dsp_key = f"eb_anc_dsp_{idx}" + grb_key = f"eb_anc_grb_{idx}" self._object_tracker(dsp_key, "oval", face_index, anc_dsp, dsp_kwargs) self._object_tracker(grb_key, "oval", face_index, anc_grb, grb_kwargs) logger.trace("Updated extract box anchor annotations") @@ -143,7 +143,8 @@ def _check_cursor_anchors(self): if tag.startswith("eb_anc_grb_") and "face_" not in tag).split("_")[-1]) - self._canvas.config(cursor="{}_{}_corner".format(*self._corner_order[corner_idx])) + pos_x, pos_y = self._corner_order[corner_idx] + self._canvas.config(cursor=f"{pos_x}_{pos_y}_corner") self._mouse_location = ("anchor", face_idx, corner_idx) return True @@ -222,14 +223,14 @@ def _drag_start(self, event): The tkinter mouse event. """ if self._mouse_location is None: - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None return self._drag_data["current_location"] = np.array((event.x, event.y)) - callback = dict(anchor=self._resize, rotate=self._rotate, box=self._move) + callback = {"anchor": self._resize, "rotate": self._rotate, "box": self._move} self._drag_callback = callback[self._mouse_location[0]] - def _drag_stop(self, event): # pylint: disable=unused-argument + def _drag_stop(self, event): # pylint:disable=unused-argument """ Trigger a viewport thumbnail update on click + drag release Parameters @@ -270,7 +271,7 @@ def _resize(self, event): The tkinter mouse event. """ face_idx = self._mouse_location[1] - face_tag = "eb_box_face_{}".format(face_idx) + face_tag = f"eb_box_face_{face_idx}" position = np.array((event.x, event.y)) box = np.array(self._canvas.coords(face_tag)) center = np.array((sum(box[0::2]) / 4, sum(box[1::2]) / 4)) @@ -365,7 +366,7 @@ def _rotate(self, event): The tkinter mouse event. """ face_idx = self._mouse_location[1] - face_tag = "eb_box_face_{}".format(face_idx) + face_tag = f"eb_box_face_{face_idx}" box = np.array(self._canvas.coords(face_tag)) position = np.array((event.x, event.y)) diff --git a/tools/manual/frameviewer/editor/landmarks.py b/tools/manual/frameviewer/editor/landmarks.py index bc2896212d..e59517e7b0 100644 --- a/tools/manual/frameviewer/editor/landmarks.py +++ b/tools/manual/frameviewer/editor/landmarks.py @@ -3,7 +3,7 @@ import gettext import numpy as np -from lib.align import AlignedFace +from lib.align import AlignedFace, LANDMARK_PARTS, LandmarkType from ._base import Editor, logger # LOCALES @@ -36,7 +36,7 @@ def __init__(self, canvas, detected_faces): super().__init__(canvas, detected_faces, control_text) # Clear selection box on an editor or frame change self._canvas._tk_action_var.trace("w", lambda *e: self._reset_selection()) - self._globals.tk_frame_index.trace("w", lambda *e: self._reset_selection()) + self._globals.var_frame_index.trace_add("write", lambda *e: self._reset_selection()) def _add_actions(self): """ Add the optional action buttons to the viewer. Current actions are Point, Select @@ -55,7 +55,7 @@ def _toggle_zoom(self, *args): # pylint:disable=unused-argument tkinter callback arguments. Required but unused. """ self._reset_selection() - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) def _reset_selection(self, event=None): # pylint:disable=unused-argument """ Reset the selection box and the selected landmark annotations. """ @@ -67,7 +67,7 @@ def _reset_selection(self, event=None): # pylint:disable=unused-argument outline="gray", state="hidden") self._canvas.coords(self._selection_box, 0, 0, 0, 0) - self._drag_data = dict() + self._drag_data = {} if event is not None: self._drag_start(event) @@ -83,7 +83,7 @@ def update_annotation(self): landmarks = aligned.landmarks + zoomed_offset # Hide all landmarks and only display selected self._canvas.itemconfig("lm_dsp", state="hidden") - self._canvas.itemconfig("lm_dsp_face_{}".format(face_index), state="normal") + self._canvas.itemconfig(f"lm_dsp_face_{face_index}", state="normal") else: landmarks = self._scale_to_display(face.landmarks_xy) for lm_idx, landmark in enumerate(landmarks): @@ -109,8 +109,8 @@ def _display_landmark(self, bounding_box, face_index, landmark_index): color = self._control_color bbox = (bounding_box[0] - radius, bounding_box[1] - radius, bounding_box[0] + radius, bounding_box[1] + radius) - key = "lm_dsp_{}".format(landmark_index) - kwargs = dict(outline=color, fill=color, width=radius) + key = f"lm_dsp_{landmark_index}" + kwargs = {"outline": color, "fill": color, "width": radius} self._object_tracker(key, "oval", face_index, bbox, kwargs) def _label_landmark(self, bounding_box, face_index, landmark_index): @@ -132,9 +132,9 @@ def _label_landmark(self, bounding_box, face_index, landmark_index): # NB The text must be visible to be able to get the bounding box, so set to hidden # after the bounding box has been retrieved - keys = ["lm_lbl_{}".format(landmark_index), "lm_lbl_bg_{}".format(landmark_index)] - text_kwargs = dict(fill="black", font=("Default", 10), text=str(landmark_index + 1)) - bg_kwargs = dict(fill="#ffffea", outline="black") + keys = [f"lm_lbl_{landmark_index}", f"lm_lbl_bg_{landmark_index}"] + text_kwargs = {"fill": "black", "font": ("Default", 10), "text": str(landmark_index + 1)} + bg_kwargs = {"fill": "#ffffea", "outline": "black"} text_id = self._object_tracker(keys[0], "text", face_index, top_left, text_kwargs) bbox = self._canvas.bbox(text_id) @@ -162,11 +162,11 @@ def _grab_landmark(self, bounding_box, face_index, landmark_index): radius = 7 bbox = (bounding_box[0] - radius, bounding_box[1] - radius, bounding_box[0] + radius, bounding_box[1] + radius) - key = "lm_grb_{}".format(landmark_index) - kwargs = dict(outline="", - fill="", - width=1, - dash=(2, 4)) + key = f"lm_grb_{landmark_index}" + kwargs = {"outline": "", + "fill": "", + "width": 1, + "dash": (2, 4)} self._object_tracker(key, "oval", face_index, bbox, kwargs) # << MOUSE HANDLING >> @@ -185,7 +185,7 @@ def _update_cursor(self, event): if self._drag_data: self._update_cursor_select_mode(event) else: - objs = self._canvas.find_withtag("lm_grb_face_{}".format(self._globals.face_index) + objs = self._canvas.find_withtag(f"lm_grb_face_{self._globals.face_index}" if self._globals.is_zoomed else "lm_grb") item_ids = set(self._canvas.find_overlapping(event.x - 6, event.y - 6, @@ -226,7 +226,7 @@ def _update_cursor_point_mode(self, item_id): self._canvas.config(cursor="none") for prefix in ("lm_lbl_", "lm_lbl_bg_"): - tag = "{}{}_face_{}".format(prefix, lm_idx, face_idx) + tag = f"{prefix}{lm_idx}_face_{face_idx}" logger.trace("Displaying: %s tag: %s", self._canvas.type(tag), tag) self._canvas.itemconfig(tag, state="normal") self._mouse_location = obj_idx @@ -271,11 +271,11 @@ def _drag_start(self, event): self._drag_data["start_location"] = (event.x, event.y) self._drag_callback = self._move_selection else: # Reset - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None self._reset_selection(event) - def _drag_stop(self, event): # pylint: disable=unused-argument + def _drag_stop(self, event): # pylint:disable=unused-argument """ In select mode, call the select mode callback. In point mode: trigger a viewport thumbnail update on click + drag release @@ -294,7 +294,7 @@ def _drag_stop(self, event): # pylint: disable=unused-argument self._det_faces.update.post_edit_trigger(self._globals.frame_index, self._mouse_location[0]) self._mouse_location = None - self._drag_data = dict() + self._drag_data = {} elif self._drag_data and self._drag_data.get("selected", False): self._drag_stop_selected() else: @@ -429,15 +429,6 @@ class Mesh(Editor): The _detected_faces data for this manual session """ def __init__(self, canvas, detected_faces): - self._landmark_mapping = dict(mouth_inner=(60, 68), - mouth_outer=(48, 60), - right_eyebrow=(17, 22), - left_eyebrow=(22, 27), - right_eye=(36, 42), - left_eye=(42, 48), - nose=(27, 36), - jaw=(0, 17), - chin=(8, 11)) super().__init__(canvas, detected_faces, None) def update_annotation(self): @@ -452,19 +443,23 @@ def update_annotation(self): centering="face", size=min(self._globals.frame_display_dims)) landmarks = aligned.landmarks + zoomed_offset + landmark_mapping = LANDMARK_PARTS[aligned.landmark_type] # Hide all meshes and only display selected self._canvas.itemconfig("Mesh", state="hidden") - self._canvas.itemconfig("Mesh_face_{}".format(face_index), state="normal") + self._canvas.itemconfig(f"Mesh_face_{face_index}", state="normal") else: landmarks = self._scale_to_display(face.landmarks_xy) + landmark_mapping = LANDMARK_PARTS[LandmarkType.from_shape(landmarks.shape)] logger.trace("Drawing Landmarks Mesh: (landmarks: %s, color: %s)", landmarks, color) - for idx, (segment, val) in enumerate(self._landmark_mapping.items()): - key = "mesh_{}".format(idx) - pts = landmarks[val[0]:val[1]].flatten() - if segment in ("right_eye", "left_eye", "mouth_inner", "mouth_outer"): - kwargs = dict(fill="", outline=color, width=1) - self._object_tracker(key, "polygon", face_index, pts, kwargs) + for idx, (start, end, fill) in enumerate(landmark_mapping.values()): + key = f"mesh_{idx}" + pts = landmarks[start:end].flatten() + if fill: + kwargs = {"fill": "", "outline": color, "width": 1} + asset = "polygon" else: - self._object_tracker(key, "line", face_index, pts, dict(fill=color, width=1)) + kwargs = {"fill": color, "width": 1} + asset = "line" + self._object_tracker(key, asset, face_index, pts, kwargs) # Place mesh as bottom annotation self._canvas.tag_raise(self.__class__.__name__, "main_image") diff --git a/tools/manual/frameviewer/editor/mask.py b/tools/manual/frameviewer/editor/mask.py index 99d3f6b09d..fec2c92d13 100644 --- a/tools/manual/frameviewer/editor/mask.py +++ b/tools/manual/frameviewer/editor/mask.py @@ -82,7 +82,9 @@ def _add_actions(self): group=None, hotkey="M") self._add_action("draw", "draw", _("Draw Tool"), group="paint", hotkey="D") self._add_action("erase", "erase", _("Erase Tool"), group="paint", hotkey="E") - self._actions["magnify"]["tk_var"].trace("w", lambda *e: self._globals.tk_update.set(True)) + self._actions["magnify"]["tk_var"].trace( + "w", + lambda *e: self._globals.var_full_update.set(True)) def _add_controls(self): """ Add the mask specific control panel controls. @@ -122,6 +124,7 @@ def _add_controls(self): default="Circle", is_radio=True, helptext=_("Select a shape for masking cursor."))) + def _set_tk_mask_change_callback(self): """ Add a trace to change the displayed mask on a mask type change. """ var = self._control_vars["display"]["MaskType"] @@ -142,21 +145,21 @@ def _on_mask_type_change(self): mask_type = self._control_vars["display"]["MaskType"].get() if mask_type == self._mask_type: return - self._meta = dict(position=self._globals.frame_index) + self._meta = {"position": self._globals.frame_index} self._mask_type = mask_type - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) def hide_annotation(self, tag=None): """ Clear the mask :attr:`_meta` dict when hiding the annotation. """ super().hide_annotation() - self._meta = dict() + self._meta = {} def update_annotation(self): """ Update the mask annotation with the latest mask. """ position = self._globals.frame_index if position != self._meta.get("position", -1): # Reset meta information when moving to a new frame - self._meta = dict(position=position) + self._meta = {"position": position} key = self.__class__.__name__ mask_type = self._control_vars["display"]["MaskType"].get().lower() color = self._control_color[1:] @@ -220,21 +223,21 @@ def _set_full_frame_meta(self, mask, mask_scale): - slices: The (`x`, `y`) slice objects required to extract the mask ROI from the full frame """ - frame_dims = self._globals.current_frame["display_dims"] + frame_dims = self._globals.current_frame.display_dims scaled_mask_roi = np.rint(mask.original_roi * - self._globals.current_frame["scale"]).astype("int32") + self._globals.current_frame.scale).astype("int32") # Scale and clip the ROI to fit within display frame boundaries clipped_roi = scaled_mask_roi.clip(min=(0, 0), max=frame_dims) # Obtain min and max points to get ROI as a rectangle - min_max = dict(min=clipped_roi.min(axis=0), max=clipped_roi.max(axis=0)) + min_max = {"min": clipped_roi.min(axis=0), "max": clipped_roi.max(axis=0)} # Create a bounding box rectangle ROI roi_dims = np.rint((min_max["max"][1] - min_max["min"][1], min_max["max"][0] - min_max["min"][0])).astype("uint16") - roi = dict(mask=np.zeros(roi_dims, dtype="uint8")[..., None], - corners=np.expand_dims(scaled_mask_roi - min_max["min"], axis=0)) + roi = {"mask": np.zeros(roi_dims, dtype="uint8")[..., None], + "corners": np.expand_dims(scaled_mask_roi - min_max["min"], axis=0)} # Block out areas outside of the actual mask ROI polygon cv2.fillPoly(roi["mask"], roi["corners"], 255) logger.trace("Setting Full Frame mask ROI. shape: %s", roi["mask"].shape) @@ -245,8 +248,8 @@ def _set_full_frame_meta(self, mask, mask_scale): # Adjust affine matrix for internal mask size and display dimensions adjustments = (np.array([[mask_scale, 0., 0.], [0., mask_scale, 0.]]), - np.array([[1 / self._globals.current_frame["scale"], 0., 0.], - [0., 1 / self._globals.current_frame["scale"], 0.], + np.array([[1 / self._globals.current_frame.scale, 0., 0.], + [0., 1 / self._globals.current_frame.scale, 0.], [0., 0., 1.]])) in_matrix = np.dot(adjustments[0], np.concatenate((mask.affine_matrix, np.array([[0., 0., 1.]])))) @@ -284,7 +287,7 @@ def _update_mask_image(self, key, face_index, rgb_color, opacity): top_left = self._zoomed_roi[:2] # Hide all masks and only display selected self._canvas.itemconfig("Mask", state="hidden") - self._canvas.itemconfig("Mask_face_{}".format(face_index), state="normal") + self._canvas.itemconfig(f"Mask_face_{face_index}", state="normal") else: display_image = self._update_mask_image_full_frame(mask, rgb_color, face_index) top_left = self._meta["top_left"][face_index] @@ -304,7 +307,7 @@ def _update_mask_image(self, key, face_index, rgb_color, opacity): "image", face_index, top_left, - dict(image=self._tk_faces[face_index], anchor=tk.NW)) + {"image": self._tk_faces[face_index], "anchor": tk.NW}) def _update_mask_image_zoomed(self, mask, rgb_color): """ Update the mask image when zoomed in. @@ -345,7 +348,7 @@ def _update_mask_image_full_frame(self, mask, rgb_color, face_index): :class: `PIL.Image` The full frame mask image formatted for display """ - frame_dims = self._globals.current_frame["display_dims"] + frame_dims = self._globals.current_frame.display_dims frame = np.zeros(frame_dims + (1, ), dtype="uint8") interpolator = self._meta["interpolator"][face_index] slices = self._meta["slices"][face_index] @@ -376,13 +379,13 @@ def _update_roi_box(self, mask, face_index, color): else: box = self._scale_to_display(mask.original_roi).flatten() top_left = box[:2] - 10 - kwargs = dict(fill=color, font=("Default", 20, "bold"), text=str(face_index)) + kwargs = {"fill": color, "font": ("Default", 20, "bold"), "text": str(face_index)} self._object_tracker("mask_text", "text", face_index, top_left, kwargs) - kwargs = dict(fill="", outline=color, width=1) + kwargs = {"fill": "", "outline": color, "width": 1} self._object_tracker("mask_roi", "polygon", face_index, box, kwargs) if self._globals.is_zoomed: # Raise box above zoomed image - self._canvas.tag_raise("mask_roi_face_{}".format(face_index)) + self._canvas.tag_raise(f"mask_roi_face_{face_index}") # << MOUSE HANDLING >> # Mouse cursor display @@ -449,7 +452,7 @@ def _drag_start(self, event, control_click=False): # pylint:disable=arguments-d """ face_idx = self._mouse_location[1] if face_idx is None: - self._drag_data = dict() + self._drag_data = {} self._drag_callback = None else: self._drag_data["starting_location"] = np.array((event.x, event.y)) @@ -531,7 +534,7 @@ def _drag_stop(self, event): if np.array_equal(self._drag_data["starting_location"], location[0]): self._get_cursor_shape_mark(self._meta["mask"][face_idx], location, face_idx) self._mask_to_alignments(face_idx) - self._drag_data = dict() + self._drag_data = {} self._update_cursor(event) def _get_cursor_shape_mark(self, img, location, face_idx): @@ -561,11 +564,10 @@ def _get_cursor_shape_mark(self, img, location, face_idx): else: cv2.circle(img, tuple(points), radius, color, thickness=-1) - def _get_cursor_shape(self, x1=0, y1=0, x2=0, y2=0, outline="black", state="hidden"): + def _get_cursor_shape(self, x_1=0, y_1=0, x_2=0, y_2=0, outline="black", state="hidden"): if self._cursor_shape_name == "Rectangle": - return self._canvas.create_rectangle(x1, y1, x2, y2, outline=outline, state=state) - else: - return self._canvas.create_oval(x1, y1, x2, y2, outline=outline, state=state) + return self._canvas.create_rectangle(x_1, y_1, x_2, y_2, outline=outline, state=state) + return self._canvas.create_oval(x_1, y_1, x_2, y_2, outline=outline, state=state) def _mask_to_alignments(self, face_index): """ Update the annotated mask to alignments. diff --git a/tools/manual/frameviewer/frame.py b/tools/manual/frameviewer/frame.py index b4495ef111..e83f3492f8 100644 --- a/tools/manual/frameviewer/frame.py +++ b/tools/manual/frameviewer/frame.py @@ -16,7 +16,7 @@ from .editor import (BoundingBox, ExtractBox, Landmarks, Mask, # noqa pylint:disable=unused-import Mesh, View) -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("tools.manual", localedir="locales", fallback=True) @@ -42,7 +42,7 @@ def __init__(self, parent, tk_globals, detected_faces): self._globals = tk_globals self._det_faces = detected_faces - self._optional_widgets = dict() + self._optional_widgets = {} self._actions_frame = ActionsFrame(self) main_frame = ttk.Frame(self) @@ -74,28 +74,28 @@ def __init__(self, parent, tk_globals, detected_faces): @property def _helptext(self): """ dict: {`name`: `help text`} Helptext lookup for navigation buttons """ - return dict( - play=_("Play/Pause (SPACE)"), - beginning=_("Go to First Frame (HOME)"), - prev=_("Go to Previous Frame (Z)"), - next=_("Go to Next Frame (X)"), - end=_("Go to Last Frame (END)"), - extract=_("Extract the faces to a folder... (Ctrl+E)"), - save=_("Save the Alignments file (Ctrl+S)"), - mode=_("Filter Frames to only those Containing the Selected Item (F)"), - distance=_("Set the distance from an 'average face' to be considered misaligned. " - "Higher distances are more restrictive")) + return { + "play": _("Play/Pause (SPACE)"), + "beginning": _("Go to First Frame (HOME)"), + "prev": _("Go to Previous Frame (Z)"), + "next": _("Go to Next Frame (X)"), + "end": _("Go to Last Frame (END)"), + "extract": _("Extract the faces to a folder... (Ctrl+E)"), + "save": _("Save the Alignments file (Ctrl+S)"), + "mode": _("Filter Frames to only those Containing the Selected Item (F)"), + "distance": _("Set the distance from an 'average face' to be considered misaligned. " + "Higher distances are more restrictive")} @property def _btn_action(self): """ dict: {`name`: `action`} Command lookup for navigation buttons """ - actions = dict(play=self._navigation.handle_play_button, - beginning=self._navigation.goto_first_frame, - prev=self._navigation.decrement_frame, - next=self._navigation.increment_frame, - end=self._navigation.goto_last_frame, - extract=self._det_faces.extract, - save=self._det_faces.save) + actions = {"play": self._navigation.handle_play_button, + "beginning": self._navigation.goto_first_frame, + "prev": self._navigation.decrement_frame, + "next": self._navigation.increment_frame, + "end": self._navigation.goto_last_frame, + "extract": self._det_faces.extract, + "save": self._det_faces.save} return actions @property @@ -146,48 +146,48 @@ def _add_nav(self): lbl_frame.pack(side=tk.RIGHT) tbox = ttk.Entry(lbl_frame, width=7, - textvariable=self._globals.tk_transport_index, + textvariable=self._globals.var_transport_index, justify=tk.RIGHT) tbox.pack(padx=0, side=tk.LEFT) - lbl = ttk.Label(lbl_frame, text="/{}".format(max_frame)) + lbl = ttk.Label(lbl_frame, text=f"/{max_frame}") lbl.pack(side=tk.RIGHT) cmd = partial(set_slider_rounding, - var=self._globals.tk_transport_index, + var=self._globals.var_transport_index, d_type=int, round_to=1, min_max=(0, max_frame)) nav = ttk.Scale(frame, - variable=self._globals.tk_transport_index, + variable=self._globals.var_transport_index, from_=0, to=max_frame, command=cmd) nav.pack(side=tk.LEFT, fill=tk.BOTH, expand=True) - self._globals.tk_transport_index.trace("w", self._set_frame_index) - return dict(entry=tbox, scale=nav, label=lbl) + self._globals.var_transport_index.trace_add("write", self._set_frame_index) + return {"entry": tbox, "scale": nav, "label": lbl} def _set_frame_index(self, *args): # pylint:disable=unused-argument """ Set the actual frame index based on current slider position and filter mode. """ try: - slider_position = self._globals.tk_transport_index.get() + slider_position = self._globals.var_transport_index.get() except TclError: # don't update the slider when the entry box has been cleared of any value return frames = self._det_faces.filter.frames_list actual_position = max(0, min(len(frames) - 1, slider_position)) if actual_position != slider_position: - self._globals.tk_transport_index.set(actual_position) + self._globals.var_transport_index.set(actual_position) frame_idx = frames[actual_position] if frames else -1 logger.trace("slider_position: %s, frame_idx: %s", actual_position, frame_idx) - self._globals.tk_frame_index.set(frame_idx) + self._globals.var_frame_index.set(frame_idx) def _add_transport(self): """ Add video transport controls """ frame = ttk.Frame(self._transport_frame) frame.pack(side=tk.BOTTOM, fill=tk.X) icons = get_images().icons - buttons = dict() + buttons = {} for action in ("play", "beginning", "prev", "next", "end", "save", "extract", "mode"): padx = (0, 6) if action in ("play", "prev", "mode") else (0, 0) side = tk.RIGHT if action in ("extract", "save", "mode") else tk.LEFT @@ -237,14 +237,14 @@ def _add_filter_mode_combo(self, frame): frame: :class:`tkinter.ttk.Frame` The Filter Frame that holds the filter combo box """ - self._globals.tk_filter_mode.set("All Frames") - self._globals.tk_filter_mode.trace("w", self._navigation.nav_scale_callback) + self._globals.var_filter_mode.set("All Frames") + self._globals.var_filter_mode.trace("w", self._navigation.nav_scale_callback) nav_frame = ttk.Frame(frame) lbl = ttk.Label(nav_frame, text="Filter:") lbl.pack(side=tk.LEFT, padx=(0, 5)) combo = ttk.Combobox( nav_frame, - textvariable=self._globals.tk_filter_mode, + textvariable=self._globals.var_filter_mode, state="readonly", values=self._filter_modes) combo.pack(side=tk.RIGHT) @@ -260,7 +260,7 @@ def _add_filter_threshold_slider(self, frame): The Filter Frame that holds the filter threshold slider """ slider_frame = ttk.Frame(frame) - tk_var = self._globals.tk_filter_distance + tk_var = self._globals.var_filter_distance min_max = (5, 20) ctl_frame = ttk.Frame(slider_frame) @@ -284,22 +284,22 @@ def _add_filter_threshold_slider(self, frame): Tooltip(item, text=self._helptext["distance"], wrap_length=200) - tk_var.trace("w", self._navigation.nav_scale_callback) + tk_var.trace_add("write", self._navigation.nav_scale_callback) self._optional_widgets["distance_slider"] = slider_frame def pack_threshold_slider(self): """ Display or hide the threshold slider depending on the current filter mode. For misaligned faces filter, display the slider. Hide for all other filters. """ - if self._globals.tk_filter_mode.get() == "Misaligned Faces": + if self._globals.var_filter_mode.get() == "Misaligned Faces": self._optional_widgets["distance_slider"].pack(side=tk.LEFT) else: self._optional_widgets["distance_slider"].pack_forget() def cycle_filter_mode(self): """ Cycle the navigation mode combo entry """ - current_mode = self._globals.filter_mode + current_mode = self._globals.var_filter_mode.get() idx = (self._filter_modes.index(current_mode) + 1) % len(self._filter_modes) - self._globals.tk_filter_mode.set(self._filter_modes[idx]) + self._globals.var_filter_mode.set(self._filter_modes[idx]) def set_action(self, key): """ Set the current action based on keyboard shortcut @@ -318,7 +318,7 @@ def _resize(self, event): framesize = (event.width, event.height) logger.trace("Resizing video frame. Framesize: %s", framesize) self._globals.set_frame_display_dims(*framesize) - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) # << TRANSPORT >> # def _play(self, *args, frame_count=None): # pylint:disable=unused-argument @@ -366,7 +366,7 @@ def __init__(self, parent): self._buttons = self._add_buttons() self._static_buttons = self._add_static_buttons() self._selected_action = self._set_selected_action_tkvar() - self._optional_buttons = dict() # Has to be set from parent after canvas is initialized + self._optional_buttons = {} # Has to be set from parent after canvas is initialized @property def actions(self): @@ -382,19 +382,19 @@ def tk_selected_action(self): def key_bindings(self): """ dict: {`key`: `action`}. The mapping of key presses to actions. Keyboard shortcut is the first letter of each action. """ - return {"F{}".format(idx + 1): action for idx, action in enumerate(self._actions)} + return {f"F{idx + 1}": action for idx, action in enumerate(self._actions)} @property def _helptext(self): """ dict: `button key`: `button helptext`. The help text to display for each button. """ inverse_keybindings = {val: key for key, val in self.key_bindings.items()} - retval = dict(View=_("View alignments"), - BoundingBox=_("Bounding box editor"), - ExtractBox=_("Location editor"), - Mask=_("Mask editor"), - Landmarks=_("Landmark point editor")) + retval = {"View": _('View alignments'), + "BoundingBox": _('Bounding box editor'), + "ExtractBox": _("Location editor"), + "Mask": _("Mask editor"), + "Landmarks": _("Landmark point editor")} for item in retval: - retval[item] += " ({})".format(inverse_keybindings[item]) + retval[item] += f" ({inverse_keybindings[item]})" return retval def _configure_styles(self): @@ -415,7 +415,7 @@ def _add_buttons(self): """ frame = ttk.Frame(self) frame.pack(side=tk.TOP, fill=tk.Y) - buttons = dict() + buttons = {} for action in self.key_bindings.values(): if action == self._initial_action: btn_style = "actions_selected.TButton" @@ -467,23 +467,24 @@ def _set_selected_action_tkvar(self): def _add_static_buttons(self): """ Add the buttons to copy alignments from previous and next frames """ - lookup = dict(copy_prev=(_("Previous"), "C"), copy_next=(_("Next"), "V"), reload=("", "R")) + lookup = {"copy_prev": (_("Previous"), "C"), + "copy_next": (_("Next"), "V"), + "reload": ("", "R")} frame = ttk.Frame(self) frame.pack(side=tk.TOP, fill=tk.Y) sep = ttk.Frame(frame, height=2, relief=tk.RIDGE) sep.pack(fill=tk.X, pady=5, side=tk.TOP) - buttons = dict() - tk_frame_index = self._globals.tk_frame_index + buttons = {} for action in ("copy_prev", "copy_next", "reload"): if action == "reload": icon = "reload3" - cmd = lambda f=tk_frame_index: self._det_faces.revert_to_saved(f.get()) # noqa + cmd = lambda f=self._globals: self._det_faces.revert_to_saved(f.frame_index) # noqa:E731,E501 # pylint:disable=line-too-long,unnecessary-lambda-assignment helptext = _("Revert to saved Alignments ({})").format(lookup[action][1]) else: icon = action direction = action.replace("copy_", "") - cmd = lambda f=tk_frame_index, d=direction: self._det_faces.update.copy( # noqa - f.get(), d) + cmd = lambda f=self._globals, d=direction: self._det_faces.update.copy( # noqa:E731,E501 # pylint:disable=line-too-long,unnecessary-lambda-assignment + f.frame_index, d) helptext = _("Copy {} Alignments ({})").format(*lookup[action]) state = ["!disabled"] if action == "copy_next" else ["disabled"] button = ttk.Button(frame, @@ -494,11 +495,11 @@ def _add_static_buttons(self): button.pack() Tooltip(button, text=helptext) buttons[action] = button - self._globals.tk_frame_index.trace("w", self._disable_enable_copy_buttons) - self._globals.tk_update.trace("w", self._disable_enable_reload_button) + self._globals.var_frame_index.trace_add("write", self._disable_enable_copy_buttons) + self._globals.var_full_update.trace_add("write", self._disable_enable_reload_button) return buttons - def _disable_enable_copy_buttons(self, *args): # pylint: disable=unused-argument + def _disable_enable_copy_buttons(self, *args): # pylint:disable=unused-argument """ Disable or enable the static buttons """ position = self._globals.frame_index face_count_per_index = self._det_faces.face_count_per_index @@ -506,12 +507,12 @@ def _disable_enable_copy_buttons(self, *args): # pylint: disable=unused-argumen for count in face_count_per_index[:position]) next_exists = position != -1 and any(count != 0 for count in face_count_per_index[position + 1:]) - states = dict(prev=["!disabled"] if prev_exists else ["disabled"], - next=["!disabled"] if next_exists else ["disabled"]) + states = {"prev": ["!disabled"] if prev_exists else ["disabled"], + "next": ["!disabled"] if next_exists else ["disabled"]} for direction in ("prev", "next"): - self._static_buttons["copy_{}".format(direction)].state(states[direction]) + self._static_buttons[f"copy_{direction}"].state(states[direction]) - def _disable_enable_reload_button(self, *args): # pylint: disable=unused-argument + def _disable_enable_reload_button(self, *args): # pylint:disable=unused-argument """ Disable or enable the static buttons """ position = self._globals.frame_index state = ["!disabled"] if (position != -1 and @@ -549,12 +550,12 @@ def add_optional_buttons(self, editors): helptext = action["helptext"] hotkey = action["hotkey"] - helptext += "" if hotkey is None else " ({})".format(hotkey.upper()) + helptext += "" if hotkey is None else f" ({hotkey.upper()})" Tooltip(button, text=helptext) self._optional_buttons.setdefault( - name, dict())[button] = dict(hotkey=hotkey, - group=group, - tk_var=action["tk_var"]) + name, {})[button] = {"hotkey": hotkey, + "group": group, + "tk_var": action["tk_var"]} self._optional_buttons[name]["frame"] = frame self._display_optional_buttons() @@ -652,9 +653,9 @@ def __init__(self, parent, tk_globals, detected_faces, actions, tk_action_var): self._actions = actions self._tk_action_var = tk_action_var self._image = BackgroundImage(self) - self._editor_globals = dict(control_tk_vars=dict(), - annotation_formats=dict(), - key_bindings=dict()) + self._editor_globals = {"control_tk_vars": {}, + "annotation_formats": {}, + "key_bindings": {}} self._max_face_count = 0 self._editors = self._get_editors() self._add_callbacks() @@ -695,17 +696,17 @@ def editors(self): @property def editor_display(self): """ dict: List of editors and any additional annotations they should display. """ - return dict(View=["BoundingBox", "ExtractBox", "Landmarks", "Mesh"], - BoundingBox=["Mesh"], - ExtractBox=["Mesh"], - Landmarks=["ExtractBox", "Mesh"], - Mask=[]) + return {"View": ["BoundingBox", "ExtractBox", "Landmarks", "Mesh"], + "BoundingBox": ["Mesh"], + "ExtractBox": ["Mesh"], + "Landmarks": ["ExtractBox", "Mesh"], + "Mask": []} @property def offset(self): """ tuple: The (`width`, `height`) offset of the canvas based on the size of the currently displayed image """ - frame_dims = self._globals.current_frame["display_dims"] + frame_dims = self._globals.current_frame.display_dims offset_x = (self._globals.frame_display_dims[0] - frame_dims[0]) / 2 offset_y = (self._globals.frame_display_dims[1] - frame_dims[1]) / 2 logger.trace("offset_x: %s, offset_y: %s", offset_x, offset_y) @@ -719,7 +720,7 @@ def _get_editors(self): dict The {`action`: :class:`Editor`} dictionary of editors for :attr:`_actions` name. """ - editors = dict() + editors = {} for editor_name in self._actions + ("Mesh", ): editor = eval(editor_name)(self, # pylint:disable=eval-used self._det_faces) @@ -731,11 +732,11 @@ def _add_callbacks(self): """ Add the callback trace functions to the :class:`tkinter.Variable` s Adds callbacks for: - :attr:`_globals.tk_update` Update the display for the current image + :attr:`_globals.var_full_update` Update the display for the current image :attr:`__tk_action_var` Update the mouse display tracking for current action """ - self._globals.tk_update.trace("w", self._update_display) - self._tk_action_var.trace("w", self._change_active_editor) + self._globals.var_full_update.trace_add("write", self._update_display) + self._tk_action_var.trace_add("write", self._change_active_editor) def _change_active_editor(self, *args): # pylint:disable=unused-argument """ Update the display for the active editor. @@ -755,7 +756,7 @@ def _change_active_editor(self, *args): # pylint:disable=unused-argument self.active_editor.bind_mouse_motion() self.active_editor.set_mouse_click_actions() - self._globals.tk_update.set(True) + self._globals.var_full_update.set(True) def _update_display(self, *args): # pylint:disable=unused-argument """ Update the display on frame cache update @@ -765,7 +766,7 @@ def _update_display(self, *args): # pylint:disable=unused-argument A little hacky, but the editors to display or hide are processed in alphabetical order, so that they are always processed in the same order (for tag lowering and raising) """ - if not self._globals.tk_update.get(): + if not self._globals.var_full_update.get(): return zoomed_centering = self.active_editor.zoomed_centering self._image.refresh(self.active_editor.view_mode) @@ -777,7 +778,7 @@ def _update_display(self, *args): # pylint:disable=unused-argument if zoomed_centering != self.active_editor.zoomed_centering: # Refresh the image if editor annotation has changed the zoom centering of the image self._image.refresh(self.active_editor.view_mode) - self._globals.tk_update.set(False) + self._globals.var_full_update.set(False) self.update_idletasks() def _hide_additional_faces(self): @@ -797,7 +798,7 @@ def _hide_additional_faces(self): self._max_face_count = current_face_count return for idx in range(current_face_count, self._max_face_count): - tag = "face_{}".format(idx) + tag = f"face_{idx}" if any(self.itemcget(item_id, "state") != "hidden" for item_id in self.find_withtag(tag)): logger.debug("Hiding face tag '%s'", tag) diff --git a/tools/manual/globals.py b/tools/manual/globals.py new file mode 100644 index 0000000000..07843f552e --- /dev/null +++ b/tools/manual/globals.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python3 +""" Holds global tkinter variables and information pertaining to the entire Manual tool """ +from __future__ import annotations + +import logging +import os +import sys +import tkinter as tk + +from dataclasses import dataclass, field + +import cv2 +import numpy as np + +from lib.gui.utils import get_config +from lib.logger import parse_class_init +from lib.utils import VIDEO_EXTENSIONS + +logger = logging.getLogger(__name__) + + +@dataclass +class CurrentFrame: + """ Dataclass for holding information about the currently displayed frame """ + image: np.ndarray = field(default_factory=lambda: np.zeros(1)) + """:class:`numpy.ndarry`: The currently displayed frame in original dimensions """ + scale: float = 1.0 + """float: The scaling factor to use to resize the image to the display window """ + interpolation: int = cv2.INTER_AREA + """int: The opencv interpolator ID to use for resizing the image to the display window """ + display_dims: tuple[int, int] = (0, 0) + """tuple[int, int]`: The size of the currently displayed frame, in the display window """ + filename: str = "" + """str: The filename of the currently displayed frame """ + + def __repr__(self) -> str: + """ Clean string representation showing numpy arrays as shape and dtype + + Returns + ------- + str + Loggable representation of the dataclass + """ + properties = [f"{k}={(v.shape, v.dtype) if isinstance(v, np.ndarray) else v}" + for k, v in self.__dict__.items()] + return f"{self.__class__.__name__} ({', '.join(properties)}" + + +@dataclass +class TKVars: + """ Holds the global TK Variables """ + frame_index: tk.IntVar + """:class:`tkinter.IntVar`: The absolute frame index of the currently displayed frame""" + transport_index: tk.IntVar + """:class:`tkinter.IntVar`: The transport index of the currently displayed frame when filters + have been applied """ + face_index: tk.IntVar + """:class:`tkinter.IntVar`: The face index of the currently selected face""" + filter_distance: tk.IntVar + """:class:`tkinter.IntVar`: The amount to filter by distance""" + + update: tk.BooleanVar + """:class:`tkinter.BooleanVar`: Whether an update has been performed """ + update_active_viewport: tk.BooleanVar + """:class:`tkinter.BooleanVar`: Whether the viewport needs updating """ + is_zoomed: tk.BooleanVar + """:class:`tkinter.BooleanVar`: Whether the main window is zoomed in to a face or out to a + full frame""" + + filter_mode: tk.StringVar + """:class:`tkinter.StringVar`: The currently selected filter mode """ + faces_size: tk.StringVar + """:class:`tkinter.StringVar`: The pixel size of faces in the viewport """ + + def __repr__(self) -> str: + """ Clean string representation showing variable type as well as their value + + Returns + ------- + str + Loggable representation of the dataclass + """ + properties = [f"{k}={v.__class__.__name__}({v.get()})" for k, v in self.__dict__.items()] + return f"{self.__class__.__name__} ({', '.join(properties)}" + + +class TkGlobals(): + """ Holds Tkinter Variables and other frame information that need to be accessible from all + areas of the GUI. + + Parameters + ---------- + input_location: str + The location of the input folder of frames or video file + """ + def __init__(self, input_location: str) -> None: + logger.debug(parse_class_init(locals())) + self._tk_vars = self._get_tk_vars() + + self._is_video = self._check_input(input_location) + self._frame_count = 0 # set by FrameLoader + self._frame_display_dims = (int(round(896 * get_config().scaling_factor)), + int(round(504 * get_config().scaling_factor))) + self._current_frame = CurrentFrame() + logger.debug("Initialized %s", self.__class__.__name__) + + @classmethod + def _get_tk_vars(cls) -> TKVars: + """ Create and initialize the tkinter variables. + + Returns + ------- + :class:`TKVars` + The global tkinter variables + """ + retval = TKVars(frame_index=tk.IntVar(value=0), + transport_index=tk.IntVar(value=0), + face_index=tk.IntVar(value=0), + filter_distance=tk.IntVar(value=10), + update=tk.BooleanVar(value=False), + update_active_viewport=tk.BooleanVar(value=False), + is_zoomed=tk.BooleanVar(value=False), + filter_mode=tk.StringVar(), + faces_size=tk.StringVar()) + logger.debug(retval) + return retval + + @property + def current_frame(self) -> CurrentFrame: + """ :class:`CurrentFrame`: The currently displayed frame in the frame viewer with it's + meta information. """ + return self._current_frame + + @property + def frame_count(self) -> int: + """ int: The total number of frames for the input location """ + return self._frame_count + + @property + def frame_display_dims(self) -> tuple[int, int]: + """ tuple: The (`width`, `height`) of the video display frame in pixels. """ + return self._frame_display_dims + + @property + def is_video(self) -> bool: + """ bool: ``True`` if the input is a video file, ``False`` if it is a folder of images. """ + return self._is_video + + # TK Variables that need to be exposed + @property + def var_full_update(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: Flag to indicate that whole GUI should be refreshed """ + return self._tk_vars.update + + @property + def var_transport_index(self) -> tk.IntVar: + """ :class:`tkinter.IntVar`: The current index of the display frame's transport slider. """ + return self._tk_vars.transport_index + + @property + def var_frame_index(self) -> tk.IntVar: + """ :class:`tkinter.IntVar`: The current absolute frame index of the currently + displayed frame. """ + return self._tk_vars.frame_index + + @property + def var_filter_distance(self) -> tk.IntVar: + """ :class:`tkinter.IntVar`: The variable holding the currently selected threshold + distance for misaligned filter mode. """ + return self._tk_vars.filter_distance + + @property + def var_filter_mode(self) -> tk.StringVar: + """ :class:`tkinter.StringVar`: The variable holding the currently selected navigation + filter mode. """ + return self._tk_vars.filter_mode + + @property + def var_faces_size(self) -> tk.StringVar: + """ :class:`tkinter..IntVar`: The variable holding the currently selected Faces Viewer + thumbnail size. """ + return self._tk_vars.faces_size + + @property + def var_update_active_viewport(self) -> tk.BooleanVar: + """ :class:`tkinter.BooleanVar`: Boolean Variable that is traced by the viewport's active + frame to update. """ + return self._tk_vars.update_active_viewport + + # Raw values returned from TK Variables + @property + def face_index(self) -> int: + """ int: The currently displayed face index when in zoomed mode. """ + return self._tk_vars.face_index.get() + + @property + def frame_index(self) -> int: + """ int: The currently displayed frame index. NB This returns -1 if there are no frames + that meet the currently selected filter criteria. """ + return self._tk_vars.frame_index.get() + + @property + def is_zoomed(self) -> bool: + """ bool: ``True`` if the frame viewer is zoomed into a face, ``False`` if the frame viewer + is displaying a full frame. """ + return self._tk_vars.is_zoomed.get() + + @staticmethod + def _check_input(frames_location: str) -> bool: + """ Check whether the input is a video + + Parameters + ---------- + frames_location: str + The input location for video or images + + Returns + ------- + bool: 'True' if input is a video 'False' if it is a folder. + """ + if os.path.isdir(frames_location): + retval = False + elif os.path.splitext(frames_location)[1].lower() in VIDEO_EXTENSIONS: + retval = True + else: + logger.error("The input location '%s' is not valid", frames_location) + sys.exit(1) + logger.debug("Input '%s' is_video: %s", frames_location, retval) + return retval + + def set_face_index(self, index: int) -> None: + """ Set the currently selected face index + + Parameters + ---------- + index: int + The currently selected face index + """ + logger.trace("Setting face index from %s to %s", # type:ignore[attr-defined] + self.face_index, index) + self._tk_vars.face_index.set(index) + + def set_frame_count(self, count: int) -> None: + """ Set the count of total number of frames to :attr:`frame_count` when the + :class:`FramesLoader` has completed loading. + + Parameters + ---------- + count: int + The number of frames that exist for this session + """ + logger.debug("Setting frame_count to : %s", count) + self._frame_count = count + + def set_current_frame(self, image: np.ndarray, filename: str) -> None: + """ Set the frame and meta information for the currently displayed frame. Populates the + attribute :attr:`current_frame` + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image used to display in the Frame Viewer + filename: str + The filename of the current frame + """ + scale = min(self.frame_display_dims[0] / image.shape[1], + self.frame_display_dims[1] / image.shape[0]) + self._current_frame.image = image + self._current_frame.filename = filename + self._current_frame.scale = scale + self._current_frame.interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + self._current_frame.display_dims = (int(round(image.shape[1] * scale)), + int(round(image.shape[0] * scale))) + logger.trace(self._current_frame) # type:ignore[attr-defined] + + def set_frame_display_dims(self, width: int, height: int) -> None: + """ Set the size, in pixels, of the video frame display window and resize the displayed + frame. + + Used on a frame resize callback, sets the :attr:frame_display_dims`. + + Parameters + ---------- + width: int + The width of the frame holding the video canvas in pixels + height: int + The height of the frame holding the video canvas in pixels + """ + self._frame_display_dims = (int(width), int(height)) + image = self._current_frame.image + scale = min(self.frame_display_dims[0] / image.shape[1], + self.frame_display_dims[1] / image.shape[0]) + self._current_frame.scale = scale + self._current_frame.interpolation = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA + self._current_frame.display_dims = (int(round(image.shape[1] * scale)), + int(round(image.shape[0] * scale))) + logger.trace(self._current_frame) # type:ignore[attr-defined] + + def set_zoomed(self, state: bool) -> None: + """ Set the current zoom state + + Parameters + ---------- + state: bool + ``True`` for zoomed ``False`` for full frame + """ + logger.trace("Setting zoom state from %s to %s", # type:ignore[attr-defined] + self.is_zoomed, state) + self._tk_vars.is_zoomed.set(state) diff --git a/tools/manual/manual.py b/tools/manual/manual.py index 360f67c9bf..2516a62b9c 100644 --- a/tools/manual/manual.py +++ b/tools/manual/manual.py @@ -1,28 +1,51 @@ #!/usr/bin/env python3 -""" The Manual Tool is a tkinter driven GUI app for editing alignments files with visual tools. -This module is the main entry point into the Manual Tool. """ +""" Main entry point for the Manual Tool. A GUI app for editing alignments files """ +from __future__ import annotations + import logging import os import sys +import typing as T import tkinter as tk from tkinter import ttk +from dataclasses import dataclass from time import sleep -import cv2 import numpy as np from lib.gui.control_helper import ControlPanel from lib.gui.utils import get_images, get_config, initialize_config, initialize_images from lib.image import SingleFrameLoader, read_image_meta +from lib.logger import parse_class_init from lib.multithreading import MultiThread -from lib.utils import _video_extensions -from plugins.extract.pipeline import Extractor, ExtractMedia +from lib.utils import handle_deprecated_cliopts +from plugins.extract import ExtractMedia, Extractor -from .detected_faces import DetectedFaces, ThumbsCreator +from .detected_faces import DetectedFaces from .faceviewer.frame import FacesFrame from .frameviewer.frame import DisplayFrame +from .globals import TkGlobals +from .thumbnails import ThumbsCreator + +if T.TYPE_CHECKING: + from argparse import Namespace + from lib.align import DetectedFace, Mask + from lib.queue_manager import EventQueue + +logger = logging.getLogger(__name__) + +TypeManualExtractor = T.Literal["FAN", "cv2-dnn", "mask"] -logger = logging.getLogger(__name__) # pylint: disable=invalid-name + +@dataclass +class _Containers: + """ Dataclass for holding the main area containers in the GUI """ + main: ttk.PanedWindow + """:class:`tkinter.ttk.PanedWindow`: The main window holding the full GUI """ + top: ttk.Frame + """:class:`tkinter.ttk.Frame: The top part (frame viewer) of the GUI""" + bottom: ttk.Frame + """:class:`tkinter.ttk.Frame: The bottom part (face viewer) of the GUI""" class Manual(tk.Tk): @@ -38,9 +61,10 @@ class Manual(tk.Tk): The :mod:`argparse` arguments as passed in from :mod:`tools.py` """ - def __init__(self, arguments): - logger.debug("Initializing %s: (arguments: '%s')", self.__class__.__name__, arguments) + def __init__(self, arguments: Namespace) -> None: + logger.debug(parse_class_init(locals())) super().__init__() + arguments = handle_deprecated_cliopts(arguments) self._validate_non_faces(arguments.frames) self._initialize_tkinter() @@ -55,24 +79,27 @@ def __init__(self, arguments): video_meta_data = self._detected_faces.video_meta_data valid_meta = all(val is not None for val in video_meta_data.values()) - loader = FrameLoader(self._globals, arguments.frames, video_meta_data) + loader = FrameLoader(self._globals, + arguments.frames, + video_meta_data, + self._detected_faces.frame_list) + if valid_meta: # Load the faces whilst other threads complete if we have valid meta data self._detected_faces.load_faces() self._containers = self._create_containers() self._wait_for_threads(extractor, loader, valid_meta) - if not valid_meta: - # Load the faces after other threads complete if meta data required updating + if not valid_meta: # If meta data needs updating, load faces after other threads self._detected_faces.load_faces() self._generate_thumbs(arguments.frames, arguments.thumb_regen, arguments.single_process) - self._display = DisplayFrame(self._containers["top"], + self._display = DisplayFrame(self._containers.top, self._globals, self._detected_faces) - _Options(self._containers["top"], self._globals, self._display) + _Options(self._containers.top, self._globals, self._display) - self._faces_frame = FacesFrame(self._containers["bottom"], + self._faces_frame = FacesFrame(self._containers.bottom, self._globals, self._detected_faces, self._display) @@ -83,7 +110,7 @@ def __init__(self, arguments): logger.debug("Initialized %s", self.__class__.__name__) @classmethod - def _validate_non_faces(cls, frames_folder): + def _validate_non_faces(cls, frames_folder: str) -> None: """ Quick check on the input to make sure that a folder of extracted faces is not being passed in. """ if not os.path.isdir(frames_folder): @@ -106,7 +133,7 @@ def _validate_non_faces(cls, frames_folder): sys.exit(1) logger.debug("Test input file '%s' does not contain Faceswap header data", test_file) - def _wait_for_threads(self, extractor, loader, valid_meta): + def _wait_for_threads(self, extractor: Aligner, loader: FrameLoader, valid_meta: bool) -> None: """ The :class:`Aligner` and :class:`FramesLoader` are launched in background threads. Wait for them to be initialized prior to proceeding. @@ -139,9 +166,10 @@ def _wait_for_threads(self, extractor, loader, valid_meta): extractor.link_faces(self._detected_faces) if not valid_meta: logger.debug("Saving video meta data to alignments file") - self._detected_faces.save_video_meta_data(**loader.video_meta_data) + self._detected_faces.save_video_meta_data( + **loader.video_meta_data) # type:ignore[arg-type] - def _generate_thumbs(self, input_location, force, single_process): + def _generate_thumbs(self, input_location: str, force: bool, single_process: bool) -> None: """ Check whether thumbnails are stored in the alignments file and if not generate them. Parameters @@ -162,7 +190,7 @@ def _generate_thumbs(self, input_location, force, single_process): thumbs.generate_cache() logger.debug("Generated thumbnails cache") - def _initialize_tkinter(self): + def _initialize_tkinter(self) -> None: """ Initialize a standalone tkinter instance. """ logger.debug("Initializing tkinter") for widget in ("TButton", "TCheckbutton", "TRadiobutton"): @@ -173,15 +201,16 @@ def _initialize_tkinter(self): self.title("Faceswap.py - Visual Alignments") logger.debug("Initialized tkinter") - def _create_containers(self): + def _create_containers(self) -> _Containers: """ Create the paned window containers for various GUI elements Returns ------- - dict: + :class:`_Containers`: The main containers of the manual tool. """ logger.debug("Creating containers") + main = ttk.PanedWindow(self, orient=tk.VERTICAL, name="pw_main") @@ -192,11 +221,13 @@ def _create_containers(self): bottom = ttk.Frame(main, name="frame_bottom") main.add(bottom) - retval = dict(main=main, top=top, bottom=bottom) + + retval = _Containers(main=main, top=top, bottom=bottom) + logger.debug("Created containers: %s", retval) return retval - def _handle_key_press(self, event): + def _handle_key_press(self, event: tk.Event) -> None: """ Keyboard shortcuts Parameters @@ -215,7 +246,7 @@ def _handle_key_press(self, event): modifiers = {0x0001: 'shift', 0x0004: 'ctrl'} - tk_pos = self._globals.tk_frame_index + globs = self._globals bindings = { "z": self._display.navigation.decrement_frame, "x": self._display.navigation.increment_frame, @@ -234,21 +265,23 @@ def _handle_key_press(self, event): "f5": lambda k=event.keysym: self._display.set_action(k), "f9": lambda k=event.keysym: self._faces_frame.set_annotation_display(k), "f10": lambda k=event.keysym: self._faces_frame.set_annotation_display(k), - "c": lambda f=tk_pos.get(), d="prev": self._detected_faces.update.copy(f, d), - "v": lambda f=tk_pos.get(), d="next": self._detected_faces.update.copy(f, d), + "c": lambda f=globs.frame_index, d="prev": self._detected_faces.update.copy(f, d), + "v": lambda f=globs.frame_index, d="next": self._detected_faces.update.copy(f, d), "ctrl_s": self._detected_faces.save, - "r": lambda f=tk_pos.get(): self._detected_faces.revert_to_saved(f)} + "r": lambda f=globs.frame_index: self._detected_faces.revert_to_saved(f)} # Allow keypad keys to be used for numbers press = event.keysym.replace("KP_", "") if event.keysym.startswith("KP_") else event.keysym + assert isinstance(event.state, int) modifier = "_".join(val for key, val in modifiers.items() if event.state & key != 0) key_press = "_".join([modifier, press]) if modifier else press if key_press.lower() in bindings: - logger.trace("key press: %s, action: %s", key_press, bindings[key_press.lower()]) + logger.trace("key press: %s, action: %s", # type:ignore[attr-defined] + key_press, bindings[key_press.lower()]) self.focus_set() bindings[key_press.lower()]() - def _set_initial_layout(self): + def _set_initial_layout(self) -> None: """ Set the favicon and the bottom frame position to correct location to display full frame window. @@ -260,12 +293,13 @@ def _set_initial_layout(self): logger.debug("Setting initial layout") self.tk.call("wm", "iconphoto", - self._w, get_images().icons["favicon"]) # pylint:disable=protected-access + self._w, # type:ignore[attr-defined] # pylint:disable=protected-access + get_images().icons["favicon"]) location = int(self.winfo_screenheight() // 1.5) - self._containers["main"].sashpos(0, location) + self._containers.main.sashpos(0, location) self.update_idletasks() - def process(self): + def process(self) -> None: """ The entry point for the Visual Alignments tool from :mod:`lib.tools.manual.cli`. Launch the tkinter Visual Alignments Window and run main loop. @@ -278,6 +312,8 @@ class _Options(ttk.Frame): # pylint:disable=too-many-ancestors """ Control panel options for currently displayed Editor. This is the right hand panel of the GUI that holds editor specific settings and annotation display settings. + Parameters + ---------- parent: :class:`tkinter.ttk.Frame` The parent frame for the control panel options tk_globals: :class:`~tools.manual.manual.TkGlobals` @@ -285,9 +321,11 @@ class _Options(ttk.Frame): # pylint:disable=too-many-ancestors display_frame: :class:`DisplayFrame` The frame that holds the editors """ - def __init__(self, parent, tk_globals, display_frame): - logger.debug("Initializing %s: (parent: %s, tk_globals: %s, display_frame: %s)", - self.__class__.__name__, parent, tk_globals, display_frame) + def __init__(self, + parent: ttk.Frame, + tk_globals: TkGlobals, + display_frame: DisplayFrame) -> None: + logger.debug(parse_class_init(locals())) super().__init__(parent) self._globals = tk_globals @@ -298,7 +336,7 @@ def __init__(self, parent, tk_globals, display_frame): self.pack(side=tk.RIGHT, fill=tk.Y) logger.debug("Initialized %s", self.__class__.__name__) - def _initialize(self): + def _initialize(self) -> dict[str, ControlPanel]: """ Initialize all of the control panels, then display the default panel. Adds the control panel to :attr:`_control_panels` and sets the traceback to update @@ -311,11 +349,16 @@ def _initialize(self): The Traceback must be set after the panel has first been packed as otherwise it interferes with the loading of the faces pane. + + Returns + ------- + dict[str, :class:`~lib.gui.control_helper.ControlPanel`] + The configured control panels """ self._initialize_face_options() frame = ttk.Frame(self) frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - panels = dict() + panels = {} for name, editor in self._display_frame.editors.items(): logger.debug("Initializing control panel for '%s' editor", name) controls = editor.controls @@ -325,14 +368,14 @@ def _initialize(self): max_columns=1, header_text=controls["header"], blank_nones=False, - label_width=18, + label_width=12, style="CPanel", scrollbar=False) panel.pack_forget() panels[name] = panel return panels - def _initialize_face_options(self): + def _initialize_face_options(self) -> None: """ Set the Face Viewer options panel, beneath the standard control options. """ frame = ttk.Frame(self) frame.pack(side=tk.BOTTOM, fill=tk.X, padx=5, pady=5) @@ -341,13 +384,13 @@ def _initialize_face_options(self): lbl = ttk.Label(size_frame, text="Face Size:") lbl.pack(side=tk.LEFT) cmb = ttk.Combobox(size_frame, - value=["Tiny", "Small", "Medium", "Large", "Extra Large"], + values=["Tiny", "Small", "Medium", "Large", "Extra Large"], state="readonly", - textvariable=self._globals.tk_faces_size) - self._globals.tk_faces_size.set("Medium") + textvariable=self._globals.var_faces_size) + self._globals.var_faces_size.set("Medium") cmb.pack(side=tk.RIGHT, padx=5) - def _set_tk_callbacks(self): + def _set_tk_callbacks(self) -> None: """ Sets the callback to change to the relevant control panel options when the selected editor is changed, and the display update on panel option change.""" self._display_frame.tk_selected_action.trace("w", self._update_options) @@ -361,9 +404,9 @@ def _set_tk_callbacks(self): logger.debug("Adding control update callback: (editor: %s, control: %s)", name, ctl.title) seen_controls.add(ctl) - ctl.tk_var.trace("w", lambda *e: self._globals.tk_update.set(True)) + ctl.tk_var.trace("w", lambda *e: self._globals.var_full_update.set(True)) - def _update_options(self, *args): # pylint:disable=unused-argument + def _update_options(self, *args) -> None: # pylint:disable=unused-argument """ Update the control panel display for the current editor. If the options have not already been set, then adds the control panel to @@ -379,7 +422,7 @@ def _update_options(self, *args): # pylint:disable=unused-argument logger.debug("Displaying control panel for editor: '%s'", editor) self._control_panels[editor].pack(expand=True, fill=tk.BOTH) - def _clear_options_frame(self): + def _clear_options_frame(self) -> None: """ Hides the currently displayed control panel """ for editor, panel in self._control_panels.items(): if panel.winfo_ismapped(): @@ -387,244 +430,6 @@ def _clear_options_frame(self): panel.pack_forget() -class TkGlobals(): - """ Holds Tkinter Variables and other frame information that need to be accessible from all - areas of the GUI. - - Parameters - ---------- - input_location: str - The location of the input folder of frames or video file - """ - def __init__(self, input_location): - logger.debug("Initializing %s: (input_location: %s)", - self.__class__.__name__, input_location) - self._tk_vars = self._get_tk_vars() - - self._is_video = self._check_input(input_location) - self._frame_count = 0 # set by FrameLoader - self._frame_display_dims = (int(round(896 * get_config().scaling_factor)), - int(round(504 * get_config().scaling_factor))) - self._current_frame = dict(image=None, - scale=None, - interpolation=None, - display_dims=None, - filename=None) - logger.debug("Initialized %s", self.__class__.__name__) - - @classmethod - def _get_tk_vars(cls): - """ Create and initialize the tkinter variables. - - Returns - ------- - dict - The variable name as key, the variable as value - """ - retval = dict() - for name in ("frame_index", "transport_index", "face_index", "filter_distance"): - var = tk.IntVar() - var.set(10 if name == "filter_distance" else 0) - retval[name] = var - for name in ("update", "update_active_viewport", "is_zoomed"): - var = tk.BooleanVar() - var.set(False) - retval[name] = var - for name in ("filter_mode", "faces_size"): - retval[name] = tk.StringVar() - return retval - - @property - def current_frame(self): - """ dict: The currently displayed frame in the frame viewer with it's meta information. Key - and Values are as follows: - - **image** (:class:`numpy.ndarry`): The currently displayed frame in original dimensions - - **scale** (`float`): The scaling factor to use to resize the image to the display - window - - **interpolation** (`int`): The opencv interpolator ID to use for resizing the image to - the display window - - **display_dims** (`tuple`): The size of the currently displayed frame, sized for the - display window - - **filename** (`str`): The filename of the currently displayed frame - """ - return self._current_frame - - @property - def frame_count(self): - """ int: The total number of frames for the input location """ - return self._frame_count - - @property - def tk_face_index(self): - """ :class:`tkinter.IntVar`: The variable that holds the face index of the selected face - within the current frame when in zoomed mode. """ - return self._tk_vars["face_index"] - - @property - def tk_update_active_viewport(self): - """ :class:`tkinter.BooleanVar`: Boolean Variable that is traced by the viewport's active - frame to update.. """ - return self._tk_vars["update_active_viewport"] - - @property - def face_index(self): - """ int: The currently displayed face index when in zoomed mode. """ - return self._tk_vars["face_index"].get() - - @property - def frame_display_dims(self): - """ tuple: The (`width`, `height`) of the video display frame in pixels. """ - return self._frame_display_dims - - @property - def frame_index(self): - """ int: The currently displayed frame index. NB This returns -1 if there are no frames - that meet the currently selected filter criteria. """ - return self._tk_vars["frame_index"].get() - - @property - def tk_frame_index(self): - """ :class:`tkinter.IntVar`: The variable holding the current frame index. """ - return self._tk_vars["frame_index"] - - @property - def filter_mode(self): - """ str: The currently selected navigation mode. """ - return self._tk_vars["filter_mode"].get() - - @property - def tk_filter_mode(self): - """ :class:`tkinter.StringVar`: The variable holding the currently selected navigation - filter mode. """ - return self._tk_vars["filter_mode"] - - @property - def tk_filter_distance(self): - """ :class:`tkinter.DoubleVar`: The variable holding the currently selected threshold - distance for misaligned filter mode. """ - return self._tk_vars["filter_distance"] - - @property - def tk_faces_size(self): - """ :class:`tkinter.StringVar`: The variable holding the currently selected Faces Viewer - thumbnail size. """ - return self._tk_vars["faces_size"] - - @property - def is_video(self): - """ bool: ``True`` if the input is a video file, ``False`` if it is a folder of images. """ - return self._is_video - - @property - def tk_is_zoomed(self): - """ :class:`tkinter.BooleanVar`: The variable holding the value indicating whether the - frame viewer is zoomed into a face or zoomed out to the full frame. """ - return self._tk_vars["is_zoomed"] - - @property - def is_zoomed(self): - """ bool: ``True`` if the frame viewer is zoomed into a face, ``False`` if the frame viewer - is displaying a full frame. """ - return self._tk_vars["is_zoomed"].get() - - @property - def tk_transport_index(self): - """ :class:`tkinter.IntVar`: The current index of the display frame's transport slider. """ - return self._tk_vars["transport_index"] - - @property - def tk_update(self): - """ :class:`tkinter.BooleanVar`: The variable holding the trigger that indicates that a - full update needs to occur. """ - return self._tk_vars["update"] - - @staticmethod - def _check_input(frames_location): - """ Check whether the input is a video - - Parameters - ---------- - frames_location: str - The input location for video or images - - Returns - ------- - bool: 'True' if input is a video 'False' if it is a folder. - """ - if os.path.isdir(frames_location): - retval = False - elif os.path.splitext(frames_location)[1].lower() in _video_extensions: - retval = True - else: - logger.error("The input location '%s' is not valid", frames_location) - sys.exit(1) - logger.debug("Input '%s' is_video: %s", frames_location, retval) - return retval - - def set_frame_count(self, count): - """ Set the count of total number of frames to :attr:`frame_count` when the - :class:`FramesLoader` has completed loading. - - Parameters - ---------- - count: int - The number of frames that exist for this session - """ - logger.debug("Setting frame_count to : %s", count) - self._frame_count = count - - def set_current_frame(self, image, filename): - """ Set the frame and meta information for the currently displayed frame. Populates the - attribute :attr:`current_frame` - - Parameters - ---------- - image: :class:`numpy.ndarray` - The image used to display in the Frame Viewer - filename: str - The filename of the current frame - """ - scale = min(self.frame_display_dims[0] / image.shape[1], - self.frame_display_dims[1] / image.shape[0]) - self._current_frame["image"] = image - self._current_frame["filename"] = filename - self._current_frame["scale"] = scale - self._current_frame["interpolation"] = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA - self._current_frame["display_dims"] = (int(round(image.shape[1] * scale)), - int(round(image.shape[0] * scale))) - logger.trace({k: v.shape if isinstance(v, np.ndarray) else v - for k, v in self._current_frame.items()}) - - def set_frame_display_dims(self, width, height): - """ Set the size, in pixels, of the video frame display window and resize the displayed - frame. - - Used on a frame resize callback, sets the :attr:frame_display_dims`. - - Parameters - ---------- - width: int - The width of the frame holding the video canvas in pixels - height: int - The height of the frame holding the video canvas in pixels - """ - self._frame_display_dims = (int(width), int(height)) - image = self._current_frame["image"] - scale = min(self.frame_display_dims[0] / image.shape[1], - self.frame_display_dims[1] / image.shape[0]) - self._current_frame["scale"] = scale - self._current_frame["interpolation"] = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA - self._current_frame["display_dims"] = (int(round(image.shape[1] * scale)), - int(round(image.shape[0] * scale))) - logger.trace({k: v.shape if isinstance(v, np.ndarray) else v - for k, v in self._current_frame.items()}) - - class Aligner(): """ The :class:`Aligner` class sets up an extraction pipeline for each of the current Faceswap Aligners, along with the Landmarks based Maskers. When new landmarks are required, the bounding @@ -639,64 +444,80 @@ class Aligner(): A list of indices correlating to connected GPUs that Tensorflow should not use. Pass ``None`` to not exclude any GPUs. """ - def __init__(self, tk_globals, exclude_gpus): + def __init__(self, tk_globals: TkGlobals, exclude_gpus: list[int] | None) -> None: logger.debug("Initializing: %s (tk_globals: %s, exclude_gpus: %s)", self.__class__.__name__, tk_globals, exclude_gpus) self._globals = tk_globals - self._aligners = {"cv2-dnn": None, "FAN": None, "mask": None} - self._aligner = "FAN" self._exclude_gpus = exclude_gpus - self._detected_faces = None - self._frame_index = None - self._face_index = None + + self._detected_faces: DetectedFaces | None = None + self._frame_index: int | None = None + self._face_index: int | None = None + + self._aligners: dict[TypeManualExtractor, Extractor | None] = {"cv2-dnn": None, + "FAN": None, + "mask": None} + self._aligner: TypeManualExtractor = "FAN" + self._init_thread = self._background_init_aligner() logger.debug("Initialized: %s", self.__class__.__name__) @property - def _in_queue(self): + def _in_queue(self) -> EventQueue: """ :class:`queue.Queue` - The input queue to the extraction pipeline. """ - return self._aligners[self._aligner].input_queue + aligner = self._aligners[self._aligner] + assert aligner is not None + return aligner.input_queue @property - def _feed_face(self): - """ :class:`plugins.extract.pipeline.ExtractMedia`: The current face for feeding into the - aligner, formatted for the pipeline """ + def _feed_face(self) -> ExtractMedia: + """ :class:`~plugins.extract.extract_media.ExtractMedia`: The current face for feeding into + the aligner, formatted for the pipeline """ + assert self._frame_index is not None + assert self._face_index is not None + assert self._detected_faces is not None face = self._detected_faces.current_faces[self._frame_index][self._face_index] return ExtractMedia( - self._globals.current_frame["filename"], - self._globals.current_frame["image"], + self._globals.current_frame.filename, + self._globals.current_frame.image, detected_faces=[face]) @property - def is_initialized(self): + def is_initialized(self) -> bool: """ bool: The Aligners are initialized in a background thread so that other tasks can be performed whilst we wait for initialization. ``True`` is returned if the aligner has completed initialization otherwise ``False``.""" thread_is_alive = self._init_thread.is_alive() if thread_is_alive: - logger.trace("Aligner not yet initialized") + logger.trace("Aligner not yet initialized") # type:ignore[attr-defined] self._init_thread.check_and_raise_error() else: - logger.trace("Aligner initialized") + logger.trace("Aligner initialized") # type:ignore[attr-defined] self._init_thread.join() return not thread_is_alive - def _background_init_aligner(self): + def _background_init_aligner(self) -> MultiThread: """ Launch the aligner in a background thread so we can run other tasks whilst - waiting for initialization """ + waiting for initialization + + Returns + ------- + :class:`lib.multithreading.MultiThread + The background aligner loader thread + """ logger.debug("Launching aligner initialization thread") thread = MultiThread(self._init_aligner, thread_count=1, - name="{}.init_aligner".format(self.__class__.__name__)) + name=f"{self.__class__.__name__}.init_aligner") thread.start() logger.debug("Launched aligner initialization thread") return thread - def _init_aligner(self): + def _init_aligner(self) -> None: """ Initialize Aligner in a background thread, and set it to :attr:`_aligner`. """ logger.debug("Initialize Aligner") # Make sure non-GPU aligner is allocated first - for model in ("mask", "cv2-dnn", "FAN"): + for model in T.get_args(TypeManualExtractor): logger.debug("Initializing aligner: %s", model) plugin = None if model == "mask" else model exclude_gpus = self._exclude_gpus if model == "FAN" else None @@ -705,14 +526,15 @@ def _init_aligner(self): ["components", "extended"], exclude_gpus=exclude_gpus, multiprocess=True, - normalize_method="hist") + normalize_method="hist", + disable_filter=True) if plugin: aligner.set_batchsize("align", 1) # Set the batchsize to 1 aligner.launch() logger.debug("Initialized %s Extractor", model) self._aligners[model] = aligner - def link_faces(self, detected_faces): + def link_faces(self, detected_faces: DetectedFaces) -> None: """ As the Aligner has the potential to take the longest to initialize, it is kicked off as early as possible. At this time :class:`~tools.manual.detected_faces.DetectedFaces` is not yet available. @@ -729,7 +551,8 @@ def link_faces(self, detected_faces): logger.debug("Linking detected_faces: %s", detected_faces) self._detected_faces = detected_faces - def get_landmarks(self, frame_index, face_index, aligner): + def get_landmarks(self, frame_index: int, face_index: int, aligner: TypeManualExtractor + ) -> np.ndarray: """ Feed the detected face into the alignment pipeline and retrieve the landmarks. The face to feed into the aligner is generated from the given frame and face indices. @@ -740,7 +563,7 @@ def get_landmarks(self, frame_index, face_index, aligner): The frame index to extract the aligned face for face_index: int The face index within the current frame to extract the face for - aligner: ["FAN", "cv2-dnn"] + aligner: Literal["FAN", "cv2-dnn"] The aligner to use to extract the face Returns @@ -748,22 +571,37 @@ def get_landmarks(self, frame_index, face_index, aligner): :class:`numpy.ndarray` The 68 point landmark alignments """ - logger.trace("frame_index: %s, face_index: %s, aligner: %s", + logger.trace("frame_index: %s, face_index: %s, aligner: %s", # type:ignore[attr-defined] frame_index, face_index, aligner) self._frame_index = frame_index self._face_index = face_index self._aligner = aligner self._in_queue.put(self._feed_face) - detected_face = next(self._aligners[aligner].detected_faces()).detected_faces[0] - logger.trace("landmarks: %s", detected_face.landmarks_xy) + extractor = self._aligners[aligner] + assert extractor is not None + detected_face = next(extractor.detected_faces()).detected_faces[0] + logger.trace("landmarks: %s", detected_face.landmarks_xy) # type:ignore[attr-defined] return detected_face.landmarks_xy - def get_masks(self, frame_index, face_index): + def _remove_nn_masks(self, detected_face: DetectedFace) -> None: + """ Remove any non-landmarks based masks on a landmark edit + + Parameters + ---------- + detected_face: + The detected face object to remove masks from + """ + del_masks = {m for m in detected_face.mask if m not in ("components", "extended")} + logger.debug("Removing masks after landmark update: %s", del_masks) + for mask in del_masks: + del detected_face.mask[mask] + + def get_masks(self, frame_index: int, face_index: int) -> dict[str, Mask]: """ Feed the aligned face into the mask pipeline and retrieve the updated masks. The face to feed into the aligner is generated from the given frame and face indices. This is to be called when a manual update is done on the landmarks, and new masks need - generating + generating. Parameters ---------- @@ -774,33 +612,38 @@ def get_masks(self, frame_index, face_index): Returns ------- - dict + dict[str, :class:`~lib.align.aligned_mask.Mask`] The updated masks """ - logger.trace("frame_index: %s, face_index: %s", frame_index, face_index) + logger.trace("frame_index: %s, face_index: %s", # type:ignore[attr-defined] + frame_index, face_index) self._frame_index = frame_index self._face_index = face_index self._aligner = "mask" self._in_queue.put(self._feed_face) + assert self._aligners["mask"] is not None detected_face = next(self._aligners["mask"].detected_faces()).detected_faces[0] + self._remove_nn_masks(detected_face) logger.debug("mask: %s", detected_face.mask) return detected_face.mask - def set_normalization_method(self, method): + def set_normalization_method(self, method: T.Literal["none", "clahe", "hist", "mean"]) -> None: """ Change the normalization method for faces fed into the aligner. The normalization method is user adjustable from the GUI. When this method is triggered the method is updated for all aligner pipelines. Parameters ---------- - method: str + method: Literal["none", "clahe", "hist", "mean"] The normalization method to use """ logger.debug("Setting normalization method to: '%s'", method) for plugin, aligner in self._aligners.items(): + assert aligner is not None if plugin == "mask": continue - aligner.set_aligner_normalization_method(method) + logger.debug("Setting to: '%s'", method) + aligner.aligner.set_normalize_method(method) class FrameLoader(): @@ -815,53 +658,93 @@ class FrameLoader(): The path to the input frames video_meta_data: dict The meta data held within the alignments file, if it exists and the input is a video + file_list: list[str] + The list of filenames that exist within the alignments file """ - def __init__(self, tk_globals, frames_location, video_meta_data): - logger.debug("Initializing %s: (tk_globals: %s, frames_location: '%s', " - "video_meta_data: %s)", self.__class__.__name__, tk_globals, frames_location, - video_meta_data) + def __init__(self, + tk_globals: TkGlobals, + frames_location: str, + video_meta_data: dict[str, list[int] | list[float] | None], + file_list: list[str]) -> None: + logger.debug(parse_class_init(locals())) self._globals = tk_globals - self._loader = None + self._loader: SingleFrameLoader | None = None self._current_idx = 0 - self._init_thread = self._background_init_frames(frames_location, video_meta_data) - self._globals.tk_frame_index.trace("w", self._set_frame) + self._init_thread = self._background_init_frames(frames_location, + video_meta_data, + file_list) + self._globals.var_frame_index.trace_add("write", self._set_frame) logger.debug("Initialized %s", self.__class__.__name__) @property - def is_initialized(self): - """ bool: ``True`` if the Frame Loader has completed initialization otherwise - ``False``. """ + def is_initialized(self) -> bool: + """ bool: ``True`` if the Frame Loader has completed initialization. """ thread_is_alive = self._init_thread.is_alive() if thread_is_alive: self._init_thread.check_and_raise_error() else: self._init_thread.join() - # Setting the initial frame cannot be done in the thread, so set when queried from main - self._set_frame(initialize=True) + self._set_frame(initialize=True) # Setting initial frame must be done from main thread return not thread_is_alive @property - def video_meta_data(self): + def video_meta_data(self) -> dict[str, list[int] | list[float] | None]: """ dict: The pts_time and key frames for the loader. """ + assert self._loader is not None return self._loader.video_meta_data - def _background_init_frames(self, frames_location, video_meta_data): + def _background_init_frames(self, + frames_location: str, + video_meta_data: dict[str, list[int] | list[float] | None], + frame_list: list[str]) -> MultiThread: """ Launch the images loader in a background thread so we can run other tasks whilst - waiting for initialization. """ + waiting for initialization. + + Parameters + ---------- + frame_location: str + The location of the source video file/frames folder + video_meta_data: dict + The meta data for video file sources + frame_list: list[str] + The list of frames that exist in the alignments file + """ thread = MultiThread(self._load_images, frames_location, video_meta_data, + frame_list, thread_count=1, - name="{}.init_frames".format(self.__class__.__name__)) + name=f"{self.__class__.__name__}.init_frames") thread.start() return thread - def _load_images(self, frames_location, video_meta_data): - """ Load the images in a background thread. """ - self._loader = SingleFrameLoader(frames_location, video_meta_data=video_meta_data) - self._globals.set_frame_count(self._loader.count) + def _load_images(self, + frames_location: str, + video_meta_data: dict[str, list[int] | list[float] | None], + frame_list: list[str]) -> None: + """ Load the images in a background thread. - def _set_frame(self, *args, initialize=False): # pylint:disable=unused-argument + Parameters + ---------- + frame_location: str + The location of the source video file/frames folder + video_meta_data: dict + The meta data for video file sources + frame_list: list[str] + The list of frames that exist in the alignments file + """ + self._loader = SingleFrameLoader(frames_location, video_meta_data=video_meta_data) + if not self._loader.is_video and len(frame_list) < self._loader.count: + files = [os.path.basename(f) for f in self._loader.file_list] + skip_list = [idx for idx, fname in enumerate(files) if fname not in frame_list] + logger.debug("Adding %s entries to skip list for images not in alignments file", + len(skip_list)) + self._loader.add_skip_list(skip_list) + self._globals.set_frame_count(self._loader.process_count) + + def _set_frame(self, # pylint:disable=unused-argument + *args, + initialize: bool = False) -> None: """ Set the currently loaded frame to :attr:`_current_frame` and trigger a full GUI update. If the loader has not been initialized, or the navigation position is the same as the @@ -877,17 +760,19 @@ def _set_frame(self, *args, initialize=False): # pylint:disable=unused-argument """ position = self._globals.frame_index if not initialize and (position == self._current_idx and not self._globals.is_zoomed): - logger.trace("Update criteria not met. Not updating: (initialize: %s, position: %s, " - "current_idx: %s, is_zoomed: %s)", initialize, position, - self._current_idx, self._globals.is_zoomed) + logger.trace("Update criteria not met. Not updating: " # type:ignore[attr-defined] + "(initialize: %s, position: %s, current_idx: %s, is_zoomed: %s)", + initialize, position, self._current_idx, self._globals.is_zoomed) return if position == -1: filename = "No Frame" frame = np.ones(self._globals.frame_display_dims + (3, ), dtype="uint8") else: + assert self._loader is not None filename, frame = self._loader.image_from_index(position) - logger.trace("filename: %s, frame: %s, position: %s", filename, frame.shape, position) + logger.trace("filename: %s, frame: %s, position: %s", # type:ignore[attr-defined] + filename, frame.shape, position) self._globals.set_current_frame(frame, filename) self._current_idx = position - self._globals.tk_update.set(True) - self._globals.tk_update_active_viewport.set(True) + self._globals.var_full_update.set(True) + self._globals.var_update_active_viewport.set(True) diff --git a/tools/manual/thumbnails.py b/tools/manual/thumbnails.py new file mode 100644 index 0000000000..617d37b4c3 --- /dev/null +++ b/tools/manual/thumbnails.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +""" Thumbnail generator for the manual tool """ +from __future__ import annotations +import logging +import typing as T +import os + +from dataclasses import dataclass +from time import sleep +from threading import Lock + +import imageio +import numpy as np + +from tqdm import tqdm +from lib.align import AlignedFace +from lib.image import SingleFrameLoader, generate_thumbnail +from lib.multithreading import MultiThread + +if T.TYPE_CHECKING: + from .detected_faces import DetectedFaces + +logger = logging.getLogger(__name__) + + +@dataclass +class ProgressBar: + """ Thread-safe progress bar for tracking thumbnail generation progress """ + pbar: tqdm | None = None + lock = Lock() + + +@dataclass +class VideoMeta: + """ Holds meta information about a video file + + Parameters + ---------- + key_frames: list[int] + List of key frame indices for the video + pts_times: list[float] + List of presentation timestams for the video + """ + key_frames: list[int] | None = None + pts_times: list[float] | None = None + + +class ThumbsCreator(): + """ Background loader to generate thumbnails for the alignments file. Generates low resolution + thumbnails in parallel threads for faster processing. + + Parameters + ---------- + detected_faces: :class:`~tool.manual.faces.DetectedFaces` + The :class:`~lib.align.DetectedFace` objects for this video + input_location: str + The location of the input folder of frames or video file + single_process: bool + ``True`` to generated thumbs in a single process otherwise ``False`` + """ + def __init__(self, + detected_faces: DetectedFaces, + input_location: str, + single_process: bool) -> None: + logger.debug("Initializing %s: (detected_faces: %s, input_location: %s, " + "single_process: %s)", self.__class__.__name__, detected_faces, + input_location, single_process) + self._size = 80 + self._pbar = ProgressBar() + self._meta = VideoMeta( + key_frames=T.cast(list[int] | None, + detected_faces.video_meta_data.get("keyframes", None)), + pts_times=T.cast(list[float] | None, + detected_faces.video_meta_data.get("pts_time", None))) + self._location = input_location + self._alignments = detected_faces._alignments + self._frame_faces = detected_faces._frame_faces + + self._is_video = self._meta.pts_times is not None and self._meta.key_frames is not None + + cpu_count = os.cpu_count() + self._num_threads = 1 if cpu_count is None or cpu_count <= 2 else cpu_count - 2 + + if self._is_video and single_process: + self._num_threads = 1 + elif self._is_video and not single_process: + assert self._meta.key_frames is not None + self._num_threads = min(self._num_threads, len(self._meta.key_frames)) + else: + self._num_threads = max(self._num_threads, 32) + self._threads: list[MultiThread] = [] + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def has_thumbs(self) -> bool: + """ bool: ``True`` if the underlying alignments file holds thumbnail images + otherwise ``False``. """ + return self._alignments.thumbnails.has_thumbnails + + def generate_cache(self) -> None: + """ Extract the face thumbnails from a video or folder of images into the + alignments file. """ + self._pbar.pbar = tqdm(desc="Caching Thumbnails", + leave=False, + total=len(self._frame_faces)) + if self._is_video: + self._launch_video() + else: + self._launch_folder() + while True: + self._check_and_raise_error() + if all(not thread.is_alive() for thread in self._threads): + break + sleep(1) + self._join_threads() + self._pbar.pbar.close() + self._alignments.save() + + # << PRIVATE METHODS >> # + def _check_and_raise_error(self) -> None: + """ Monitor the loading threads for errors and raise if any occur. """ + for thread in self._threads: + thread.check_and_raise_error() + + def _join_threads(self) -> None: + """ Join the loading threads """ + logger.debug("Joining face viewer loading threads") + for thread in self._threads: + thread.join() + + def _launch_video(self) -> None: + """ Launch multiple :class:`lib.multithreading.MultiThread` objects to load faces from + a video file. + + Splits the video into segments and passes each of these segments to separate background + threads for some speed up. + """ + key_frames = self._meta.key_frames + pts_times = self._meta.pts_times + assert key_frames is not None and pts_times is not None + key_frame_split = len(key_frames) // self._num_threads + for idx in range(self._num_threads): + is_final = idx == self._num_threads - 1 + start_idx: int = idx * key_frame_split + keyframe_idx = len(key_frames) - 1 if is_final else start_idx + key_frame_split + end_idx = key_frames[keyframe_idx] + start_pts = pts_times[key_frames[start_idx]] + end_pts = False if idx + 1 == self._num_threads else pts_times[end_idx] + starting_index = pts_times.index(start_pts) + if end_pts: + segment_count = len(pts_times[key_frames[start_idx]:end_idx]) + else: + segment_count = len(pts_times[key_frames[start_idx]:]) + logger.debug("thread index: %s, start_idx: %s, end_idx: %s, start_pts: %s, " + "end_pts: %s, starting_index: %s, segment_count: %s", idx, start_idx, + end_idx, start_pts, end_pts, starting_index, segment_count) + thread = MultiThread(self._load_from_video, + start_pts, + end_pts, + starting_index, + segment_count) + thread.start() + self._threads.append(thread) + + def _launch_folder(self) -> None: + """ Launch :class:`lib.multithreading.MultiThread` to retrieve faces from a + folder of images. + + Goes through the file list one at a time, passing each file to a separate background + thread for some speed up. + """ + reader = SingleFrameLoader(self._location) + num_threads = min(reader.count, self._num_threads) + frame_split = reader.count // self._num_threads + logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s", + reader.count, num_threads, frame_split) + for idx in range(num_threads): + is_final = idx == num_threads - 1 + start_idx = idx * frame_split + end_idx = reader.count if is_final else start_idx + frame_split + thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx) + thread.start() + self._threads.append(thread) + + def _load_from_video(self, + pts_start: float, + pts_end: float, + start_index: int, + segment_count: int) -> None: + """ Loads faces from video for the given segment of the source video. + + Each segment of the video is extracted from in a different background thread. + + Parameters + ---------- + pts_start: float + The start time to cut the segment out of the video + pts_end: float + The end time to cut the segment out of the video + start_index: int + The frame index that this segment starts from. Used for calculating the actual frame + index of each frame extracted + segment_count: int + The number of frames that appear in this segment. Used for ending early in case more + frames come out of the segment than should appear (sometimes more frames are picked up + at the end of the segment, so these are discarded) + """ + logger.debug("pts_start: %s, pts_end: %s, start_index: %s, segment_count: %s", + pts_start, pts_end, start_index, segment_count) + reader = self._get_reader(pts_start, pts_end) + idx = 0 + sample_filename, ext = os.path.splitext(next(fname for fname in self._alignments.data)) + vidname = sample_filename[:sample_filename.rfind("_")] + for idx, frame in enumerate(reader): + frame_idx = idx + start_index + filename = f"{vidname}_{frame_idx + 1:06d}{ext}" + self._set_thumbail(filename, frame[..., ::-1], frame_idx) + if idx == segment_count - 1: + # Sometimes extra frames are picked up at the end of a segment, so stop + # processing when segment frame count has been hit. + break + reader.close() + logger.debug("Segment complete: (starting_frame_index: %s, processed_count: %s)", + start_index, idx) + + def _get_reader(self, pts_start: float, pts_end: float): + """ Get an imageio iterator for this thread's segment. + + Parameters + ---------- + pts_start: float + The start time to cut the segment out of the video + pts_end: float + The end time to cut the segment out of the video + + Returns + ------- + :class:`imageio.Reader` + A reader iterator for the requested segment of video + """ + input_params = ["-ss", str(pts_start)] + if pts_end: + input_params.extend(["-to", str(pts_end)]) + logger.debug("pts_start: %s, pts_end: %s, input_params: %s", + pts_start, pts_end, input_params) + return imageio.get_reader(self._location, + "ffmpeg", # type:ignore[arg-type] + input_params=input_params) + + def _load_from_folder(self, + reader: SingleFrameLoader, + start_index: int, + end_index: int) -> None: + """ Loads faces from the given range of frame indices from a folder of images. + + Each frame range is extracted in a different background thread. + + Parameters + ---------- + reader: :class:`lib.image.SingleFrameLoader` + The reader that is used to retrieve the requested frame + start_index: int + The starting frame index for the images to extract faces from + end_index: int + The end frame index for the images to extract faces from + """ + logger.debug("reader: %s, start_index: %s, end_index: %s", + reader, start_index, end_index) + for frame_index in range(start_index, end_index): + filename, frame = reader.image_from_index(frame_index) + self._set_thumbail(filename, frame, frame_index) + logger.debug("Segment complete: (start_index: %s, processed_count: %s)", + start_index, end_index - start_index) + + def _set_thumbail(self, filename: str, frame: np.ndarray, frame_index: int) -> None: + """ Extracts the faces from the frame and adds to alignments file + + Parameters + ---------- + filename: str + The filename of the frame within the alignments file + frame: :class:`numpy.ndarray` + The frame that contains the faces + frame_index: int + The frame index of this frame in the :attr:`_frame_faces` + """ + for face_idx, face in enumerate(self._frame_faces[frame_index]): + aligned = AlignedFace(face.landmarks_xy, + image=frame, + centering="head", + size=96) + face.thumbnail = generate_thumbnail(aligned.face, size=96) + assert face.thumbnail is not None + self._alignments.thumbnails.add_thumbnail(filename, face_idx, face.thumbnail) + with self._pbar.lock: + assert self._pbar.pbar is not None + self._pbar.pbar.update(1) diff --git a/tools/mask/cli.py b/tools/mask/cli.py index ac0d8be662..cc14bb1b9d 100644 --- a/tools/mask/cli.py +++ b/tools/mask/cli.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ +import argparse import gettext from lib.cli.args import FaceSwapArgs @@ -11,8 +12,8 @@ _LANG = gettext.translation("tools.mask.cli", localedir="locales", fallback=True) _ = _LANG.gettext - -_HELPTEXT = _("This command lets you generate masks for existing alignments.") +_HELPTEXT = _("This tool allows you to generate, import, export or preview masks for existing " + "alignments.") class MaskArgs(FaceSwapArgs): @@ -21,130 +22,224 @@ class MaskArgs(FaceSwapArgs): @staticmethod def get_info(): """ Return command information """ - return _("Mask tool\nGenerate masks for existing alignments files.") - - def get_argument_list(self): - argument_list = list() - argument_list.append(dict( - opts=("-a", "--alignments"), - action=FileFullPaths, - type=str, - group=_("data"), - required=True, - filetypes="alignments", - help=_("Full path to the alignments file to add the mask to. NB: if the mask already " - "exists in the alignments file it will be overwritten."))) - argument_list.append(dict( - opts=("-i", "--input"), - action=DirOrFileFullPaths, - type=str, - group=_("data"), - filetypes="video", - required=True, - help=_("Directory containing extracted faces, source frames, or a video file."))) - argument_list.append(dict( - opts=("-it", "--input-type"), - action=Radio, - type=str.lower, - choices=("faces", "frames"), - dest="input_type", - group=_("data"), - default="frames", - help=_("R|Whether the `input` is a folder of faces or a folder frames/video" - "\nL|faces: The input is a folder containing extracted faces." - "\nL|frames: The input is a folder containing frames or is a video"))) - argument_list.append(dict( - opts=("-M", "--masker"), - action=Radio, - type=str.lower, - choices=PluginLoader.get_available_extractors("mask"), - default="extended", - group=_("process"), - help=_("R|Masker to use." - "\nL|bisenet-fp: Relatively lightweight NN based mask that provides more " - "refined control over the area to be masked including full head masking " - "(configurable in mask settings)." - "\nL|components: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks to create a mask." - "\nL|extended: Mask designed to provide facial segmentation based on the " - "positioning of landmark locations. A convex hull is constructed around the " - "exterior of the landmarks and the mask is extended upwards onto the forehead." - "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " - "faces clear of obstructions. Profile faces and obstructions may result in " - "sub-par performance." - "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " - "frontal faces. The mask model has been specifically trained to recognize " - "some facial obstructions (hands and eyeglasses). Profile faces may result in " - "sub-par performance." - "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " - "faces. The mask model has been trained by community members and will need " - "testing for further description. Profile faces may result in sub-par " - "performance."))) - argument_list.append(dict( - opts=("-p", "--processing"), - action=Radio, - type=str.lower, - choices=("all", "missing", "output"), - default="missing", - group=_("process"), - help=_("R|Whether to update all masks in the alignments files, only those faces " - "that do not already have a mask of the given `mask type` or just to output " - "the masks to the `output` location." - "\nL|all: Update the mask for all faces in the alignments file." - "\nL|missing: Create a mask for all faces in the alignments file where a mask " - "does not previously exist." - "\nL|output: Don't update the masks, just output them for review in the given " - "output folder."))) - argument_list.append(dict( - opts=("-o", "--output-folder"), - action=DirFullPaths, - dest="output", - type=str, - group=_("output"), - help=_("Optional output location. If provided, a preview of the masks created will " - "be output in the given folder."))) - argument_list.append(dict( - opts=("-b", "--blur_kernel"), - action=Slider, - type=int, - group=_("output"), - min_max=(0, 9), - default=3, - rounding=1, - help=_("Apply gaussian blur to the mask output. Has the effect of smoothing the " - "edges of the mask giving less of a hard edge. the size is in pixels. This " - "value should be odd, if an even number is passed in then it will be rounded " - "to the next odd number. NB: Only effects the output preview. Set to 0 for " - "off"))) - argument_list.append(dict( - opts=("-t", "--threshold"), - action=Slider, - type=int, - group=_("output"), - min_max=(0, 50), - default=4, - rounding=1, - help=_("Helps reduce 'blotchiness' on some masks by making light shades white " - "and dark shades black. Higher values will impact more of the mask. NB: " - "Only effects the output preview. Set to 0 for off"))) - argument_list.append(dict( - opts=("-ot", "--output-type"), - action=Radio, - type=str.lower, - choices=("combined", "masked", "mask"), - default="combined", - group=_("output"), - help=_("R|How to format the output when processing is set to 'output'." - "\nL|combined: The image contains the face/frame, face mask and masked face." - "\nL|masked: Output the face/frame as rgba image with the face masked." - "\nL|mask: Only output the mask as a single channel image."))) - argument_list.append(dict( - opts=("-f", "--full-frame"), - action="store_true", - default=False, - group=_("output"), - help=_("R|Whether to output the whole frame or only the face box when using " - "output processing. Only has an effect when using frames as input."))) + return _("Mask tool\nGenerate, import, export or preview masks for existing alignments " + "files.") + @staticmethod + def get_argument_list(): + argument_list = [] + argument_list.append({ + "opts": ("-a", "--alignments"), + "action": FileFullPaths, + "type": str, + "group": _("data"), + "required": False, + "filetypes": "alignments", + "help": _( + "Full path to the alignments file that contains the masks if not at the " + "default location. NB: If the input-type is faces and you wish to update the " + "corresponding alignments file, then you must provide a value here as the " + "location cannot be automatically detected.")}) + argument_list.append({ + "opts": ("-i", "--input"), + "action": DirOrFileFullPaths, + "type": str, + "group": _("data"), + "filetypes": "video", + "required": True, + "help": _( + "Directory containing extracted faces, source frames, or a video file.")}) + argument_list.append({ + "opts": ("-I", "--input-type"), + "action": Radio, + "type": str.lower, + "choices": ("faces", "frames"), + "dest": "input_type", + "group": _("data"), + "default": "frames", + "help": _( + "R|Whether the `input` is a folder of faces or a folder frames/video" + "\nL|faces: The input is a folder containing extracted faces." + "\nL|frames: The input is a folder containing frames or is a video")}) + argument_list.append({ + "opts": ("-B", "--batch-mode"), + "action": "store_true", + "dest": "batch_mode", + "default": False, + "group": _("data"), + "help": _( + "R|Run the mask tool on multiple sources. If selected then the other options " + "should be set as follows:" + "\nL|input: A parent folder containing either all of the video files to be " + "processed, or containing sub-folders of frames/faces." + "\nL|output-folder: If provided, then sub-folders will be created within the " + "given location to hold the previews for each input." + "\nL|alignments: Alignments field will be ignored for batch processing. The " + "alignments files must exist at the default location (for frames). For batch " + "processing of masks with 'faces' as the input type, then only the PNG header " + "within the extracted faces will be updated.")}) + argument_list.append({ + "opts": ("-M", "--masker"), + "action": Radio, + "type": str.lower, + "choices": PluginLoader.get_available_extractors("mask"), + "default": "extended", + "group": _("process"), + "help": _( + "R|Masker to use." + "\nL|bisenet-fp: Relatively lightweight NN based mask that provides more " + "refined control over the area to be masked including full head masking " + "(configurable in mask settings)." + "\nL|components: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks to create a mask." + "\nL|custom: A dummy mask that fills the mask area with all 1s or 0s " + "(configurable in settings). This is only required if you intend to manually " + "edit the custom masks yourself in the manual tool. This mask does not use the " + "GPU." + "\nL|extended: Mask designed to provide facial segmentation based on the " + "positioning of landmark locations. A convex hull is constructed around the " + "exterior of the landmarks and the mask is extended upwards onto the forehead." + "\nL|vgg-clear: Mask designed to provide smart segmentation of mostly frontal " + "faces clear of obstructions. Profile faces and obstructions may result in " + "sub-par performance." + "\nL|vgg-obstructed: Mask designed to provide smart segmentation of mostly " + "frontal faces. The mask model has been specifically trained to recognize " + "some facial obstructions (hands and eyeglasses). Profile faces may result in " + "sub-par performance." + "\nL|unet-dfl: Mask designed to provide smart segmentation of mostly frontal " + "faces. The mask model has been trained by community members. Profile faces " + "may result in sub-par performance.")}) + argument_list.append({ + "opts": ("-p", "--processing"), + "action": Radio, + "type": str.lower, + "choices": ("all", "missing", "output", "import"), + "default": "all", + "group": _("process"), + "help": _( + "R|The Mask tool process to perform." + "\nL|all: Update the mask for all faces in the alignments file for the selected " + "'masker'." + "\nL|missing: Create a mask for all faces in the alignments file where a mask " + "does not previously exist for the selected 'masker'." + "\nL|output: Don't update the masks, just output the selected 'masker' for " + "review/editing in external tools to the given output folder." + "\nL|import: Import masks that have been edited outside of faceswap into the " + "alignments file. Note: 'custom' must be the selected 'masker' and the masks must " + "be in the same format as the 'input-type' (frames or faces)")}) + argument_list.append({ + "opts": ("-m", "--mask-path"), + "action": DirFullPaths, + "type": str, + "group": _("import"), + "help": _( + "R|Import only. The path to the folder that contains masks to be imported." + "\nL|How the masks are provided is not important, but they will be stored, " + "internally, as 8-bit grayscale images." + "\nL|If the input are images, then the masks must be named exactly the same as " + "input frames/faces (excluding the file extension)." + "\nL|If the input is a video file, then the filename of the masks is not " + "important but should contain the frame number at the end of the filename (but " + "before the file extension). The frame number can be separated from the rest of " + "the filename by any non-numeric character and can be padded by any number of " + "zeros. The frame number must correspond correctly to the frame number in the " + "original video (starting from frame 1).")}) + argument_list.append({ + "opts": ("-c", "--centering"), + "action": Radio, + "type": str.lower, + "choices": ("face", "head", "legacy"), + "default": "face", + "group": _("import"), + "help": _( + "R|Import/Output only. When importing masks, this is the centering to use. For " + "output this is only used for outputting custom imported masks, and should " + "correspond to the centering used when importing the mask. Note: For any job " + "other than 'import' and 'output' this option is ignored as mask centering is " + "handled internally." + "\nL|face: Centers the mask on the center of the face, adjusting for " + "pitch and yaw. Outside of requirements for full head masking/training, this " + "is likely to be the best choice." + "\nL|head: Centers the mask on the center of the head, adjusting for " + "pitch and yaw. Note: You should only select head centering if you intend to " + "include the full head (including hair) within the mask and are looking to " + "train a full head model." + "\nL|legacy: The 'original' extraction technique. Centers the mask near the " + " of the nose with and crops closely to the face. Can result in the edges of " + "the mask appearing outside of the training area.")}) + argument_list.append({ + "opts": ("-s", "--storage-size"), + "dest": "storage_size", + "action": Slider, + "type": int, + "group": _("import"), + "min_max": (64, 1024), + "default": 128, + "rounding": 64, + "help": _( + "Import only. The size, in pixels to internally store the mask at.\nThe default " + "is 128 which is fine for nearly all usecases. Larger sizes will result in larger " + "alignments files and longer processing.")}) + argument_list.append({ + "opts": ("-o", "--output-folder"), + "action": DirFullPaths, + "dest": "output", + "type": str, + "group": _("output"), + "help": _( + "Optional output location. If provided, a preview of the masks created will " + "be output in the given folder.")}) + argument_list.append({ + "opts": ("-b", "--blur_kernel"), + "action": Slider, + "type": int, + "group": _("output"), + "min_max": (0, 9), + "default": 0, + "rounding": 1, + "help": _( + "Apply gaussian blur to the mask output. Has the effect of smoothing the " + "edges of the mask giving less of a hard edge. the size is in pixels. This " + "value should be odd, if an even number is passed in then it will be rounded " + "to the next odd number. NB: Only effects the output preview. Set to 0 for " + "off")}) + argument_list.append({ + "opts": ("-t", "--threshold"), + "action": Slider, + "type": int, + "group": _("output"), + "min_max": (0, 50), + "default": 0, + "rounding": 1, + "help": _( + "Helps reduce 'blotchiness' on some masks by making light shades white " + "and dark shades black. Higher values will impact more of the mask. NB: " + "Only effects the output preview. Set to 0 for off")}) + argument_list.append({ + "opts": ("-O", "--output-type"), + "action": Radio, + "type": str.lower, + "choices": ("combined", "masked", "mask"), + "default": "combined", + "group": _("output"), + "help": _( + "R|How to format the output when processing is set to 'output'." + "\nL|combined: The image contains the face/frame, face mask and masked face." + "\nL|masked: Output the face/frame as rgba image with the face masked." + "\nL|mask: Only output the mask as a single channel image.")}) + argument_list.append({ + "opts": ("-f", "--full-frame"), + "action": "store_true", + "default": False, + "group": _("output"), + "help": _( + "R|Whether to output the whole frame or only the face box when using " + "output processing. Only has an effect when using frames as input.")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-it", ), + "type": str, + "dest": "depr_input-type_it_I", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/mask/loader.py b/tools/mask/loader.py new file mode 100644 index 0000000000..8f50d81c48 --- /dev/null +++ b/tools/mask/loader.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +""" Handles loading of faces/frames from source locations and pairing with alignments +information """ +from __future__ import annotations + +import logging +import os +import typing as T + +import numpy as np +from tqdm import tqdm + +from lib.align import DetectedFace, update_legacy_png_header +from lib.align.alignments import AlignmentFileDict +from lib.image import FacesLoader, ImagesLoader +from plugins.extract import ExtractMedia + +if T.TYPE_CHECKING: + from lib.align import Alignments + from lib.align.alignments import PNGHeaderDict +logger = logging.getLogger(__name__) + + +class Loader: + """ Loader for reading source data from disk, and yielding the output paired with alignment + information + + Parameters + ---------- + location: str + Full path to the source files location + is_faces: bool + ``True`` if the source is a folder of faceswap extracted faces + """ + def __init__(self, location: str, is_faces: bool) -> None: + logger.debug("Initializing %s (location: %s, is_faces: %s)", + self.__class__.__name__, location, is_faces) + + self._is_faces = is_faces + self._loader = FacesLoader(location) if is_faces else ImagesLoader(location) + self._alignments: Alignments | None = None + self._skip_count = 0 + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def file_list(self) -> list[str]: + """list[str]: Full file list of source files to be loaded """ + return self._loader.file_list + + @property + def is_video(self) -> bool: + """bool: ``True`` if the source is a video file otherwise ``False`` """ + return self._loader.is_video + + @property + def location(self) -> str: + """str: Full path to the source folder/video file location """ + return self._loader.location + + @property + def skip_count(self) -> int: + """int: The number of faces/frames that have been skipped due to no match in alignments + file """ + return self._skip_count + + def add_alignments(self, alignments: Alignments | None) -> None: + """ Add the loaded alignments to :attr:`_alignments` for content matching + + Parameters + ---------- + alignments: :class:`~lib.align.Alignments` | None + The alignments file object or ``None`` if not provided + """ + logger.debug("Adding alignments to loader: %s", alignments) + self._alignments = alignments + + @classmethod + def _get_detected_face(cls, alignment: AlignmentFileDict) -> DetectedFace: + """ Convert an alignment dict item to a detected_face object + + Parameters + ---------- + alignment: :class:`lib.align.alignments.AlignmentFileDict` + The alignment dict for a face + + Returns + ------- + :class:`~lib.align.detected_face.DetectedFace`: + The corresponding detected_face object for the alignment + """ + detected_face = DetectedFace() + detected_face.from_alignment(alignment) + return detected_face + + def _process_face(self, + filename: str, + image: np.ndarray, + metadata: PNGHeaderDict) -> ExtractMedia | None: + """ Process a single face when masking from face images + + Parameters + ---------- + filename: str + the filename currently being processed + image: :class:`numpy.ndarray` + The current face being processed + metadata: dict + The source frame metadata from the PNG header + + Returns + ------- + :class:`plugins.pipeline.ExtractMedia` | None + the extract media object for the processed face or ``None`` if alignment information + could not be found + """ + frame_name = metadata["source"]["source_filename"] + face_index = metadata["source"]["face_index"] + + if self._alignments is None: # mask from PNG header + lookup_index = 0 + alignments = [T.cast(AlignmentFileDict, metadata["alignments"])] + else: # mask from Alignments file + lookup_index = face_index + alignments = self._alignments.get_faces_in_frame(frame_name) + if not alignments or face_index > len(alignments) - 1: + self._skip_count += 1 + logger.warning("Skipping Face not found in alignments file: '%s'", filename) + return None + + alignment = alignments[lookup_index] + detected_face = self._get_detected_face(alignment) + + retval = ExtractMedia(filename, image, detected_faces=[detected_face], is_aligned=True) + retval.add_frame_metadata(metadata["source"]) + return retval + + def _from_faces(self) -> T.Generator[ExtractMedia, None, None]: + """ Load content from pre-aligned faces and pair with corresponding metadata + + Yields + ------ + :class:`plugins.pipeline.ExtractMedia` + the extract media object for the processed face + """ + log_once = False + for filename, image, metadata in tqdm(self._loader.load(), total=self._loader.count): + if not metadata: # Legacy faces. Update the headers + if self._alignments is None: + logger.error("Legacy faces have been discovered, but no alignments file " + "provided. You must provide an alignments file for this face set") + break + + if not log_once: + logger.warning("Legacy faces discovered. These faces will be updated") + log_once = True + + metadata = update_legacy_png_header(filename, self._alignments) + if not metadata: # Face not found + self._skip_count += 1 + logger.warning("Legacy face not found in alignments file. This face has not " + "been updated: '%s'", filename) + continue + + if "source_frame_dims" not in metadata.get("source", {}): + logger.error("The faces need to be re-extracted as at least some of them do not " + "contain information required to correctly generate masks.") + logger.error("You can re-extract the face-set by using the Alignments Tool's " + "Extract job.") + break + + retval = self._process_face(filename, image, metadata) + if retval is None: + continue + + yield retval + + def _from_frames(self) -> T.Generator[ExtractMedia, None, None]: + """ Load content from frames and and pair with corresponding metadata + + Yields + ------ + :class:`plugins.pipeline.ExtractMedia` + the extract media object for the processed face + """ + assert self._alignments is not None + for filename, image in tqdm(self._loader.load(), total=self._loader.count): + frame = os.path.basename(filename) + + if not self._alignments.frame_exists(frame): + self._skip_count += 1 + logger.warning("Skipping frame not in alignments file: '%s'", frame) + continue + + if not self._alignments.frame_has_faces(frame): + logger.debug("Skipping frame with no faces: '%s'", frame) + continue + + faces_in_frame = self._alignments.get_faces_in_frame(frame) + detected_faces = [self._get_detected_face(alignment) for alignment in faces_in_frame] + retval = ExtractMedia(filename, image, detected_faces=detected_faces) + yield retval + + def load(self) -> T.Generator[ExtractMedia, None, None]: + """ Load content from source and pair with corresponding alignment data + + Yields + ------ + :class:`plugins.pipeline.ExtractMedia` + the extract media object for the processed face + """ + if self._is_faces: + iterator = self._from_faces + else: + iterator = self._from_frames + + for media in iterator(): + yield media + + if self._skip_count > 0: + logger.warning("%s face(s) skipped due to not existing in the alignments file", + self._skip_count) diff --git a/tools/mask/mask.py b/tools/mask/mask.py index 6dbb2ac197..25e129bda2 100644 --- a/tools/mask/mask.py +++ b/tools/mask/mask.py @@ -1,467 +1,307 @@ #!/usr/bin/env python3 """ Tool to generate masks and previews of masks for existing alignments file """ +from __future__ import annotations import logging import os import sys -import cv2 -import numpy as np -from tqdm import tqdm +from argparse import Namespace +from multiprocessing import Process -from lib.align import Alignments, AlignedFace, DetectedFace, update_legacy_png_header -from lib.image import FacesLoader, ImagesLoader, ImagesSaver, encode_image +from lib.align import Alignments -from lib.multithreading import MultiThread -from lib.utils import get_folder -from plugins.extract.pipeline import Extractor, ExtractMedia +from lib.utils import handle_deprecated_cliopts, VIDEO_EXTENSIONS +from plugins.extract import ExtractMedia +from .loader import Loader +from .mask_import import Import +from .mask_generate import MaskGenerator +from .mask_output import Output -logger = logging.getLogger(__name__) # pylint:disable=invalid-name +logger = logging.getLogger(__name__) -class Mask(): # pylint:disable=too-few-public-methods + +class Mask: """ This tool is part of the Faceswap Tools suite and should be called from ``python tools.py mask`` command. Faceswap Masks tool. Generate masks from existing alignments files, and output masks for preview. + Wrapper for the mask process to run in either batch mode or single use mode + Parameters ---------- arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` """ - def __init__(self, arguments): + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s: (arguments: %s", self.__class__.__name__, arguments) - self._update_type = arguments.processing - self._input_is_faces = arguments.input_type == "faces" - self._mask_type = arguments.masker - self._output = dict(opts=dict(blur_kernel=arguments.blur_kernel, - threshold=arguments.threshold), - type=arguments.output_type, - full_frame=arguments.full_frame, - suffix=self._get_output_suffix(arguments)) - self._counts = dict(face=0, skip=0, update=0) + if arguments.batch_mode and arguments.processing == "import": + logger.error("Batch mode is not supported for 'import' processing") + sys.exit(0) - self._check_input(arguments.input) - self._saver = self._set_saver(arguments) - loader = FacesLoader if self._input_is_faces else ImagesLoader - self._loader = loader(arguments.input) - self._faces_saver = None + self._args = arguments + self._input_locations = self._get_input_locations() - self._alignments = Alignments(os.path.dirname(arguments.alignments), - filename=os.path.basename(arguments.alignments)) + def _get_input_locations(self) -> list[str]: + """ Obtain the full path to input locations. Will be a list of locations if batch mode is + selected, or containing a single location if batch mode is not selected. - self._extractor = self._get_extractor(arguments.exclude_gpus) - self._extractor_input_thread = self._feed_extractor() + Returns + ------- + list: + The list of input location paths + """ + if not self._args.batch_mode: + return [self._args.input] + + if not os.path.isdir(self._args.input): + logger.error("Batch mode is selected but input '%s' is not a folder", self._args.input) + sys.exit(1) + + retval = [os.path.join(self._args.input, fname) + for fname in os.listdir(self._args.input) + if os.path.isdir(os.path.join(self._args.input, fname)) + or os.path.splitext(fname)[-1].lower() in VIDEO_EXTENSIONS] + logger.info("Batch mode selected. Processing locations: %s", retval) + return retval - logger.debug("Initialized %s", self.__class__.__name__) + def _get_output_location(self, input_location: str) -> str: + """ Obtain the path to an output folder for faces for a given input location. - def _check_input(self, mask_input): - """ Check the input is valid. If it isn't exit with a logged error + A sub-folder within the user supplied output location will be returned based on + the input filename Parameters ---------- - mask_input: str - Path to the input folder/video + input_location: str + The full path to an input video or folder of images """ - if not os.path.exists(mask_input): - logger.error("Location cannot be found: '%s'", mask_input) - sys.exit(0) - if os.path.isfile(mask_input) and self._input_is_faces: - logger.error("Input type 'faces' was selected but input is not a folder: '%s'", - mask_input) - sys.exit(0) - logger.debug("input '%s' is valid", mask_input) + retval = os.path.join(self._args.output, + os.path.splitext(os.path.basename(input_location))[0]) + logger.debug("Returning output: '%s' for input: '%s'", retval, input_location) + return retval - def _set_saver(self, arguments): - """ set the saver in a background thread + @staticmethod + def _run_mask_process(arguments: Namespace) -> None: + """ The mask process to be run in a spawned process. + + In some instances, batch-mode memory leaks. Launching each job in a separate process + prevents this leak. Parameters ---------- arguments: :class:`argparse.Namespace` - The :mod:`argparse` arguments as passed in from :mod:`tools.py` - - Returns - ------- - ``None`` or :class:`lib.image.ImagesSaver`: - If output is requested, returns a :class:`lib.image.ImagesSaver` otherwise - returns ``None`` + The :mod:`argparse` arguments to be used for the given job """ - if not hasattr(arguments, "output") or arguments.output is None or not arguments.output: - if self._update_type == "output": - logger.error("Processing set as 'output' but no output folder provided.") - sys.exit(0) - logger.debug("No output provided. Not creating saver") - return None - output_dir = get_folder(arguments.output, make_folder=True) - logger.info("Saving preview masks to: '%s'", output_dir) - saver = ImagesSaver(output_dir) - logger.debug(saver) - return saver + logger.debug("Starting process: (arguments: %s)", arguments) + mask = _Mask(arguments) + mask.process() + logger.debug("Finished process: (arguments: %s)", arguments) - def _get_extractor(self, exclude_gpus): - """ Obtain a Mask extractor plugin and launch it + def process(self) -> None: + """ The entry point for triggering the Extraction Process. - Parameters - ---------- - exclude_gpus: list or ``None`` - A list of indices correlating to connected GPUs that Tensorflow should not use. Pass - ``None`` to not exclude any GPUs. - - Returns - ------- - :class:`plugins.extract.pipeline.Extractor`: - The launched Extractor + Should only be called from :class:`lib.cli.launcher.ScriptExecutor` """ - if self._update_type == "output": - logger.debug("Update type `output` selected. Not launching extractor") - return None - logger.debug("masker: %s", self._mask_type) - extractor = Extractor(None, None, self._mask_type, - exclude_gpus=exclude_gpus, - image_is_aligned=self._input_is_faces) - extractor.launch() - logger.debug(extractor) - return extractor - - def _feed_extractor(self): - """ Feed the input queue to the Extractor from a faces folder or from source frames in a - background thread - - Returns - ------- - :class:`lib.multithreading.Multithread`: - The thread that is feeding the extractor. - """ - masker_input = getattr(self, - "_input_{}".format("faces" if self._input_is_faces else "frames")) - logger.debug("masker_input: %s", masker_input) + for idx, location in enumerate(self._input_locations): + if self._args.batch_mode: + logger.info("Processing job %s of %s: %s", + idx + 1, len(self._input_locations), location) + arguments = Namespace(**self._args.__dict__) + arguments.input = location + # Due to differences in how alignments are handled for frames/faces, only default + # locations allowed + arguments.alignments = None + if self._args.output: + arguments.output = self._get_output_location(location) + else: + arguments = self._args - args = tuple() if self._update_type == "output" else (self._extractor.input_queue, ) - input_thread = MultiThread(masker_input, *args, thread_count=1) - input_thread.start() - logger.debug(input_thread) - return input_thread + if len(self._input_locations) > 1: + proc = Process(target=self._run_mask_process, args=(arguments, )) + proc.start() + proc.join() + else: + self._run_mask_process(arguments) - def _input_faces(self, *args): - """ Input pre-aligned faces to the Extractor plugin inside a thread - Parameters - ---------- - args: tuple - The arguments that are to be loaded inside this thread. Contains the queue that the - faces should be put to - """ - log_once = False - logger.debug("args: %s", args) - if self._update_type != "output": - queue = args[0] - for filename, image, metadata in tqdm(self._loader.load(), total=self._loader.count): - if not metadata: # Legacy faces. Update the headers - if not log_once: - logger.warning("Legacy faces discovered. These faces will be updated") - log_once = True - metadata = update_legacy_png_header(filename, self._alignments) - if not metadata: # Face not found - self._counts["skip"] += 1 - logger.warning("Legacy face not found in alignments file. This face has not " - "been updated: '%s'", filename) - continue - if "source_frame_dims" not in metadata["source"]: - logger.error("The faces need to be re-extracted as at least some of them do not " - "contain information required to correctly generate masks.") - logger.error("You can re-extract the face-set by using the Alignments Tool's " - "Extract job.") - break - frame_name = metadata["source"]["source_filename"] - face_index = metadata["source"]["face_index"] - alignment = self._alignments.get_faces_in_frame(frame_name) - if not alignment or face_index > len(alignment) - 1: - self._counts["skip"] += 1 - logger.warning("Skipping Face not found in alignments file: '%s'", filename) - continue - alignment = alignment[face_index] - self._counts["face"] += 1 - - if self._check_for_missing(frame_name, face_index, alignment): - continue - - detected_face = self._get_detected_face(alignment) - if self._update_type == "output": - detected_face.image = image - self._save(frame_name, face_index, detected_face) - else: - media = ExtractMedia(filename, image, detected_faces=[detected_face]) - # Hacky overload of ExtractMedia's shape parameter to apply the actual original - # frame dimension - media._image_shape = (*metadata["source"]["source_frame_dims"], 3) - setattr(media, "mask_tool_face_info", metadata["source"]) # TODO formalize - queue.put(media) - self._counts["update"] += 1 - if self._update_type != "output": - queue.put("EOF") - - def _input_frames(self, *args): - """ Input frames to the Extractor plugin inside a thread +class _Mask: + """ This tool is part of the Faceswap Tools suite and should be called from + ``python tools.py mask`` command. - Parameters - ---------- - args: tuple - The arguments that are to be loaded inside this thread. Contains the queue that the - faces should be put to - """ - logger.debug("args: %s", args) - if self._update_type != "output": - queue = args[0] - for filename, image in tqdm(self._loader.load(), total=self._loader.count): - frame = os.path.basename(filename) - if not self._alignments.frame_exists(frame): - self._counts["skip"] += 1 - logger.warning("Skipping frame not in alignments file: '%s'", frame) - continue - if not self._alignments.frame_has_faces(frame): - logger.debug("Skipping frame with no faces: '%s'", frame) - continue - - faces_in_frame = self._alignments.get_faces_in_frame(frame) - self._counts["face"] += len(faces_in_frame) - - # To keep face indexes correct/cover off where only one face in an image is missing a - # mask where there are multiple faces we process all faces again for any frames which - # have missing masks. - if all(self._check_for_missing(frame, idx, alignment) - for idx, alignment in enumerate(faces_in_frame)): - continue - - detected_faces = [self._get_detected_face(alignment) for alignment in faces_in_frame] - if self._update_type == "output": - for idx, detected_face in enumerate(detected_faces): - detected_face.image = image - self._save(frame, idx, detected_face) - else: - self._counts["update"] += len(detected_faces) - queue.put(ExtractMedia(filename, image, detected_faces=detected_faces)) - if self._update_type != "output": - queue.put("EOF") + Faceswap Masks tool. Generate masks from existing alignments files, and output masks + for preview. - def _check_for_missing(self, frame, idx, alignment): - """ Check if the alignment is missing the requested mask_type + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The :mod:`argparse` arguments as passed in from :mod:`tools.py` + """ + def __init__(self, arguments: Namespace) -> None: + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + arguments = handle_deprecated_cliopts(arguments) + self._update_type = arguments.processing + self._input_is_faces = arguments.input_type == "faces" + self._check_input(arguments.input) - Parameters - ---------- - frame: str - The frame name in the alignments file - idx: int - The index of the face for this frame in the alignments file - alignment: dict - The alignment for a face + self._loader = Loader(arguments.input, self._input_is_faces) + self._alignments = self._get_alignments(arguments.alignments, arguments.input) + + if self._loader.is_video and self._alignments is not None: + self._alignments.update_legacy_has_source(os.path.basename(self._loader.location)) + + self._loader.add_alignments(self._alignments) + + self._output = Output(arguments, self._alignments, self._loader.file_list) + + self._import = None + if self._update_type == "import": + self._import = Import(arguments.mask_path, + arguments.centering, + arguments.storage_size, + self._input_is_faces, + self._loader, + self._alignments, + arguments.input, + arguments.masker) + + self._mask_gen: MaskGenerator | None = None + if self._update_type in ("all", "missing"): + self._mask_gen = MaskGenerator(arguments.masker, + self._update_type == "all", + self._input_is_faces, + self._loader, + self._alignments, + arguments.input, + arguments.exclude_gpus) - Returns - ------- - bool: - ``True`` if the update_type is "missing" and the mask does not exist in the alignments - file otherwise ``False`` - """ - retval = (self._update_type == "missing" and - alignment.get("mask", None) is not None and - alignment["mask"].get(self._mask_type, None) is not None) - if retval: - logger.debug("Mask pre-exists for face: '%s' - %s", frame, idx) - return retval + logger.debug("Initialized %s", self.__class__.__name__) - def _get_output_suffix(self, arguments): - """ The filename suffix, based on selected output options. + def _check_input(self, mask_input: str) -> None: + """ Check the input is valid. If it isn't exit with a logged error Parameters ---------- - arguments: :class:`argparse.Namespace` - The command line arguments for the mask tool - - Returns - ------- - str: - The suffix to be appended to the output filename + mask_input: str + Path to the input folder/video """ - sfx = "mask_preview_" - sfx += "face_" if not arguments.full_frame or self._input_is_faces else "frame_" - sfx += "{}.png".format(arguments.output_type) - return sfx + if not os.path.exists(mask_input): + logger.error("Location cannot be found: '%s'", mask_input) + sys.exit(0) + if os.path.isfile(mask_input) and self._input_is_faces: + logger.error("Input type 'faces' was selected but input is not a folder: '%s'", + mask_input) + sys.exit(0) + logger.debug("input '%s' is valid", mask_input) - @staticmethod - def _get_detected_face(alignment): - """ Convert an alignment dict item to a detected_face object + def _get_alignments(self, alignments: str | None, input_location: str) -> Alignments | None: + """ Obtain the alignments from either the given alignments location or the default + location. Parameters ---------- - alignment: dict - The alignment dict for a face + alignments: str | None + Full path to the alignemnts file if provided or ``None`` if not + input_location: str + Full path to the source files to be used by the mask tool Returns ------- - :class:`lib.FacesDetect.detected_face`: - The corresponding detected_face object for the alignment - """ - detected_face = DetectedFace() - detected_face.from_alignment(alignment) - return detected_face - - def process(self): - """ The entry point for the Mask tool from :file:`lib.tools.cli`. Runs the Mask process """ - logger.debug("Starting masker process") - updater = getattr(self, "_update_{}".format("faces" if self._input_is_faces else "frames")) - if self._update_type != "output": - if self._input_is_faces: - self._faces_saver = ImagesSaver(self._loader.location, as_bytes=True) - for extractor_output in self._extractor.detected_faces(): - self._extractor_input_thread.check_and_raise_error() - updater(extractor_output) - if self._counts["update"] != 0: - self._alignments.backup() - self._alignments.save() - if self._input_is_faces: - self._faces_saver.close() - - self._extractor_input_thread.join() - if self._saver is not None: - self._saver.close() - - if self._counts["skip"] != 0: - logger.warning("%s face(s) skipped due to not existing in the alignments file", - self._counts["skip"]) - if self._update_type != "output": - if self._counts["update"] == 0: - logger.warning("No masks were updated of the %s faces seen", self._counts["face"]) - else: - logger.info("Updated masks for %s faces of %s", - self._counts["update"], self._counts["face"]) - logger.debug("Completed masker process") - - def _update_faces(self, extractor_output): - """ Update alignments for the mask if the input type is a faces folder - - If an output location has been indicated, then puts the mask preview to the save queue - - Parameters - ---------- - extractor_output: dict - The output from the :class:`plugins.extract.pipeline.Extractor` object + ``None`` or :class:`~lib.align.alignments.Alignments`: + If output is requested, returns a :class:`~lib.align.alignments.Alignments` otherwise + returns ``None`` """ - for face in extractor_output.detected_faces: - frame_name = extractor_output.mask_tool_face_info["source_filename"] - face_index = extractor_output.mask_tool_face_info["face_index"] - logger.trace("Saving face: (frame: %s, face index: %s)", frame_name, face_index) - - self._alignments.update_face(frame_name, face_index, face.to_alignment()) - metadata = dict(alignments=face.to_png_meta(), - source=extractor_output.mask_tool_face_info) - self._faces_saver.save(extractor_output.filename, - encode_image(extractor_output.image, ".png", metadata=metadata)) + if alignments: + logger.debug("Alignments location provided: %s", alignments) + return Alignments(os.path.dirname(alignments), + filename=os.path.basename(alignments)) + if self._input_is_faces and self._update_type == "output": + logger.debug("No alignments file provided for faces. Using PNG Header for output") + return None + if self._input_is_faces: + logger.warning("Faces input selected without an alignments file. Masks wil only " + "be updated in the faces' PNG Header") + return None - if self._saver is not None: - face.image = extractor_output.image - self._save(frame_name, face_index, face) + folder = input_location + if self._loader.is_video: + logger.debug("Alignments from Video File: '%s'", folder) + folder, filename = os.path.split(folder) + filename = f"{os.path.splitext(filename)[0]}_alignments.fsa" + else: + logger.debug("Alignments from Input Folder: '%s'", folder) + filename = "alignments" - def _update_frames(self, extractor_output): - """ Update alignments for the mask if the input type is a frames folder or video + retval = Alignments(folder, filename=filename) + return retval - If an output location has been indicated, then puts the mask preview to the save queue + def _save_output(self, media: ExtractMedia) -> None: + """ Output masks to disk Parameters ---------- - extractor_output: dict - The output from the :class:`plugins.extract.pipeline.Extractor` object + media: :class:`~plugins.extract.extract_media.ExtractMedia` + The extract media holding the faces to output """ - frame = os.path.basename(extractor_output.filename) - for idx, face in enumerate(extractor_output.detected_faces): - self._alignments.update_face(frame, idx, face.to_alignment()) - if self._saver is not None: - face.image = extractor_output.image - self._save(frame, idx, face) + filename = os.path.basename(media.frame_metadata["source_filename"] + if self._input_is_faces else media.filename) + dims = media.frame_metadata["source_frame_dims"] if self._input_is_faces else None + for idx, face in enumerate(media.detected_faces): + face_idx = media.frame_metadata["face_index"] if self._input_is_faces else idx + face.image = media.image + self._output.save(filename, face_idx, face, frame_dims=dims) + + def _generate_masks(self) -> None: + """ Generate masks from a mask plugin """ + assert self._mask_gen is not None + + logger.info("Generating masks") + + for media in self._mask_gen.process(): + if self._output.should_save: + self._save_output(media) + + def _import_masks(self) -> None: + """ Import masks that have been generated outside of faceswap """ + assert self._import is not None + logger.info("Importing masks") + + for media in self._loader.load(): + self._import.import_mask(media) + if self._output.should_save: + self._save_output(media) + + if self._alignments is not None and self._import.update_count > 0: + self._alignments.backup() + self._alignments.save() + + if self._import.skip_count > 0: + logger.warning("No masks were found for %s item(s), so these have not been imported", + self._import.skip_count) + + logger.info("Imported masks for %s faces of %s", + self._import.update_count, self._import.update_count + self._import.skip_count) + + def _output_masks(self) -> None: + """ Output masks to selected output folder """ + for media in self._loader.load(): + self._save_output(media) + + def process(self) -> None: + """ The entry point for the Mask tool from :file:`lib.tools.cli`. Runs the Mask process """ + logger.debug("Starting masker process") - def _save(self, frame, idx, detected_face): - """ Build the mask preview image and save + if self._update_type in ("all", "missing"): + self._generate_masks() - Parameters - ---------- - frame: str - The frame name in the alignments file - idx: int - The index of the face for this frame in the alignments file - detected_face: `lib.FacesDetect.detected_face` - A detected_face object for a face - """ - if self._mask_type == "bisenet-fp": - mask_types = [f"{self._mask_type}_{area}" for area in ("face", "head")] - else: - mask_types = [self._mask_type] - - if detected_face.mask is None or not any(mask in detected_face.mask - for mask in mask_types): - logger.warning("Mask type '%s' does not exist for frame '%s' index %s. Skipping", - self._mask_type, frame, idx) - return - - for mask_type in mask_types: - if mask_type not in detected_face.mask: - # If extracting bisenet mask, then skip versions which don't exist - continue - filename = os.path.join(self._saver.location, "{}_{}_{}".format( - os.path.splitext(frame)[0], - idx, - f"{mask_type}_{self._output['suffix']}")) - image = self._create_image(detected_face, mask_type) - logger.trace("filename: '%s', image_shape: %s", filename, image.shape) - self._saver.save(filename, image) - - def _create_image(self, detected_face, mask_type): - """ Create a mask preview image for saving out to disk + if self._update_type == "import": + self._import_masks() - Parameters - ---------- - detected_face: `lib.FacesDetect.detected_face` - A detected_face object for a face - mask_type: str - The stored mask type name to create the image for + if self._update_type == "output": + self._output_masks() - Returns - ------- - :class:`numpy.ndarray`: - A preview image depending on the output type in one of the following forms: - - Containing 3 sub images: The original face, the masked face and the mask - - The mask only - - The masked face - """ - mask = detected_face.mask[mask_type] - mask.set_blur_and_threshold(**self._output["opts"]) - if not self._output["full_frame"] or self._input_is_faces: - if self._input_is_faces: - face = AlignedFace(detected_face.landmarks_xy, - image=detected_face.image, - centering=mask.stored_centering, - size=detected_face.image.shape[0], - is_aligned=True).face - else: - centering = "legacy" if self._alignments.version == 1.0 else mask.stored_centering - detected_face.load_aligned(detected_face.image, centering=centering, force=True) - face = detected_face.aligned.face - mask = cv2.resize(detected_face.mask[mask_type].mask, - (face.shape[1], face.shape[0]), - interpolation=cv2.INTER_CUBIC)[..., None] - else: - face = np.array(detected_face.image) # cv2 fails if this comes as imageio.core.Array - mask = mask.get_full_frame_mask(face.shape[1], face.shape[0]) - mask = np.expand_dims(mask, -1) - - height, width = face.shape[:2] - if self._output["type"] == "combined": - masked = (face.astype("float32") * mask.astype("float32") / 255.).astype("uint8") - mask = np.tile(mask, 3) - for img in (face, masked, mask): - cv2.rectangle(img, (0, 0), (width - 1, height - 1), (255, 255, 255), 1) - out_image = np.concatenate((face, masked, mask), axis=1) - elif self._output["type"] == "mask": - out_image = mask - elif self._output["type"] == "masked": - out_image = np.concatenate([face, mask], axis=-1) - return out_image + self._output.close() + logger.debug("Completed masker process") diff --git a/tools/mask/mask_generate.py b/tools/mask/mask_generate.py new file mode 100644 index 0000000000..eb3cd6af44 --- /dev/null +++ b/tools/mask/mask_generate.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +""" Handles the generation of masks from faceswap for upating into an alignments file """ +from __future__ import annotations + +import logging +import os +import typing as T + +from lib.image import encode_image, ImagesSaver +from lib.multithreading import MultiThread +from plugins.extract import Extractor + +if T.TYPE_CHECKING: + from lib.align import Alignments, DetectedFace + from lib.align.alignments import PNGHeaderDict + from lib.queue_manager import EventQueue + from plugins.extract import ExtractMedia + from .loader import Loader + + +logger = logging.getLogger(__name__) + + +class MaskGenerator: + """ Uses faceswap's extract pipeline to generate masks and update them into the alignments file + and/or extracted face PNG Headers + + Parameters + ---------- + mask_type: str + The mask type to generate + update_all: bool + ``True`` to update all faces, ``False`` to only update faces missing masks + input_is_faces: bool + ``True`` if the input are faceswap extracted faces otherwise ``False`` + exclude_gpus: list[int] + List of any GPU IDs that should be excluded + loader: :class:`tools.mask.loader.Loader` + The loader for loading source images/video from disk + """ + def __init__(self, + mask_type: str, + update_all: bool, + input_is_faces: bool, + loader: Loader, + alignments: Alignments | None, + input_location: str, + exclude_gpus: list[int]) -> None: + logger.debug("Initializing %s (mask_type: %s, update_all: %s, input_is_faces: %s, " + "loader: %s, alignments: %s, input_location: %s, exclude_gpus: %s)", + self.__class__.__name__, mask_type, update_all, input_is_faces, loader, + alignments, input_location, exclude_gpus) + + self._update_all = update_all + self._is_faces = input_is_faces + self._alignments = alignments + + self._extractor = self._get_extractor(mask_type, exclude_gpus) + self._mask_type = self._set_correct_mask_type(mask_type) + self._input_thread = self._set_loader_thread(loader) + self._saver = ImagesSaver(input_location, as_bytes=True) if input_is_faces else None + + self._counts: dict[T.Literal["face", "update"], int] = {"face": 0, "update": 0} + + logger.debug("Initialized %s", self.__class__.__name__) + + def _get_extractor(self, mask_type, exclude_gpus: list[int]) -> Extractor: + """ Obtain a Mask extractor plugin and launch it + + Parameters + ---------- + mask_type: str + The mask type to generate + exclude_gpus: list or ``None`` + A list of indices correlating to connected GPUs that Tensorflow should not use. Pass + ``None`` to not exclude any GPUs. + + Returns + ------- + :class:`plugins.extract.pipeline.Extractor`: + The launched Extractor + """ + logger.debug("masker: %s", mask_type) + extractor = Extractor(None, None, mask_type, exclude_gpus=exclude_gpus) + extractor.launch() + logger.debug(extractor) + return extractor + + def _set_correct_mask_type(self, mask_type: str) -> str: + """ Some masks have multiple variants that they can be saved as depending on config options + + Parameters + ---------- + mask_type: str + The mask type to generate + + Returns + ------- + str + The actual mask variant to update + """ + if mask_type != "bisenet-fp": + return mask_type + + # Hacky look up into masker to get the type of mask + mask_plugin = self._extractor._mask[0] # pylint:disable=protected-access + assert mask_plugin is not None + mtype = "head" if mask_plugin.config.get("include_hair", False) else "face" + new_type = f"{mask_type}_{mtype}" + logger.debug("Updating '%s' to '%s'", mask_type, new_type) + return new_type + + def _needs_update(self, frame: str, idx: int, face: DetectedFace) -> bool: + """ Check if the mask for the current alignment needs updating for the requested mask_type + + Parameters + ---------- + frame: str + The frame name in the alignments file + idx: int + The index of the face for this frame in the alignments file + face: :class:`~lib.align.DetectedFace` + The dected face object to check + + Returns + ------- + bool: + ``True`` if the mask needs to be updated otherwise ``False`` + """ + if self._update_all: + return True + + retval = not face.mask or face.mask.get(self._mask_type, None) is None + + logger.trace("Needs updating: %s, '%s' - %s", # type:ignore[attr-defined] + retval, frame, idx) + return retval + + def _feed_extractor(self, loader: Loader, extract_queue: EventQueue) -> None: + """ Process to feed the extractor from inside a thread + + Parameters + ---------- + loader: class:`tools.mask.loader.Loader` + The loader for loading source images/video from disk + extract_queue: :class:`lib.queue_manager.EventQueue` + The input queue to the extraction pipeline + """ + for media in loader.load(): + self._counts["face"] += len(media.detected_faces) + + if self._is_faces: + assert len(media.detected_faces) == 1 + needs_update = self._needs_update(media.frame_metadata["source_filename"], + media.frame_metadata["face_index"], + media.detected_faces[0]) + else: + # To keep face indexes correct/cover off where only one face in an image is missing + # a mask where there are multiple faces we process all faces again for any frames + # which have missing masks. + needs_update = any(self._needs_update(media.filename, idx, detected_face) + for idx, detected_face in enumerate(media.detected_faces)) + + if not needs_update: + logger.trace("No masks need updating in '%s'", # type:ignore[attr-defined] + media.filename) + continue + + logger.trace("Passing to extractor: '%s'", media.filename) # type:ignore[attr-defined] + extract_queue.put(media) + + logger.debug("Terminating loader thread") + extract_queue.put("EOF") + + def _set_loader_thread(self, loader: Loader) -> MultiThread: + """ Set the iterator to load ExtractMedia objects into the mask extraction pipeline + so we can just iterate through the output masks + + Parameters + ---------- + loader: class:`tools.mask.loader.Loader` + The loader for loading source images/video from disk + """ + in_queue = self._extractor.input_queue + logger.debug("Starting load thread: (loader: %s, queue: %s)", loader, in_queue) + in_thread = MultiThread(self._feed_extractor, loader, in_queue, thread_count=1) + in_thread.start() + logger.debug("Started load thread: %s", in_thread) + return in_thread + + def _update_from_face(self, media: ExtractMedia) -> None: + """ Update the alignments file and/or the extracted face + + Parameters + ---------- + media: :class:`~lib.extract.pipeline.ExtractMedia` + The ExtractMedia object with updated masks + """ + assert len(media.detected_faces) == 1 + assert self._saver is not None + + fname = media.frame_metadata["source_filename"] + idx = media.frame_metadata["face_index"] + face = media.detected_faces[0] + + if self._alignments is not None: + logger.trace("Updating face %s in frame '%s'", idx, fname) # type:ignore[attr-defined] + self._alignments.update_face(fname, idx, face.to_alignment()) + + logger.trace("Updating extracted face: '%s'", media.filename) # type:ignore[attr-defined] + meta: PNGHeaderDict = {"alignments": face.to_png_meta(), "source": media.frame_metadata} + self._saver.save(media.filename, encode_image(media.image, ".png", metadata=meta)) + + def _update_from_frame(self, media: ExtractMedia) -> None: + """ Update the alignments file + + Parameters + ---------- + media: :class:`~lib.extract.pipeline.ExtractMedia` + The ExtractMedia object with updated masks + """ + assert self._alignments is not None + fname = os.path.basename(media.filename) + logger.trace("Updating %s faces in frame '%s'", # type:ignore[attr-defined] + len(media.detected_faces), fname) + for idx, face in enumerate(media.detected_faces): + self._alignments.update_face(fname, idx, face.to_alignment()) + + def _finalize(self) -> None: + """ Close thread and save alignments on completion """ + logger.debug("Finalizing MaskGenerator") + self._input_thread.join() + + if self._counts["update"] > 0 and self._alignments is not None: + logger.debug("Saving alignments") + self._alignments.backup() + self._alignments.save() + + if self._saver is not None: + logger.debug("Closing face saver") + self._saver.close() + + if self._counts["update"] == 0: + logger.warning("No masks were updated of the %s faces seen", self._counts["face"]) + else: + logger.info("Updated masks for %s faces of %s", + self._counts["update"], self._counts["face"]) + + def process(self) -> T.Generator[ExtractMedia, None, None]: + """ Process the output from the extractor pipeline + + Yields + ------ + :class:`~lib.extract.pipeline.ExtractMedia` + The ExtractMedia object with updated masks + """ + for media in self._extractor.detected_faces(): + self._input_thread.check_and_raise_error() + self._counts["update"] += len(media.detected_faces) + + if self._is_faces: + self._update_from_face(media) + else: + self._update_from_frame(media) + + yield media + + self._finalize() + logger.debug("Completed MaskGenerator process") diff --git a/tools/mask/mask_import.py b/tools/mask/mask_import.py new file mode 100644 index 0000000000..4192ce03ab --- /dev/null +++ b/tools/mask/mask_import.py @@ -0,0 +1,406 @@ +#!/usr/bin/env python3 +""" Import mask processing for faceswap's mask tool """ +from __future__ import annotations + +import logging +import os +import re +import sys +import typing as T + +import cv2 +from tqdm import tqdm + +from lib.align import AlignedFace +from lib.image import encode_image, ImagesSaver +from lib.utils import get_image_paths + +if T.TYPE_CHECKING: + import numpy as np + from .loader import Loader + from plugins.extract import ExtractMedia + from lib.align import Alignments, DetectedFace + from lib.align.alignments import PNGHeaderDict + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +class Import: + """ Import masks from disk into an Alignments file + + Parameters + ---------- + import_path: str + The path to the input images + centering: Literal["face", "head", "legacy"] + The centering to store the mask at + storage_size: int + The size to store the mask at + input_is_faces: bool + ``True`` if the input is aligned faces otherwise ``False`` + loader: :class:`~tools.mask.loader.Loader` + The source file loader object + alignments: :class:`~lib.align.alignments.Alignments` | None + The alignments file object for the faces, if provided + mask_type: str + The mask type to update to + """ + def __init__(self, + import_path: str, + centering: CenteringType, + storage_size: int, + input_is_faces: bool, + loader: Loader, + alignments: Alignments | None, + input_location: str, + mask_type: str) -> None: + logger.debug("Initializing %s (import_path: %s, centering: %s, storage_size: %s, " + "input_is_faces: %s, loader: %s, alignments: %s, input_location: %s, " + "mask_type: %s)", self.__class__.__name__, import_path, centering, + storage_size, input_is_faces, loader, alignments, input_location, mask_type) + + self._validate_mask_type(mask_type) + + self._centering = centering + self._size = storage_size + self._is_faces = input_is_faces + self._alignments = alignments + self._re_frame_num = re.compile(r"\d+$") + self._mapping = self._generate_mapping(import_path, loader) + + self._saver = ImagesSaver(input_location, as_bytes=True) if input_is_faces else None + self._counts: dict[T.Literal["skip", "update"], int] = {"skip": 0, "update": 0} + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def skip_count(self) -> int: + """ int: Number of masks that were skipped as they do not exist for given faces """ + return self._counts["skip"] + + @property + def update_count(self) -> int: + """ int: Number of masks that were skipped as they do not exist for given faces """ + return self._counts["update"] + + @classmethod + def _validate_mask_type(cls, mask_type: str) -> None: + """ Validate that the mask type is 'custom' to ensure user does not accidentally overwrite + existing masks they may have editted + + Parameters + ---------- + mask_type: str + The mask type that has been selected + """ + if mask_type == "custom": + return + + logger.error("Masker 'custom' must be selected for importing masks") + sys.exit(1) + + @classmethod + def _get_file_list(cls, path: str) -> list[str]: + """ Check the nask folder exists and obtain the list of images + + Parameters + ---------- + path: str + Full path to the location of mask images to be imported + + Returns + ------- + list[str] + list of full paths to all of the images in the mask folder + """ + if not os.path.isdir(path): + logger.error("Mask path: '%s' is not a folder", path) + sys.exit(1) + paths = get_image_paths(path) + if not paths: + logger.error("Mask path '%s' contains no images", path) + sys.exit(1) + return paths + + def _warn_extra_masks(self, file_list: list[str]) -> None: + """ Generate a warning for each mask that exists that does not correspond to a match in the + source input + + Parameters + ---------- + file_list: list[str] + List of mask files that could not be mapped to a source image + """ + if not file_list: + logger.debug("All masks exist in the source data") + return + + for fname in file_list: + logger.warning("Extra mask file found: '%s'", os.path.basename(fname)) + + logger.warning("%s mask file(s) do not exist in the source data so will not be imported " + "(see above)", len(file_list)) + + def _file_list_to_frame_number(self, file_list: list[str]) -> dict[int, str]: + """ Extract frame numbers from mask file names and return as a dictionary + + Parameters + ---------- + file_list: list[str] + List of full paths to masks to extract frame number from + + Returns + ------- + dict[int, str] + Dictionary of frame numbers to filenames + """ + retval: dict[int, str] = {} + for filename in file_list: + frame_num = self._re_frame_num.findall(os.path.splitext(os.path.basename(filename))[0]) + + if not frame_num or len(frame_num) > 1: + logger.error("Could not detect frame number from mask file '%s'. " + "Check your filenames", os.path.basename(filename)) + sys.exit(1) + + fnum = int(frame_num[0]) + + if fnum in retval: + logger.error("Frame number %s for mask file '%s' already exists from file: '%s'. " + "Check your filenames", + fnum, os.path.basename(filename), os.path.basename(retval[fnum])) + sys.exit(1) + + retval[fnum] = filename + + logger.debug("Files: %s, frame_numbers: %s", len(file_list), len(retval)) + + return retval + + def _map_video(self, file_list: list[str], source_files: list[str]) -> dict[str, str]: + """ Generate the mapping between the source data and the masks to be imported for + video sources + + Parameters + ---------- + file_list: list[str] + List of full paths to masks to be imported + source_files: list[str] + list of filenames withing the source file + + Returns + ------- + dict[str, str] + Source filenames mapped to full path location of mask to be imported + """ + retval = {} + unmapped = [] + mask_frames = self._file_list_to_frame_number(file_list) + for filename in tqdm(source_files, desc="Mapping masks to input", leave=False): + src_idx = int(os.path.splitext(filename)[0].rsplit("_", maxsplit=1)[-1]) + mapped = mask_frames.pop(src_idx, "") + if not mapped: + unmapped.append(filename) + continue + retval[os.path.basename(filename)] = mapped + + if len(unmapped) == len(source_files): + logger.error("No masks map between the source data and the mask folder. " + "Check your filenames") + sys.exit(1) + + self._warn_extra_masks(list(mask_frames.values())) + logger.debug("Source: %s, Mask: %s, Mapped: %s", + len(source_files), len(file_list), len(retval)) + return retval + + def _map_images(self, file_list: list[str], source_files: list[str]) -> dict[str, str]: + """ Generate the mapping between the source data and the masks to be imported for + folder of image sources + + Parameters + ---------- + file_list: list[str] + List of full paths to masks to be imported + source_files: list[str] + list of filenames withing the source file + + Returns + ------- + dict[str, str] + Source filenames mapped to full path location of mask to be imported + """ + mask_count = len(file_list) + retval = {} + unmapped = [] + for filename in tqdm(source_files, desc="Mapping masks to input", leave=False): + fname = os.path.splitext(os.path.basename(filename))[0] + mapped = next((f for f in file_list + if os.path.splitext(os.path.basename(f))[0] == fname), "") + if not mapped: + unmapped.append(filename) + continue + retval[os.path.basename(filename)] = file_list.pop(file_list.index(mapped)) + + if len(unmapped) == len(source_files): + logger.error("No masks map between the source data and the mask folder. " + "Check your filenames") + sys.exit(1) + + self._warn_extra_masks(file_list) + + logger.debug("Source: %s, Mask: %s, Mapped: %s", + len(source_files), mask_count, len(retval)) + return retval + + def _generate_mapping(self, import_path: str, loader: Loader) -> dict[str, str]: + """ Generate the mapping between the source data and the masks to be imported + + Parameters + ---------- + import_path: str + The path to the input images + loader: :class:`~tools.mask.loader.Loader` + The source file loader object + + Returns + ------- + dict[str, str] + Source filenames mapped to full path location of mask to be imported + """ + file_list = self._get_file_list(import_path) + if loader.is_video: + retval = self._map_video(file_list, loader.file_list) + else: + retval = self._map_images(file_list, loader.file_list) + + return retval + + def _store_mask(self, face: DetectedFace, mask: np.ndarray) -> None: + """ Store the mask to the given DetectedFace object + + Parameters + ---------- + face: :class:`~lib.align.detected_face.DetectedFace` + The detected face object to store the mask to + mask: :class:`numpy.ndarray` + The mask to store + """ + aligned = AlignedFace(face.landmarks_xy, + mask[..., None] if self._is_faces else mask, + centering=self._centering, + size=self._size, + is_aligned=self._is_faces, + dtype="float32") + assert aligned.face is not None + face.add_mask(f"custom_{self._centering}", + aligned.face / 255., + aligned.adjusted_matrix, + aligned.interpolators[1], + storage_size=self._size, + storage_centering=self._centering) + + def _store_mask_face(self, media: ExtractMedia, mask: np.ndarray) -> None: + """ Store the mask when the input is aligned faceswap faces + + Parameters + ---------- + media: :class:`~plugins.extract.extract_media.ExtractMedia` + The extract media object containing the face(s) to import the mask for + + mask: :class:`numpy.ndarray` + The mask loaded from disk + """ + assert self._saver is not None + assert len(media.detected_faces) == 1 + + logger.trace("Adding mask for '%s'", media.filename) # type:ignore[attr-defined] + + face = media.detected_faces[0] + self._store_mask(face, mask) + + if self._alignments is not None: + idx = media.frame_metadata["source_filename"] + fname = media.frame_metadata["face_index"] + logger.trace("Updating face %s in frame '%s'", idx, fname) # type:ignore[attr-defined] + self._alignments.update_face(idx, + fname, + face.to_alignment()) + + logger.trace("Updating extracted face: '%s'", media.filename) # type:ignore[attr-defined] + meta: PNGHeaderDict = {"alignments": face.to_png_meta(), "source": media.frame_metadata} + self._saver.save(media.filename, encode_image(media.image, ".png", metadata=meta)) + + @classmethod + def _resize_mask(cls, mask: np.ndarray, dims: tuple[int, int]) -> np.ndarray: + """ Resize a mask to the given dimensions + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask to resize + dims: tuple[int, int] + The (height, width) target size + + Returns + ------- + :class:`numpy.ndarray` + The resized mask, or the original mask if no resizing required + """ + if mask.shape[:2] == dims: + return mask + logger.trace("Resizing mask from %s to %s", mask.shape, dims) # type:ignore[attr-defined] + interp = cv2.INTER_AREA if mask.shape[0] > dims[0] else cv2.INTER_CUBIC + + mask = cv2.resize(mask, tuple(reversed(dims)), interpolation=interp) + return mask + + def _store_mask_frame(self, media: ExtractMedia, mask: np.ndarray) -> None: + """ Store the mask when the input is frames + + Parameters + ---------- + media: :class:`~plugins.extract.extract_media.ExtractMedia` + The extract media object containing the face(s) to import the mask for + + mask: :class:`numpy.ndarray` + The mask loaded from disk + """ + assert self._alignments is not None + logger.trace("Adding %s mask(s) for '%s'", # type:ignore[attr-defined] + len(media.detected_faces), media.filename) + + mask = self._resize_mask(mask, media.image_size) + + for idx, face in enumerate(media.detected_faces): + self._store_mask(face, mask) + self._alignments.update_face(os.path.basename(media.filename), + idx, + face.to_alignment()) + + def import_mask(self, media: ExtractMedia) -> None: + """ Import the mask for the given Extract Media object + + Parameters + ---------- + media: :class:`~plugins.extract.extract_media.ExtractMedia` + The extract media object containing the face(s) to import the mask for + """ + mask_file = self._mapping.get(os.path.basename(media.filename)) + if not mask_file: + self._counts["skip"] += 1 + logger.warning("No mask file found for: '%s'", os.path.basename(media.filename)) + return + + mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE) + + logger.trace("Loaded mask for frame '%s': %s", # type:ignore[attr-defined] + os.path.basename(mask_file), mask.shape) + + self._counts["update"] += len(media.detected_faces) + + if self._is_faces: + self._store_mask_face(media, mask) + else: + self._store_mask_frame(media, mask) diff --git a/tools/mask/mask_output.py b/tools/mask/mask_output.py new file mode 100644 index 0000000000..79ffb809e3 --- /dev/null +++ b/tools/mask/mask_output.py @@ -0,0 +1,519 @@ +#!/usr/bin/env python3 +""" Output processing for faceswap's mask tool """ +from __future__ import annotations + +import logging +import os +import sys +import typing as T +from argparse import Namespace + +import cv2 +import numpy as np +from tqdm import tqdm + +from lib.align import AlignedFace +from lib.align.alignments import AlignmentDict + +from lib.image import ImagesSaver, read_image_meta_batch +from lib.utils import get_folder +from scripts.fsmedia import Alignments as ExtractAlignments + +if T.TYPE_CHECKING: + from lib.align import Alignments, DetectedFace + from lib.align.aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +class Output: + """ Handles outputting of masks for preview/editting to disk + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments that the mask tool was called with + alignments: :class:~`lib.align.alignments.Alignments` | None + The alignments file object (or ``None`` if not provided and input is faces) + file_list: list[str] + Full file list for the loader. Used for extracting alignments from faces + """ + def __init__(self, arguments: Namespace, + alignments: Alignments | None, + file_list: list[str]) -> None: + logger.debug("Initializing %s (arguments: %s, alignments: %s, file_list: %s)", + self.__class__.__name__, arguments, alignments, len(file_list)) + + self._blur_kernel: int = arguments.blur_kernel + self._threshold: int = arguments.threshold + self._type: T.Literal["combined", "masked", "mask"] = arguments.output_type + self._full_frame: bool = arguments.full_frame + self._mask_type = arguments.masker + self._centering: CenteringType = arguments.centering + + self._input_is_faces = arguments.input_type == "faces" + self._saver = self._set_saver(arguments.output, arguments.processing) + self._alignments = self._get_alignments(alignments, file_list) + + self._full_frame_cache: dict[str, list[tuple[int, DetectedFace]]] = {} + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def should_save(self) -> bool: + """bool: ``True`` if mask images should be output otherwise ``False`` """ + return self._saver is not None + + def _get_subfolder(self, output: str) -> str: + """ Obtain a subfolder within the output folder to save the output based on selected + output options. + + Parameters + ---------- + output: str + Full path to the root output folder + + Returns + ------- + str: + The full path to where masks should be saved + """ + out_type = "frame" if self._full_frame else "face" + retval = os.path.join(output, + f"{self._mask_type}_{out_type}_{self._type}") + logger.info("Saving masks to '%s'", retval) + return retval + + def _set_saver(self, output: str | None, processing: str) -> ImagesSaver | None: + """ set the saver in a background thread + + Parameters + ---------- + output: str + Full path to the root output folder if provided + processing: str + The processing that has been selected + + Returns + ------- + ``None`` or :class:`lib.image.ImagesSaver`: + If output is requested, returns a :class:`lib.image.ImagesSaver` otherwise + returns ``None`` + """ + if output is None or not output: + if processing == "output": + logger.error("Processing set as 'output' but no output folder provided.") + sys.exit(0) + logger.debug("No output provided. Not creating saver") + return None + output_dir = get_folder(self._get_subfolder(output), make_folder=True) + retval = ImagesSaver(output_dir) + logger.debug(retval) + return retval + + def _get_alignments(self, + alignments: Alignments | None, + file_list: list[str]) -> Alignments | None: + """ Obtain the alignments file. If input is faces and full frame output is requested then + the file needs to be generated from the input faces, if not provided + + Parameters + ---------- + alignments: :class:~`lib.align.alignments.Alignments` | None + The alignments file object (or ``None`` if not provided and input is faces) + file_list: list[str] + Full paths to ihe mask tool input files + + Returns + ------- + :class:~`lib.align.alignments.Alignments` | None + The alignments file if provided and/or is required otherwise ``None`` + """ + if alignments is not None or not self._full_frame: + return alignments + logger.debug("Generating alignments from faces") + + data = T.cast(dict[str, AlignmentDict], {}) + for _, meta in tqdm(read_image_meta_batch(file_list), + desc="Reading alignments from faces", + total=len(file_list), + leave=False): + fname = meta["itxt"]["source"]["source_filename"] + aln = meta["itxt"]["alignments"] + data.setdefault(fname, {}).setdefault("faces", # type:ignore[typeddict-item] + []).append(aln) + + dummy_args = Namespace(alignments_path="/dummy/alignments.fsa") + retval = ExtractAlignments(dummy_args, is_extract=True) + retval.update_from_dict(data) + return retval + + def _get_background_frame(self, detected_faces: list[DetectedFace], frame_dims: tuple[int, int] + ) -> np.ndarray: + """ Obtain the background image when final output is in full frame format. There will only + ever be one background, even when there are multiple faces + + The output image will depend on the requested output type and whether the input is faces + or frames + + Parameters + ---------- + detected_faces: list[:class:`~lib.align.detected_face.DetectedFace`] + Detected face objects for the output image + frame_dims: tuple[int, int] + The size of the original frame + + Returns + ------- + :class:`numpy.ndarray` + The full frame background image for applying masks to + """ + if self._type == "mask": + return np.zeros(frame_dims, dtype="uint8") + + if not self._input_is_faces: # Frame is in the detected faces object + assert detected_faces[0].image is not None + return np.ascontiguousarray(detected_faces[0].image) + + # Outputting to frames, but input is faces. Apply the face patches to an empty canvas + retval = np.zeros((*frame_dims, 3), dtype="uint8") + for detected_face in detected_faces: + assert detected_face.image is not None + face = AlignedFace(detected_face.landmarks_xy, + image=detected_face.image, + centering="head", + size=detected_face.image.shape[0], + is_aligned=True) + border = cv2.BORDER_TRANSPARENT if len(detected_faces) > 1 else cv2.BORDER_CONSTANT + assert face.face is not None + cv2.warpAffine(face.face, + face.adjusted_matrix, + tuple(reversed(frame_dims)), + retval, + flags=cv2.WARP_INVERSE_MAP | face.interpolators[1], + borderMode=border) + return retval + + def _get_background_face(self, + detected_face: DetectedFace, + mask_centering: CenteringType, + mask_size: int) -> np.ndarray: + """ Obtain the background images when the output is faces + + The output image will depend on the requested output type and whether the input is faces + or frames + + Parameters + ---------- + detected_face: :class:`~lib.align.detected_face.DetectedFace` + Detected face object for the output image + mask_centering: Literal["face", "head", "legacy"] + The centering of the stored mask + mask_size: int + The pixel size of the stored mask + + Returns + ------- + list[]:class:`numpy.ndarray`] + The face background image for applying masks to for each detected face object + """ + if self._type == "mask": + return np.zeros((mask_size, mask_size), dtype="uint8") + + assert detected_face.image is not None + + if self._input_is_faces: + retval = AlignedFace(detected_face.landmarks_xy, + image=detected_face.image, + centering=mask_centering, + size=mask_size, + is_aligned=True).face + else: + centering: CenteringType = ("legacy" if self._alignments is not None and + self._alignments.version == 1.0 + else mask_centering) + detected_face.load_aligned(detected_face.image, + size=mask_size, + centering=centering, + force=True) + retval = detected_face.aligned.face + + assert retval is not None + return retval + + def _get_background(self, + detected_faces: list[DetectedFace], + frame_dims: tuple[int, int], + mask_centering: CenteringType, + mask_size: int) -> np.ndarray: + """ Obtain the background image that the final outut will be placed on + + Parameters + ---------- + detected_faces: list[:class:`~lib.align.detected_face.DetectedFace`] + Detected face objects for the output image + frame_dims: tuple[int, int] + The size of the original frame + mask_centering: Literal["face", "head", "legacy"] + The centering of the stored mask + mask_size: int + The pixel size of the stored mask + + Returns + ------- + :class:`numpy.ndarray` + The background image for the mask output + """ + if self._full_frame: + retval = self._get_background_frame(detected_faces, frame_dims) + else: + assert len(detected_faces) == 1 # If outputting faces, we should only receive 1 face + retval = self._get_background_face(detected_faces[0], mask_centering, mask_size) + + logger.trace("Background image (size: %s, dtype: %s)", # type:ignore[attr-defined] + retval.shape, retval.dtype) + return retval + + def _get_mask(self, + detected_faces: list[DetectedFace], + mask_type: str, + mask_dims: tuple[int, int]) -> np.ndarray: + """ Generate the mask to be applied to the final output frame + + Parameters + ---------- + detected_faces: list[:class:`~lib.align.detected_face.DetectedFace`] + Detected face objects to generate the masks from + mask_type: str + The mask-type to use + mask_dims : tuple[int, int] + The size of the mask to output + + Returns + ------- + :class:`numpy.ndarray` + The final mask to apply to the output image + """ + retval = np.zeros(mask_dims, dtype="uint8") + for face in detected_faces: + mask_object = face.mask[mask_type] + mask_object.set_blur_and_threshold(blur_kernel=self._blur_kernel, + threshold=self._threshold) + if self._full_frame: + mask = mask_object.get_full_frame_mask(*reversed(mask_dims)) + else: + mask = mask_object.mask[..., 0] + np.maximum(retval, mask, out=retval) + logger.trace("Final mask (shape: %s, dtype: %s)", # type:ignore[attr-defined] + retval.shape, retval.dtype) + return retval + + def _build_output_image(self, background: np.ndarray, mask: np.ndarray) -> np.ndarray: + """ Collate the mask and images for the final output image, depending on selected output + type + + Parameters + ---------- + background: :class:`numpy.ndarray` + The image that the mask will be applied to + mask: :class:`numpy.ndarray` + The mask to output + + Returns + ------- + :class:`numpy.ndarray` + The final output image + """ + if self._type == "mask": + return mask + + mask = mask[..., None] + if self._type == "masked": + return np.concatenate([background, mask], axis=-1) + + height, width = background.shape[:2] + masked = (background.astype("float32") * mask.astype("float32") / 255.).astype("uint8") + mask = np.tile(mask, 3) + for img in (background, masked, mask): + cv2.rectangle(img, (0, 0), (width - 1, height - 1), (255, 255, 255), 1) + axis = 0 if background.shape[0] < background.shape[1] else 1 + retval = np.concatenate((background, masked, mask), axis=axis) + + return retval + + def _create_image(self, + detected_faces: list[DetectedFace], + mask_type: str, + frame_dims: tuple[int, int] | None) -> np.ndarray: + """ Create a mask preview image for saving out to disk + + Parameters + ---------- + detected_faces: list[:class:`~lib.align.detected_face.DetectedFace`] + Detected face objects for the output image + mask_type: str + The mask_type to process + frame_dims: tuple[int, int] | None + The size of the original frame, if input is faces otherwise ``None`` + + Returns + ------- + :class:`numpy.ndarray`: + A preview image depending on the output type in one of the following forms: + - Containing 3 sub images: The original face, the masked face and the mask + - The mask only + - The masked face + """ + assert detected_faces[0].image is not None + dims = T.cast(tuple[int, int], + frame_dims if self._input_is_faces else detected_faces[0].image.shape[:2]) + assert dims is not None and len(dims) == 2 + + mask_centering = detected_faces[0].mask[mask_type].stored_centering + mask_size = detected_faces[0].mask[mask_type].stored_size + + background = self._get_background(detected_faces, dims, mask_centering, mask_size) + mask = self._get_mask(detected_faces, + mask_type, + dims if self._full_frame else (mask_size, mask_size)) + retval = self._build_output_image(background, mask) + + logger.trace("Output image (shape: %s, dtype: %s)", # type:ignore[attr-defined] + retval.shape, retval.dtype) + return retval + + def _handle_cache(self, + frame: str, + idx: int, + detected_face: DetectedFace) -> list[tuple[int, DetectedFace]]: + """ For full frame output, cache any faces until all detected faces have been seen. For + face output, just return the detected_face object inside a list + + Parameters + ---------- + frame: str + The frame name in the alignments file + idx: int + The index of the face for this frame in the alignments file + detected_face: :class:`~lib.align.detected_face.DetectedFace` + A detected_face object for a face + + Returns + ------- + list[tuple[int, :class:`~lib.align.detected_face.DetectedFace`]] + Face index and detected face objects to be processed for this output, if any + """ + if not self._full_frame: + return [(idx, detected_face)] + + assert self._alignments is not None + faces_in_frame = self._alignments.count_faces_in_frame(frame) + if faces_in_frame == 1: + return [(idx, detected_face)] + + self._full_frame_cache.setdefault(frame, []).append((idx, detected_face)) + + if len(self._full_frame_cache[frame]) != faces_in_frame: + logger.trace("Caching face for frame '%s'", frame) # type:ignore[attr-defined] + return [] + + retval = self._full_frame_cache.pop(frame) + logger.trace("Processing '%s' from cache: %s", frame, retval) # type:ignore[attr-defined] + return retval + + def _get_mask_types(self, + frame: str, + detected_faces: list[tuple[int, DetectedFace]]) -> list[str]: + """ Get the mask type names for the select mask type. Remove any detected faces where + the selected mask does not exist + + Parameters + ---------- + frame: str + The frame name in the alignments file + idx: int + The index of the face for this frame in the alignments file + detected_face: list[tuple[int, :class:`~lib.align.detected_face.DetectedFace`] + The face index and detected_face object for output + + Returns + ------- + list[str] + List of mask type names to be processed + """ + if self._mask_type == "bisenet-fp": + mask_types = [f"{self._mask_type}_{area}" for area in ("face", "head")] + else: + mask_types = [self._mask_type] + + if self._mask_type == "custom": + mask_types.append(f"{self._mask_type}_{self._centering}") + + final_masks = set() + for idx in reversed(range(len(detected_faces))): + face_idx, detected_face = detected_faces[idx] + if detected_face.mask is None or not any(mask in detected_face.mask + for mask in mask_types): + logger.warning("Mask type '%s' does not exist for frame '%s' index %s. Skipping", + self._mask_type, frame, face_idx) + del detected_faces[idx] + continue + final_masks.update([m for m in detected_face.mask if m in mask_types]) + + retval = list(final_masks) + logger.trace("Handling mask types: %s", retval) # type:ignore[attr-defined] + return retval + + def save(self, + frame: str, + idx: int, + detected_face: DetectedFace, + frame_dims: tuple[int, int] | None = None) -> None: + """ Build the mask preview image and save + + Parameters + ---------- + frame: str + The frame name in the alignments file + idx: int + The index of the face for this frame in the alignments file + detected_face: :class:`~lib.align.detected_face.DetectedFace` + A detected_face object for a face + frame_dims: tuple[int, int] | None, optional + The size of the original frame, if input is faces otherwise ``None``. Default: ``None`` + """ + assert self._saver is not None + + faces = self._handle_cache(frame, idx, detected_face) + if not faces: + return + + mask_types = self._get_mask_types(frame, faces) + if not faces or not mask_types: + logger.debug("No valid faces/masks to process for '%s'", frame) + return + + for mask_type in mask_types: + detected_faces = [f[1] for f in faces if mask_type in f[1].mask] + if not detected_face: + logger.warning("No '%s' masks to output for '%s'", mask_type, frame) + continue + if len(detected_faces) != len(faces): + logger.warning("Some '%s' masks are missing for '%s'", mask_type, frame) + + image = self._create_image(detected_faces, mask_type, frame_dims) + filename = os.path.splitext(frame)[0] + if len(mask_types) > 1: + filename += f"_{mask_type}" + if not self._full_frame: + filename += f"_{idx}" + filename = os.path.join(self._saver.location, f"{filename}.png") + logger.trace("filename: '%s', image_shape: %s", filename, image.shape) # type: ignore + self._saver.save(filename, image) + + def close(self) -> None: + """ Shut down the image saver if it is open """ + if self._saver is None: + return + logger.debug("Shutting down saver") + self._saver.close() diff --git a/tools/model/__init__.py b/tools/model/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/model/cli.py b/tools/model/cli.py new file mode 100644 index 0000000000..68d1e8e455 --- /dev/null +++ b/tools/model/cli.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +""" Command Line Arguments for tools """ +import gettext +import typing as T + +from lib.cli.args import FaceSwapArgs +from lib.cli.actions import DirFullPaths, Radio + +# LOCALES +_LANG = gettext.translation("tools.restore.cli", localedir="locales", fallback=True) +_ = _LANG.gettext + +_HELPTEXT = _("This tool lets you perform actions on saved Faceswap models.") + + +class ModelArgs(FaceSwapArgs): + """ Class to perform actions on model files """ + + @staticmethod + def get_info() -> str: + """ Return command information """ + return _("A tool for performing actions on Faceswap trained model files") + + @staticmethod + def get_argument_list() -> list[dict[str, T.Any]]: + """ Put the arguments in a list so that they are accessible from both argparse and gui """ + argument_list = [] + argument_list.append({ + "opts": ("-m", "--model-dir"), + "action": DirFullPaths, + "dest": "model_dir", + "required": True, + "help": _( + "Model directory. A directory containing the model you wish to perform an action " + "on.")}) + argument_list.append({ + "opts": ("-j", "--job"), + "action": Radio, + "type": str, + "choices": ("inference", "nan-scan", "restore"), + "required": True, + "help": _( + "R|Choose which action you want to perform." + "\nL|'inference' - Create an inference only copy of the model. Strips any layers " + "from the model which are only required for training. NB: This is for exporting " + "the model for use in external applications. Inference generated models cannot be " + "used within Faceswap. See the 'format' option for specifying the model output " + "format." + "\nL|'nan-scan' - Scan the model file for NaNs or Infs (invalid data)." + "\nL|'restore' - Restore a model from backup.")}) + argument_list.append({ + "opts": ("-f", "--format"), + "action": Radio, + "type": str, + "choices": ("h5", "saved-model"), + "default": "h5", + "group": _("inference"), + "help": _( + "R|The format to save the model as. Note: Only used for 'inference' job." + "\nL|'h5' - Standard Keras H5 format. Does not store any custom layer " + "information. Layers will need to be loaded from Faceswap to use." + "\nL|'saved-model' - Tensorflow's Saved Model format. Contains all information " + "required to load the model outside of Faceswap.")}) + argument_list.append({ + "opts": ("-s", "--swap-model"), + "action": "store_true", + "dest": "swap_model", + "default": False, + "group": _("inference"), + "help": _( + "Only used for 'inference' job. Generate the inference model for B -> A instead " + "of A -> B.")}) + return argument_list diff --git a/tools/model/model.py b/tools/model/model.py new file mode 100644 index 0000000000..0cb3a033b3 --- /dev/null +++ b/tools/model/model.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python3 +""" Tool to restore models from backup """ +from __future__ import annotations +import logging +import os +import sys +import typing as T + +import numpy as np +import tensorflow as tf +from tensorflow import keras + +from lib.model.backup_restore import Backup + +# Import the following libs for custom objects +from lib.model import initializers, layers, normalization # noqa # pylint:disable=unused-import +from plugins.train.model._base.model import _Inference + + +if T.TYPE_CHECKING: + import argparse + +logger = logging.getLogger(__name__) + + +class Model(): + """ Tool to perform actions on a model file. + + Parameters + ---------- + :class:`argparse.Namespace` + The command line arguments calling the model tool + """ + def __init__(self, arguments: argparse.Namespace) -> None: + logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) + self._configure_tensorflow() + self._model_dir = self._check_folder(arguments.model_dir) + self._job = self._get_job(arguments) + + @classmethod + def _configure_tensorflow(cls) -> None: + """ Disable eager execution and force Tensorflow into CPU mode. """ + tf.config.set_visible_devices([], device_type="GPU") + tf.compat.v1.disable_eager_execution() + + @classmethod + def _get_job(cls, arguments: argparse.Namespace) -> T.Any: + """ Get the correct object that holds the selected job. + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments received for the Model tool which will be used to initiate + the selected job + + Returns + ------- + class + The object that will perform the selected job + """ + jobs = {"inference": Inference, + "nan-scan": NaNScan, + "restore": Restore} + return jobs[arguments.job](arguments) + + @classmethod + def _check_folder(cls, model_dir: str) -> str: + """ Check that the passed in model folder exists and contains a valid model. + + If the passed in value fails any checks, process exits. + + Parameters + ---------- + model_dir: str + The model folder to be checked + + Returns + ------- + str + The confirmed location of the model folder. + """ + if not os.path.exists(model_dir): + logger.error("Model folder does not exist: '%s'", model_dir) + sys.exit(1) + + chkfiles = [fname + for fname in os.listdir(model_dir) + if fname.endswith(".h5") + and not os.path.splitext(fname)[0].endswith("_inference")] + + if not chkfiles: + logger.error("Could not find a model in the supplied folder: '%s'", model_dir) + sys.exit(1) + + if len(chkfiles) > 1: + logger.error("More than one model file found in the model folder: '%s'", model_dir) + sys.exit(1) + + model_name = os.path.splitext(chkfiles[0])[0].title() + logger.info("%s Model found", model_name) + return model_dir + + def process(self) -> None: + """ Call the selected model job.""" + self._job.process() + + +class Inference(): + """ Save an inference model from a trained Faceswap model. + + Parameters + ---------- + :class:`argparse.Namespace` + The command line arguments calling the model tool + """ + def __init__(self, arguments: argparse.Namespace) -> None: + self._switch = arguments.swap_model + self._format = arguments.format + self._input_file, self._output_file = self._get_output_file(arguments.model_dir) + + def _get_output_file(self, model_dir: str) -> tuple[str, str]: + """ Obtain the full path for the output model file/folder + + Parameters + ---------- + model_dir: str + The full path to the folder containing the Faceswap trained model .h5 file + + Returns + ------- + str + The full path to the source model file + str + The full path to the inference model save location + """ + model_name = next(fname for fname in os.listdir(model_dir) if fname.endswith(".h5")) + in_path = os.path.join(model_dir, model_name) + logger.debug("Model input path: '%s'", in_path) + + model_name = f"{os.path.splitext(model_name)[0]}_inference" + model_name = f"{model_name}.h5" if self._format == "h5" else model_name + out_path = os.path.join(model_dir, model_name) + logger.debug("Inference output path: '%s'", out_path) + return in_path, out_path + + def process(self) -> None: + """ Run the inference model creation process. """ + logger.info("Loading model '%s'", self._input_file) + model = keras.models.load_model(self._input_file, compile=False) + logger.info("Creating inference model...") + inference = _Inference(model, self._switch).model + logger.info("Saving to: '%s'", self._output_file) + inference.save(self._output_file) + + +class NaNScan(): + """ Tool to scan for NaN and Infs in model weights. + + Parameters + ---------- + :class:`argparse.Namespace` + The command line arguments calling the model tool + """ + def __init__(self, arguments: argparse.Namespace) -> None: + logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) + self._model_file = self._get_model_filename(arguments.model_dir) + + @classmethod + def _get_model_filename(cls, model_dir: str) -> str: + """ Obtain the full path the model's .h5 file. + + Parameters + ---------- + model_dir: str + The full path to the folder containing the model file + + Returns + ------- + str + The full path to the saved model file + """ + model_file = next(fname for fname in os.listdir(model_dir) if fname.endswith(".h5")) + return os.path.join(model_dir, model_file) + + def _parse_weights(self, + layer: keras.models.Model | keras.layers.Layer) -> dict: + """ Recursively pass through sub-models to scan layer weights""" + weights = layer.get_weights() + logger.debug("Processing weights for layer '%s', length: '%s'", + layer.name, len(weights)) + + if not weights: + logger.debug("Skipping layer with no weights: %s", layer.name) + return {} + + if hasattr(layer, "layers"): # Must be a submodel + retval = {} + for lyr in layer.layers: + info = self._parse_weights(lyr) + if not info: + continue + retval[lyr.name] = info + return retval + + nans = sum(np.count_nonzero(np.isnan(w)) for w in weights) + infs = sum(np.count_nonzero(np.isinf(w)) for w in weights) + + if nans + infs == 0: + return {} + return {"nans": nans, "infs": infs} + + def _parse_output(self, errors: dict, indent: int = 0) -> None: + """ Parse the output of the errors dictionary and print a pretty summary. + + Parameters + ---------- + errors: dict + The nested dictionary of errors found when parsing the weights + + indent: int, optional + How far should the current printed line be indented. Default: `0` + """ + for key, val in errors.items(): + logline = f"|{'--' * indent} " + logline += key.ljust(50 - len(logline)) + if isinstance(val, dict) and "nans" not in val: + logger.info(logline) + self._parse_output(val, indent + 1) + elif isinstance(val, dict) and "nans" in val: + logline += f"nans: {val['nans']}, infs: {val['infs']}" + logger.info(logline.ljust(30)) + + def process(self) -> None: + """ Scan the loaded model for NaNs and Infs and output summary. """ + logger.info("Loading model...") + model = keras.models.load_model(self._model_file, compile=False) + logger.info("Parsing weights for invalid values...") + errors = self._parse_weights(model) + + if not errors: + logger.info("No invalid values found in model: '%s'", self._model_file) + sys.exit(1) + + logger.info("Invalid values found in model: %s", self._model_file) + self._parse_output(errors) + + +class Restore(): + """ Restore a model from backup. + + Parameters + ---------- + :class:`argparse.Namespace` + The command line arguments calling the model tool + """ + def __init__(self, arguments: argparse.Namespace) -> None: + logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) + self._model_dir = arguments.model_dir + self._model_name = self._get_model_name() + + def process(self) -> None: + """ Perform the Restore process """ + logger.info("Starting Model Restore...") + backup = Backup(self._model_dir, self._model_name) + backup.restore() + logger.info("Completed Model Restore") + + def _get_model_name(self) -> str: + """ Additional checks to make sure that a backup exists in the model location. """ + bkfiles = [fname for fname in os.listdir(self._model_dir) if fname.endswith(".bk")] + if not bkfiles: + logger.error("Could not find any backup files in the supplied folder: '%s'", + self._model_dir) + sys.exit(1) + logger.verbose("Backup files: %s)", bkfiles) # type:ignore + + model_name = next(fname for fname in bkfiles if fname.endswith(".h5.bk")) + return model_name[:-6] diff --git a/tools/preview/cli.py b/tools/preview/cli.py index 3ee985d30d..147f0449cc 100644 --- a/tools/preview/cli.py +++ b/tools/preview/cli.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ +import argparse import gettext +import typing as T from lib.cli.args import FaceSwapArgs from lib.cli.actions import DirOrFileFullPaths, DirFullPaths, FileFullPaths @@ -17,42 +19,63 @@ class PreviewArgs(FaceSwapArgs): """ Class to parse the command line arguments for Preview (Convert Settings) tool """ @staticmethod - def get_info(): - """ Return command information """ + def get_info() -> str: + """ Return command information + + Returns + ------- + str + Top line information about the Preview tool + """ return _("Preview tool\nAllows you to configure your convert settings with a live preview") - def get_argument_list(self): - - argument_list = list() - argument_list.append(dict( - opts=("-i", "--input-dir"), - action=DirOrFileFullPaths, - filetypes="video", - dest="input_dir", - group=_("data"), - required=True, - help=_("Input directory or video. Either a directory containing the image files you " - "wish to process or path to a video file."))) - argument_list.append(dict( - opts=("-al", "--alignments"), - action=FileFullPaths, - filetypes="alignments", - type=str, - group=_("data"), - dest="alignments_path", - help=_("Path to the alignments file for the input, if not at the default location"))) - argument_list.append(dict( - opts=("-m", "--model-dir"), - action=DirFullPaths, - dest="model_dir", - group=_("data"), - required=True, - help=_("Model directory. A directory containing the trained model you wish to " - "process."))) - argument_list.append(dict( - opts=("-s", "--swap-model"), - action="store_true", - dest="swap_model", - default=False, - help=_("Swap the model. Instead of A -> B, swap B -> A"))) + @staticmethod + def get_argument_list() -> list[dict[str, T.Any]]: + """ Put the arguments in a list so that they are accessible from both argparse and gui + + Returns + ------- + list[dict[str, Any]] + Top command line options for the preview tool + """ + argument_list = [] + argument_list.append({ + "opts": ("-i", "--input-dir"), + "action": DirOrFileFullPaths, + "filetypes": "video", + "dest": "input_dir", + "group": _("data"), + "required": True, + "help": _( + "Input directory or video. Either a directory containing the image files you wish " + "to process or path to a video file.")}) + argument_list.append({ + "opts": ("-a", "--alignments"), + "action": FileFullPaths, + "filetypes": "alignments", + "type": str, + "group": _("data"), + "dest": "alignments_path", + "help": _( + "Path to the alignments file for the input, if not at the default location")}) + argument_list.append({ + "opts": ("-m", "--model-dir"), + "action": DirFullPaths, + "dest": "model_dir", + "group": _("data"), + "required": True, + "help": _( + "Model directory. A directory containing the trained model you wish to process.")}) + argument_list.append({ + "opts": ("-s", "--swap-model"), + "action": "store_true", + "dest": "swap_model", + "default": False, + "help": _("Swap the model. Instead of A -> B, swap B -> A")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-al", ), + "type": str, + "dest": "depr_alignments_al_a", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/preview/control_panels.py b/tools/preview/control_panels.py new file mode 100644 index 0000000000..3dc55ba2e2 --- /dev/null +++ b/tools/preview/control_panels.py @@ -0,0 +1,681 @@ +#!/usr/bin/env python3 +""" Manages the widgets that hold the bottom 'control' area of the preview tool """ +from __future__ import annotations +import gettext +import logging +import typing as T + +import tkinter as tk + +from tkinter import ttk +from configparser import ConfigParser + +from lib.gui.custom_widgets import Tooltip +from lib.gui.control_helper import ControlPanel, ControlPanelOption +from lib.gui.utils import get_images +from plugins.plugin_loader import PluginLoader +from plugins.convert._config import Config + +if T.TYPE_CHECKING: + from collections.abc import Callable + from .preview import Preview + +logger = logging.getLogger(__name__) + +# LOCALES +_LANG = gettext.translation("tools.preview", localedir="locales", fallback=True) +_ = _LANG.gettext + + +class ConfigTools(): + """ Tools for loading, saving, setting and retrieving configuration file values. + + Attributes + ---------- + tk_vars: dict + Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` + """ + def __init__(self) -> None: + self._config = Config(None) + self.tk_vars: dict[str, dict[str, tk.BooleanVar | tk.StringVar | tk.IntVar | tk.DoubleVar] + ] = {} + self._config_dicts = self._get_config_dicts() # Holds currently saved config + + @property + def config(self) -> Config: + """ :class:`plugins.convert._config.Config` The convert configuration """ + return self._config + + @property + def config_dicts(self) -> dict[str, T.Any]: + """ dict: The convert configuration options in dictionary form.""" + return self._config_dicts + + @property + def sections(self) -> list[str]: + """ list: The sorted section names that exist within the convert Configuration options. """ + return sorted(set(plugin.split(".")[0] for plugin in self._config.config.sections() + if plugin.split(".")[0] != "writer")) + + @property + def plugins_dict(self) -> dict[str, list[str]]: + """ dict: Dictionary of configuration option sections as key with a list of containing + plugins as the value """ + return {section: sorted([plugin.split(".")[1] for plugin in self._config.config.sections() + if plugin.split(".")[0] == section]) + for section in self.sections} + + def update_config(self) -> None: + """ Update :attr:`config` with the currently selected values from the GUI. """ + for section, items in self.tk_vars.items(): + for item, value in items.items(): + try: + new_value = str(value.get()) + except tk.TclError as err: + # When manually filling in text fields, blank values will + # raise an error on numeric data types so return 0 + logger.debug("Error getting value. Defaulting to 0. Error: %s", str(err)) + new_value = str(0) + old_value = self._config.config[section][item] + if new_value != old_value: + logger.trace("Updating config: %s, %s from %s to %s", # type: ignore + section, item, old_value, new_value) + self._config.config[section][item] = new_value + + def _get_config_dicts(self) -> dict[str, dict[str, T.Any]]: + """ Obtain a custom configuration dictionary for convert configuration items in use + by the preview tool formatted for control helper. + + Returns + ------- + dict + Each configuration section as keys, with the values as a dict of option: + :class:`lib.gui.control_helper.ControlOption` pairs. """ + logger.debug("Formatting Config for GUI") + config_dicts: dict[str, dict[str, T.Any]] = {} + for section in self._config.config.sections(): + if section.startswith("writer."): + continue + for key, val in self._config.defaults[section].items.items(): + if key == "helptext": + config_dicts.setdefault(section, {})[key] = val + continue + cp_option = ControlPanelOption(title=key, + dtype=val.datatype, + group=val.group, + default=val.default, + initial_value=self._config.get(section, key), + choices=val.choices, + is_radio=val.gui_radio, + rounding=val.rounding, + min_max=val.min_max, + helptext=val.helptext) + self.tk_vars.setdefault(section, {})[key] = cp_option.tk_var + config_dicts.setdefault(section, {})[key] = cp_option + logger.debug("Formatted Config for GUI: %s", config_dicts) + return config_dicts + + def reset_config_to_saved(self, section: str | None = None) -> None: + """ Reset the GUI parameters to their saved values within the configuration file. + + Parameters + ---------- + section: str, optional + The configuration section to reset the values for, If ``None`` provided then all + sections are reset. Default: ``None`` + """ + logger.debug("Resetting to saved config: %s", section) + sections = [section] if section is not None else list(self.tk_vars.keys()) + for config_section in sections: + for item, options in self._config_dicts[config_section].items(): + if item == "helptext": + continue + val = options.value + if val != self.tk_vars[config_section][item].get(): + self.tk_vars[config_section][item].set(val) + logger.debug("Setting %s - %s to saved value %s", config_section, item, val) + logger.debug("Reset to saved config: %s", section) + + def reset_config_to_default(self, section: str | None = None) -> None: + """ Reset the GUI parameters to their default configuration values. + + Parameters + ---------- + section: str, optional + The configuration section to reset the values for, If ``None`` provided then all + sections are reset. Default: ``None`` + """ + logger.debug("Resetting to default: %s", section) + sections = [section] if section is not None else list(self.tk_vars.keys()) + for config_section in sections: + for item, options in self._config_dicts[config_section].items(): + if item == "helptext": + continue + default = options.default + if default != self.tk_vars[config_section][item].get(): + self.tk_vars[config_section][item].set(default) + logger.debug("Setting %s - %s to default value %s", + config_section, item, default) + logger.debug("Reset to default: %s", section) + + def save_config(self, section: str | None = None) -> None: + """ Save the configuration ``.ini`` file with the currently stored values. + + Notes + ----- + We cannot edit the existing saved config as comments tend to get removed, so we create + a new config and populate that. + + Parameters + ---------- + section: str, optional + The configuration section to save, If ``None`` provided then all sections are saved. + Default: ``None`` + """ + logger.debug("Saving %s config", section) + + new_config = ConfigParser(allow_no_value=True) + + for section_name, sect in self._config.defaults.items(): + logger.debug("Adding section: '%s')", section_name) + self._config.insert_config_section(section_name, + sect.helptext, + config=new_config) + for item, options in sect.items.items(): + if item == "helptext": + continue # helptext already written at top + if ((section is not None and section_name != section) + or section_name not in self.tk_vars): + # retain saved values that have not been updated + new_opt = self._config.get(section_name, item) + logger.debug("Retaining option: (item: '%s', value: '%s')", item, new_opt) + else: + new_opt = self.tk_vars[section_name][item].get() + logger.debug("Setting option: (item: '%s', value: '%s')", item, new_opt) + + # Set config_dicts value to new saved value + self._config_dicts[section_name][item].set_initial_value(new_opt) + + helptext = self._config.format_help(options.helptext, is_section=False) + new_config.set(section_name, helptext) + new_config.set(section_name, item, str(new_opt)) + + self._config.config = new_config + self._config.save_config() + logger.info("Saved config: '%s'", self._config.configfile) + + +class BusyProgressBar(): + """ An infinite progress bar for when a thread is running to swap/patch a group of samples """ + def __init__(self, parent: ttk.Frame) -> None: + self._progress_bar = self._add_busy_indicator(parent) + + def _add_busy_indicator(self, parent: ttk.Frame) -> ttk.Progressbar: + """ Place progress bar into bottom bar to indicate when processing. + + Parameters + ---------- + parent: tkinter object + The tkinter object that holds the busy indicator + + Returns + ------- + ttk.Progressbar + A Progress bar to indicate that the Preview tool is busy + """ + logger.debug("Placing busy indicator") + pbar = ttk.Progressbar(parent, mode="indeterminate") + pbar.pack(side=tk.LEFT) + pbar.pack_forget() + return pbar + + def stop(self) -> None: + """ Stop and hide progress bar """ + logger.debug("Stopping busy indicator") + if not self._progress_bar.winfo_ismapped(): + logger.debug("busy indicator already hidden") + return + self._progress_bar.stop() + self._progress_bar.pack_forget() + + def start(self) -> None: + """ Start and display progress bar """ + logger.debug("Starting busy indicator") + if self._progress_bar.winfo_ismapped(): + logger.debug("busy indicator already started") + return + + self._progress_bar.pack(side=tk.LEFT, padx=5, pady=(5, 10), fill=tk.X, expand=True) + self._progress_bar.start(25) + + +class ActionFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ Frame that holds the left hand side options panel containing the command line options. + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + parent: tkinter object + The parent tkinter object that holds the Action Frame + """ + def __init__(self, app: Preview, parent: ttk.Frame) -> None: + logger.debug("Initializing %s: (app: %s, parent: %s)", + self.__class__.__name__, app, parent) + self._app = app + + super().__init__(parent) + self.pack(side=tk.LEFT, anchor=tk.N, fill=tk.Y) + self._tk_vars: dict[str, tk.StringVar] = {} + + self._options = { + "color": app._patch.converter.cli_arguments.color_adjustment.replace("-", "_"), + "mask_type": app._patch.converter.cli_arguments.mask_type.replace("-", "_"), + "face_scale": app._patch.converter.cli_arguments.face_scale} + defaults = {opt: self._format_to_display(val) if opt != "face_scale" else val + for opt, val in self._options.items()} + self._busy_bar = self._build_frame(defaults, + app._samples.generate, + app._refresh, + app._samples.available_masks, + app._samples.predictor.has_predicted_mask) + + @property + def convert_args(self) -> dict[str, T.Any]: + """ dict: Currently selected Command line arguments from the :class:`ActionFrame`. """ + retval = {opt if opt != "color" else "color_adjustment": + self._format_from_display(self._tk_vars[opt].get()) + for opt in self._options if opt != "face_scale"} + retval["face_scale"] = self._tk_vars["face_scale"].get() + return retval + + @property + def busy_progress_bar(self) -> BusyProgressBar: + """ :class:`BusyProgressBar`: The progress bar that appears on the left hand side whilst a + swap/patch is being applied """ + return self._busy_bar + + @staticmethod + def _format_from_display(var: str) -> str: + """ Format a variable from the display version to the command line action version. + + Parameters + ---------- + var: str + The variable name to format + + Returns + ------- + str + The formatted variable name + """ + return var.replace(" ", "_").lower() + + @staticmethod + def _format_to_display(var: str) -> str: + """ Format a variable from the command line action version to the display version. + Parameters + ---------- + var: str + The variable name to format + + Returns + ------- + str + The formatted variable name + """ + return var.replace("_", " ").replace("-", " ").title() + + def _build_frame(self, + defaults: dict[str, T.Any], + refresh_callback: Callable[[], None], + patch_callback: Callable[[], None], + available_masks: list[str], + has_predicted_mask: bool) -> BusyProgressBar: + """ Build the :class:`ActionFrame`. + + Parameters + ---------- + defaults: dict + The default command line options + patch_callback: python function + The function to execute when a patch callback is received + refresh_callback: python function + The function to execute when a refresh callback is received + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + ttk.Progressbar + A Progress bar to indicate that the Preview tool is busy + """ + logger.debug("Building Action frame") + + bottom_frame = ttk.Frame(self) + bottom_frame.pack(side=tk.BOTTOM, fill=tk.X, anchor=tk.S) + top_frame = ttk.Frame(self) + top_frame.pack(side=tk.TOP, fill=tk.BOTH, anchor=tk.N, expand=True) + + self._add_cli_choices(top_frame, defaults, available_masks, has_predicted_mask) + + busy_indicator = BusyProgressBar(bottom_frame) + self._add_refresh_button(bottom_frame, refresh_callback) + self._add_patch_callback(patch_callback) + self._add_actions(bottom_frame) + logger.debug("Built Action frame") + return busy_indicator + + def _add_cli_choices(self, + parent: ttk.Frame, + defaults: dict[str, T.Any], + available_masks: list[str], + has_predicted_mask: bool) -> None: + """ Create :class:`lib.gui.control_helper.ControlPanel` object for the command + line options. + + parent: :class:`ttk.Frame` + The frame to hold the command line choices + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + """ + cp_options = self._get_control_panel_options(defaults, available_masks, has_predicted_mask) + panel_kwargs = {"blank_nones": False, "label_width": 10, "style": "CPanel"} + ControlPanel(parent, cp_options, header_text=None, **panel_kwargs) + + def _get_control_panel_options(self, + defaults: dict[str, T.Any], + available_masks: list[str], + has_predicted_mask: bool) -> list[ControlPanelOption]: + """ Create :class:`lib.gui.control_helper.ControlPanelOption` objects for the command + line options. + + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + list + The list of `lib.gui.control_helper.ControlPanelOption` objects for the Action Frame + """ + cp_options: list[ControlPanelOption] = [] + for opt in self._options: + if opt == "face_scale": + cp_option = ControlPanelOption(title=opt, + dtype=float, + default=0.0, + rounding=2, + min_max=(-10., 10.), + group="Command Line Choices") + else: + if opt == "mask_type": + choices = self._create_mask_choices(defaults, + available_masks, + has_predicted_mask) + else: + choices = PluginLoader.get_available_convert_plugins(opt, True) + cp_option = ControlPanelOption(title=opt, + dtype=str, + default=defaults[opt], + initial_value=defaults[opt], + choices=choices, + group="Command Line Choices", + is_radio=False) + self._tk_vars[opt] = cp_option.tk_var + cp_options.append(cp_option) + return cp_options + + def _create_mask_choices(self, + defaults: dict[str, T.Any], + available_masks: list[str], + has_predicted_mask: bool) -> list[str]: + """ Set the mask choices and default mask based on available masks. + + Parameters + ---------- + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + list + The masks that are available to use from the alignments file + """ + logger.debug("Initial mask choices: %s", available_masks) + if has_predicted_mask: + available_masks += ["predicted"] + if "none" not in available_masks: + available_masks += ["none"] + if self._format_from_display(defaults["mask_type"]) not in available_masks: + logger.debug("Setting default mask to first available: %s", available_masks[0]) + defaults["mask_type"] = available_masks[0] + logger.debug("Final mask choices: %s", available_masks) + return available_masks + + @classmethod + def _add_refresh_button(cls, + parent: ttk.Frame, + refresh_callback: Callable[[], None]) -> None: + """ Add a button to refresh the images. + + Parameters + ---------- + refresh_callback: python function + The function to execute when the refresh button is pressed + """ + btn = ttk.Button(parent, text="Update Samples", command=refresh_callback) + btn.pack(padx=5, pady=5, side=tk.TOP, fill=tk.X, anchor=tk.N) + + def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: + """ Add callback to re-patch images on action option change. + + Parameters + ---------- + patch_callback: python function + The function to execute when the images require patching + """ + for tk_var in self._tk_vars.values(): + tk_var.trace("w", patch_callback) + + def _add_actions(self, parent: ttk.Frame) -> None: + """ Add Action Buttons to the :class:`ActionFrame` + + Parameters + ---------- + parent: tkinter object + The tkinter object that holds the action buttons + """ + logger.debug("Adding util buttons") + frame = ttk.Frame(parent) + frame.pack(padx=5, pady=(5, 10), side=tk.RIGHT, fill=tk.X, anchor=tk.E) + + for utl in ("save", "clear", "reload"): + logger.debug("Adding button: '%s'", utl) + img = get_images().icons[utl] + if utl == "save": + text = _("Save full config") + action = self._app.config_tools.save_config + elif utl == "clear": + text = _("Reset full config to default values") + action = self._app.config_tools.reset_config_to_default + elif utl == "reload": + text = _("Reset full config to saved values") + action = self._app.config_tools.reset_config_to_saved + + btnutl = ttk.Button(frame, + image=img, + command=action) + btnutl.pack(padx=2, side=tk.RIGHT) + Tooltip(btnutl, text=text, wrap_length=200) + logger.debug("Added util buttons") + + +class OptionsBook(ttk.Notebook): # pylint:disable=too-many-ancestors + """ The notebook that holds the Convert configuration options. + + Parameters + ---------- + parent: tkinter object + The parent tkinter object that holds the Options book + config_tools: :class:`ConfigTools` + Tools for loading and saving configuration files + patch_callback: python function + The function to execute when a patch callback is received + + Attributes + ---------- + config_tools: :class:`ConfigTools` + Tools for loading and saving configuration files + """ + def __init__(self, + parent: ttk.Frame, + config_tools: ConfigTools, + patch_callback: Callable[[], None]) -> None: + logger.debug("Initializing %s: (parent: %s, config: %s)", + self.__class__.__name__, parent, config_tools) + super().__init__(parent) + self.pack(side=tk.RIGHT, anchor=tk.N, fill=tk.BOTH, expand=True) + self.config_tools = config_tools + + self._tabs: dict[str, dict[str, ttk.Notebook | ConfigFrame]] = {} + self._build_tabs() + self._build_sub_tabs() + self._add_patch_callback(patch_callback) + logger.debug("Initialized %s", self.__class__.__name__) + + def _build_tabs(self) -> None: + """ Build the notebook tabs for the each configuration section. """ + logger.debug("Build Tabs") + for section in self.config_tools.sections: + tab = ttk.Notebook(self) + self._tabs[section] = {"tab": tab} + self.add(tab, text=section.replace("_", " ").title()) + + def _build_sub_tabs(self) -> None: + """ Build the notebook sub tabs for each convert section's plugin. """ + for section, plugins in self.config_tools.plugins_dict.items(): + for plugin in plugins: + config_key = ".".join((section, plugin)) + config_dict = self.config_tools.config_dicts[config_key] + tab = ConfigFrame(self, config_key, config_dict) + self._tabs[section][plugin] = tab + text = plugin.replace("_", " ").title() + T.cast(ttk.Notebook, self._tabs[section]["tab"]).add(tab, text=text) + + def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: + """ Add callback to re-patch images on configuration option change. + + Parameters + ---------- + patch_callback: python function + The function to execute when the images require patching + """ + for plugins in self.config_tools.tk_vars.values(): + for tk_var in plugins.values(): + tk_var.trace("w", patch_callback) + + +class ConfigFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ Holds the configuration options for a convert plugin inside the :class:`OptionsBook`. + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + options: dict + The options for this section/plugin + """ + + def __init__(self, + parent: OptionsBook, + config_key: str, + options: dict[str, T.Any]): + logger.debug("Initializing %s", self.__class__.__name__) + super().__init__(parent) + self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + + self._options = options + + self._action_frame = ttk.Frame(self) + self._action_frame.pack(padx=0, pady=(0, 5), side=tk.BOTTOM, fill=tk.X, anchor=tk.E) + self._add_frame_separator() + + self._build_frame(parent, config_key) + logger.debug("Initialized %s", self.__class__.__name__) + + def _build_frame(self, parent: OptionsBook, config_key: str) -> None: + """ Build the options frame for this command + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + """ + logger.debug("Add Config Frame") + panel_kwargs = {"columns": 2, "option_columns": 2, "blank_nones": False, "style": "CPanel"} + frame = ttk.Frame(self) + frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + cp_options = [opt for key, opt in self._options.items() if key != "helptext"] + ControlPanel(frame, cp_options, header_text=None, **panel_kwargs) + self._add_actions(parent, config_key) + logger.debug("Added Config Frame") + + def _add_frame_separator(self) -> None: + """ Add a separator between top and bottom frames. """ + logger.debug("Add frame seperator") + sep = ttk.Frame(self._action_frame, height=2, relief=tk.RIDGE) + sep.pack(fill=tk.X, pady=5, side=tk.TOP) + logger.debug("Added frame seperator") + + def _add_actions(self, parent: OptionsBook, config_key: str) -> None: + """ Add Action Buttons. + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + """ + logger.debug("Adding util buttons") + + title = config_key.split(".")[1].replace("_", " ").title() + btn_frame = ttk.Frame(self._action_frame) + btn_frame.pack(padx=5, side=tk.BOTTOM, fill=tk.X) + for utl in ("save", "clear", "reload"): + logger.debug("Adding button: '%s'", utl) + img = get_images().icons[utl] + if utl == "save": + text = _(f"Save {title} config") + action = parent.config_tools.save_config + elif utl == "clear": + text = _(f"Reset {title} config to default values") + action = parent.config_tools.reset_config_to_default + elif utl == "reload": + text = _(f"Reset {title} config to saved values") + action = parent.config_tools.reset_config_to_saved + + btnutl = ttk.Button(btn_frame, + image=img, + command=lambda cmd=action: cmd(config_key)) # type: ignore + btnutl.pack(padx=2, side=tk.RIGHT) + Tooltip(btnutl, text=text, wrap_length=200) + logger.debug("Added util buttons") diff --git a/tools/preview/preview.py b/tools/preview/preview.py index 1efda36b87..4ac81b3627 100644 --- a/tools/preview/preview.py +++ b/tools/preview/preview.py @@ -1,43 +1,47 @@ #!/usr/bin/env python3 """ Tool to preview swaps and tweak configuration prior to running a convert """ - +from __future__ import annotations import gettext import logging import random import tkinter as tk +import typing as T + from tkinter import ttk import os import sys -from configparser import ConfigParser from threading import Event, Lock, Thread -import cv2 import numpy as np -from PIL import Image, ImageTk -from lib.align import DetectedFace, transform_image -from lib.cli.args import ConvertArgs +from lib.align import DetectedFace +from lib.cli.args_extract_convert import ConvertArgs from lib.gui.utils import get_images, get_config, initialize_config, initialize_images -from lib.gui.custom_widgets import Tooltip -from lib.gui.control_helper import ControlPanel, ControlPanelOption from lib.convert import Converter -from lib.utils import FaceswapError +from lib.utils import FaceswapError, handle_deprecated_cliopts from lib.queue_manager import queue_manager from scripts.fsmedia import Alignments, Images -from scripts.convert import Predict +from scripts.convert import Predict, ConvertItem + +from plugins.extract import ExtractMedia -from plugins.plugin_loader import PluginLoader -from plugins.convert._config import Config +from .control_panels import ActionFrame, ConfigTools, OptionsBook +from .viewer import FacesDisplay, ImagesCanvas -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +if T.TYPE_CHECKING: + from argparse import Namespace + from lib.queue_manager import EventQueue + from .control_panels import BusyProgressBar + +logger = logging.getLogger(__name__) # LOCALES _LANG = gettext.translation("tools.preview", localedir="locales", fallback=True) _ = _LANG.gettext -class Preview(tk.Tk): # pylint:disable=too-few-public-methods +class Preview(tk.Tk): """ This tool is part of the Faceswap Tools suite and should be called from ``python tools.py preview`` command. @@ -50,44 +54,63 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` """ + _w: str - def __init__(self, arguments): + def __init__(self, arguments: Namespace) -> None: logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) super().__init__() + arguments = handle_deprecated_cliopts(arguments) self._config_tools = ConfigTools() self._lock = Lock() - - self._tk_vars = dict(refresh=tk.BooleanVar(), busy=tk.BooleanVar()) - for val in self._tk_vars.values(): - val.set(False) - self._display = FacesDisplay(256, 64, self._tk_vars) - - trigger_patch = Event() - self._samples = Samples(arguments, 5, self._display, self._lock, trigger_patch) - self._patch = Patch(arguments, - self._available_masks, - self._samples, - self._display, - self._lock, - trigger_patch, - self._config_tools, - self._tk_vars) + self._dispatcher = Dispatcher(self) + self._display = FacesDisplay(self, 256, 64) + self._samples = Samples(self, arguments, 5) + self._patch = Patch(self, arguments) self._initialize_tkinter() - self._image_canvas = None - self._opts_book = None - self._cli_frame = None # cli frame holds cli options + self._image_canvas: ImagesCanvas | None = None + self._opts_book: OptionsBook | None = None + self._cli_frame: ActionFrame | None = None # cli frame holds cli options logger.debug("Initialized %s", self.__class__.__name__) @property - def _available_masks(self): - """ list: The mask names that are available for every face in the alignments file """ - retval = [key - for key, val in self._samples.alignments.mask_summary.items() - if val == self._samples.alignments.faces_count] - return retval + def config_tools(self) -> "ConfigTools": + """ :class:`ConfigTools`: The object responsible for parsing configuration options and + updating to/from the GUI """ + return self._config_tools + + @property + def dispatcher(self) -> "Dispatcher": + """ :class:`Dispatcher`: The object responsible for triggering events and variables and + handling global GUI state """ + return self._dispatcher + + @property + def display(self) -> FacesDisplay: + """ :class:`~tools.preview.viewer.FacesDisplay`: The object that holds the sample, + converted and patched faces """ + return self._display + + @property + def lock(self) -> Lock: + """ :class:`threading.Lock`: The threading lock object for the Preview GUI """ + return self._lock + + @property + def progress_bar(self) -> BusyProgressBar: + """ :class:`~tools.preview.control_panels.BusyProgressBar`: The progress bar that indicates + a swap/patch thread is running """ + assert self._cli_frame is not None + return self._cli_frame.busy_progress_bar + + def update_display(self): + """ Update the images in the canvas and redraw """ + if not hasattr(self, "_image_canvas"): # On first call object not yet created + return + assert self._image_canvas is not None + self._image_canvas.reload() - def _initialize_tkinter(self): + def _initialize_tkinter(self) -> None: """ Initialize a standalone tkinter instance. """ logger.debug("Initializing tkinter") initialize_config(self, None, None) @@ -97,10 +120,11 @@ def _initialize_tkinter(self): self.tk.call( "wm", "iconphoto", - self._w, get_images().icons["favicon"]) # pylint:disable=protected-access + self._w, + get_images().icons["favicon"]) # pylint:disable=protected-access logger.debug("Initialized tkinter") - def process(self): + def process(self) -> None: """ The entry point for the Preview tool from :file:`lib.tools.cli`. Launch the tkinter preview Window and run main loop. @@ -108,43 +132,35 @@ def process(self): self._build_ui() self.mainloop() - def _refresh(self, *args): - """ Load new faces to display in preview. + def _refresh(self, *args) -> None: + """ Patch faces with current convert settings. Parameters ---------- *args: tuple Unused, but required for tkinter callback. """ - logger.trace("Refreshing swapped faces. args: %s", args) - self._tk_vars["busy"].set(True) + logger.debug("Patching swapped faces. args: %s", args) + self._dispatcher.set_busy() self._config_tools.update_config() with self._lock: + assert self._cli_frame is not None self._patch.converter_arguments = self._cli_frame.convert_args - self._patch.current_config = self._config_tools.config - self._patch.trigger.set() - logger.trace("Refreshed swapped faces") - def _build_ui(self): + self._dispatcher.set_needs_patch() + logger.debug("Patched swapped faces") + + def _build_ui(self) -> None: """ Build the elements for displaying preview images and options panels. """ container = ttk.PanedWindow(self, orient=tk.VERTICAL) container.pack(fill=tk.BOTH, expand=True) - container.preview_display = self._display - self._image_canvas = ImagesCanvas(container, self._tk_vars) + setattr(container, "preview_display", self._display) # TODO subclass not setattr + self._image_canvas = ImagesCanvas(self, container) container.add(self._image_canvas, weight=3) options_frame = ttk.Frame(container) - self._cli_frame = ActionFrame( - options_frame, - self._available_masks, - self._samples.predictor.has_predicted_mask, - self._patch.converter.cli_arguments.color_adjustment.replace("-", "_"), - self._patch.converter.cli_arguments.mask_type.replace("-", "_"), - self._config_tools, - self._refresh, - self._samples.generate, - self._tk_vars) + self._cli_frame = ActionFrame(self, options_frame) self._opts_book = OptionsBook(options_frame, self._config_tools, self._refresh) @@ -153,6 +169,90 @@ def _build_ui(self): container.sashpos(0, int(400 * get_config().scaling_factor)) +class Dispatcher(): + """ Handles the app level tk.Variables and the threading events. Dispatches events to the + correct location and handles GUI state whilst events are handled + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + """ + def __init__(self, app: Preview): + logger.debug("Initializing %s: (app: %s)", self.__class__.__name__, app) + self._app = app + self._tk_busy = tk.BooleanVar(value=False) + self._evnt_needs_patch = Event() + self._is_updating = False + self._stacked_event = False + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def needs_patch(self) -> Event: + """:class:`threading.Event`. Set by the parent and cleared by the child. Informs the child + patching thread that a run needs to be processed """ + return self._evnt_needs_patch + + # TKInter Variables + def set_busy(self) -> None: + """ Set the tkinter busy variable to ``True`` and display the busy progress bar """ + if self._tk_busy.get(): + logger.debug("Busy event is already set. Doing nothing") + return + if not hasattr(self._app, "progress_bar"): + logger.debug("Not setting busy during initial startup") + return + + logger.debug("Setting busy event to True") + self._tk_busy.set(True) + self._app.progress_bar.start() + self._app.update_idletasks() + + def _unset_busy(self) -> None: + """ Set the tkinter busy variable to ``False`` and hide the busy progress bar """ + self._is_updating = False + if not self._tk_busy.get(): + logger.debug("busy unset when already unset. Doing nothing") + return + logger.debug("Setting busy event to False") + self._tk_busy.set(False) + self._app.progress_bar.stop() + self._app.update_idletasks() + + # Threading Events + def _wait_for_patch(self) -> None: + """ Wait for a patch thread to complete before triggering a display refresh and unsetting + the busy indicators """ + logger.debug("Checking for patch completion...") + if self._evnt_needs_patch.is_set(): + logger.debug("Samples not patched. Waiting...") + self._app.after(1000, self._wait_for_patch) + return + + logger.debug("Patch completion detected") + self._app.update_display() + self._unset_busy() + + if self._stacked_event: + logger.debug("Processing last stacked event") + self.set_busy() + self._stacked_event = False + self.set_needs_patch() + return + + def set_needs_patch(self) -> None: + """ Sends a trigger to the patching thread that it needs to be run. Waits for the patching + to complete prior to triggering a display refresh and unsetting the busy indicators """ + if self._is_updating: + logger.debug("Request to run patch when it is already running. Adding stacked event.") + self._stacked_event = True + return + self._is_updating = True + logger.debug("Triggering patch") + self._evnt_needs_patch.set() + self._wait_for_patch() + + class Samples(): """ The display samples. @@ -165,28 +265,21 @@ class Samples(): Parameters ---------- + app: :class:`Preview` + The main tkinter Preview app arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` sample_size: int The number of samples to take from the input video/images - display: :class:`FacesDisplay` - The display section of the Preview GUI. - lock: :class:`threading.Lock` - A threading lock to prevent multiple GUI updates at the same time. - trigger_patch: :class:`threading.Event` - An event to indicate that a converter patch should be run """ - def __init__(self, arguments, sample_size, display, lock, trigger_patch): - logger.debug("Initializing %s: (arguments: '%s', sample_size: %s, display: %s, lock: %s, " - "trigger_patch: %s)", self.__class__.__name__, arguments, sample_size, - display, lock, trigger_patch) + def __init__(self, app: Preview, arguments: Namespace, sample_size: int) -> None: + logger.debug("Initializing %s: (app: %s, arguments: '%s', sample_size: %s)", + self.__class__.__name__, app, arguments, sample_size) self._sample_size = sample_size - self._display = display - self._lock = lock - self._trigger_patch = trigger_patch - self._input_images = [] - self._predicted_images = [] + self._app = app + self._input_images: list[ConvertItem] = [] + self._predicted_images: list[tuple[ConvertItem, np.ndarray]] = [] self._images = Images(arguments) self._alignments = Alignments(arguments, @@ -197,48 +290,61 @@ def __init__(self, arguments, sample_size, display, lock, trigger_patch): "file was generated. You need to update the file to proceed.") logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.") sys.exit(1) + if not self._alignments.have_alignments_file: logger.error("Alignments file not found at: '%s'", self._alignments.file) sys.exit(1) + + if self._images.is_video: + assert isinstance(self._images.input_images, str) + self._alignments.update_legacy_has_source(os.path.basename(self._images.input_images)) + self._filelist = self._get_filelist() self._indices = self._get_indices() - self._predictor = Predict(queue_manager.get_queue("preview_predict_in"), - sample_size, - arguments) - self._display.set_centering(self._predictor.centering) + self._predictor = Predict(self._sample_size, arguments) + self._predictor.launch(queue_manager.get_queue("preview_predict_in")) + self._app._display.set_centering(self._predictor.centering) self.generate() logger.debug("Initialized %s", self.__class__.__name__) @property - def sample_size(self): + def available_masks(self) -> list[str]: + """ list: The mask names that are available for every face in the alignments file """ + retval = [key + for key, val in self.alignments.mask_summary.items() + if val == self.alignments.faces_count] + return retval + + @property + def sample_size(self) -> int: """ int: The number of samples to take from the input video/images """ return self._sample_size @property - def predicted_images(self): + def predicted_images(self) -> list[tuple[ConvertItem, np.ndarray]]: """ list: The predicted faces output from the Faceswap model """ return self._predicted_images @property - def alignments(self): + def alignments(self) -> Alignments: """ :class:`~lib.align.Alignments`: The alignments for the preview faces """ return self._alignments @property - def predictor(self): + def predictor(self) -> Predict: """ :class:`~scripts.convert.Predict`: The Predictor for the Faceswap model """ return self._predictor @property - def _random_choice(self): + def _random_choice(self) -> list[int]: """ list: Random indices from the :attr:`_indices` group """ retval = [random.choice(indices) for indices in self._indices] logger.debug(retval) return retval - def _get_filelist(self): + def _get_filelist(self) -> list[str]: """ Get a list of files for the input, filtering out those frames which do not contain faces. @@ -248,8 +354,9 @@ def _get_filelist(self): A list of filenames of frames that contain faces. """ logger.debug("Filtering file list to frames with faces") - if self._images.is_video: - filelist = [f"{os.path.splitext(self._images.input_images)[0]}_{frame_no:06d}.png" + if isinstance(self._images.input_images, str): + vid_name, ext = os.path.splitext(self._images.input_images) + filelist = [f"{vid_name}_{frame_no:06d}{ext}" for frame_no in range(1, self._images.images_found + 1)] else: filelist = self._images.input_images @@ -266,7 +373,7 @@ def _get_filelist(self): raise FaceswapError(msg) from err return retval - def _get_indices(self): + def _get_indices(self) -> list[list[int]]: """ Get indices for each sample group. Obtain :attr:`self.sample_size` evenly sized groups of indices @@ -279,6 +386,7 @@ def _get_indices(self): """ # Remove start and end values to get a list divisible by self.sample_size no_files = len(self._filelist) + self._sample_size = min(self._sample_size, no_files) crop = no_files % self._sample_size top_tail = list(range(no_files))[ crop // 2:no_files - (crop - (crop // 2))] @@ -291,17 +399,20 @@ def _get_indices(self): for idx, pool in enumerate(retval)]) return retval - def generate(self): + def generate(self) -> None: """ Generate a sample set. Selects :attr:`sample_size` random faces. Runs them through prediction to obtain the swap, then trigger the patch event to run the faces through patching. """ + logger.debug("Generating new random samples") + self._app.dispatcher.set_busy() self._load_frames() self._predict() - self._trigger_patch.set() + self._app.dispatcher.set_needs_patch() + logger.debug("Generated new random samples") - def _load_frames(self): + def _load_frames(self) -> None: """ Load a sample of random frames. * Picks a random face from each indices group. @@ -320,27 +431,28 @@ def _load_frames(self): face = self._alignments.get_faces_in_frame(filename)[0] detected_face = DetectedFace() detected_face.from_alignment(face, image=image) - self._input_images.append({"filename": filename, - "image": image, - "detected_faces": [detected_face]}) - self._display.source = self._input_images - self._display.update_source = True - logger.debug("Selected frames: %s", [frame["filename"] for frame in self._input_images]) - - def _predict(self): + inbound = ExtractMedia(filename=filename, image=image, detected_faces=[detected_face]) + self._input_images.append(ConvertItem(inbound=inbound)) + self._app.display.source = self._input_images + self._app.display.update_source = True + logger.debug("Selected frames: %s", + [frame.inbound.filename for frame in self._input_images]) + + def _predict(self) -> None: """ Predict from the loaded frames. With a threading lock (to prevent stacking), run the selected faces through the Faceswap model predict function and add the output to :attr:`predicted` """ - with self._lock: + with self._app.lock: self._predicted_images = [] for frame in self._input_images: self._predictor.in_queue.put(frame) idx = 0 while idx < self._sample_size: logger.debug("Predicting face %s of %s", idx + 1, self._sample_size) - items = self._predictor.out_queue.get() + items: (T.Literal["EOF"] | + list[tuple[ConvertItem, np.ndarray]]) = self._predictor.out_queue.get() if items == "EOF": logger.debug("Received EOF") break @@ -359,79 +471,50 @@ class Patch(): Parameters ---------- + app: :class:`Preview` + The main tkinter Preview app arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` - available_masks: list - The masks that are available for convert - samples: :class:`Samples` - The Samples for display. - display: :class:`FacesDisplay` - The display section of the Preview GUI. - lock: :class:`threading.Lock` - A threading lock to prevent multiple GUI updates at the same time. - trigger: :class:`threading.Event` - An event to indicate that a converter patch should be run - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` Attributes ---------- converter_arguments: dict The currently selected converter command line arguments for the patch queue - current_config::class:`lib.config.FaceswapConfig` - The currently set configuration for the patch queue """ - def __init__(self, arguments, available_masks, samples, - display, lock, trigger, config_tools, tk_vars): - logger.debug("Initializing %s: (arguments: '%s', available_masks: %s, samples: %s, " - "display: %s, lock: %s, trigger: %s, config_tools: %s, tk_vars %s)", - self.__class__.__name__, arguments, available_masks, samples, display, lock, - trigger, config_tools, tk_vars) - self._samples = samples + def __init__(self, app: Preview, arguments: Namespace) -> None: + logger.debug("Initializing %s: (app: %s, arguments: '%s')", + self.__class__.__name__, app, arguments) + self._app = app self._queue_patch_in = queue_manager.get_queue("preview_patch_in") - self._display = display - self._lock = lock - self._trigger = trigger - self.current_config = config_tools.config - self.converter_arguments = None # Updated converter arguments dict + self.converter_arguments: dict[str, T.Any] | None = None # Updated converter args configfile = arguments.configfile if hasattr(arguments, "configfile") else None - self._converter = Converter(output_size=self._samples.predictor.output_size, - coverage_ratio=self._samples.predictor.coverage_ratio, - centering=self._samples.predictor.centering, + self._converter = Converter(output_size=app._samples.predictor.output_size, + coverage_ratio=app._samples.predictor.coverage_ratio, + centering=app._samples.predictor.centering, draw_transparent=False, pre_encode=None, - arguments=self._generate_converter_arguments(arguments, - available_masks), + arguments=self._generate_converter_arguments( + arguments, + app._samples.available_masks), configfile=configfile) - self._shutdown = Event() - self._thread = Thread(target=self._process, name="patch_thread", - args=(self._trigger, - self._shutdown, - self._queue_patch_in, - self._samples, - tk_vars), + args=(self._queue_patch_in, + self._app.dispatcher.needs_patch, + app._samples), daemon=True) self._thread.start() logger.debug("Initializing %s", self.__class__.__name__) @property - def trigger(self): - """ :class:`threading.Event`: The trigger to indicate that a patching run should - commence. """ - return self._trigger - - @property - def converter(self): + def converter(self) -> Converter: """ :class:`lib.convert.Converter`: The converter to use for patching the images. """ return self._converter @staticmethod - def _generate_converter_arguments(arguments, available_masks): + def _generate_converter_arguments(arguments: Namespace, + available_masks: list[str]) -> Namespace: """ Add the default converter arguments to the initial arguments. Ensure the mask selection is available. @@ -448,7 +531,7 @@ def _generate_converter_arguments(arguments, available_masks): arguments added """ valid_masks = available_masks + ["none"] - converter_arguments = ConvertArgs(None, "convert").get_optional_arguments() + converter_arguments = ConvertArgs(None, "convert").get_optional_arguments() # type: ignore for item in converter_arguments: value = item.get("default", None) # Skip options without a default value @@ -466,7 +549,10 @@ def _generate_converter_arguments(arguments, available_masks): logger.debug(arguments) return arguments - def _process(self, trigger_event, shutdown_event, patch_queue_in, samples, tk_vars): + def _process(self, + patch_queue_in: EventQueue, + trigger_event: Event, + samples: Samples) -> None: """ The face patching process. Runs in a thread, and waits for an event to be set. Once triggered, runs a patching @@ -474,44 +560,36 @@ def _process(self, trigger_event, shutdown_event, patch_queue_in, samples, tk_va Parameters ---------- - trigger_event: :class:`threading.Event` - Set by parent process when a patching run should be executed - shutdown_event :class:`threading.Event` - Set by parent process if a shutdown has been requested - patch_queue_in: :class:`queue.Queue` + patch_queue_in: :class:`~lib.queue_manager.EventQueue` The input queue for the patching process + trigger_event: :class:`threading.Event` + The event that indicates a patching run needs to be processed samples: :class:`Samples` The Samples for display. - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` """ - logger.debug("Launching patch process thread: (trigger_event: %s, shutdown_event: %s, " - "patch_queue_in: %s, samples: %s, tk_vars: %s)", trigger_event, - shutdown_event, patch_queue_in, samples, tk_vars) + logger.debug("Launching patch process thread: (patch_queue_in: %s, trigger_event: %s, " + "samples: %s)", patch_queue_in, trigger_event, samples) patch_queue_out = queue_manager.get_queue("preview_patch_out") while True: trigger = trigger_event.wait(1) - if shutdown_event.is_set(): - logger.debug("Shutdown received") - break if not trigger: continue - # Clear trigger so calling process can set it during this run - trigger_event.clear() + logger.debug("Patch Triggered") queue_manager.flush_queue("preview_patch_in") self._feed_swapped_faces(patch_queue_in, samples) - with self._lock: + with self._app.lock: self._update_converter_arguments() - self._converter.reinitialize(config=self.current_config) + self._converter.reinitialize(config=self._app.config_tools.config) swapped = self._patch_faces(patch_queue_in, patch_queue_out, samples.sample_size) - with self._lock: - self._display.destination = swapped - tk_vars["refresh"].set(True) - tk_vars["busy"].set(False) + with self._app.lock: + self._app.display.destination = swapped + + logger.debug("Patch complete") + trigger_event.clear() logger.debug("Closed patch process thread") - def _update_converter_arguments(self): + def _update_converter_arguments(self) -> None: """ Update the converter arguments to the currently selected values. """ logger.debug("Updating Converter cli arguments") if self.converter_arguments is None: @@ -523,31 +601,35 @@ def _update_converter_arguments(self): logger.debug("Updated Converter cli arguments") @staticmethod - def _feed_swapped_faces(patch_queue_in, samples): + def _feed_swapped_faces(patch_queue_in: EventQueue, samples: Samples) -> None: """ Feed swapped faces to the converter's in-queue. Parameters ---------- - patch_queue_in: :class:`queue.Queue` + patch_queue_in: :class:`~lib.queue_manager.EventQueue` The input queue for the patching process samples: :class:`Samples` The Samples for display. """ - logger.trace("feeding swapped faces to converter") + logger.debug("feeding swapped faces to converter") for item in samples.predicted_images: patch_queue_in.put(item) - logger.trace("fed %s swapped faces to converter", len(samples.predicted_images)) - logger.trace("Putting EOF to converter") + logger.debug("fed %s swapped faces to converter", + len(samples.predicted_images)) + logger.debug("Putting EOF to converter") patch_queue_in.put("EOF") - def _patch_faces(self, queue_in, queue_out, sample_size): + def _patch_faces(self, + queue_in: EventQueue, + queue_out: EventQueue, + sample_size: int) -> list[np.ndarray]: """ Patch faces. Run the convert process on the swapped faces and return the patched faces. - patch_queue_in: :class:`queue.Queue` + patch_queue_in: :class:`~lib.queue_manager.EventQueue` The input queue for the patching process - queue_out: :class:`queue.Queue` + queue_out: :class:`~lib.queue_manager.EventQueue` The output queue from the patching process sample_size: int The number of samples to be displayed @@ -557,921 +639,15 @@ def _patch_faces(self, queue_in, queue_out, sample_size): list The swapped faces patched with the selected convert settings """ - logger.trace("Patching faces") + logger.debug("Patching faces") self._converter.process(queue_in, queue_out) swapped = [] idx = 0 while idx < sample_size: - logger.trace("Patching image %s of %s", idx + 1, sample_size) + logger.debug("Patching image %s of %s", idx + 1, sample_size) item = queue_out.get() swapped.append(item[1]) - logger.trace("Patched image %s of %s", idx + 1, sample_size) + logger.debug("Patched image %s of %s", idx + 1, sample_size) idx += 1 - logger.trace("Patched faces") + logger.debug("Patched faces") return swapped - - -class FacesDisplay(): - """ Compiles the 2 rows of sample faces (original and swapped) into a single image - - Parameters - ---------- - size: int - The size of each individual face sample in pixels - padding: int - The amount of extra padding to apply to the outside of the face - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - - Attributes - ---------- - update_source: bool - Flag to indicate that the source images for the preview have been updated, so the preview - should be recompiled. - source: list - The list of :class:`numpy.ndarray` source preview images for top row of display - destination: list - The list of :class:`numpy.ndarray` swapped and patched preview images for bottom row of - display - """ - def __init__(self, size, padding, tk_vars): - logger.trace("Initializing %s: (size: %s, padding: %s, tk_vars: %s)", - self.__class__.__name__, size, padding, tk_vars) - self._size = size - self._display_dims = (1, 1) - self._tk_vars = tk_vars - self._padding = padding - - self._faces = {} - self._centering = None - self._faces_source = None - self._faces_dest = None - self._tk_image = None - - # Set from Samples - self.update_source = False - self.source = [] # Source images, filenames + detected faces - # Set from Patch - self.destination = [] # Swapped + patched images - - logger.trace("Initialized %s", self.__class__.__name__) - - @property - def tk_image(self): - """ :class:`PIL.ImageTk.PhotoImage`: The compiled preview display in tkinter display - format """ - return self._tk_image - - @property - def _total_columns(self): - """ Return the total number of images that are being displayed """ - return len(self.source) - - def set_centering(self, centering): - """ The centering that the model uses is not known at initialization time. - Set :attr:`_centering` when the model has been loaded. - - Parameters - ---------- - centering: str - The centering that the model was trained on - """ - self._centering = centering - - def set_display_dimensions(self, dimensions): - """ Adjust the size of the frame that will hold the preview samples. - - Parameters - ---------- - dimensions: tuple - The (`width`, `height`) of the frame that holds the preview - """ - self._display_dims = dimensions - - def update_tk_image(self): - """ Build the full preview images and compile :attr:`tk_image` for display. """ - logger.trace("Updating tk image") - self._build_faces_image() - img = np.vstack((self._faces_source, self._faces_dest)) - size = self._get_scale_size(img) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - img = Image.fromarray(img) - img = img.resize(size, Image.ANTIALIAS) - self._tk_image = ImageTk.PhotoImage(img) - self._tk_vars["refresh"].set(False) - logger.trace("Updated tk image") - - def _get_scale_size(self, image): - """ Get the size that the full preview image should be resized to fit in the - display window. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The full sized compiled preview image - - Returns - ------- - tuple - The (`width`, `height`) that the display image should be sized to fit in the display - window - """ - frameratio = float(self._display_dims[0]) / float(self._display_dims[1]) - imgratio = float(image.shape[1]) / float(image.shape[0]) - - if frameratio <= imgratio: - scale = self._display_dims[0] / float(image.shape[1]) - size = (self._display_dims[0], max(1, int(image.shape[0] * scale))) - else: - scale = self._display_dims[1] / float(image.shape[0]) - size = (max(1, int(image.shape[1] * scale)), self._display_dims[1]) - logger.trace("scale: %s, size: %s", scale, size) - return size - - def _build_faces_image(self): - """ Compile the source and destination rows of the preview image. """ - logger.trace("Building Faces Image") - update_all = self.update_source - self._faces_from_frames() - if update_all: - header = self._header_text() - source = np.hstack([self._draw_rect(face) for face in self._faces["src"]]) - self._faces_source = np.vstack((header, source)) - self._faces_dest = np.hstack([self._draw_rect(face) for face in self._faces["dst"]]) - logger.debug("source row shape: %s, swapped row shape: %s", - self._faces_dest.shape, self._faces_source.shape) - - def _faces_from_frames(self): - """ Extract the preview faces from the source frames and apply the requisite padding. """ - logger.debug("Extracting faces from frames: Number images: %s", len(self.source)) - if self.update_source: - self._crop_source_faces() - self._crop_destination_faces() - logger.debug("Extracted faces from frames: %s", - {k: len(v) for k, v in self._faces.items()}) - - def _crop_source_faces(self): - """ Extract the source faces from the source frames, along with their filenames and the - transformation matrix used to extract the faces. """ - logger.debug("Updating source faces") - self._faces = {} - for image in self.source: - detected_face = image["detected_faces"][0] - src_img = image["image"] - detected_face.load_aligned(src_img, size=self._size, centering=self._centering) - matrix = detected_face.aligned.matrix - self._faces.setdefault("filenames", - []).append(os.path.splitext(image["filename"])[0]) - self._faces.setdefault("matrix", []).append(matrix) - self._faces.setdefault("src", []).append(transform_image(src_img, - matrix, - self._size, - self._padding)) - self.update_source = False - logger.debug("Updated source faces") - - def _crop_destination_faces(self): - """ Extract the swapped faces from the swapped frames using the source face destination - matrices. """ - logger.debug("Updating destination faces") - self._faces["dst"] = [] - destination = self.destination if self.destination else [np.ones_like(src["image"]) - for src in self.source] - for idx, image in enumerate(destination): - self._faces["dst"].append(transform_image(image, - self._faces["matrix"][idx], - self._size, - self._padding)) - logger.debug("Updated destination faces") - - def _header_text(self): - """ Create the header text displaying the frame name for each preview column. - - Returns - ------- - :class:`numpy.ndarray` - The header row of the preview image containing the frame names for each column - """ - font_scale = self._size / 640 - height = self._size // 8 - font = cv2.FONT_HERSHEY_SIMPLEX - # Get size of placed text for positioning - text_sizes = [cv2.getTextSize(self._faces["filenames"][idx], - font, - font_scale, - 1)[0] - for idx in range(self._total_columns)] - # Get X and Y co-ordinates for each text item - text_y = int((height + text_sizes[0][1]) / 2) - text_x = [int((self._size - text_sizes[idx][0]) / 2) + self._size * idx - for idx in range(self._total_columns)] - logger.debug("filenames: %s, text_sizes: %s, text_x: %s, text_y: %s", - self._faces["filenames"], text_sizes, text_x, text_y) - header_box = np.ones((height, self._size * self._total_columns, 3), np.uint8) * 255 - for idx, text in enumerate(self._faces["filenames"]): - cv2.putText(header_box, - text, - (text_x[idx], text_y), - font, - font_scale, - (0, 0, 0), - 1, - lineType=cv2.LINE_AA) - logger.debug("header_box.shape: %s", header_box.shape) - return header_box - - def _draw_rect(self, image): - """ Place a white border around a given image. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The image to place a border on to - Returns - ------- - :class:`numpy.ndarray` - The given image with a border drawn around the outside - """ - cv2.rectangle(image, (0, 0), (self._size - 1, self._size - 1), (255, 255, 255), 1) - image = np.clip(image, 0.0, 255.0) - return image.astype("uint8") - - -class ConfigTools(): - """ Tools for loading, saving, setting and retrieving configuration file values. - - Attributes - ---------- - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self): - self._config = Config(None) - self.tk_vars = {} - self._config_dicts = self._get_config_dicts() # Holds currently saved config - - @property - def config(self): - """ :class:`plugins.convert._config.Config` The convert configuration """ - return self._config - - @property - def config_dicts(self): - """ dict: The convert configuration options in dictionary form.""" - return self._config_dicts - - @property - def sections(self): - """ list: The sorted section names that exist within the convert Configuration options. """ - return sorted(set(plugin.split(".")[0] for plugin in self._config.config.sections() - if plugin.split(".")[0] != "writer")) - - @property - def plugins_dict(self): - """ dict: Dictionary of configuration option sections as key with a list of containing - plugins as the value """ - return {section: sorted([plugin.split(".")[1] for plugin in self._config.config.sections() - if plugin.split(".")[0] == section]) - for section in self.sections} - - def update_config(self): - """ Update :attr:`config` with the currently selected values from the GUI. """ - for section, items in self.tk_vars.items(): - for item, value in items.items(): - try: - new_value = str(value.get()) - except tk.TclError as err: - # When manually filling in text fields, blank values will - # raise an error on numeric data types so return 0 - logger.debug("Error getting value. Defaulting to 0. Error: %s", str(err)) - new_value = str(0) - old_value = self._config.config[section][item] - if new_value != old_value: - logger.trace("Updating config: %s, %s from %s to %s", - section, item, old_value, new_value) - self._config.config[section][item] = new_value - - def _get_config_dicts(self): - """ Obtain a custom configuration dictionary for convert configuration items in use - by the preview tool formatted for control helper. - - Returns - ------- - dict - Each configuration section as keys, with the values as a dict of option: - :class:`lib.gui.control_helper.ControlOption` pairs. """ - logger.debug("Formatting Config for GUI") - config_dicts = {} - for section in self._config.config.sections(): - if section.startswith("writer."): - continue - for key, val in self._config.defaults[section].items(): - if key == "helptext": - config_dicts.setdefault(section, {})[key] = val - continue - cp_option = ControlPanelOption(title=key, - dtype=val["type"], - group=val["group"], - default=val["default"], - initial_value=self._config.get(section, key), - choices=val["choices"], - is_radio=val["gui_radio"], - rounding=val["rounding"], - min_max=val["min_max"], - helptext=val["helptext"]) - self.tk_vars.setdefault(section, {})[key] = cp_option.tk_var - config_dicts.setdefault(section, {})[key] = cp_option - logger.debug("Formatted Config for GUI: %s", config_dicts) - return config_dicts - - def reset_config_to_saved(self, section=None): - """ Reset the GUI parameters to their saved values within the configuration file. - - Parameters - ---------- - section: str, optional - The configuration section to reset the values for, If ``None`` provided then all - sections are reset. Default: ``None`` - """ - logger.debug("Resetting to saved config: %s", section) - sections = [section] if section is not None else list(self.tk_vars.keys()) - for config_section in sections: - for item, options in self._config_dicts[config_section].items(): - if item == "helptext": - continue - val = options.value - if val != self.tk_vars[config_section][item].get(): - self.tk_vars[config_section][item].set(val) - logger.debug("Setting %s - %s to saved value %s", config_section, item, val) - logger.debug("Reset to saved config: %s", section) - - def reset_config_to_default(self, section=None): - """ Reset the GUI parameters to their default configuration values. - - Parameters - ---------- - section: str, optional - The configuration section to reset the values for, If ``None`` provided then all - sections are reset. Default: ``None`` - """ - logger.debug("Resetting to default: %s", section) - sections = [section] if section is not None else list(self.tk_vars.keys()) - for config_section in sections: - for item, options in self._config_dicts[config_section].items(): - if item == "helptext": - continue - default = options.default - if default != self.tk_vars[config_section][item].get(): - self.tk_vars[config_section][item].set(default) - logger.debug("Setting %s - %s to default value %s", - config_section, item, default) - logger.debug("Reset to default: %s", section) - - def save_config(self, section=None): - """ Save the configuration ``.ini`` file with the currently stored values. - - Notes - ----- - We cannot edit the existing saved config as comments tend to get removed, so we create - a new config and populate that. - - Parameters - ---------- - section: str, optional - The configuration section to save, If ``None`` provided then all sections are saved. - Default: ``None`` - """ - logger.debug("Saving %s config", section) - - new_config = ConfigParser(allow_no_value=True) - - for config_section, items in self._config.defaults.items(): - logger.debug("Adding section: '%s')", config_section) - self._config.insert_config_section(config_section, - items["helptext"], - config=new_config) - for item, options in items.items(): - if item == "helptext": - continue # helptext already written at top - if ((section is not None and config_section != section) - or config_section not in self.tk_vars): - # retain saved values that have not been updated - new_opt = self._config.get(config_section, item) - logger.debug("Retaining option: (item: '%s', value: '%s')", item, new_opt) - else: - new_opt = self.tk_vars[config_section][item].get() - logger.debug("Setting option: (item: '%s', value: '%s')", item, new_opt) - - # Set config_dicts value to new saved value - self._config_dicts[config_section][item].set_initial_value(new_opt) - - helptext = self._config.format_help(options["helptext"], is_section=False) - new_config.set(config_section, helptext) - new_config.set(config_section, item, str(new_opt)) - - self._config.config = new_config - self._config.save_config() - logger.info("Saved config: '%s'", self._config.configfile) - - -class ImagesCanvas(ttk.Frame): # pylint:disable=too-many-ancestors - """ tkinter Canvas that holds the preview images. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the canvas - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self, parent, tk_vars): - logger.debug("Initializing %s: (parent: %s, tk_vars: %s)", - self.__class__.__name__, parent, tk_vars) - super().__init__(parent) - self.pack(expand=True, fill=tk.BOTH, padx=2, pady=2) - - self._refresh_display_trigger = tk_vars["refresh"] - self._refresh_display_trigger.trace("w", self._refresh_display_callback) - self._display = parent.preview_display - self._canvas = tk.Canvas(self, bd=0, highlightthickness=0) - self._canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - self._displaycanvas = self._canvas.create_image(0, 0, - image=self._display.tk_image, - anchor=tk.NW) - self.bind("", self._resize) - logger.debug("Initialized %s", self.__class__.__name__) - - def _refresh_display_callback(self, *args): - """ Add a trace to refresh display on callback """ - if not self._refresh_display_trigger.get(): - return - logger.trace("Refresh display trigger received: %s", args) - self._reload() - - def _resize(self, event): - """ Resize the image to fit the frame, maintaining aspect ratio """ - logger.trace("Resizing preview image") - framesize = (event.width, event.height) - self._display.set_display_dimensions(framesize) - self._reload() - - def _reload(self): - """ Reload the preview image """ - logger.trace("Reloading preview image") - self._display.update_tk_image() - self._canvas.itemconfig(self._displaycanvas, image=self._display.tk_image) - - -class ActionFrame(ttk.Frame): # pylint: disable=too-many-ancestors - """ Frame that holds the left hand side options panel containing the command line options. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the Action Frame - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - selected_color: str - The selected color adjustment type - selected_mask_type: str - The selected mask type - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - patch_callback: python function - The function to execute when a patch callback is received - refresh_callback: python function - The function to execute when a refresh callback is received - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self, parent, available_masks, has_predicted_mask, selected_color, - selected_mask_type, config_tools, patch_callback, refresh_callback, tk_vars): - logger.debug("Initializing %s: (available_masks: %s, has_predicted_mask: %s, " - "selected_color: %s, selected_mask_type: %s, patch_callback: %s, " - "refresh_callback: %s, tk_vars: %s)", - self.__class__.__name__, available_masks, has_predicted_mask, selected_color, - selected_mask_type, patch_callback, refresh_callback, tk_vars) - self._config_tools = config_tools - - super().__init__(parent) - self.pack(side=tk.LEFT, anchor=tk.N, fill=tk.Y) - self._options = ["color", "mask_type"] - self._busy_tkvar = tk_vars["busy"] - self._tk_vars = {} - - d_locals = locals() - defaults = {opt: self._format_to_display(d_locals[f"selected_{opt}"]) - for opt in self._options} - self._busy_indicator = self._build_frame(defaults, - refresh_callback, - patch_callback, - available_masks, - has_predicted_mask) - - @property - def convert_args(self): - """ dict: Currently selected Command line arguments from the :class:`ActionFrame`. """ - return {opt if opt != "color" else "color_adjustment": - self._format_from_display(self._tk_vars[opt].get()) - for opt in self._options} - - @staticmethod - def _format_from_display(var): - """ Format a variable from the display version to the command line action version. - - Parameters - ---------- - var: str - The variable name to format - - Returns - ------- - str - The formatted variable name - """ - return var.replace(" ", "_").lower() - - @staticmethod - def _format_to_display(var): - """ Format a variable from the command line action version to the display version. - Parameters - ---------- - var: str - The variable name to format - - Returns - ------- - str - The formatted variable name - """ - return var.replace("_", " ").replace("-", " ").title() - - def _build_frame(self, defaults, refresh_callback, patch_callback, - available_masks, has_predicted_mask): - """ Build the :class:`ActionFrame`. - - Parameters - ---------- - defaults: dict - The default command line options - patch_callback: python function - The function to execute when a patch callback is received - refresh_callback: python function - The function to execute when a refresh callback is received - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - ttk.Progressbar - A Progress bar to indicate that the Preview tool is busy - """ - logger.debug("Building Action frame") - - bottom_frame = ttk.Frame(self) - bottom_frame.pack(side=tk.BOTTOM, fill=tk.X, anchor=tk.S) - top_frame = ttk.Frame(self) - top_frame.pack(side=tk.TOP, fill=tk.BOTH, anchor=tk.N, expand=True) - - self._add_cli_choices(top_frame, defaults, available_masks, has_predicted_mask) - - busy_indicator = self._add_busy_indicator(bottom_frame) - self._add_refresh_button(bottom_frame, refresh_callback) - self._add_patch_callback(patch_callback) - self._add_actions(bottom_frame) - logger.debug("Built Action frame") - return busy_indicator - - def _add_cli_choices(self, parent, defaults, available_masks, has_predicted_mask): - """ Create :class:`lib.gui.control_helper.ControlPanel` object for the command - line options. - - parent: :class:`ttk.Frame` - The frame to hold the command line choices - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - """ - cp_options = self._get_control_panel_options(defaults, available_masks, has_predicted_mask) - panel_kwargs = dict(blank_nones=False, label_width=10, style="CPanel") - ControlPanel(parent, cp_options, header_text=None, **panel_kwargs) - - def _get_control_panel_options(self, defaults, available_masks, has_predicted_mask): - """ Create :class:`lib.gui.control_helper.ControlPanelOption` objects for the command - line options. - - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - list - The list of `lib.gui.control_helper.ControlPanelOption` objects for the Action Frame - """ - cp_options = [] - for opt in self._options: - if opt == "mask_type": - choices = self._create_mask_choices(defaults, available_masks, has_predicted_mask) - else: - choices = PluginLoader.get_available_convert_plugins(opt, True) - cp_option = ControlPanelOption(title=opt, - dtype=str, - default=defaults[opt], - initial_value=defaults[opt], - choices=choices, - group="Command Line Choices", - is_radio=False) - self._tk_vars[opt] = cp_option.tk_var - cp_options.append(cp_option) - return cp_options - - @staticmethod - def _create_mask_choices(defaults, available_masks, has_predicted_mask): - """ Set the mask choices and default mask based on available masks. - - Parameters - ---------- - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - list - The masks that are available to use from the alignments file - """ - logger.debug("Initial mask choices: %s", available_masks) - if has_predicted_mask: - available_masks += ["predicted"] - if "none" not in available_masks: - available_masks += ["none"] - if defaults["mask_type"] not in available_masks: - logger.debug("Setting default mask to first available: %s", available_masks[0]) - defaults["mask_type"] = available_masks[0] - logger.debug("Final mask choices: %s", available_masks) - return available_masks - - @staticmethod - def _add_refresh_button(parent, refresh_callback): - """ Add a button to refresh the images. - - Parameters - ---------- - refresh_callback: python function - The function to execute when the refresh button is pressed - """ - btn = ttk.Button(parent, text="Update Samples", command=refresh_callback) - btn.pack(padx=5, pady=5, side=tk.TOP, fill=tk.X, anchor=tk.N) - - def _add_patch_callback(self, patch_callback): - """ Add callback to re-patch images on action option change. - - Parameters - ---------- - patch_callback: python function - The function to execute when the images require patching - """ - for tk_var in self._tk_vars.values(): - tk_var.trace("w", patch_callback) - - def _add_busy_indicator(self, parent): - """ Place progress bar into bottom bar to indicate when processing. - - Parameters - ---------- - parent: tkinter object - The tkinter object that holds the busy indicator - - Returns - ------- - ttk.Progressbar - A Progress bar to indicate that the Preview tool is busy - """ - logger.debug("Placing busy indicator") - pbar = ttk.Progressbar(parent, mode="indeterminate") - pbar.pack(side=tk.LEFT) - pbar.pack_forget() - self._busy_tkvar.trace("w", self._busy_indicator_trace) - return pbar - - def _busy_indicator_trace(self, *args): - """ Show or hide busy indicator based on whether the preview is updating. - - Parameters - ---------- - args: unused - Required for tkinter event, but unused - """ - logger.trace("Busy indicator trace: %s", args) - if self._busy_tkvar.get(): - self._start_busy_indicator() - else: - self._stop_busy_indicator() - - def _stop_busy_indicator(self): - """ Stop and hide progress bar """ - logger.debug("Stopping busy indicator") - self._busy_indicator.stop() - self._busy_indicator.pack_forget() - - def _start_busy_indicator(self): - """ Start and display progress bar """ - logger.debug("Starting busy indicator") - self._busy_indicator.pack(side=tk.LEFT, padx=5, pady=(5, 10), fill=tk.X, expand=True) - self._busy_indicator.start() - - def _add_actions(self, parent): - """ Add Action Buttons to the :class:`ActionFrame` - - Parameters - ---------- - parent: tkinter object - The tkinter object that holds the action buttons - """ - logger.debug("Adding util buttons") - frame = ttk.Frame(parent) - frame.pack(padx=5, pady=(5, 10), side=tk.RIGHT, fill=tk.X, anchor=tk.E) - - for utl in ("save", "clear", "reload"): - logger.debug("Adding button: '%s'", utl) - img = get_images().icons[utl] - if utl == "save": - text = _("Save full config") - action = self._config_tools.save_config - elif utl == "clear": - text = _("Reset full config to default values") - action = self._config_tools.reset_config_to_default - elif utl == "reload": - text = _("Reset full config to saved values") - action = self._config_tools.reset_config_to_saved - - btnutl = ttk.Button(frame, - image=img, - command=action) - btnutl.pack(padx=2, side=tk.RIGHT) - Tooltip(btnutl, text=text, wrap_length=200) - logger.debug("Added util buttons") - - -class OptionsBook(ttk.Notebook): # pylint:disable=too-many-ancestors - """ The notebook that holds the Convert configuration options. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the Options book - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - patch_callback: python function - The function to execute when a patch callback is received - - Attributes - ---------- - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - """ - def __init__(self, parent, config_tools, patch_callback): - logger.debug("Initializing %s: (parent: %s, config: %s)", - self.__class__.__name__, parent, config_tools) - super().__init__(parent) - self.pack(side=tk.RIGHT, anchor=tk.N, fill=tk.BOTH, expand=True) - self.config_tools = config_tools - - self._tabs = {} - self._build_tabs() - self._build_sub_tabs() - self._add_patch_callback(patch_callback) - logger.debug("Initialized %s", self.__class__.__name__) - - def _build_tabs(self): - """ Build the notebook tabs for the each configuration section. """ - logger.debug("Build Tabs") - for section in self.config_tools.sections: - tab = ttk.Notebook(self) - self._tabs[section] = {"tab": tab} - self.add(tab, text=section.replace("_", " ").title()) - - def _build_sub_tabs(self): - """ Build the notebook sub tabs for each convert section's plugin. """ - for section, plugins in self.config_tools.plugins_dict.items(): - for plugin in plugins: - config_key = ".".join((section, plugin)) - config_dict = self.config_tools.config_dicts[config_key] - tab = ConfigFrame(self, config_key, config_dict) - self._tabs[section][plugin] = tab - self._tabs[section]["tab"].add(tab, text=plugin.replace("_", " ").title()) - - def _add_patch_callback(self, patch_callback): - """ Add callback to re-patch images on configuration option change. - - Parameters - ---------- - patch_callback: python function - The function to execute when the images require patching - """ - for plugins in self.config_tools.tk_vars.values(): - for tk_var in plugins.values(): - tk_var.trace("w", patch_callback) - - -class ConfigFrame(ttk.Frame): # pylint: disable=too-many-ancestors - """ Holds the configuration options for a convert plugin inside the :class:`OptionsBook`. - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - options: dict - The options for this section/plugin - """ - - def __init__(self, parent, config_key, options): - logger.debug("Initializing %s", self.__class__.__name__) - super().__init__(parent) - self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - - self._options = options - - self._action_frame = ttk.Frame(self) - self._action_frame.pack(padx=0, pady=(0, 5), side=tk.BOTTOM, fill=tk.X, anchor=tk.E) - self._add_frame_separator() - - self._build_frame(parent, config_key) - logger.debug("Initialized %s", self.__class__.__name__) - - def _build_frame(self, parent, config_key): - """ Build the options frame for this command - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - """ - logger.debug("Add Config Frame") - panel_kwargs = dict(columns=2, option_columns=2, blank_nones=False, style="CPanel") - frame = ttk.Frame(self) - frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - cp_options = [opt for key, opt in self._options.items() if key != "helptext"] - ControlPanel(frame, cp_options, header_text=None, **panel_kwargs) - self._add_actions(parent, config_key) - logger.debug("Added Config Frame") - - def _add_frame_separator(self): - """ Add a separator between top and bottom frames. """ - logger.debug("Add frame seperator") - sep = ttk.Frame(self._action_frame, height=2, relief=tk.RIDGE) - sep.pack(fill=tk.X, pady=5, side=tk.TOP) - logger.debug("Added frame seperator") - - def _add_actions(self, parent, config_key): - """ Add Action Buttons. - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - """ - logger.debug("Adding util buttons") - - title = config_key.split(".")[1].replace("_", " ").title() - btn_frame = ttk.Frame(self._action_frame) - btn_frame.pack(padx=5, side=tk.BOTTOM, fill=tk.X) - for utl in ("save", "clear", "reload"): - logger.debug("Adding button: '%s'", utl) - img = get_images().icons[utl] - if utl == "save": - text = _(f"Save {title} config") - action = parent.config_tools.save_config - elif utl == "clear": - text = _(f"Reset {title} config to default values") - action = parent.config_tools.reset_config_to_default - elif utl == "reload": - text = _(f"Reset {title} config to saved values") - action = parent.config_tools.reset_config_to_saved - - btnutl = ttk.Button(btn_frame, - image=img, - command=lambda cmd=action: cmd(config_key)) - btnutl.pack(padx=2, side=tk.RIGHT) - Tooltip(btnutl, text=text, wrap_length=200) - logger.debug("Added util buttons") diff --git a/tools/preview/viewer.py b/tools/preview/viewer.py new file mode 100644 index 0000000000..7abe11b96d --- /dev/null +++ b/tools/preview/viewer.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python3 +""" Manages the widgets that hold the top 'viewer' area of the preview tool """ +from __future__ import annotations +import logging +import os +import tkinter as tk +import typing as T + +from tkinter import ttk +from dataclasses import dataclass, field + +import cv2 +import numpy as np +from PIL import Image, ImageTk + +from lib.align import transform_image +from lib.align.aligned_face import CenteringType +from scripts.convert import ConvertItem + + +if T.TYPE_CHECKING: + from .preview import Preview + +logger = logging.getLogger(__name__) + + +@dataclass +class _Faces: + """ Dataclass for holding faces """ + filenames: list[str] = field(default_factory=list) + matrix: list[np.ndarray] = field(default_factory=list) + src: list[np.ndarray] = field(default_factory=list) + dst: list[np.ndarray] = field(default_factory=list) + + +class FacesDisplay(): + """ Compiles the 2 rows of sample faces (original and swapped) into a single image + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + size: int + The size of each individual face sample in pixels + padding: int + The amount of extra padding to apply to the outside of the face + + Attributes + ---------- + update_source: bool + Flag to indicate that the source images for the preview have been updated, so the preview + should be recompiled. + source: list + The list of :class:`numpy.ndarray` source preview images for top row of display + destination: list + The list of :class:`numpy.ndarray` swapped and patched preview images for bottom row of + display + """ + def __init__(self, app: Preview, size: int, padding: int) -> None: + logger.trace("Initializing %s: (app: %s, size: %s, padding: %s)", # type: ignore + self.__class__.__name__, app, size, padding) + self._size = size + self._display_dims = (1, 1) + self._app = app + self._padding = padding + + self._faces = _Faces() + self._centering: CenteringType | None = None + self._faces_source: np.ndarray = np.array([]) + self._faces_dest: np.ndarray = np.array([]) + self._tk_image: ImageTk.PhotoImage | None = None + + # Set from Samples + self.update_source = False + self.source: list[ConvertItem] = [] # Source images, filenames + detected faces + # Set from Patch + self.destination: list[np.ndarray] = [] # Swapped + patched images + + logger.trace("Initialized %s", self.__class__.__name__) # type: ignore + + @property + def tk_image(self) -> ImageTk.PhotoImage | None: + """ :class:`PIL.ImageTk.PhotoImage`: The compiled preview display in tkinter display + format """ + return self._tk_image + + @property + def _total_columns(self) -> int: + """ int: The total number of images that are being displayed """ + return len(self.source) + + def set_centering(self, centering: CenteringType) -> None: + """ The centering that the model uses is not known at initialization time. + Set :attr:`_centering` when the model has been loaded. + + Parameters + ---------- + centering: str + The centering that the model was trained on + """ + self._centering = centering + + def set_display_dimensions(self, dimensions: tuple[int, int]) -> None: + """ Adjust the size of the frame that will hold the preview samples. + + Parameters + ---------- + dimensions: tuple + The (`width`, `height`) of the frame that holds the preview + """ + self._display_dims = dimensions + + def update_tk_image(self) -> None: + """ Build the full preview images and compile :attr:`tk_image` for display. """ + logger.trace("Updating tk image") # type: ignore + self._build_faces_image() + img = np.vstack((self._faces_source, self._faces_dest)) + size = self._get_scale_size(img) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + pilimg = Image.fromarray(img) + pilimg = pilimg.resize(size, Image.ANTIALIAS) + self._tk_image = ImageTk.PhotoImage(pilimg) + logger.trace("Updated tk image") # type: ignore + + def _get_scale_size(self, image: np.ndarray) -> tuple[int, int]: + """ Get the size that the full preview image should be resized to fit in the + display window. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The full sized compiled preview image + + Returns + ------- + tuple + The (`width`, `height`) that the display image should be sized to fit in the display + window + """ + frameratio = float(self._display_dims[0]) / float(self._display_dims[1]) + imgratio = float(image.shape[1]) / float(image.shape[0]) + + if frameratio <= imgratio: + scale = self._display_dims[0] / float(image.shape[1]) + size = (self._display_dims[0], max(1, int(image.shape[0] * scale))) + else: + scale = self._display_dims[1] / float(image.shape[0]) + size = (max(1, int(image.shape[1] * scale)), self._display_dims[1]) + logger.trace("scale: %s, size: %s", scale, size) # type: ignore + return size + + def _build_faces_image(self) -> None: + """ Compile the source and destination rows of the preview image. """ + logger.trace("Building Faces Image") # type: ignore + update_all = self.update_source + self._faces_from_frames() + if update_all: + header = self._header_text() + source = np.hstack([self._draw_rect(face) for face in self._faces.src]) + self._faces_source = np.vstack((header, source)) + self._faces_dest = np.hstack([self._draw_rect(face) for face in self._faces.dst]) + logger.debug("source row shape: %s, swapped row shape: %s", + self._faces_dest.shape, self._faces_source.shape) + + def _faces_from_frames(self) -> None: + """ Extract the preview faces from the source frames and apply the requisite padding. """ + logger.debug("Extracting faces from frames: Number images: %s", len(self.source)) + if self.update_source: + self._crop_source_faces() + self._crop_destination_faces() + logger.debug("Extracted faces from frames: %s", + {k: len(v) for k, v in self._faces.__dict__.items()}) + + def _crop_source_faces(self) -> None: + """ Extract the source faces from the source frames, along with their filenames and the + transformation matrix used to extract the faces. """ + logger.debug("Updating source faces") + self._faces = _Faces() # Init new class + for item in self.source: + detected_face = item.inbound.detected_faces[0] + src_img = item.inbound.image + detected_face.load_aligned(src_img, + size=self._size, + centering=T.cast(CenteringType, self._centering)) + matrix = detected_face.aligned.matrix + self._faces.filenames.append(os.path.splitext(item.inbound.filename)[0]) + self._faces.matrix.append(matrix) + self._faces.src.append(transform_image(src_img, matrix, self._size, self._padding)) + self.update_source = False + logger.debug("Updated source faces") + + def _crop_destination_faces(self) -> None: + """ Extract the swapped faces from the swapped frames using the source face destination + matrices. """ + logger.debug("Updating destination faces") + self._faces.dst = [] + destination = self.destination if self.destination else [np.ones_like(src.inbound.image) + for src in self.source] + for idx, image in enumerate(destination): + self._faces.dst.append(transform_image(image, + self._faces.matrix[idx], + self._size, + self._padding)) + logger.debug("Updated destination faces") + + def _header_text(self) -> np.ndarray: + """ Create the header text displaying the frame name for each preview column. + + Returns + ------- + :class:`numpy.ndarray` + The header row of the preview image containing the frame names for each column + """ + font_scale = self._size / 640 + height = self._size // 8 + font = cv2.FONT_HERSHEY_SIMPLEX + # Get size of placed text for positioning + text_sizes = [cv2.getTextSize(self._faces.filenames[idx], + font, + font_scale, + 1)[0] + for idx in range(self._total_columns)] + # Get X and Y co-ordinates for each text item + text_y = int((height + text_sizes[0][1]) / 2) + text_x = [int((self._size - text_sizes[idx][0]) / 2) + self._size * idx + for idx in range(self._total_columns)] + logger.debug("filenames: %s, text_sizes: %s, text_x: %s, text_y: %s", + self._faces.filenames, text_sizes, text_x, text_y) + header_box = np.ones((height, self._size * self._total_columns, 3), np.uint8) * 255 + for idx, text in enumerate(self._faces.filenames): + cv2.putText(header_box, + text, + (text_x[idx], text_y), + font, + font_scale, + (0, 0, 0), + 1, + lineType=cv2.LINE_AA) + logger.debug("header_box.shape: %s", header_box.shape) + return header_box + + def _draw_rect(self, image: np.ndarray) -> np.ndarray: + """ Place a white border around a given image. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to place a border on to + Returns + ------- + :class:`numpy.ndarray` + The given image with a border drawn around the outside + """ + cv2.rectangle(image, (0, 0), (self._size - 1, self._size - 1), (255, 255, 255), 1) + image = np.clip(image, 0.0, 255.0) + return image.astype("uint8") + + +class ImagesCanvas(ttk.Frame): # pylint:disable=too-many-ancestors + """ tkinter Canvas that holds the preview images. + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + parent: tkinter object + The parent tkinter object that holds the canvas + """ + def __init__(self, app: Preview, parent: ttk.PanedWindow) -> None: + logger.debug("Initializing %s: (app: %s, parent: %s)", + self.__class__.__name__, app, parent) + super().__init__(parent) + self.pack(expand=True, fill=tk.BOTH, padx=2, pady=2) + + self._display: FacesDisplay = parent.preview_display # type: ignore + self._canvas = tk.Canvas(self, bd=0, highlightthickness=0) + self._canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + self._displaycanvas = self._canvas.create_image(0, 0, + image=self._display.tk_image, + anchor=tk.NW) + self.bind("", self._resize) + logger.debug("Initialized %s", self.__class__.__name__) + + def _resize(self, event: tk.Event) -> None: + """ Resize the image to fit the frame, maintaining aspect ratio """ + logger.debug("Resizing preview image") + framesize = (event.width, event.height) + self._display.set_display_dimensions(framesize) + self.reload() + + def reload(self) -> None: + """ Update the images in the canvas and redraw """ + logger.debug("Reloading preview image") + self._display.update_tk_image() + self._canvas.itemconfig(self._displaycanvas, image=self._display.tk_image) + logger.debug("Reloaded preview image") diff --git a/tools/restore/cli.py b/tools/restore/cli.py deleted file mode 100644 index fe3ca9aa77..0000000000 --- a/tools/restore/cli.py +++ /dev/null @@ -1,35 +0,0 @@ -#!/usr/bin/env python3 -""" Command Line Arguments for tools """ -import gettext - -from lib.cli.args import FaceSwapArgs -from lib.cli.actions import DirFullPaths - - -# LOCALES -_LANG = gettext.translation("tools.restore.cli", localedir="locales", fallback=True) -_ = _LANG.gettext - -_HELPTEXT = _("This command lets you restore models from backup.") - - -class RestoreArgs(FaceSwapArgs): - """ Class to restore model files from backup """ - - @staticmethod - def get_info(): - """ Return command information """ - return _("A tool for restoring models from backup (.bk) files") - - @staticmethod - def get_argument_list(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ - argument_list = list() - argument_list.append(dict( - opts=("-m", "--model-dir"), - action=DirFullPaths, - dest="model_dir", - required=True, - help=_("Model directory. A directory containing the model you wish to restore from " - "backup."))) - return argument_list diff --git a/tools/restore/restore.py b/tools/restore/restore.py deleted file mode 100644 index 497a254d51..0000000000 --- a/tools/restore/restore.py +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env python3 -""" Tool to restore models from backup """ - -import logging -import os -import sys - -from lib.model.backup_restore import Backup - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - - -class Restore(): - """ Restore a model from backup """ - - def __init__(self, arguments): - logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments) - self.model_dir = arguments.model_dir - self.model_name = None - - def process(self): - """ Perform the Restore process """ - logger.info("Starting Model Restore...") - self.validate() - backup = Backup(self.model_dir, self.model_name) - backup.restore() - logger.info("Completed Model Restore") - - def validate(self): - """ Make sure there is only one model in the target folder """ - if not os.path.exists(self.model_dir): - logger.error("Folder does not exist: '%s'", self.model_dir) - sys.exit(1) - chkfiles = [fname for fname in os.listdir(self.model_dir) if fname.endswith("_state.json")] - bkfiles = [fname for fname in os.listdir(self.model_dir) if fname.endswith(".bk")] - if not chkfiles: - logger.error("Could not find a model in the supplied folder: '%s'", self.model_dir) - sys.exit(1) - if len(chkfiles) > 1: - logger.error("More than one model found in the supplied folder: '%s'", self.model_dir) - sys.exit(1) - if not bkfiles: - logger.error("Could not find any backup files in the supplied folder: '%s'", - self.model_dir) - sys.exit(1) - self.model_name = chkfiles[0].replace("_state.json", "") - logger.info("%s Model found", self.model_name.title()) - logger.verbose("Backup files: %s)", bkfiles) diff --git a/tools/sort/cli.py b/tools/sort/cli.py index bde6b9ad8a..d85b9637a5 100644 --- a/tools/sort/cli.py +++ b/tools/sort/cli.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ +import argparse import gettext from lib.cli.args import FaceSwapArgs @@ -12,6 +13,63 @@ _HELPTEXT = _("This command lets you sort images using various methods.") +_SORT_METHODS = ( + "none", "blur", "blur-fft", "distance", "face", "face-cnn", "face-cnn-dissim", + "yaw", "pitch", "roll", "hist", "hist-dissim", "color-black", "color-gray", "color-luma", + "color-green", "color-orange", "size") + +_GPTHRESHOLD = _(" Adjust the '-t' ('--threshold') parameter to control the strength of grouping.") +_GPCOLOR = _(" Adjust the '-b' ('--bins') parameter to control the number of bins for grouping. " + "Each image is allocated to a bin by the percentage of color pixels that appear in " + "the image.") +_GPDEGREES = _(" Adjust the '-b' ('--bins') parameter to control the number of bins for grouping. " + "Each image is allocated to a bin by the number of degrees the face is orientated " + "from center.") +_GPLINEAR = _(" Adjust the '-b' ('--bins') parameter to control the number of bins for grouping. " + "The minimum and maximum values are taken for the chosen sort metric. The bins " + "are then populated with the results from the group sorting.") +_METHOD_TEXT = { + "blur": _("faces by blurriness."), + "blur-fft": _("faces by fft filtered blurriness."), + "distance": _("faces by the estimated distance of the alignments from an 'average' face. This " + "can be useful for eliminating misaligned faces. Sorts from most like an " + "average face to least like an average face."), + "face": _("faces using VGG Face2 by face similarity. This uses a pairwise clustering " + "algorithm to check the distances between 512 features on every face in your set " + "and order them appropriately."), + "face-cnn": _("faces by their landmarks."), + "face-cnn-dissim": _("Like 'face-cnn' but sorts by dissimilarity."), + "yaw": _("faces by Yaw (rotation left to right)."), + "pitch": _("faces by Pitch (rotation up and down)."), + "roll": _("faces by Roll (rotation). Aligned faces should have a roll value close to zero. " + "The further the Roll value from zero the higher liklihood the face is misaligned."), + "hist": _("faces by their color histogram."), + "hist-dissim": _("Like 'hist' but sorts by dissimilarity."), + "color-gray": _("images by the average intensity of the converted grayscale color channel."), + "color-black": _("images by their number of black pixels. Useful when faces are near borders " + "and a large part of the image is black."), + "color-luma": _("images by the average intensity of the converted Y color channel. Bright " + "lighting and oversaturated images will be ranked first."), + "color-green": _("images by the average intensity of the converted Cg color channel. Green " + "images will be ranked first and red images will be last."), + "color-orange": _("images by the average intensity of the converted Co color channel. Orange " + "images will be ranked first and blue images will be last."), + "size": _("images by their size in the original frame. Faces further from the camera and from " + "lower resolution sources will be sorted first, whilst faces closer to the camera " + "and from higher resolution sources will be sorted last.")} + +_BIN_TYPES = [ + (("face", "face-cnn", "face-cnn-dissim", "hist", "hist-dissim"), _GPTHRESHOLD), + (("color-black", "color-gray", "color-luma", "color-green", "color-orange"), _GPCOLOR), + (("yaw", "pitch", "roll"), _GPDEGREES), + (("blur", "blur-fft", "distance", "size"), _GPLINEAR)] +_SORT_HELP = "" +_GROUP_HELP = "" + +for method in sorted(_METHOD_TEXT): + _SORT_HELP += f"\nL|{method}: {_('Sort')} {_METHOD_TEXT[method]}" + _GROUP_HELP += (f"\nL|{method}: {_('Group')} {_METHOD_TEXT[method]} " + f"{next((x[1] for x in _BIN_TYPES if method in x[0]), '')}") class SortArgs(FaceSwapArgs): @@ -25,147 +83,148 @@ def get_info(): @staticmethod def get_argument_list(): """ Put the arguments in a list so that they are accessible from both argparse and gui """ - argument_list = list() - argument_list.append(dict( - opts=('-i', '--input'), - action=DirFullPaths, - dest="input_dir", - group=_("data"), - help=_("Input directory of aligned faces."), - required=True)) - argument_list.append(dict( - opts=('-o', '--output'), - action=DirFullPaths, - dest="output_dir", - group=_("data"), - help=_("Output directory for sorted aligned faces."))) - argument_list.append(dict( - opts=('-s', '--sort-by'), - action=Radio, - type=str, - choices=("blur", "blur-fft", "distance", "face", "face-cnn", "face-cnn-dissim", - "face-yaw", "hist", "hist-dissim", "color-gray", "color-luma", "color-green", - "color-orange", "size", "black-pixels"), - dest='sort_method', - group=_("sort settings"), - default="face", - help=_("R|Sort by method. Choose how images are sorted. " - "\nL|'blur': Sort faces by blurriness." - "\nL|'blur-fft': Sort faces by fft filtered blurriness." - "\nL|'distance' Sort faces by the estimated distance of the alignments from an " - "'average' face. This can be useful for eliminating misaligned faces." - "\nL|'face': Use VGG Face to sort by face similarity. This uses a pairwise " - "clustering algorithm to check the distances between 512 features on every " - "face in your set and order them appropriately." - "\nL|'face-cnn': Sort faces by their landmarks. You can adjust the threshold " - "with the '-t' (--ref_threshold) option." - "\nL|'face-cnn-dissim': Like 'face-cnn' but sorts by dissimilarity." - "\nL|'face-yaw': Sort faces by Yaw (rotation left to right)." - "\nL|'hist': Sort faces by their color histogram. You can adjust the threshold " - "with the '-t' (--ref_threshold) option." - "\nL|'hist-dissim': Like 'hist' but sorts by dissimilarity." - "\nL|'color-gray': Sort images by the average intensity of the converted " - "grayscale color channel." - "\nL|'color-luma': Sort images by the average intensity of the converted Y " - "color channel. Bright lighting and oversaturated images will be ranked first." - "\nL|'color-green': Sort images by the average intensity of the converted Cg " - "color channel. Green images will be ranked first and red images will be last." - "\nL|'color-orange': Sort images by the average intensity of the converted Co " - "color channel. Orange images will be ranked first and blue images will be " - "last." - "\nL|'size': Sort images by their size in the original frame. Faces closer to " - "the camera and from higher resolution sources will be sorted first, whilst " - "faces further from the camera and from lower resolution sources will be " - "sorted last." - "\nL|'black-pixels': Sort images by their number of black pixels. Useful when " - "faces are near borders and a large part of the image is black." - "\nDefault: face"))) - argument_list.append(dict( - opts=('-k', '--keep'), - action='store_true', - dest='keep_original', - default=False, - group=_("output"), - help=_("Keeps the original files in the input directory. Be careful when using this " - "with rename grouping and no specified output directory as this would keep the " - "original and renamed files in the same directory."))) - argument_list.append(dict( - opts=('-t', '--ref_threshold'), - action=Slider, - min_max=(-1.0, 10.0), - rounding=2, - type=float, - dest='min_threshold', - group=_("sort settings"), - default=-1.0, - help=_("Float value. Minimum threshold to use for grouping comparison with 'face-cnn' " - "and 'hist' methods. The lower the value the more discriminating the grouping " - "is. Leaving -1.0 will allow the program set the default value automatically. " - "For face-cnn 7.2 should be enough, with 4 being very discriminating. For hist " - "0.3 should be enough, with 0.2 being very discriminating. Be careful setting " - "a value that's too low in a directory with many images, as this could result " - "in a lot of directories being created. Defaults: face-cnn 7.2, hist 0.3"))) - argument_list.append(dict( - opts=('-fp', '--final-process'), - action=Radio, - type=str, - choices=("folders", "rename"), - dest='final_process', - default="rename", - group=_("output"), - help=_("R|Default: rename." - "\nL|'folders': files are sorted using the -s/--sort-by method, then they are " - "organized into folders using the -g/--group-by grouping method." - "\nL|'rename': files are sorted using the -s/--sort-by then they are " - "renamed."))) - argument_list.append(dict( - opts=('-g', '--group-by'), - action=Radio, - type=str, - choices=("blur", "blur-fft", "face-cnn", "face-yaw", "hist", "black-pixels"), - dest='group_method', - group=_("output"), - default="hist", - help=_("Group by method. When -fp/--final-processing by folders choose the how the " - "images are grouped after sorting. Default: hist"))) - argument_list.append(dict( - opts=('-b', '--bins'), - action=Slider, - min_max=(1, 100), - rounding=1, - type=int, - dest='num_bins', - group=_("output"), - default=5, - help=_("Integer value. Number of folders that will be used to group by blur, " - "face-yaw and black-pixels. For blur folder 0 will be the least blurry, while " - "the last folder will be the blurriest. For face-yaw the number of bins is by " - "how much 180 degrees is divided. So if you use 18, then each folder will be " - "a 10 degree increment. Folder 0 will contain faces looking the most to the " - "left whereas the last folder will contain the faces looking the most to the " - "right. If the number of images doesn't divide evenly into the number of " - "bins, the remaining images get put in the last bin. For black-pixels it " - "represents the divider of the percentage of black pixels. For 10, first " - "folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, " - "etc. Default value: 5"))) - argument_list.append(dict( - opts=('-l', '--log-changes'), - action='store_true', - group=_("settings"), - default=False, - help=_("Logs file renaming changes if grouping by renaming, or it logs the file " - "copying/movement if grouping by folders. If no log file is specified with " - "'--log-file', then a 'sort_log.json' file will be created in the input " - "directory."))) - argument_list.append(dict( - opts=('-lf', '--log-file'), - action=SaveFileFullPaths, - filetypes="alignments", - group=_("settings"), - dest='log_file_path', - default='sort_log.json', - help=_("Specify a log file to use for saving the renaming or grouping information. If " - "specified extension isn't 'json' or 'yaml', then json will be used as the " - "serializer, with the supplied filename. Default: sort_log.json"))) - + argument_list = [] + argument_list.append({ + "opts": ('-i', '--input'), + "action": DirFullPaths, + "dest": "input_dir", + "group": _("data"), + "help": _("Input directory of aligned faces."), + "required": True}) + argument_list.append({ + "opts": ('-o', '--output'), + "action": DirFullPaths, + "dest": "output_dir", + "group": _("data"), + "help": _( + "Output directory for sorted aligned faces. If not provided and 'keep' is " + "selected then a new folder called 'sorted' will be created within the input " + "folder to house the output. If not provided and 'keep' is not selected then the " + "images will be sorted in-place, overwriting the original contents of the " + "'input_dir'")}) + argument_list.append({ + "opts": ("-B", "--batch-mode"), + "action": "store_true", + "dest": "batch_mode", + "default": False, + "group": _("data"), + "help": _( + "R|If selected then the input_dir should be a parent folder containing multiple " + "folders of faces you wish to sort. The faces will be output to separate sub-" + "folders in the output_dir")}) + argument_list.append({ + "opts": ('-s', '--sort-by'), + "action": Radio, + "type": str, + "choices": _SORT_METHODS, + "dest": 'sort_method', + "group": _("sort settings"), + "default": "face", + "help": _( + "R|Choose how images are sorted. Selecting a sort method gives the images a new " + "filename based on the order the image appears within the given method." + "\nL|'none': Don't sort the images. When a 'group-by' method is selected, " + "selecting 'none' means that the files will be moved/copied into their respective " + "bins, but the files will keep their original filenames. Selecting 'none' for " + "both 'sort-by' and 'group-by' will do nothing" + _SORT_HELP + "\nDefault: face")}) + argument_list.append({ + "opts": ('-g', '--group-by'), + "action": Radio, + "type": str, + "choices": _SORT_METHODS, + "dest": 'group_method', + "group": _("group settings"), + "default": "none", + "help": _( + "R|Selecting a group by method will move/copy files into numbered bins based on " + "the selected method." + "\nL|'none': Don't bin the images. Folders will be sorted by the selected 'sort-" + "by' but will not be binned, instead they will be sorted into a single folder. " + "Selecting 'none' for both 'sort-by' and 'group-by' will do nothing" + + _GROUP_HELP + "\nDefault: none")}) + argument_list.append({ + "opts": ('-k', '--keep'), + "action": 'store_true', + "dest": 'keep_original', + "default": False, + "group": _("data"), + "help": _( + "Whether to keep the original files in their original location. Choosing a 'sort-" + "by' method means that the files have to be renamed. Selecting 'keep' means that " + "the original files will be kept, and the renamed files will be created in the " + "specified output folder. Unselecting keep means that the original files will be " + "moved and renamed based on the selected sort/group criteria.")}) + argument_list.append({ + "opts": ('-t', '--threshold'), + "action": Slider, + "min_max": (-1.0, 10.0), + "rounding": 2, + "type": float, + "dest": 'threshold', + "group": _("group settings"), + "default": -1.0, + "help": _( + "R|Float value. Minimum threshold to use for grouping comparison with 'face-cnn' " + "'hist' and 'face' methods." + "\nThe lower the value the more discriminating the grouping is. Leaving -1.0 will " + "allow Faceswap to choose the default value." + "\nL|For 'face-cnn' 7.2 should be enough, with 4 being very discriminating. " + "\nL|For 'hist' 0.3 should be enough, with 0.2 being very discriminating. " + "\nL|For 'face' between 0.1 (more bins) to 0.5 (fewer bins) should be about right." + "\nBe careful setting a value that's too extrene in a directory with many images, " + "as this could result in a lot of folders being created. Defaults: face-cnn 7.2, " + "hist 0.3, face 0.25")}) + argument_list.append({ + "opts": ('-b', '--bins'), + "action": Slider, + "min_max": (1, 100), + "rounding": 1, + "type": int, + "dest": 'num_bins', + "group": _("group settings"), + "default": 5, + "help": _( + "R|Integer value. Used to control the number of bins created for grouping by: any " + "'blur' methods, 'color' methods or 'face metric' methods ('distance', 'size') " + "and 'orientation; methods ('yaw', 'pitch'). For any other grouping " + "methods see the '-t' ('--threshold') option." + "\nL|For 'face metric' methods the bins are filled, according the the " + "distribution of faces between the minimum and maximum chosen metric." + "\nL|For 'color' methods the number of bins represents the divider of the " + "percentage of colored pixels. Eg. For a bin number of '5': The first folder will " + "have the faces with 0%% to 20%% colored pixels, second 21%% to 40%%, etc. Any " + "empty bins will be deleted, so you may end up with fewer bins than selected." + "\nL|For 'blur' methods folder 0 will be the least blurry, while the last folder " + "will be the blurriest." + "\nL|For 'orientation' methods the number of bins is dictated by how much 180 " + "degrees is divided. Eg. If 18 is selected, then each folder will be a 10 degree " + "increment. Folder 0 will contain faces looking the most to the left/down whereas " + "the last folder will contain the faces looking the most to the right/up. NB: " + "Some bins may be empty if faces do not fit the criteria. \nDefault value: 5")}) + argument_list.append({ + "opts": ('-l', '--log-changes'), + "action": 'store_true', + "group": _("settings"), + "default": False, + "help": _( + "Logs file renaming changes if grouping by renaming, or it logs the file copying/" + "movement if grouping by folders. If no log file is specified with '--log-file', " + "then a 'sort_log.json' file will be created in the input directory.")}) + argument_list.append({ + "opts": ('-f', '--log-file'), + "action": SaveFileFullPaths, + "filetypes": "alignments", + "group": _("settings"), + "dest": 'log_file_path', + "default": 'sort_log.json', + "help": _( + "Specify a log file to use for saving the renaming or grouping information. If " + "specified extension isn't 'json' or 'yaml', then json will be used as the " + "serializer, with the supplied filename. Default: sort_log.json")}) + # Deprecated multi-character switches + argument_list.append({ + "opts": ("-lf", ), + "type": str, + "dest": "depr_log-file_lf_f", + "help": argparse.SUPPRESS}) return argument_list diff --git a/tools/sort/sort.py b/tools/sort/sort.py index 47b7ea2155..c963f9af3b 100644 --- a/tools/sort/sort.py +++ b/tools/sort/sort.py @@ -2,1002 +2,330 @@ """ A tool that allows for sorting and grouping images in different ways. """ +from __future__ import annotations import logging import os import sys -import operator -from concurrent import futures -from shutil import copyfile +import typing as T + +from argparse import Namespace +from shutil import copyfile, rmtree -import numpy as np -import cv2 from tqdm import tqdm # faceswap imports -from lib.serializer import get_serializer_from_filename -from lib.align import AlignedFace, DetectedFace -from lib.image import FacesLoader, read_image, read_image_meta_batch -from lib.utils import FaceswapError -from plugins.extract.recognition.vgg_face2_keras import VGGFace2 as VGGFace -from plugins.extract.pipeline import Extractor, ExtractMedia +from lib.serializer import Serializer, get_serializer_from_filename +from lib.utils import handle_deprecated_cliopts -logger = logging.getLogger(__name__) # pylint: disable=invalid-name +from .sort_methods import SortBlur, SortColor, SortFace, SortHistogram, SortMultiMethod +from .sort_methods_aligned import SortDistance, SortFaceCNN, SortPitch, SortSize, SortYaw, SortRoll +if T.TYPE_CHECKING: + from .sort_methods import SortMethod -class Sort(): - """ Sorts folders of faces based on input criteria """ - # pylint: disable=no-member - - def __init__(self, arguments): - self._args = arguments - self.changes = None - self.serializer = None - self._vgg_face = None - self._loader = FacesLoader(self._args.input_dir) - - def process(self): - """ Main processing function of the sort tool """ - - # Setting default argument values that cannot be set by argparse - - # Set output folder to the same value as input folder - # if the user didn't specify it. - if self._args.output_dir is None: - logger.verbose("No output directory provided. Using input folder as output folder.") - self._args.output_dir = self._args.input_dir - - # Assigning default threshold values based on grouping method - if (self._args.final_process == "folders" - and self._args.min_threshold < 0.0): - method = self._args.group_method.lower() - if method == 'face-cnn': - self._args.min_threshold = 7.2 - elif method == 'hist': - self._args.min_threshold = 0.3 - - # Load VGG Face if sorting by face - if self._args.sort_method.lower() == "face": - self._vgg_face = VGGFace(exclude_gpus=self._args.exclude_gpus) - self._vgg_face.init_model() - - # If logging is enabled, prepare container - if self._args.log_changes: - self.changes = dict() +logger = logging.getLogger(__name__) - # Assign default sort_log.json value if user didn't specify one - if self._args.log_file_path == 'sort_log.json': - self._args.log_file_path = os.path.join(self._args.input_dir, - 'sort_log.json') - # Set serializer based on log file extension - self.serializer = get_serializer_from_filename(self._args.log_file_path) +class Sort(): + """ Sorts folders of faces based on input criteria + + Wrapper for the sort process to run in either batch mode or single use mode + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The arguments to be passed to the extraction process as generated from Faceswap's command + line arguments + """ + def __init__(self, arguments: Namespace) -> None: + logger.debug("Initializing: %s (args: %s)", self.__class__.__name__, arguments) + self._args = handle_deprecated_cliopts(arguments) + self._input_locations = self._get_input_locations() + logger.debug("Initialized: %s", self.__class__.__name__) + + def _get_input_locations(self) -> list[str]: + """ Obtain the full path to input locations. Will be a list of locations if batch mode is + selected, or a containing a single location if batch mode is not selected. - # Prepare sort, group and final process method names - _sort = "sort_" + self._args.sort_method.lower() - _group = "group_" + self._args.group_method.lower() - _final = "final_process_" + self._args.final_process.lower() - if _sort.startswith('sort_color-'): - self._args.color_method = _sort.replace('sort_color-', '') - _sort = _sort[:10] - self._args.sort_method = _sort.replace('-', '_') - self._args.group_method = _group.replace('-', '_') - self._args.final_process = _final.replace('-', '_') - - self.sort_process() - - def launch_aligner(self): - """ Load the aligner plugin to retrieve landmarks """ - extractor = Extractor(None, "fan", None, - normalize_method="hist", exclude_gpus=self._args.exclude_gpus) - extractor.set_batchsize("align", 1) - extractor.launch() - return extractor - - @staticmethod - def alignment_dict(filename, image): - """ Set the image to an ExtractMedia object for alignment """ - height, width = image.shape[:2] - face = DetectedFace(x=0, w=width, y=0, h=height) - return ExtractMedia(filename, image, detected_faces=[face]) - - def _get_landmarks(self): - """ Multi-threaded, parallel and sequentially ordered landmark loader """ - extractor = self.launch_aligner() - filename_list, image_list = self._get_images() - feed_list = list(map(Sort.alignment_dict, filename_list, image_list)) - landmarks = np.zeros((len(feed_list), 68, 2), dtype='float32') - - logger.info("Finding landmarks in images...") - # TODO thread the put to queue so we don't have to put and get at the same time - # Or even better, set up a proper background loader from disk (i.e. use lib.image.ImageIO) - for idx, feed in enumerate(tqdm(feed_list, desc="Aligning", file=sys.stdout)): - extractor.input_queue.put(feed) - landmarks[idx] = next(extractor.detected_faces()).detected_faces[0].landmarks_xy - - return filename_list, image_list, landmarks - - def _get_images(self): - """ Multi-threaded, parallel and sequentially ordered image loader """ - logger.info("Loading images...") - filename_list = self.find_images(self._args.input_dir) - with futures.ThreadPoolExecutor() as executor: - image_list = list(tqdm(executor.map(read_image, filename_list), - desc="Loading Images", - file=sys.stdout, - total=len(filename_list))) - - return filename_list, image_list - - def sort_process(self): - """ - This method dynamically assigns the functions that will be used to run - the core process of sorting, optionally grouping, renaming/moving into - folders. After the functions are assigned they are executed. + Returns + ------- + list: + The list of input location paths """ - sort_method = self._args.sort_method.lower() - group_method = self._args.group_method.lower() - final_method = self._args.final_process.lower() + if not self._args.batch_mode: + return [self._args.input_dir] - img_list = getattr(self, sort_method)() - if "folders" in final_method: - # Check if non-dissimilarity sort method and group method are not the same - if group_method.replace('group_', '') not in sort_method: - img_list = self.reload_images(group_method, img_list) - img_list = getattr(self, group_method)(img_list) - else: - img_list = getattr(self, group_method)(img_list) + retval = [os.path.join(self._args.input_dir, fname) + for fname in os.listdir(self._args.input_dir) + if os.path.isdir(os.path.join(self._args.input_dir, fname))] + logger.debug("Input locations: %s", retval) + return retval - getattr(self, final_method)(img_list) + def _output_for_input(self, input_location: str) -> str: + """ Obtain the path to an output folder for faces for a given input location. - logger.info("Done.") + If not running in batch mode, then the user supplied output location will be returned, + otherwise a sub-folder within the user supplied output location will be returned based on + the input filename - # Methods for sorting - def sort_distance(self): - """ Sort by comparison of face landmark points to mean face by average distance of core - landmarks. """ - logger.info("Sorting by average distance of landmarks...") - filenames = [] - distances = [] - filelist = [os.path.join(self._loader.location, fname) - for fname in os.listdir(self._loader.location) - if os.path.splitext(fname)[-1] == ".png"] - for filename, metadata in tqdm(read_image_meta_batch(filelist), - total=len(filelist), - desc="Calculating Distances"): - if not metadata: - msg = ("The images to be sorted do not contain alignment data. Images must have " - "been generated by Faceswap's Extract process.\nIf you are sorting an " - "older faceset, then you should re-extract the faces from your source " - "alignments file to generate this data.") - raise FaceswapError(msg) - alignments = metadata["itxt"]["alignments"] - aligned_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32")) - filenames.append(filename) - distances.append(aligned_face.average_distance) - - logger.info("Sorting...") - matched_list = list(zip(filenames, distances)) - img_list = sorted(matched_list, key=operator.itemgetter(1)) - return img_list - - def sort_blur(self): - """ Sort by blur amount """ - logger.info("Sorting by estimated image blur...") - - blurs = [(filename, self.estimate_blur(image, metadata)) - for filename, image, metadata in tqdm(self._loader.load(), - desc="Estimating blur", - total=self._loader.count, - leave=False)] - logger.info("Sorting...") - return sorted(blurs, key=lambda x: x[1], reverse=True) - - def sort_blur_fft(self): - """ Sort by fft filtered blur amount with fft""" - logger.info("Sorting by estimated fft filtered image blur...") - - fft_blurs = [(filename, self.estimate_blur_fft(image, metadata)) - for filename, image, metadata in tqdm(self._loader.load(), - desc="Estimating fft blur score", - total=self._loader.count, - leave=False)] - logger.info("Sorting...") - return sorted(fft_blurs, key=lambda x: x[1], reverse=True) - - def sort_color(self): - """ Score by channel average intensity """ - logger.info("Sorting by channel average intensity...") - desired_channel = {'gray': 0, 'luma': 0, 'orange': 1, 'green': 2} - method = self._args.color_method - channel_to_sort = next(v for (k, v) in desired_channel.items() if method.endswith(k)) - filename_list, image_list = self._get_images() - - logger.info("Converting to appropriate colorspace...") - same_size = all(img.size == image_list[0].size for img in image_list) - images = np.array(image_list, dtype='float32')[None, ...] if same_size else image_list - converted_images = self._convert_color(images, same_size, method) - - logger.info("Scoring each image...") - if same_size: - scores = np.average(converted_images[0], axis=(1, 2)) - else: - progress_bar = tqdm(converted_images, desc="Scoring", file=sys.stdout) - scores = np.array([np.average(image, axis=(0, 1)) for image in progress_bar]) - - logger.info("Sorting...") - matched_list = list(zip(filename_list, scores[:, channel_to_sort])) - sorted_file_img_list = sorted(matched_list, key=operator.itemgetter(1), reverse=True) - return sorted_file_img_list - - def sort_face(self): - """ Sort by identity similarity """ - logger.info("Sorting by identity similarity...") - filenames = [] - preds = [] - for filename, image, metadata in tqdm(self._loader.load(), - desc="Classifying Faces", - total=self._loader.count, - leave=False): - if not metadata: - msg = ("The images to be sorted do not contain alignment data. Images must have " - "been generated by Faceswap's Extract process.\nIf you are sorting an " - "older faceset, then you should re-extract the faces from your source " - "alignments file to generate this data.") - raise FaceswapError(msg) - alignments = metadata["alignments"] - face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), - image=image, - centering="legacy", - size=self._vgg_face.input_size, - is_aligned=True).face - filenames.append(filename) - preds.append(self._vgg_face.predict(face)) - - logger.info("Sorting by ward linkage...") - - indices = self._vgg_face.sorted_similarity(np.array(preds), method="ward") - img_list = np.array(filenames)[indices] - return img_list - - def sort_face_cnn(self): - """ Sort by landmark similarity """ - logger.info("Sorting by landmark similarity...") - filename_list, _, landmarks = self._get_landmarks() - img_list = list(zip(filename_list, landmarks)) - - logger.info("Comparing landmarks and sorting...") - img_list_len = len(img_list) - for i in tqdm(range(0, img_list_len - 1), desc="Comparing", file=sys.stdout): - min_score = float("inf") - j_min_score = i + 1 - for j in range(i + 1, img_list_len): - fl1 = img_list[i][1] - fl2 = img_list[j][1] - score = np.sum(np.absolute((fl2 - fl1).flatten())) - if score < min_score: - min_score = score - j_min_score = j - (img_list[i + 1], img_list[j_min_score]) = (img_list[j_min_score], img_list[i + 1]) - return img_list - - def sort_face_cnn_dissim(self): - """ Sort by landmark dissimilarity """ - logger.info("Sorting by landmark dissimilarity...") - filename_list, _, landmarks = self._get_landmarks() - scores = np.zeros(len(filename_list), dtype='float32') - img_list = list(list(items) for items in zip(filename_list, landmarks, scores)) - - logger.info("Comparing landmarks...") - img_list_len = len(img_list) - for i in tqdm(range(0, img_list_len - 1), desc="Comparing", file=sys.stdout): - score_total = 0 - for j in range(i + 1, img_list_len): - if i == j: - continue - fl1 = img_list[i][1] - fl2 = img_list[j][1] - score_total += np.sum(np.absolute((fl2 - fl1).flatten())) - img_list[i][2] = score_total - - logger.info("Sorting...") - img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True) - return img_list - - def sort_face_yaw(self): - """ Sort by estimated face yaw angle """ - logger.info("Sorting by estimated face yaw angle..") - filenames = [] - yaws = [] - for filename, image, metadata in tqdm(self._loader.load(), - desc="Classifying Faces", - total=self._loader.count, - leave=False): - if not metadata: - msg = ("The images to be sorted do not contain alignment data. Images must have " - "been generated by Faceswap's Extract process.\nIf you are sorting an " - "older faceset, then you should re-extract the faces from your source " - "alignments file to generate this data.") - raise FaceswapError(msg) - alignments = metadata["alignments"] - aligned_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), - image=image, - centering="legacy", - is_aligned=True) - filenames.append(filename) - yaws.append(aligned_face.pose.yaw) - - logger.info("Sorting...") - matched_list = list(zip(filenames, yaws)) - img_list = sorted(matched_list, key=operator.itemgetter(1), reverse=True) - return img_list - - def sort_hist(self): - """ Sort by image histogram similarity """ - logger.info("Sorting by histogram similarity...") - - # TODO We have metadata here, so we can mask the face for hist sorting - img_list = [(filename, cv2.calcHist([image], [0], None, [256], [0, 256])) - for filename, image, _ in tqdm(self._loader.load(), - desc="Calculating histograms", - total=self._loader.count, - leave=False)] - - logger.info("Comparing histograms and sorting...") - img_list_len = len(img_list) - for i in tqdm(range(0, img_list_len - 1), desc="Comparing histograms", file=sys.stdout): - min_score = float("inf") - j_min_score = i + 1 - for j in range(i + 1, img_list_len): - score = cv2.compareHist(img_list[i][1], img_list[j][1], cv2.HISTCMP_BHATTACHARYYA) - if score < min_score: - min_score = score - j_min_score = j - (img_list[i + 1], img_list[j_min_score]) = (img_list[j_min_score], img_list[i + 1]) - return img_list - - def sort_hist_dissim(self): - """ Sort by image histogram dissimilarity """ - logger.info("Sorting by histogram dissimilarity...") - - # TODO We have metadata here, so we can mask the face for hist sorting - img_list = [[filename, cv2.calcHist([image], [0], None, [256], [0, 256]), 0.0] - for filename, image, _ in tqdm(self._loader.load(), - desc="Calculating histograms", - total=self._loader.count, - leave=False)] - - img_list_len = len(img_list) - for i in tqdm(range(0, img_list_len), desc="Comparing histograms", file=sys.stdout): - score_total = 0 - for j in range(0, img_list_len): - if i == j: - continue - score_total += cv2.compareHist(img_list[i][1], - img_list[j][1], - cv2.HISTCMP_BHATTACHARYYA) - img_list[i][2] = score_total - - logger.info("Sorting...") - return sorted(img_list, key=lambda x: x[2], reverse=True) - - def sort_size(self): - """ Sort the faces by largest face (in original frame) to smallest """ - logger.info("Sorting by original face size...") - img_list = [] - for filename, image, metadata in tqdm(self._loader.load(), - desc="Calculating face sizes", - total=self._loader.count, - leave=False): - if not metadata: - msg = ("The images to be sorted do not contain alignment data. Images must have " - "been generated by Faceswap's Extract process.\nIf you are sorting an " - "older faceset, then you should re-extract the faces from your source " - "alignments file to generate this data.") - raise FaceswapError(msg) - alignments = metadata["alignments"] - aligned_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), - image=image, - centering="legacy", - is_aligned=True) - roi = aligned_face.original_roi - size = ((roi[1][0] - roi[0][0]) ** 2 + # pylint:disable=unsubscriptable-object - (roi[1][1] - roi[0][1]) ** 2) ** 0.5 # pylint:disable=unsubscriptable-object - img_list.append((filename, size)) - - logger.info("Sorting...") - return sorted(img_list, key=lambda x: x[1], reverse=True) - - def sort_black_pixels(self): - """ Sort by percentage of black pixels - - Calculates the sum of black pixels, get the percentage X 3 channels + Parameters + ---------- + input_location: str + The full path to an input video or folder of images """ - logger.info("Sorting by percentage of black pixels...") - img_list = [(filename, np.ndarray.all(image == [0, 0, 0], axis=2).sum()/image.size*100*3) - for filename, image, _ in tqdm(self._loader.load(), - desc="Calculating black pixels", - total=self._loader.count, - leave=False)] - img_list_len = len(img_list) - for i in tqdm(range(0, img_list_len - 1), desc="Comparing black pixels", file=sys.stdout): - for j in range(0, img_list_len-i-1): - if img_list[j][1] > img_list[j+1][1]: - temp = img_list[j] - img_list[j] = img_list[j+1] - img_list[j+1] = temp - return img_list - - # Methods for grouping - def group_blur(self, img_list): - """ Group into bins by blur """ - # Starting the binning process - num_bins = self._args.num_bins - - # The last bin will get all extra images if it's - # not possible to distribute them evenly - num_per_bin = len(img_list) // num_bins - remainder = len(img_list) % num_bins - - logger.info("Grouping by blur...") - bins = [[] for _ in range(num_bins)] - idx = 0 - for i in range(num_bins): - for _ in range(num_per_bin): - bins[i].append(img_list[idx][0]) - idx += 1 - - # If remainder is 0, nothing gets added to the last bin. - for i in range(1, remainder + 1): - bins[-1].append(img_list[-i][0]) + if not self._args.batch_mode or self._args.output_dir is None: + return self._args.output_dir - return bins + retval = os.path.join(self._args.output_dir, os.path.basename(input_location)) + logger.debug("Returning output: '%s' for input: '%s'", retval, input_location) + return retval - def group_blur_fft(self, img_list): - """ Group into bins by fft blur score""" - # Starting the binning process - num_bins = self._args.num_bins + def process(self) -> None: + """ The entry point for triggering the Sort Process. - # The last bin will get all extra images if it's - # not possible to distribute them evenly - num_per_bin = len(img_list) // num_bins - remainder = len(img_list) % num_bins - - logger.info("Grouping by fft blur score...") - bins = [[] for _ in range(num_bins)] - idx = 0 - for i in range(num_bins): - for _ in range(num_per_bin): - bins[i].append(img_list[idx][0]) - idx += 1 - - # If remainder is 0, nothing gets added to the last bin. - for i in range(1, remainder + 1): - bins[-1].append(img_list[-i][0]) - - return bins - - def group_face_cnn(self, img_list): - """ Group into bins by CNN face similarity """ - logger.info("Grouping by face-cnn similarity...") - - # Groups are of the form: group_num -> reference faces - reference_groups = dict() - - # Bins array, where index is the group number and value is - # an array containing the file paths to the images in that group. - bins = [] - - # Comparison threshold used to decide how similar - # faces have to be to be grouped together. - # It is multiplied by 1000 here to allow the cli option to use smaller - # numbers. - min_threshold = self._args.min_threshold * 1000 - - img_list_len = len(img_list) + Should only be called from :class:`lib.cli.launcher.ScriptExecutor` + """ + logger.info('Starting, this may take a while...') + inputs = self._input_locations + if self._args.batch_mode: + logger.info("Batch mode selected processing: %s", self._input_locations) + for job_no, location in enumerate(self._input_locations): + if self._args.batch_mode: + logger.info("Processing job %s of %s: '%s'", job_no + 1, len(inputs), location) + arguments = Namespace(**self._args.__dict__) + arguments.input_dir = location + arguments.output_dir = self._output_for_input(location) + else: + arguments = self._args + sort = _Sort(arguments) + sort.process() - for i in tqdm(range(0, img_list_len - 1), - desc="Grouping", - file=sys.stdout): - fl1 = img_list[i][1] - current_best = [-1, float("inf")] +class _Sort(): + """ Sorts folders of faces based on input criteria """ + def __init__(self, arguments: Namespace) -> None: + logger.debug("Initializing %s: arguments: %s", self.__class__.__name__, arguments) + self._processes = {"blur": SortBlur, + "blur_fft": SortBlur, + "distance": SortDistance, + "yaw": SortYaw, + "pitch": SortPitch, + "roll": SortRoll, + "size": SortSize, + "face": SortFace, + "face_cnn": SortFaceCNN, + "face_cnn_dissim": SortFaceCNN, + "hist": SortHistogram, + "hist_dissim": SortHistogram, + "color_black": SortColor, + "color_gray": SortColor, + "color_luma": SortColor, + "color_green": SortColor, + "color_orange": SortColor} + + self._args = self._parse_arguments(arguments) + self._changes: dict[str, str] = {} + self.serializer: Serializer | None = None + + if arguments.log_changes: + self.serializer = get_serializer_from_filename(arguments.log_file_path) + + self._sorter = self._get_sorter() + logger.debug("Initialized %s", self.__class__.__name__) + + def _set_output_folder(self, arguments): + """ Set the output folder correctly if it has not been provided + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process - for key, references in reference_groups.items(): - try: - score = self.get_avg_score_faces_cnn(fl1, references) - except TypeError: - score = float("inf") - except ZeroDivisionError: - score = float("inf") - if score < current_best[1]: - current_best[0], current_best[1] = key, score + Returns + ------- + :class:`argparse.Namespace` + The command line arguments with output folder correctly set + """ + logger.debug("setting output folder: %s", arguments.output_dir) + input_dir = arguments.input_dir + output_dir = arguments.output_dir + sort_method = arguments.sort_method + group_method = arguments.group_method + + needs_rename = sort_method != "none" and group_method == "none" + + if needs_rename and arguments.keep_original and (not output_dir or + output_dir == input_dir): + output_dir = os.path.join(input_dir, "sorted") + logger.warning("No output folder selected, but files need renaming. " + "Outputting to: '%s'", output_dir) + elif not output_dir: + output_dir = input_dir + logger.warning("No output folder selected, files will be sorted in place in: '%s'", + output_dir) + + arguments.output_dir = output_dir + logger.debug("Set output folder: %s", arguments.output_dir) + return arguments + + def _parse_arguments(self, arguments): + """ Parse the arguments and update/format relevant choices - if current_best[1] < min_threshold: - reference_groups[current_best[0]].append(fl1[0]) - bins[current_best[0]].append(img_list[i][0]) - else: - reference_groups[len(reference_groups)] = [img_list[i][1]] - bins.append([img_list[i][0]]) + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process - return bins + Returns + ------- + :class:`argparse.Namespace` + The formatted command line arguments + """ + logger.debug("Cleaning arguments: %s", arguments) + if arguments.sort_method == "none" and arguments.group_method == "none": + logger.error("Both sort-by and group-by are 'None'. Nothing to do.") + sys.exit(1) - def group_face_yaw(self, img_list): - """ Group into bins by yaw of face """ - # Starting the binning process - num_bins = self._args.num_bins + # Prepare sort, group and final process method names + arguments.sort_method = arguments.sort_method.lower().replace("-", "_") + arguments.group_method = arguments.group_method.lower().replace("-", "_") - # The last bin will get all extra images if it's - # not possible to distribute them evenly - num_per_bin = len(img_list) // num_bins - remainder = len(img_list) % num_bins + arguments = self._set_output_folder(arguments) - logger.info("Grouping by face-yaw...") - bins = [[] for _ in range(num_bins)] - idx = 0 - for i in range(num_bins): - for _ in range(num_per_bin): - bins[i].append(img_list[idx][0]) - idx += 1 + if arguments.log_changes and arguments.log_file_path == "sort_log.json": + # Assign default sort_log.json value if user didn't specify one + arguments.log_file_path = os.path.join(self._args.input_dir, 'sort_log.json') - # If remainder is 0, nothing gets added to the last bin. - for i in range(1, remainder + 1): - bins[-1].append(img_list[-i][0]) + logger.debug("Cleaned arguments: %s", arguments) + return arguments - return bins + def _get_sorter(self) -> SortMethod: + """ Obtain a sorter/grouper combo for the selected sort/group by options - def group_black_pixels(self, img_list): - """ Group into bins by percentage of black pixels - :type img_list: (str, float) + Returns + ------- + :class:`SortMethod` + The sorter or combined sorter for sorting and grouping based on user selections """ - logger.info("Grouping by percentage of black pixels...") - - # Starting the binning process - bins = [[] for _ in range(self._args.num_bins)] - # Get edges of bins from 0 to 100 - bins_edges = self._near_split(100, self._args.num_bins) - # Get the proper bin number for each img order - img_bins = np.digitize([x[1] for x in img_list], bins_edges, right=True) - - # Place imgs in bins - for idx, _bin in enumerate(img_bins): - bins[_bin].append(img_list[idx][0]) - - return bins - - def group_hist(self, img_list): - """ Group into bins by histogram """ - logger.info("Grouping by histogram...") - - # Groups are of the form: group_num -> reference histogram - reference_groups = dict() - - # Bins array, where index is the group number and value is - # an array containing the file paths to the images in that group - bins = [] - - min_threshold = self._args.min_threshold - - img_list_len = len(img_list) - reference_groups[0] = [img_list[0][1]] - bins.append([img_list[0][0]]) - - for i in tqdm(range(1, img_list_len), - desc="Grouping", - file=sys.stdout): - current_best = [-1, float("inf")] - for key, value in reference_groups.items(): - score = self.get_avg_score_hist(img_list[i][1], value) - if score < current_best[1]: - current_best[0], current_best[1] = key, score - - if current_best[1] < min_threshold: - reference_groups[current_best[0]].append(img_list[i][1]) - bins[current_best[0]].append(img_list[i][0]) - else: - reference_groups[len(reference_groups)] = [img_list[i][1]] - bins.append([img_list[i][0]]) + sort_method = self._args.sort_method + group_method = self._args.group_method - return bins + sort_method = group_method if sort_method == "none" else sort_method + sorter = self._processes[sort_method](self._args, + is_group=self._args.sort_method == "none") - # Final process methods - def final_process_rename(self, img_list): - """ Rename the files """ - output_dir = self._args.output_dir - - process_file = self.set_process_file_method(self._args.log_changes, - self._args.keep_original) - - # Make sure output directory exists - if not os.path.exists(output_dir): - os.makedirs(output_dir) - - description = ( - "Copying and Renaming" if self._args.keep_original - else "Moving and Renaming" - ) - - for i in tqdm(range(0, len(img_list)), - desc=description, - leave=False, - file=sys.stdout): - src = img_list[i] if isinstance(img_list[i], str) else img_list[i][0] - src_basename = os.path.basename(src) - - dst = os.path.join(output_dir, '{:05d}_{}'.format(i, src_basename)) - try: - process_file(src, dst, self.changes) - except FileNotFoundError as err: - logger.error(err) - logger.error('fail to rename %s', src) - - for i in tqdm(range(0, len(img_list)), - desc=description, - file=sys.stdout): - renaming = self.set_renaming_method(self._args.log_changes) - fname = img_list[i] if isinstance(img_list[i], str) else img_list[i][0] - src, dst = renaming(fname, output_dir, i, self.changes) - - try: - os.rename(src, dst) - except FileNotFoundError as err: - logger.error(err) - logger.error('fail to rename %s', format(src)) + if sort_method != "none" and group_method != "none" and group_method != sort_method: + grouper = self._processes[group_method](self._args, is_group=True) + retval = SortMultiMethod(self._args, sorter, grouper) + logger.debug("Got sorter + grouper: %s (%s, %s)", retval, sorter, grouper) - if self._args.log_changes: - self.write_to_log(self.changes) - - def final_process_folders(self, bins): - """ Move the files to folders """ - output_dir = self._args.output_dir + else: - process_file = self.set_process_file_method(self._args.log_changes, - self._args.keep_original) - - # First create new directories to avoid checking - # for directory existence in the moving loop - logger.info("Creating group directories.") - for i in range(len(bins)): - directory = os.path.join(output_dir, str(i)) - if not os.path.exists(directory): - os.makedirs(directory) - - description = ( - "Copying into Groups" if self._args.keep_original - else "Moving into Groups" - ) - - logger.info("Total groups found: %s", len(bins)) - for i in tqdm(range(len(bins)), desc=description, file=sys.stdout): - for j in range(len(bins[i])): - src = bins[i][j] - src_basename = os.path.basename(src) - - dst = os.path.join(output_dir, str(i), src_basename) - try: - process_file(src, dst, self.changes) - except FileNotFoundError as err: - logger.error(err) - logger.error("Failed to move '%s' to '%s'", src, dst) + retval = sorter - if self._args.log_changes: - self.write_to_log(self.changes) + logger.debug("Final sorter: %s", retval) + return retval - # Various helper methods - def write_to_log(self, changes): + def _write_to_log(self, changes): """ Write the changes to log file """ logger.info("Writing sort log to: '%s'", self._args.log_file_path) self.serializer.save(self._args.log_file_path, changes) - def reload_images(self, group_method, img_list): - """ - Reloads the image list by replacing the comparative values with those - that the chosen grouping method expects. - :param group_method: str name of the grouping method that will be used. - :param img_list: image list that has been sorted by one of the sort - methods. - :return: img_list but with the comparative values that the chosen - grouping method expects. + def process(self) -> None: + """ Main processing function of the sort tool + + This method dynamically assigns the functions that will be used to run + the core process of sorting, optionally grouping, renaming/moving into + folders. After the functions are assigned they are executed. """ - logger.info("Preparing to group...") - if group_method == 'group_blur': - filename_list, image_list = self._get_images() - blurs = [self.estimate_blur(img) for img in image_list] - temp_list = list(zip(filename_list, blurs)) - elif group_method == 'group_blur_fft': - filename_list, image_list = self._get_images() - fft_blurs = [self.estimate_blur_fft(img) for img in image_list] - temp_list = list(zip(filename_list, fft_blurs)) - elif group_method == 'group_face_cnn': - filename_list, image_list, landmarks = self._get_landmarks() - temp_list = list(zip(filename_list, landmarks)) - elif group_method == 'group_face_yaw': - filename_list, image_list, landmarks = self._get_landmarks() - yaws = [self.calc_landmarks_face_yaw(mark) for mark in landmarks] - temp_list = list(zip(filename_list, yaws)) - elif group_method == 'group_hist': - filename_list, image_list = self._get_images() - histograms = [cv2.calcHist([img], [0], None, [256], [0, 256]) for img in image_list] - temp_list = list(zip(filename_list, histograms)) - elif group_method == 'group_black_pixels': - filename_list, image_list = self._get_images() - black_pixels = [np.ndarray.all(img == [0, 0, 0], axis=2).sum()/img.size*100*3 - for img in image_list] - temp_list = list(zip(filename_list, black_pixels)) + if self._args.group_method != "none": + # Check if non-dissimilarity sort method and group method are not the same + self._output_groups() else: - raise ValueError("{} group_method not found.".format(group_method)) + self._output_non_grouped() + + if self._args.log_changes: + self._write_to_log(self._changes) - return self.splice_lists(img_list, temp_list) + logger.info("Done.") - @staticmethod - def _near_split(bin_range, num_bins): - """ Obtain the split for the given number of bins for the given range + def _sort_file(self, source: str, destination: str) -> None: + """ Copy or move a file based on whether 'keep original' has been selected and log changes + if required. Parameters ---------- - bin_range: int - The range of data to separate into bins - num_bins: int - The number of bins to create - - Returns - ------- - list - The split dividers for the given number of bins for the given range + source: str + The full path to the source file that is being sorted + destination: str + The full path to where the source file should be moved/renamed """ - quotient, remainder = divmod(bin_range, num_bins) - seps = [quotient + 1] * remainder + [quotient] * (num_bins - remainder) - uplimit = 0 - bins = [0] - for sep in seps: - bins.append(uplimit + sep) - uplimit += sep - return bins - - @staticmethod - def _convert_color(imgs, same_size, method): - """ Helper function to convert color spaces """ - - if method.endswith('gray'): - conversion = np.array([[0.0722], [0.7152], [0.2126]]) - else: - conversion = np.array([[0.25, 0.5, 0.25], [-0.5, 0.0, 0.5], [-0.25, 0.5, -0.25]]) - - if same_size: - path = 'greedy' - operation = 'bijk, kl -> bijl' if method.endswith('gray') else 'bijl, kl -> bijk' - else: - operation = 'ijk, kl -> ijl' if method.endswith('gray') else 'ijl, kl -> ijk' - path = np.einsum_path(operation, imgs[0][..., :3], conversion, optimize='optimal')[0] - - progress_bar = tqdm(imgs, desc="Converting", file=sys.stdout) - images = [np.einsum(operation, img[..., :3], conversion, optimize=path).astype('float32') - for img in progress_bar] - return images + try: + if self._args.keep_original: + copyfile(source, destination) + else: + os.rename(source, destination) + except FileNotFoundError as err: + logger.error("Failed to sort '%s' to '%s'. Original error: %s", + source, destination, str(err)) - @staticmethod - def splice_lists(sorted_list, new_vals_list): - """ - This method replaces the value at index 1 in each sub-list in the - sorted_list with the value that is calculated for the same img_path, - but found in new_vals_list. - - Format of lists: [[img_path, value], [img_path2, value2], ...] - - :param sorted_list: list that has been sorted by one of the sort - methods. - :param new_vals_list: list that has been loaded by a different method - than the sorted_list. - :return: list that is sorted in the same way as the input sorted list - but the values corresponding to each image are from new_vals_list. - """ - new_list = [] - # Make new list of just image paths to serve as an index - val_index_list = [i[0] for i in new_vals_list] - for i in tqdm(range(len(sorted_list)), desc="Splicing", file=sys.stdout): - current_img = sorted_list[i] if isinstance(sorted_list[i], str) else sorted_list[i][0] - new_val_index = val_index_list.index(current_img) - new_list.append([current_img, new_vals_list[new_val_index][1]]) - - return new_list - - @staticmethod - def find_images(input_dir): - """ Return list of images at specified location """ - result = [] - extensions = [".jpg", ".png", ".jpeg"] - for root, _, files in os.walk(input_dir): - for file in files: - if os.path.splitext(file)[1].lower() in extensions: - result.append(os.path.join(root, file)) - break - return result - - @classmethod - def estimate_blur(cls, image, metadata=None): - """ Estimate the amount of blur an image has with the variance of the Laplacian. - Normalize by pixel number to offset the effect of image size on pixel gradients & variance. + if self._args.log_changes: + self._changes[source] = destination - Parameters - ---------- - image: :class:`numpy.ndarray` - The face image to calculate blur for - metadata: dict, optional - The metadata for the face image or ``None`` if no metadata is available. If metadata is - provided the face will be masked by the "components" mask prior to calculating blur. - Default:``None`` + def _output_groups(self) -> None: + """ Move the files to folders. - Returns - ------- - float - The estimated blur score for the face + Obtains the bins and original filenames from :attr:`_sorter` and outputs into appropriate + bins in the output location """ - if metadata is not None: - alignments = metadata["alignments"] - det_face = DetectedFace() - det_face.from_png_meta(alignments) - aln_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), - image=image, - centering="legacy", - size=256, - is_aligned=True) - mask = det_face.mask["components"] - mask.set_sub_crop(aln_face.pose.offset[mask.stored_centering] * -1, centering="legacy") - mask = cv2.resize(mask.mask, (256, 256), interpolation=cv2.INTER_CUBIC)[..., None] - image = np.minimum(aln_face.face, mask) - if image.ndim == 3: - image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - blur_map = cv2.Laplacian(image, cv2.CV_32F) - score = np.var(blur_map) / np.sqrt(image.shape[0] * image.shape[1]) - return score - - @classmethod - def estimate_blur_fft(cls, image, metadata=None): - """ Estimate the amount of blur a fft filtered image has. + is_rename = self._args.sort_method != "none" + + logger.info("Creating %s group folders in '%s'.", + len(self._sorter.binned), self._args.output_dir) + bin_names = [f"_{b}" for b in self._sorter.bin_names] + if is_rename: + bin_names = [f"{name}_by_{self._args.sort_method}" for name in bin_names] + for name in bin_names: + folder = os.path.join(self._args.output_dir, name) + if os.path.exists(folder): + rmtree(folder) + os.makedirs(folder) + + description = f"{'Copying' if self._args.keep_original else 'Moving'} into groups" + description += " and renaming" if is_rename else "" + + pbar = tqdm(range(len(self._sorter.sorted_filelist)), + desc=description, + file=sys.stdout, + leave=False) + idx = 0 + for bin_id, bin_ in enumerate(self._sorter.binned): + pbar.set_description(f"{description}: Bin {bin_id + 1} of {len(self._sorter.binned)}") + output_path = os.path.join(self._args.output_dir, bin_names[bin_id]) + if not bin_: + logger.debug("Removing empty bin: %s", output_path) + os.rmdir(output_path) + for source in bin_: + basename = os.path.basename(source) + dst_name = f"{idx:06d}_{basename}" if is_rename else basename + dest = os.path.join(output_path, dst_name) + self._sort_file(source, dest) + idx += 1 + pbar.update(1) - Parameters - ---------- - image: :class:`numpy.ndarray` - Use Fourier Transform to analyze the frequency characteristics of the masked - face using 2D Discrete Fourier Transform (DFT) filter to find the frequency domain. - A mean value is assigned to the magnitude spectrum and returns a blur score. - Adapted from https://www.pyimagesearch.com/2020/06/15/ - opencv-fast-fourier-transform-fft-for-blur-detection-in-images-and-video-streams/ - metadata: dict, optional - The metadata for the face image or ``None`` if no metadata is available. If metadata is - provided the face will be masked by the "components" mask prior to calculating blur. - Default:``None`` + # Output methods + def _output_non_grouped(self) -> None: + """ Output non-grouped files. - Returns - ------- - float - The estimated fft blur score for the face - """ - if metadata is not None: - alignments = metadata["alignments"] - det_face = DetectedFace() - det_face.from_png_meta(alignments) - aln_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), - image=image, - centering="legacy", - size=256, - is_aligned=True) - mask = det_face.mask["components"] - mask.set_sub_crop(aln_face.pose.offset[mask.stored_centering] * -1, centering="legacy") - mask = cv2.resize(mask.mask, (256, 256), interpolation=cv2.INTER_CUBIC)[..., None] - image = np.minimum(aln_face.face, mask) - if image.ndim == 3: - image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - height, width = image.shape - c_height, c_width = (int(height / 2.0), int(width / 2.0)) - fft = np.fft.fft2(image) - fft_shift = np.fft.fftshift(fft) - fft_shift[c_height - 75:c_height + 75, c_width - 75:c_width + 75] = 0 - ifft_shift = np.fft.ifftshift(fft_shift) - shift_back = np.fft.ifft2(ifft_shift) - magnitude = np.log(np.abs(shift_back)) - score = np.mean(magnitude) - return score - - @staticmethod - def calc_landmarks_face_pitch(flm): - """ UNUSED - Calculate the amount of pitch in a face """ - var_t = ((flm[6][1] - flm[8][1]) + (flm[10][1] - flm[8][1])) / 2.0 - var_b = flm[8][1] - return var_b - var_t - - @staticmethod - def calc_landmarks_face_yaw(flm): - """ Calculate the amount of yaw in a face """ - var_l = ((flm[27][0] - flm[0][0]) - + (flm[28][0] - flm[1][0]) - + (flm[29][0] - flm[2][0])) / 3.0 - var_r = ((flm[16][0] - flm[27][0]) - + (flm[15][0] - flm[28][0]) - + (flm[14][0] - flm[29][0])) / 3.0 - return var_r - var_l - - @staticmethod - def set_process_file_method(log_changes, keep_original): - """ - Assigns the final file processing method based on whether changes are - being logged and whether the original files are being kept in the - input directory. - Relevant cli arguments: -k, -l - :return: function reference + These are files which are sorted but not binned, so just the filename gets updated """ - if log_changes: - if keep_original: - def process_file(src, dst, changes): - """ Process file method if logging changes - and keeping original """ - copyfile(src, dst) - changes[src] = dst - - else: - def process_file(src, dst, changes): - """ Process file method if logging changes - and not keeping original """ - os.rename(src, dst) - changes[src] = dst + output_dir = self._args.output_dir + os.makedirs(output_dir, exist_ok=True) - else: - if keep_original: - def process_file(src, dst, changes): # pylint: disable=unused-argument - """ Process file method if not logging changes - and keeping original """ - copyfile(src, dst) + description = f"{'Copying' if self._args.keep_original else 'Moving'} and renaming" + for idx, source in enumerate(tqdm(self._sorter.sorted_filelist, + desc=description, + file=sys.stdout, + leave=False)): + dest = os.path.join(output_dir, f"{idx:06d}_{os.path.basename(source)}") - else: - def process_file(src, dst, changes): # pylint: disable=unused-argument - """ Process file method if not logging changes - and not keeping original """ - os.rename(src, dst) - return process_file - - @staticmethod - def set_renaming_method(log_changes): - """ Set the method for renaming files """ - if log_changes: - def renaming(src, output_dir, i, changes): - """ Rename files method if logging changes """ - src_basename = os.path.basename(src) - - __src = os.path.join(output_dir, - '{:05d}_{}'.format(i, src_basename)) - dst = os.path.join( - output_dir, - '{:05d}{}'.format(i, os.path.splitext(src_basename)[1])) - changes[src] = dst - return __src, dst - else: - def renaming(src, output_dir, i, changes): # pylint: disable=unused-argument - """ Rename files method if not logging changes """ - src_basename = os.path.basename(src) - - src = os.path.join(output_dir, - '{:05d}_{}'.format(i, src_basename)) - dst = os.path.join( - output_dir, - '{:05d}{}'.format(i, os.path.splitext(src_basename)[1])) - return src, dst - return renaming - - @staticmethod - def get_avg_score_hist(img1, references): - """ Return the average histogram score between a face and - reference image """ - scores = [] - for img2 in references: - score = cv2.compareHist(img1, img2, cv2.HISTCMP_BHATTACHARYYA) - scores.append(score) - return sum(scores) / len(scores) - - @staticmethod - def get_avg_score_faces_cnn(fl1, references): - """ Return the average CNN similarity score - between a face and reference image """ - scores = [] - for fl2 in references: - score = np.sum(np.absolute((fl2 - fl1).flatten())) - scores.append(score) - return sum(scores) / len(scores) + self._sort_file(source, dest) diff --git a/tools/sort/sort_methods.py b/tools/sort/sort_methods.py new file mode 100644 index 0000000000..f2a4b29526 --- /dev/null +++ b/tools/sort/sort_methods.py @@ -0,0 +1,1110 @@ +#!/usr/bin/env python3 +""" Sorting methods for the sorting tool. + +All sorting methods inherit from :class:`SortMethod` and control functions for scorting one item, +sorting a full list of scores and binning based on those sorted scores. +""" +from __future__ import annotations +import logging +import operator +import sys +import typing as T + +from collections.abc import Generator + +import cv2 +import numpy as np +from tqdm import tqdm + +from lib.align import AlignedFace, DetectedFace, LandmarkType +from lib.image import FacesLoader, ImagesLoader, read_image_meta_batch, update_existing_metadata +from lib.utils import FaceswapError +from plugins.extract.recognition.vgg_face2 import Cluster, Recognition as VGGFace + +if T.TYPE_CHECKING: + from argparse import Namespace + from lib.align.alignments import PNGHeaderAlignmentsDict, PNGHeaderSourceDict + +logger = logging.getLogger(__name__) + + +ImgMetaType: T.TypeAlias = Generator[tuple[str, + np.ndarray | None, + T.Union["PNGHeaderAlignmentsDict", None]], None, None] + + +class InfoLoader(): + """ Loads aligned faces and/or face metadata + + Parameters + ---------- + input_dir: str + Full path to containing folder of faces to be supported + loader_type: ["face", "meta", "all"] + Dictates the type of iterator that will be used. "face" just loads the image with the + filename, "meta" just loads the image alignment data with the filename. "all" loads + the image and the alignment data with the filename + """ + def __init__(self, + input_dir: str, + info_type: T.Literal["face", "meta", "all"]) -> None: + logger.debug("Initializing: %s (input_dir: %s, info_type: %s)", + self.__class__.__name__, input_dir, info_type) + self._info_type = info_type + self._iterator = None + self._description = "Reading image statistics..." + self._loader = ImagesLoader(input_dir) if info_type == "face" else FacesLoader(input_dir) + self._cached_source_data: dict[str, PNGHeaderSourceDict] = {} + if self._loader.count == 0: + logger.error("No images to process in location: '%s'", input_dir) + sys.exit(1) + + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def filelist_count(self) -> int: + """ int: The number of files to be processed """ + return len(self._loader.file_list) + + def _get_iterator(self) -> ImgMetaType: + """ Obtain the iterator for the selected :attr:`info_type`. + + Returns + ------- + generator + The correct generator for the given info_type + """ + if self._info_type == "all": + return self._full_data_reader() + if self._info_type == "meta": + return self._metadata_reader() + return self._image_data_reader() + + def __call__(self) -> ImgMetaType: + """ Return the selected iterator + + The resulting generator: + + Yields + ------ + filename: str + The filename that has been read + image: :class:`numpy.ndarray or ``None`` + The aligned face image loaded from disk for 'face' and 'all' info_types + otherwise ``None`` + alignments: dict or ``None`` + The alignments dict for 'all' and 'meta' infor_types otherwise ``None`` + """ + iterator = self._get_iterator() + return iterator + + def _get_alignments(self, + filename: str, + metadata: dict[str, T.Any]) -> PNGHeaderAlignmentsDict | None: + """ Obtain the alignments from a PNG Header. + + The other image metadata is cached locally in case a sort method needs to write back to the + PNG header + + Parameters + ---------- + filename: str + Full path to the image PNG file + metadata: dict + The header data from a PNG file + + Returns + ------- + dict or ``None`` + The alignments dictionary from the PNG header, if it exists, otherwise ``None`` + """ + if not metadata or not metadata.get("alignments") or not metadata.get("source"): + return None + self._cached_source_data[filename] = metadata["source"] + return metadata["alignments"] + + def _metadata_reader(self) -> ImgMetaType: + """ Load metadata from saved aligned faces + + Yields + ------ + filename: str + The filename that has been read + image: None + This will always be ``None`` with the metadata reader + alignments: dict or ``None`` + The alignment data for the given face or ``None`` if no alignments found + """ + for filename, metadata in tqdm(read_image_meta_batch(self._loader.file_list), + total=self._loader.count, + desc=self._description, + leave=False): + alignments = self._get_alignments(filename, metadata.get("itxt", {})) + yield filename, None, alignments + + def _full_data_reader(self) -> ImgMetaType: + """ Load the image and metadata from a folder of aligned faces + + Yields + ------ + filename: str + The filename that has been read + image: :class:`numpy.ndarray + The aligned face image loaded from disk + alignments: dict or ``None`` + The alignment data for the given face or ``None`` if no alignments found + """ + for filename, image, metadata in tqdm(self._loader.load(), + desc=self._description, + total=self._loader.count, + leave=False): + alignments = self._get_alignments(filename, metadata) + yield filename, image, alignments + + def _image_data_reader(self) -> ImgMetaType: + """ Just loads the images with their filenames + + Yields + ------ + filename: str + The filename that has been read + image: :class:`numpy.ndarray + The aligned face image loaded from disk + alignments: ``None`` + Alignments will always be ``None`` with the image data reader + """ + for filename, image in tqdm(self._loader.load(), + desc=self._description, + total=self._loader.count, + leave=False): + yield filename, image, None + + def update_png_header(self, filename: str, alignments: PNGHeaderAlignmentsDict) -> None: + """ Update the PNG header of the given file with the given alignments. + + NB: Header information can only be updated if the face is already on at least alignment + version 2.2. If below this version, then the header is not updated + + + Parameters + ---------- + filename: str + Full path to the PNG file to update + alignments: dict + The alignments to update into the PNG header + """ + vers = self._cached_source_data[filename]["alignments_version"] + if vers < 2.2: + return + + self._cached_source_data[filename]["alignments_version"] = 2.3 if vers == 2.2 else vers + header = {"alignments": alignments, "source": self._cached_source_data[filename]} + update_existing_metadata(filename, header) + + +class SortMethod(): + """ Parent class for sort methods. All sort methods should inherit from this class + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + loader_type: ["face", "meta", "all"] + The type of image loader to use. "face" just loads the image with the filename, "meta" + just loads the image alignment data with the filename. "all" loads the image and the + alignment data with the filename + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + _log_mask_once = False + + def __init__(self, + arguments: Namespace, + loader_type: T.Literal["face", "meta", "all"] = "meta", + is_group: bool = False) -> None: + logger.debug("Initializing %s: loader_type: '%s' is_group: %s, arguments: %s", + self.__class__.__name__, loader_type, is_group, arguments) + self._is_group = is_group + self._log_once = True + self._method = arguments.group_method if self._is_group else arguments.sort_method + + self._num_bins: int = arguments.num_bins + self._bin_names: list[str] = [] + + self._loader_type = loader_type + self._iterator = self._get_file_iterator(arguments.input_dir) + + self._result: list[tuple[str, float | np.ndarray]] = [] + self._binned: list[list[str]] = [] + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def loader_type(self) -> T.Literal["face", "meta", "all"]: + """ ["face", "meta", "all"]: The loader that this sorter uses """ + return self._loader_type + + @property + def binned(self) -> list[list[str]]: + """ list: List of bins (list) containing the filenames belonging to the bin. The binning + process is called when this property is first accessed""" + if not self._binned: + self._binned = self._binning() + logger.debug({f"bin_{idx}": len(bin_) for idx, bin_ in enumerate(self._binned)}) + return self._binned + + @property + def sorted_filelist(self) -> list[str]: + """ list: List of sorted filenames for given sorter in a single list. The sort process is + called when this property is first accessed """ + if not self._result: + self._sort_filelist() + retval = [item[0] for item in self._result] + logger.debug(retval) + else: + retval = [item[0] for item in self._result] + return retval + + @property + def bin_names(self) -> list[str]: + """ list: The name of each created bin, if they exist, otherwise an empty list """ + return self._bin_names + + def _get_file_iterator(self, input_dir: str) -> InfoLoader: + """ Override for method specific iterators. + + Parameters + ---------- + input_dir: str + Full path to containing folder of faces to be supported + + Returns + ------- + :class:`InfoLoader` + The correct InfoLoader iterator for the current sort method + """ + return InfoLoader(input_dir, self.loader_type) + + def _sort_filelist(self) -> None: + """ Call the sort method's logic to populate the :attr:`_results` attribute. + + Put logic for scoring an individual frame in in :attr:`score_image` of the child + + Returns + ------- + list + The sorted file. A list of tuples with the filename in the first position and score in + the second position + """ + for filename, image, alignments in self._iterator(): + self.score_image(filename, image, alignments) + + self.sort() + logger.debug("sorted list: %s", + [r[0] if isinstance(r, (tuple, list)) else r for r in self._result]) + + @classmethod + def _get_unique_labels(cls, numbers: np.ndarray) -> list[str]: + """ For a list of threshold values for displaying in the bin name, get the lowest number of + decimal figures (down to int) required to have a unique set of folder names and return the + formatted numbers. + + Parameters + ---------- + numbers: :class:`numpy.ndarray` + The list of floating point threshold numbers being used as boundary points + + Returns + ------- + list[str] + The string formatted numbers at the lowest precision possible to represent them + uniquely + """ + i = 0 + while True: + rounded = [round(n, i) for n in numbers] + if len(set(rounded)) == len(numbers): + break + i += 1 + + if i == 0: + retval = [str(int(n)) for n in rounded] + else: + pre, post = zip(*[str(r).split(".") for r in rounded]) + rpad = max(len(x) for x in post) + retval = [f"{str(int(left))}.{str(int(right)).ljust(rpad, '0')}" + for left, right in zip(pre, post)] + logger.debug("rounded values: %s, formatted labels: %s", rounded, retval) + return retval + + def _binning_linear_threshold(self, units: str = "", multiplier: int = 1) -> list[list[str]]: + """ Standard linear binning method for binning by threshold. + + The minimum and maximum result from :attr:`_result` are taken, A range is created between + these min and max values and is divided to get the number of bins to hold the data + + Parameters + ---------- + units, str, optional + The units to use for the bin name for displaying the threshold values. This this should + correspond the value in position 1 of :attr:`_result`. + Default: "" (no units) + multiplier: int, optional + The amount to multiply the contents in position 1 of :attr:`_results` for displaying in + the bin folder name + + Returns + ------- + list + List of bins of filenames + """ + sizes = np.array([i[1] for i in self._result]) + thresholds = np.linspace(sizes.min(), sizes.max(), self._num_bins + 1) + labels = self._get_unique_labels(thresholds * multiplier) + + self._bin_names = [f"{self._method}_{idx:03d}_" + f"{labels[idx]}{units}_to_{labels[idx + 1]}{units}" + for idx in range(self._num_bins)] + + bins: list[list[str]] = [[] for _ in range(self._num_bins)] + for filename, result in self._result: + bin_idx = next(bin_id for bin_id, thresh in enumerate(thresholds) + if result <= thresh) - 1 + bins[bin_idx].append(filename) + + return bins + + def _binning(self) -> list[list[str]]: + """ Called when :attr:`binning` is first accessed. Checks if sorting has been done, if not + triggers it, then does binning + + Returns + ------- + list + List of bins of filenames + """ + if not self._result: + self._sort_filelist() + retval = self.binning() + + if not self._bin_names: + self._bin_names = [f"{self._method}_{i:03d}" for i in range(len(retval))] + + logger.debug({bin_name: len(bin_) for bin_name, bin_ in zip(self._bin_names, retval)}) + + return retval + + def sort(self) -> None: + """ Override for method specific logic for sorting the loaded statistics + + The scored list :attr:`_result` should be sorted in place + """ + raise NotImplementedError() + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Override for sort method's specificic logic. This method should be executed to get a + single score from a single image and add the result to :attr:`_result` + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` or ``None`` + A face image loaded from disk or ``None`` + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + raise NotImplementedError() + + def binning(self) -> list[list[str]]: + """ Group into bins by their sorted score. Override for method specific binning techniques. + + Binning takes the results from :attr:`_result` compiled during :func:`_sort_filelist` and + organizes into bins for output. + + Returns + ------- + list + List of bins of filenames + """ + raise NotImplementedError() + + @classmethod + def _mask_face(cls, image: np.ndarray, alignments: PNGHeaderAlignmentsDict) -> np.ndarray: + """ Function for applying the mask to an aligned face if both the face image and alignment + data are available. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The aligned face image loaded from disk + alignments: Dict + The alignments data corresponding to the loaded image + + Returns + ------- + :class:`numpy.ndarray` + The original image with the mask applied + """ + det_face = DetectedFace() + det_face.from_png_meta(alignments) + aln_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), + image=image, + centering="legacy", + size=256, + is_aligned=True) + assert aln_face.face is not None + + mask = det_face.mask.get("components", det_face.mask.get("extended", None)) + + if mask is None and not cls._log_mask_once: + logger.warning("No masks are available for masking the data. Results are likely to be " + "sub-standard") + cls._log_mask_once = True + + if mask is None: + return aln_face.face + + mask.set_sub_crop(aln_face.pose.offset[mask.stored_centering], + aln_face.pose.offset["legacy"], + centering="legacy") + nmask = cv2.resize(mask.mask, (256, 256), interpolation=cv2.INTER_CUBIC)[..., None] + return np.minimum(aln_face.face, nmask) + + +class SortMultiMethod(SortMethod): + """ A Parent sort method that runs 2 different underlying methods (one for sorting one for + binning) in instances where grouping has been requested, but the sort method is different from + the group method + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + sort_method: :class:`SortMethod` + A sort method object for sorting the images + group_method: :class:`SortMethod` + A sort method object used for sorting and binning the images + """ + def __init__(self, + arguments: Namespace, + sort_method: SortMethod, + group_method: SortMethod) -> None: + self._sorter = sort_method + self._grouper = group_method + self._is_built = False + super().__init__(arguments) + + def _get_file_iterator(self, input_dir: str) -> InfoLoader: + """ Override to get a group specific iterator. If the sorter and grouper use the same kind + of iterator, use that. Otherwise return the 'all' iterator, as which ever way it is cut all + outputs will be required. Monkey patch the actual loader used into the children in case of + any callbacks. + + Parameters + ---------- + input_dir: str + Full path to containing folder of faces to be supported + + Returns + ------- + :class:`InfoLoader` + The correct InfoLoader iterator for the current sort method + """ + if self._sorter.loader_type == self._grouper.loader_type: + retval = InfoLoader(input_dir, self._sorter.loader_type) + else: + retval = InfoLoader(input_dir, "all") + self._sorter._iterator = retval # pylint:disable=protected-access + self._grouper._iterator = retval # pylint:disable=protected-access + return retval + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Score a single image for sort method: "distance", "yaw" "pitch" or "size" and add the + result to :attr:`_result` + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` or ``None`` + A face image loaded from disk or ``None`` + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + self._sorter.score_image(filename, image, alignments) + self._grouper.score_image(filename, image, alignments) + + def sort(self) -> None: + """ Sort the sorter and grouper methods """ + logger.debug("Sorting") + self._sorter.sort() + self._result = self._sorter.sorted_filelist # type:ignore + self._grouper.sort() + self._binned = self._grouper.binned + self._bin_names = self._grouper.bin_names + logger.debug("Sorted") + + def binning(self) -> list[list[str]]: + """ Override standard binning, to bin by the group-by method and sort by the sorting + method. + + Go through the grouped binned results, and reorder each bin contents based on the + sorted list + + Returns + ------- + list + List of bins of filenames + """ + sorted_ = self._result + output: list[list[str]] = [] + for bin_ in tqdm(self._binned, desc="Binning and sorting", file=sys.stdout, leave=False): + indices: dict[int, str] = {} + for filename in bin_: + indices[sorted_.index(filename)] = filename + output.append([indices[idx] for idx in sorted(indices)]) + return output + + +class SortBlur(SortMethod): + """ Sort images by blur or blur-fft amount + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + def __init__(self, arguments: Namespace, is_group: bool = False) -> None: + super().__init__(arguments, loader_type="all", is_group=is_group) + method = arguments.group_method if self._is_group else arguments.sort_method + self._use_fft = method == "blur_fft" + + def estimate_blur(self, image: np.ndarray, alignments=None) -> float: + """ Estimate the amount of blur an image has with the variance of the Laplacian. + Normalize by pixel number to offset the effect of image size on pixel gradients & variance. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The face image to calculate blur for + alignments: dict, optional + The metadata for the face image or ``None`` if no metadata is available. If metadata is + provided the face will be masked by the "components" mask prior to calculating blur. + Default:``None`` + + Returns + ------- + float + The estimated blur score for the face + """ + if alignments is not None: + image = self._mask_face(image, alignments) + if image.ndim == 3: + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + blur_map = cv2.Laplacian(image, cv2.CV_32F) + score = np.var(blur_map) / np.sqrt(image.shape[0] * image.shape[1]) + return score + + def estimate_blur_fft(self, + image: np.ndarray, + alignments: PNGHeaderAlignmentsDict | None = None) -> float: + """ Estimate the amount of blur a fft filtered image has. + + Parameters + ---------- + image: :class:`numpy.ndarray` + Use Fourier Transform to analyze the frequency characteristics of the masked + face using 2D Discrete Fourier Transform (DFT) filter to find the frequency domain. + A mean value is assigned to the magnitude spectrum and returns a blur score. + Adapted from https://www.pyimagesearch.com/2020/06/15/ + opencv-fast-fourier-transform-fft-for-blur-detection-in-images-and-video-streams/ + alignments: dict, optional + The metadata for the face image or ``None`` if no metadata is available. If metadata is + provided the face will be masked by the "components" mask prior to calculating blur. + Default:``None`` + + Returns + ------- + float + The estimated fft blur score for the face + """ + if alignments is not None: + image = self._mask_face(image, alignments) + + if image.ndim == 3: + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + height, width = image.shape + c_height, c_width = (int(height / 2.0), int(width / 2.0)) + fft = np.fft.fft2(image) + fft_shift = np.fft.fftshift(fft) + fft_shift[c_height - 75:c_height + 75, c_width - 75:c_width + 75] = 0 + ifft_shift = np.fft.ifftshift(fft_shift) + shift_back = np.fft.ifft2(ifft_shift) + magnitude = np.log(np.abs(shift_back)) + score = np.mean(magnitude) + + return score + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Score a single image for blur or blur-fft and add the result to :attr:`_result` + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` + A face image loaded from disk + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + assert image is not None + if self._log_once: + msg = "Grouping" if self._is_group else "Sorting" + inf = "fft_filtered " if self._use_fft else " " + logger.info("%s by estimated %simage blur...", msg, inf) + self._log_once = False + + estimator = self.estimate_blur_fft if self._use_fft else self.estimate_blur + self._result.append((filename, estimator(image, alignments))) + + def sort(self) -> None: + """ Sort by metric score. Order in reverse for distance sort. """ + logger.info("Sorting...") + self._result = sorted(self._result, key=operator.itemgetter(1), reverse=True) + + def binning(self) -> list[list[str]]: + """ Create bins to split linearly from the lowest to the highest sample value + + Returns + ------- + list + List of bins of filenames + """ + return self._binning_linear_threshold(multiplier=100) + + +class SortColor(SortMethod): + """ Score by channel average intensity or black pixels. + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + def __init__(self, arguments: Namespace, is_group: bool = False) -> None: + super().__init__(arguments, loader_type="face", is_group=is_group) + self._desired_channel = {'gray': 0, 'luma': 0, 'orange': 1, 'green': 2} + + method = arguments.group_method if self._is_group else arguments.sort_method + self._method = method.replace("color_", "") + + def _convert_color(self, image: np.ndarray) -> np.ndarray: + """ Helper function to convert color spaces + + Parameters + ---------- + image: :class:`numpy.ndarray` + The original image to convert color space for + + Returns + ------- + :class:`numpy.ndarray` + The color converted image + """ + if self._method == 'gray': + conversion = np.array([[0.0722], [0.7152], [0.2126]]) + else: + conversion = np.array([[0.25, 0.5, 0.25], [-0.5, 0.0, 0.5], [-0.25, 0.5, -0.25]]) + + operation = 'ijk, kl -> ijl' if self._method == "gray" else 'ijl, kl -> ijk' + path = np.einsum_path(operation, image[..., :3], conversion, optimize='optimal')[0] + return np.einsum(operation, image[..., :3], conversion, optimize=path).astype('float32') + + def _near_split(self, bin_range: int) -> list[int]: + """ Obtain the split for the given number of bins for the given range + + Parameters + ---------- + bin_range: int + The range of data to separate into bins + + Returns + ------- + list + The split dividers for the given number of bins for the given range + """ + quotient, remainder = divmod(bin_range, self._num_bins) + seps = [quotient + 1] * remainder + [quotient] * (self._num_bins - remainder) + uplimit = 0 + bins = [0] + for sep in seps: + bins.append(uplimit + sep) + uplimit += sep + return bins + + def binning(self) -> list[list[str]]: + """ Group into bins by percentage of black pixels """ + # TODO. Only grouped by black pixels. Check color + + logger.info("Grouping by percentage of %s...", self._method) + + # Starting the binning process + bins: list[list[str]] = [[] for _ in range(self._num_bins)] + # Get edges of bins from 0 to 100 + bins_edges = self._near_split(100) + # Get the proper bin number for each img order + img_bins = np.digitize([float(x[1]) for x in self._result], bins_edges, right=True) + + # Place imgs in bins + for idx, _bin in enumerate(img_bins): + bins[_bin].append(self._result[idx][0]) + + retval = [b for b in bins if b] + return retval + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Score a single image for color + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` + A face image loaded from disk + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + if self._log_once: + msg = "Grouping" if self._is_group else "Sorting" + if self._method == "black": + logger.info("%s by percentage of black pixels...", msg) + else: + logger.info("%s by channel average intensity...", msg) + self._log_once = False + + assert image is not None + if self._method == "black": + score = np.ndarray.all(image == [0, 0, 0], axis=2).sum()/image.size*100*3 + else: + channel_to_sort = self._desired_channel[self._method] + score = np.average(self._convert_color(image), axis=(0, 1))[channel_to_sort] + self._result.append((filename, score)) + + def sort(self) -> None: + """ Sort by metric score. Order in reverse for distance sort. """ + if self._method == "black": + self._sort_black_pixels() + return + self._result = sorted(self._result, key=operator.itemgetter(1), reverse=True) + + def _sort_black_pixels(self) -> None: + """ Sort by percentage of black pixels + + Calculates the sum of black pixels, gets the percentage X 3 channels + """ + img_list_len = len(self._result) + for i in tqdm(range(0, img_list_len - 1), + desc="Comparing black pixels", file=sys.stdout, + leave=False): + for j in range(0, img_list_len-i-1): + if self._result[j][1] > self._result[j+1][1]: + temp = self._result[j] + self._result[j] = self._result[j+1] + self._result[j+1] = temp + + +class SortFace(SortMethod): + """ Sort by identity similarity using VGG Face 2 + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + + _logged_lm_count_once = False + _warning = ("Extracted faces do not contain facial landmark data. Results sorted by this " + "method are likely to be sub-standard.") + + def __init__(self, arguments: Namespace, is_group: bool = False) -> None: + super().__init__(arguments, loader_type="all", is_group=is_group) + self._vgg_face = VGGFace(exclude_gpus=arguments.exclude_gpus) + self._vgg_face.init_model() + threshold = arguments.threshold + self._output_update_info = True + self._threshold: float | None = 0.25 if threshold < 0 else threshold + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Processing logic for sort by face method. + + Reads header information from the PNG file to look for VGGFace2 embedding. If it does not + exist, the embedding is obtained and added back into the PNG Header. + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` + A face image loaded from disk + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + if not alignments: + msg = ("The images to be sorted do not contain alignment data. Images must have " + "been generated by Faceswap's Extract process.\nIf you are sorting an " + "older faceset, then you should re-extract the faces from your source " + "alignments file to generate this data.") + raise FaceswapError(msg) + + if self._log_once: + msg = "Grouping" if self._is_group else "Sorting" + logger.info("%s by identity similarity...", msg) + self._log_once = False + + if alignments.get("identity", {}).get("vggface2"): + embedding = np.array(alignments["identity"]["vggface2"], dtype="float32") + + if not self._logged_lm_count_once and len(alignments["landmarks_xy"]) == 4: + logger.warning(self._warning) + self._logged_lm_count_once = True + + self._result.append((filename, embedding)) + return + + if self._output_update_info: + logger.info("VGG Face2 Embeddings are being written to the image header. " + "Sorting by this method will be quicker next time") + self._output_update_info = False + + a_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), + image=image, + centering="legacy", + size=self._vgg_face.input_size, + is_aligned=True) + + if a_face.landmark_type == LandmarkType.LM_2D_4 and not self._logged_lm_count_once: + logger.warning(self._warning) + self._logged_lm_count_once = True + + face = a_face.face + assert face is not None + embedding = self._vgg_face.predict(face[None, ...])[0] + alignments.setdefault("identity", {})["vggface2"] = embedding.tolist() + self._iterator.update_png_header(filename, alignments) + self._result.append((filename, embedding)) + + def sort(self) -> None: + """ Sort by dendogram. + + Parameters + ---------- + matched_list: list + The list of tuples with filename in first position and face encoding in the 2nd + + Returns + ------- + list + The original list, sorted for this metric + """ + logger.info("Sorting by ward linkage. This may take some time...") + preds = np.array([item[1] for item in self._result]) + indices = Cluster(np.array(preds), "ward", threshold=self._threshold)() + self._result = [(self._result[idx][0], float(score)) for idx, score in indices] + + def binning(self) -> list[list[str]]: + """ Group into bins by their sorted score + + The bin ID has been output in the 2nd column of :attr:`_result` so use that for binnin + + Returns + ------- + list + List of bins of filenames + """ + num_bins = len(set(int(i[1]) for i in self._result)) + logger.info("Grouping by %s...", self.__class__.__name__.replace("Sort", "")) + bins: list[list[str]] = [[] for _ in range(num_bins)] + + for filename, bin_id in self._result: + bins[int(bin_id)].append(filename) + + return bins + + +class SortHistogram(SortMethod): + """ Sort by image histogram similarity or dissimilarity + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + def __init__(self, arguments: Namespace, is_group: bool = False) -> None: + super().__init__(arguments, loader_type="all", is_group=is_group) + method = arguments.group_method if self._is_group else arguments.sort_method + self._is_dissim = method == "hist-dissim" + self._threshold: float = 0.3 if arguments.threshold < 0.0 else arguments.threshold + + def _calc_histogram(self, + image: np.ndarray, + alignments: PNGHeaderAlignmentsDict | None) -> np.ndarray: + if alignments: + image = self._mask_face(image, alignments) + return cv2.calcHist([image], [0], None, [256], [0, 256]) + + def _sort_dissim(self) -> None: + """ Sort histograms by dissimilarity """ + img_list_len = len(self._result) + for i in tqdm(range(0, img_list_len), + desc="Comparing histograms", + file=sys.stdout, + leave=False): + score_total = 0 + for j in range(0, img_list_len): + if i == j: + continue + score_total += cv2.compareHist(self._result[i][1], + self._result[j][1], + cv2.HISTCMP_BHATTACHARYYA) + self._result[i][2] = score_total + + self._result = sorted(self._result, key=operator.itemgetter(2), reverse=True) + + def _sort_sim(self) -> None: + """ Sort histograms by similarity """ + img_list_len = len(self._result) + for i in tqdm(range(0, img_list_len - 1), + desc="Comparing histograms", + file=sys.stdout, + leave=False): + min_score = float("inf") + j_min_score = i + 1 + for j in range(i + 1, img_list_len): + score = cv2.compareHist(self._result[i][1], + self._result[j][1], + cv2.HISTCMP_BHATTACHARYYA) + if score < min_score: + min_score = score + j_min_score = j + (self._result[i + 1], self._result[j_min_score]) = (self._result[j_min_score], + self._result[i + 1]) + + @classmethod + def _get_avg_score(cls, image: np.ndarray, references: list[np.ndarray]) -> float: + """ Return the average histogram score between a face and reference images + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to test + references: list + List of reference images to test the original image against + + Returns + ------- + float + The average score between the histograms + """ + scores = [] + for img2 in references: + score = cv2.compareHist(image, img2, cv2.HISTCMP_BHATTACHARYYA) + scores.append(score) + return sum(scores) / len(scores) + + def binning(self) -> list[list[str]]: + """ Group into bins by histogram """ + msg = "dissimilarity" if self._is_dissim else "similarity" + logger.info("Grouping by %s...", msg) + + # Groups are of the form: group_num -> reference histogram + reference_groups: dict[int, list[np.ndarray]] = {} + + # Bins array, where index is the group number and value is + # an array containing the file paths to the images in that group + bins: list[list[str]] = [] + + threshold = self._threshold + + img_list_len = len(self._result) + reference_groups[0] = [T.cast(np.ndarray, self._result[0][1])] + bins.append([self._result[0][0]]) + + for i in tqdm(range(1, img_list_len), + desc="Grouping", + file=sys.stdout, + leave=False): + current_key = -1 + current_score = float("inf") + for key, value in reference_groups.items(): + score = self._get_avg_score(self._result[i][1], value) + if score < current_score: + current_key, current_score = key, score + + if current_score < threshold: + reference_groups[T.cast(int, current_key)].append(self._result[i][1]) + bins[current_key].append(self._result[i][0]) + else: + reference_groups[len(reference_groups)] = [self._result[i][1]] + bins.append([self._result[i][0]]) + + return bins + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Collect the histogram for the given face + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` + A face image loaded from disk + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + if self._log_once: + msg = "Grouping" if self._is_group else "Sorting" + logger.info("%s by histogram similarity...", msg) + self._log_once = False + + assert image is not None + self._result.append((filename, self._calc_histogram(image, alignments))) + + def sort(self) -> None: + """ Sort by histogram. """ + logger.info("Comparing histograms and sorting...") + if self._is_dissim: + self._sort_dissim() + return + self._sort_sim() diff --git a/tools/sort/sort_methods_aligned.py b/tools/sort/sort_methods_aligned.py new file mode 100644 index 0000000000..8f0ff0b8ea --- /dev/null +++ b/tools/sort/sort_methods_aligned.py @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 +""" Sorting methods that use the properties of a :class:`lib.align.AlignedFace` object to obtain +their sorting metrics. +""" +from __future__ import annotations +import logging +import operator +import sys +import typing as T + +import numpy as np +from tqdm import tqdm + +from lib.align import AlignedFace, LandmarkType +from lib.utils import FaceswapError +from .sort_methods import SortMethod + +if T.TYPE_CHECKING: + from argparse import Namespace + from lib.align.alignments import PNGHeaderAlignmentsDict + +logger = logging.getLogger(__name__) + + +class SortAlignedMetric(SortMethod): + """ Sort by comparison of metrics stored in an Aligned Face objects. This is a parent class + for sort by aligned metrics methods. Individual methods should inherit from this class + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + sort_reverse: bool, optional + ``True`` if the sorted results should be in reverse order. Default: ``True`` + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + + _logged_lm_count_once: bool = False + + def _get_metric(self, aligned_face: AlignedFace) -> np.ndarray | float: + """ Obtain the correct metric for the given sort method" + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float or :class:`numpy.ndarray` + The metric for the current face based on chosen sort method + """ + raise NotImplementedError + + def sort(self) -> None: + """ Sort by metric score. Order in reverse for distance sort. """ + logger.info("Sorting...") + self._result = sorted(self._result, key=operator.itemgetter(1), reverse=True) + + def score_image(self, + filename: str, + image: np.ndarray | None, + alignments: PNGHeaderAlignmentsDict | None) -> None: + """ Score a single image for sort method: "distance", "yaw", "pitch" or "size" and add the + result to :attr:`_result` + + Parameters + ---------- + filename: str + The filename of the currently processing image + image: :class:`np.ndarray` or ``None`` + A face image loaded from disk or ``None`` + alignments: dict or ``None`` + The alignments dictionary for the aligned face or ``None`` + """ + if self._log_once: + msg = "Grouping" if self._is_group else "Sorting" + logger.info("%s by %s...", msg, self._method) + self._log_once = False + + if not alignments: + msg = ("The images to be sorted do not contain alignment data. Images must have " + "been generated by Faceswap's Extract process.\nIf you are sorting an " + "older faceset, then you should re-extract the faces from your source " + "alignments file to generate this data.") + raise FaceswapError(msg) + + face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32")) + if (not self._logged_lm_count_once + and face.landmark_type == LandmarkType.LM_2D_4 + and self.__class__.__name__ != "SortSize"): + logger.warning("You have selected to sort by an aligned metric, but at least one face " + "does not contain facial landmark data. This probably won't work") + self._logged_lm_count_once = True + self._result.append((filename, self._get_metric(face))) + + +class SortDistance(SortAlignedMetric): + """ Sorting mechanism for sorting faces from small to large """ + def _get_metric(self, aligned_face: AlignedFace) -> float: + """ Obtain the distance from mean face metric for the given face + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The distance metric for the current face + """ + return aligned_face.average_distance + + def sort(self) -> None: + """ Override default sort to sort in ascending order. """ + logger.info("Sorting...") + self._result = sorted(self._result, key=operator.itemgetter(1), reverse=False) + + def binning(self) -> list[list[str]]: + """ Create bins to split linearly from the lowest to the highest sample value + + Returns + ------- + list + List of bins of filenames + """ + return self._binning_linear_threshold(multiplier=100) + + +class SortPitch(SortAlignedMetric): + """ Sorting mechansim for sorting a face by pitch (down to up) """ + def _get_metric(self, aligned_face: AlignedFace) -> float: + """ Obtain the pitch metric for the given face + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The pitch metric for the current face + """ + return aligned_face.pose.pitch + + def binning(self) -> list[list[str]]: + """ Create bins from 0 degrees to 180 degrees based on number of bins + + Allocate item to bin when it is in range of one of the pre-allocated bins + + Returns + ------- + list + List of bins of filenames + """ + thresholds = np.linspace(90, -90, self._num_bins + 1) + + # Start bin names from 0 for more intuitive experience + names = np.flip(thresholds.astype("int")) + 90 + self._bin_names = [f"{self._method}_" + f"{idx:03d}_{int(names[idx])}" + f"degs_to_{int(names[idx + 1])}degs" + for idx in range(self._num_bins)] + + bins: list[list[str]] = [[] for _ in range(self._num_bins)] + for filename, result in self._result: + result = np.clip(result, -90.0, 90.0) + bin_idx = next(bin_id for bin_id, thresh in enumerate(thresholds) + if result >= thresh) - 1 + bins[bin_idx].append(filename) + return bins + + +class SortYaw(SortPitch): + """ Sorting mechansim for sorting a face by yaw (left to right). Same logic as sort pitch, but + with different metric """ + def _get_metric(self, aligned_face: AlignedFace) -> float: + """ Obtain the yaw metric for the given face + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The yaw metric for the current face + """ + return aligned_face.pose.yaw + + +class SortRoll(SortPitch): + """ Sorting mechansim for sorting a face by roll (rotation). Same logic as sort pitch, but + with different metric """ + def _get_metric(self, aligned_face: AlignedFace) -> float: + """ Obtain the roll metric for the given face + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The yaw metric for the current face + """ + return aligned_face.pose.roll + + +class SortSize(SortAlignedMetric): + """ Sorting mechanism for sorting faces from small to large """ + def _get_metric(self, aligned_face: AlignedFace) -> float: + """ Obtain the size metric for the given face + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The size metric for the current face + """ + roi = aligned_face.original_roi + size = ((roi[1][0] - roi[0][0]) ** 2 + (roi[1][1] - roi[0][1]) ** 2) ** 0.5 + return size + + def binning(self) -> list[list[str]]: + """ Create bins to split linearly from the lowest to the highest sample value + + Allocate item to bin when it is in range of one of the pre-allocated bins + + Returns + ------- + list + List of bins of filenames + """ + return self._binning_linear_threshold(units="px") + + +class SortFaceCNN(SortAlignedMetric): + """ Sort by landmark similarity or dissimilarity + + Parameters + ---------- + arguments: :class:`argparse.Namespace` + The command line arguments passed to the sort process + is_group: bool, optional + Set to ``True`` if this class is going to be called exclusively for binning. + Default: ``False`` + """ + def __init__(self, arguments: Namespace, is_group: bool = False) -> None: + super().__init__(arguments, is_group=is_group) + self._is_dissim = self._method == "face-cnn-dissim" + self._threshold: float = 7.2 if arguments.threshold < 1.0 else arguments.threshold + + def _get_metric(self, aligned_face: AlignedFace) -> np.ndarray: + """ Obtain the xy aligned landmarks for the face" + + Parameters + ---------- + aligned_face: :class:`lib.align.AlignedFace` + The aligned face to extract the metric from + + Returns + ------- + float + The metric for the current face based on chosen sort method + """ + return aligned_face.landmarks + + def sort(self) -> None: + """ Sort by landmarks. """ + logger.info("Comparing landmarks and sorting...") + if self._is_dissim: + self._sort_landmarks_dissim() + return + self._sort_landmarks_ssim() + + def _sort_landmarks_ssim(self) -> None: + """ Sort landmarks by similarity """ + img_list_len = len(self._result) + for i in tqdm(range(0, img_list_len - 1), desc="Comparing", file=sys.stdout, leave=False): + min_score = float("inf") + j_min_score = i + 1 + for j in range(i + 1, img_list_len): + fl1 = self._result[i][1] + fl2 = self._result[j][1] + score = np.sum(np.absolute((fl2 - fl1).flatten())) + if score < min_score: + min_score = score + j_min_score = j + (self._result[i + 1], self._result[j_min_score]) = (self._result[j_min_score], + self._result[i + 1]) + + def _sort_landmarks_dissim(self) -> None: + """ Sort landmarks by dissimilarity """ + logger.info("Comparing landmarks...") + img_list_len = len(self._result) + for i in tqdm(range(0, img_list_len - 1), desc="Comparing", file=sys.stdout, leave=False): + score_total = 0 + for j in range(i + 1, img_list_len): + if i == j: + continue + fl1 = self._result[i][1] + fl2 = self._result[j][1] + score_total += np.sum(np.absolute((fl2 - fl1).flatten())) + self._result[i][2] = score_total + + logger.info("Sorting...") + self._result = sorted(self._result, key=operator.itemgetter(2), reverse=True) + + def binning(self) -> list[list[str]]: + """ Group into bins by CNN face similarity + + Returns + ------- + list + List of bins of filenames + """ + msg = "dissimilarity" if self._is_dissim else "similarity" + logger.info("Grouping by face-cnn %s...", msg) + + # Groups are of the form: group_num -> reference faces + reference_groups: dict[int, list[np.ndarray]] = {} + + # Bins array, where index is the group number and value is + # an array containing the file paths to the images in that group. + bins: list[list[str]] = [] + + # Comparison threshold used to decide how similar + # faces have to be to be grouped together. + # It is multiplied by 1000 here to allow the cli option to use smaller + # numbers. + threshold = self._threshold * 1000 + img_list_len = len(self._result) + + for i in tqdm(range(0, img_list_len - 1), + desc="Grouping", + file=sys.stdout, + leave=False): + fl1 = self._result[i][1] + + current_key = -1 + current_score = float("inf") + + for key, references in reference_groups.items(): + try: + score = self._get_avg_score(fl1, references) + except TypeError: + score = float("inf") + except ZeroDivisionError: + score = float("inf") + if score < current_score: + current_key, current_score = key, score + + if current_score < threshold: + reference_groups[current_key].append(fl1[0]) + bins[current_key].append(self._result[i][0]) + else: + reference_groups[len(reference_groups)] = [self._result[i][1]] + bins.append([self._result[i][0]]) + + return bins + + @classmethod + def _get_avg_score(cls, face: np.ndarray, references: list[np.ndarray]) -> float: + """ Return the average CNN similarity score between a face and reference images + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face to check against reference images + references: list + List of reference arrays to compare the face against + + Returns + ------- + float + The average score between the face and the references + """ + scores = [] + for ref in references: + score = np.sum(np.absolute((ref - face).flatten())) + scores.append(score) + return sum(scores) / len(scores) diff --git a/update_deps.py b/update_deps.py index b4a49452ab..8065de564b 100644 --- a/update_deps.py +++ b/update_deps.py @@ -3,30 +3,32 @@ Checks for installed Conda / Pip packages and updates accordingly """ +import logging +import os +import sys -from setup import Environment, Install, Output +from lib.logger import log_setup +from setup import Environment, Install -_LOGGER = None +logger = logging.getLogger(__name__) -def output(msg): - """ Output to print or logger """ - if _LOGGER is not None: - _LOGGER.info(msg) - else: - Output().info(msg) +def main(is_gui=False) -> None: + """ Check for and update dependencies - -def main(logger=None): - """ Check for and update dependencies """ - if logger is not None: - global _LOGGER # pylint:disable=global-statement - _LOGGER = logger - output("Updating dependencies...") - update = Environment(logger=logger, updater=True) - Install(update) - output("Dependencies updated") + Parameters + ---------- + is_gui: bool, optional + ``True`` if being called by the GUI. Prevents the updater from outputting progress bars + which get scrambled in the GUI + """ + logger.info("Updating dependencies...") + update = Environment(updater=True) + Install(update, is_gui=is_gui) + logger.info("Dependencies updated") if __name__ == "__main__": + logfile = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "faceswap_update.log") + log_setup("INFO", logfile, "setup") main()