waifu2x-caffe denoiser script

denoise.py

import glob
import argparse
import subprocess
import shutil
import sys
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Denoise images with waifu2x.")
    parser.add_argument("images", nargs="+", type=Path, help="The images to process.")
    parser.add_argument(
        "-o",
        "--out",
        type=Path,
        required=False,
        help="The path to write output to. Defaults to source directory.",
    )
    parser.add_argument(
        "-e",
        "--end",
        default="_denoised",
        help="The string to append to the end of the output file's name. Defaults to `_denoised`.",
    )
    parser.add_argument(
        "-f",
        "--format",
        default="png",
        help="The format to save output as. Defaults to `png`.",
    )
    parser.add_argument(
        "-q",
        "--quality",
        type=int,
        default=95,
        help="The quality percentage for lossly compression formats [0-100], Defaults to `95`.",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    if args.quality not in range(0, 101):
        raise ValueError("Quality must be in the range [0-100].")

    waifu_dir = shutil.which("waifu2x-caffe-cui")
    if waifu_dir is None:
        raise SystemError("Could not find waifu2x-caffe. Please add it to the PATH.")
    waifu_dir = Path(waifu_dir).parent

    for image_path in (
        Path(img_path)
        for img_glob in args.images
        for img_path in glob.glob(str(img_glob))
    ):
        image_path = image_path.resolve()

        process_args = [
            "waifu2x-caffe-cui",
            "-m",
            "noise",
            "-n",
            "3",
            "-p",
            "cudnn",
            "--model-dir",
            str(waifu_dir / "models" / "cunet"),
            "-i",
            str(image_path),
            "-e",
            args.format,
        ]

        out_path = None
        if args.out is not None:
            out_path = args.out.resolve()
            if not out_path.is_absolute():
                out_path = Path.cwd() / out_path
            if out_path.is_dir():
                out_path = out_path / Path(
                    image_path.stem + args.end + "." + args.format
                )
        else:
            out_path = image_path.parent / Path(
                image_path.stem + args.end + "." + args.format
            )
        process_args.extend(("-o", str(out_path)))
        print(out_path)

        retcode = subprocess.run(process_args, cwd=waifu_dir).returncode

        if retcode != 0:
            print("Error")
            sys.exit(retcode)

    print("Done")


if __name__ == "__main__":
    main()