OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-29 16:09:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

added model arg

Browse Source
This commit is contained in:
K4YT3X 2018-04-20 13:00:51 -04:00
parent a858fa8e9b
commit 254c761546
1 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
video2x.py
View File

@ -31,6 +31,7 @@ Version 1.1 alpha
from ffmpeg import FFMPEG from ffmpeg import FFMPEG
from fractions import Fraction from fractions import Fraction
from waifu2x import WAIFU2X from waifu2x import WAIFU2X
import avalon_framework as avalon
import argparse import argparse
import json import json
import os import os
@ -58,6 +59,7 @@ def processArguments():
control_group.add_argument("-f", "--factor", help="Factor to enlarge video by", action="store", default=2) control_group.add_argument("-f", "--factor", help="Factor to enlarge video by", action="store", default=2)
control_group.add_argument("-v", "--video", help="Specify video file", action="store", default=False) control_group.add_argument("-v", "--video", help="Specify video file", action="store", default=False)
control_group.add_argument("-o", "--output", help="Specify output file", action="store", default=False) control_group.add_argument("-o", "--output", help="Specify output file", action="store", default=False)
control_group.add_argument("-y", "--model_type", help="Specify model to use", action="store", default="anime_style_art_rgb")
control_group.add_argument("--cpu", help="Use CPU for enlarging", action="store_true", default=False) control_group.add_argument("--cpu", help="Use CPU for enlarging", action="store_true", default=False)
control_group.add_argument("--gpu", help="Use GPU for enlarging", action="store_true", default=False) control_group.add_argument("--gpu", help="Use GPU for enlarging", action="store_true", default=False)
control_group.add_argument("--cudnn", help="Use CUDNN for enlarging", action="store_true", default=False) control_group.add_argument("--cudnn", help="Use CUDNN for enlarging", action="store_true", default=False)
@ -82,12 +84,26 @@ def get_vid_info():
return json.loads(json_str) return json.loads(json_str)
def main(): def check_model_type(args):
models_available = ["upconv_7_anime_style_art_rgb", "upconv_7_photo",
"anime_style_art_rgb", "photo", "anime_style_art_y"]
if args.model_type not in models_available:
avalon.error('Specified model type not found!')
avalon.info("Available models:")
for model in models_available:
print(model)
exit(1)
def video2x():
"""Main controller for Video2X """Main controller for Video2X
This function controls the flow of video conversion This function controls the flow of video conversion
and handles all necessary functions. and handles all necessary functions.
""" """
check_model_type(args)
if args.cpu: if args.cpu:
method = "cpu" method = "cpu"
elif args.gpu: elif args.gpu:
@ -112,7 +128,7 @@ def main():
# Upscale Frames # Upscale Frames
if not os.path.isdir(FOLDEROUT): if not os.path.isdir(FOLDEROUT):
os.mkdir(FOLDEROUT) os.mkdir(FOLDEROUT)
w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height) w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height, args.model_type)
# Frames to Video # Frames to Video
fm.to_vid(framerate, final_resolution, FOLDEROUT) fm.to_vid(framerate, final_resolution, FOLDEROUT)
@ -138,7 +154,7 @@ elif not args.cpu and not args.gpu and not args.cudnn:
if __name__ == '__main__': if __name__ == '__main__':
try: try:
main() video2x()
except Exception as e: except Exception as e:
# This code block is reserved for future # This code block is reserved for future
# fail-safe handlers # fail-safe handlers