mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 06:38:22 +00:00
moved main to __main__
This commit is contained in:
parent
e01d24c164
commit
862b811517
@ -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
|
||||
|
@ -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())
|
||||
|
Loading…
Reference in New Issue
Block a user