waifu2x batch operation wrapper

waifu.sh

#!/bin/bash
# example method for best quality output:
#waifu2x -scale 2 -method scale -tta 1 -tta_level 8 -model_dir /usr/share/waifu2x/models/photo -i ~/test/test_1280.png -o ~/test/test_waifu2x_tta8.png

success(){
    # play a sound on successful file processing
    paplay ~/Music/sfx/185845__lloydevans09__incorrect-function.wav
    return 0
}
failure(){
    # play a sound on failed processing
    paplay ~/Music/sfx/342756__rhodesmas__failure-01.wav
    return 1
}

waifu () {
    if [[ $# -eq 0 ]]; then
        cat << EOF
Usage: waifu [OPTIONS] FILE1 [FILE2 ...]

# Doesn't support -l (list of filenames).
# Can scale above 2x
# Default sane arguments provided are: -s2 -tta8

Options:
    -d                  dry run, just print out the commands to exec.
    -n<1|2|3>           noise reduction level
    -s<NUM>             scale ratio; default: -s2
    -j<NUM>             numbers of threads launching at the same time; default: -j2
    -tta<NUM>           test-time augmentation averages the upscaling results of the following 8 augmented inputs to reduce artifacts; default: 8 (slow and expensive!)
    -model_dir_<PATH>   path to model directory [/usr/share/waifu2x/models/photo]
    --                  terminate options list
EOF
        #waifu2x --version
        #waifu2x --list-processor
        return 0
    fi

    local DRYRUN=''
    declare -i local SCALE=''
    local NOISE=''
    local JOBS=''
    local TTA=''
    local MDIR=''
    local METHOD=''

    while [[ $# -gt 0 ]]; do
        opt="$1"

        case $opt in
            --) # terminate options list
                shift
                break
            ;;
            -d) # dry run
                DRYRUN="yes"
            ;;
            -s*) # scale ratio
                SCALE=${opt:2}
            ;;
            -n*) # noise reduction level
                NOISE=${opt:2}
            ;;
            -j*) # concurrent jobs
                JOBS=${opt:2}
            ;;
            -tta*) # better quality at 8, worse at 0
                TTA=${opt:4}
            ;;
            -model_dir_*) # path to model dir
                MDIR=${opt:11}
            ;;
            *) # no more options
                break
            ;;
        esac
        shift
    done

    if [[ $# -eq 0 ]]; then
        echo no input file was specified, exiting.
        return 1
    fi

    local CONFIG=""
    local POSTFIX="_waifu"

    if [[ -n "$SCALE" ]] && [[ -n "$NOISE" ]]; then
        METHOD="noise_scale"
    elif [[ -n "$SCALE" ]] && [[ ! -n "$NOISE" ]]; then
        METHOD="scale"
    elif [[ ! -n "$SCALE" ]] && [[ -n "$NOISE" ]]; then
        METHOD="noise"
    # "user" method not implemented
    fi

    if [[ -n "$METHOD" ]]; then
        CONFIG+=" -m $METHOD"
    fi

    if [[ -n "$SCALE" ]] && [[ $SCALE -le 2 ]]; then
        CONFIG+=" -scale $SCALE"
        POSTFIX+="_s$SCALE"
    else
        POSTFIX+="_s$SCALE"
    fi

    if [[ -n "$MDIR" ]]; then
        CONFIG+=" -model_dir $MDIR"
    else
        CONFIG+=" -model_dir "/usr/share/waifu2x/models/photo""
    fi

    if [[ -n "$NOISE" ]]; then
        CONFIG+=" --noise_level $NOISE"
        POSTFIX+="_n$NOISE"
    else
        POSTFIX+=""
    fi

    if [[ -n "${TTA}" ]]; then
        CONFIG+=" -tta 1 -tta_level $TTA"
        POSTFIX+="_tta${TTA}"
    else
        POSTFIX+=""
    fi

    if [[ -n "$JOBS" ]]; then
        CONFIG+=" -thread $JOBS"
    else
        # defaults to 2 jobs
        CONFIG+=" -thread 2"
    fi

    for f in "$@"; do
        local NAME=$(basename -- "$f")
        local EXT="${NAME##*.}"
        local NAME="${NAME%.*}"
        
    if [[ -n "$DRYRUN" ]]; then
            echo "waifu2x${CONFIG} -i "$f" -o "$NAME$POSTFIX.$EXT""

    # user asked for more than scale 2, need several passes
    elif [[ "$SCALE" -gt 2 ]]; then
        local TEMP_NAME="${f}_temp_upscale"

        # first pass
        waifu2x $CONFIG -i "$f" -o "$TEMP_NAME"
        SCALE=${SCALE}-2

        # subsequent passes
        while [[ "$SCALE" -gt 0 ]];
        do
            if [[ $SCALE -eq 2 ]]; then # last pass
            waifu2x $CONFIG -i "$TEMP_NAME" -o "$NAME$POSTFIX.$EXT"
            break
        else
            waifu2x $CONFIG -i "$TEMP_NAME" -o "$TEMP_NAME"
            SCALE=${SCALE}-2
        fi
        done

        rm --interactive=never -v "$TEMP_NAME"

        # if the expected output file exists, we consider it a success
        if [[ -f "$NAME$POSTFIX.$EXT" ]]; then
            success
        else
            failure
        fi
    else
        waifu2x $CONFIG -i "$f" -o "$NAME$POSTFIX.$EXT"
        # if the exit code is 0, we consider it a success
        if [[ $? -eq 0 ]]; then
            success
        else
            failure
        fi
        fi
    done
}

# if only one argument and path exists, call with default sane values
if [[ $# -eq 1 ]] && [[ -f "$1" ]] ; then
    waifu -s2 -tta8 "$@"
elif [[ $# -eq 2 ]] && [[ "$1" = "-d" ]] && [[ -f "$2" ]]; then
    waifu -s2 -tta8 "$@"
else
    waifu "$@"
fi