moved main to __main__

This commit is contained in:
k4yt3x 2022-05-01 08:34:00 +00:00
parent e01d24c164
commit 862b811517
2 changed files with 211 additions and 2 deletions

View File

@ -31,4 +31,4 @@ __version__ = "5.0.0-beta5"
# generated by the following lines
from .interpolator import Interpolator
from .upscaler import Upscaler
from .video2x import Video2X, main
from .video2x import Video2X

View File

@ -22,9 +22,218 @@ Date Created: July 3, 2021
Last Modified: February 26, 2022
"""
import argparse
import os
import pathlib
import sys
from .video2x import main
from loguru import logger
from rich import print as rich_print
from . import __version__
from .video2x import LOGURU_FORMAT, Video2X
LEGAL_INFO = f"""Video2X\t\t{__version__}
Author:\t\tK4YT3X
License:\tGNU AGPL v3
Github Page:\thttps://github.com/k4yt3x/video2x
Contact:\ti@k4yt3x.com"""
# algorithms available for upscaling tasks
UPSCALING_ALGORITHMS = [
"waifu2x",
"srmd",
"realsr",
"realcugan",
]
# algorithms available for frame interpolation tasks
INTERPOLATION_ALGORITHMS = ["rife"]
def parse_arguments() -> argparse.Namespace:
"""
parse command line arguments
:rtype argparse.Namespace: command parsing results
"""
parser = argparse.ArgumentParser(
prog="video2x",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--version", help="show version information and exit", action="store_true"
)
parser.add_argument(
"-i",
"--input",
type=pathlib.Path,
help="input file/directory path",
required=True,
)
parser.add_argument(
"-o",
"--output",
type=pathlib.Path,
help="output file/directory path",
required=True,
)
parser.add_argument(
"-p", "--processes", type=int, help="number of processes to launch", default=1
)
parser.add_argument(
"-l",
"--loglevel",
choices=["trace", "debug", "info", "success", "warning", "error", "critical"],
default="info",
)
# upscaler arguments
action = parser.add_subparsers(
help="action to perform", dest="action", required=True
)
upscale = action.add_parser(
"upscale",
help="upscale a file",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
add_help=False,
)
upscale.add_argument(
"--help", action="help", help="show this help message and exit"
)
upscale.add_argument("-w", "--width", type=int, help="output width")
upscale.add_argument("-h", "--height", type=int, help="output height")
upscale.add_argument("-n", "--noise", type=int, help="denoise level", default=3)
upscale.add_argument(
"-a",
"--algorithm",
choices=UPSCALING_ALGORITHMS,
help="algorithm to use for upscaling",
default=UPSCALING_ALGORITHMS[0],
)
upscale.add_argument(
"-t",
"--threshold",
type=float,
help=(
"skip if the percent difference between two adjacent frames is below this"
" value; set to 0 to process all frames"
),
default=0,
)
# interpolator arguments
interpolate = action.add_parser(
"interpolate",
help="interpolate frames for file",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
add_help=False,
)
interpolate.add_argument(
"--help", action="help", help="show this help message and exit"
)
interpolate.add_argument(
"-a",
"--algorithm",
choices=UPSCALING_ALGORITHMS,
help="algorithm to use for upscaling",
default=INTERPOLATION_ALGORITHMS[0],
)
interpolate.add_argument(
"-t",
"--threshold",
type=float,
help=(
"skip if the percent difference between two adjacent frames exceeds this"
" value; set to 100 to interpolate all frames"
),
default=10,
)
return parser.parse_args()
def main() -> int:
"""
command line entrypoint for direct CLI invocation
:rtype int: 0 if completed successfully, else other int
"""
try:
# display version and lawful informaition
if "--version" in sys.argv:
rich_print(LEGAL_INFO)
return 0
# parse command line arguments
args = parse_arguments()
# check input/output file paths
if not args.input.exists():
logger.critical(f"Cannot find input file: {args.input}")
return 1
if not args.input.is_file():
logger.critical("Input path is not a file")
return 1
if not args.output.parent.exists():
logger.critical(f"Output directory does not exist: {args.output.parent}")
return 1
# set logger level
if os.environ.get("LOGURU_LEVEL") is None:
os.environ["LOGURU_LEVEL"] = args.loglevel.upper()
# remove default handler
logger.remove()
# add new sink with custom handler
logger.add(sys.stderr, colorize=True, format=LOGURU_FORMAT)
# print package version and copyright notice
logger.opt(colors=True).info(f"<magenta>Video2X {__version__}</magenta>")
logger.opt(colors=True).info(
"<magenta>Copyright (C) 2018-2022 K4YT3X and contributors.</magenta>"
)
# initialize video2x object
video2x = Video2X()
if args.action == "upscale":
video2x.upscale(
args.input,
args.output,
args.width,
args.height,
args.noise,
args.processes,
args.threshold,
args.algorithm,
)
elif args.action == "interpolate":
video2x.interpolate(
args.input,
args.output,
args.processes,
args.threshold,
args.algorithm,
)
# don't print the traceback for manual terminations
except KeyboardInterrupt:
return 2
except Exception as error:
logger.exception(error)
return 1
# if no exceptions were produced
else:
logger.success("Processing completed successfully")
return 0
if __name__ == "__main__":
sys.exit(main())