updates processing unit selection option

This commit is contained in:
K4YT3X 2018-02-24 16:55:26 -05:00
parent 57733d602b
commit 15780d5355
2 changed files with 25 additions and 3 deletions

View File

@ -36,6 +36,9 @@ def processArguments():
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("-o", "--output", help="Specify output file", action="store", 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("--cudnn", help="Use CUDNN for enlarging", action="store_true", default=False)
args = parser.parse_args()
@ -57,8 +60,15 @@ def main():
"""Main flow control function for video2x.
This function takes care of the order of operation.
"""
if args.cpu:
method = "cpu"
elif args.gpu:
method = "gpu"
elif args.cudnn:
method = "cudnn"
fm = FFMPEG("\"" + FFMPEG_PATH + "ffmpeg.exe\"", args.output)
w2 = WAIFU2X(WAIFU2X_PATH)
w2 = WAIFU2X(WAIFU2X_PATH, method)
# Extract Frames
if not os.path.isdir(FOLDERIN):
@ -84,4 +94,15 @@ def main():
processArguments()
if not args.video:
print("Error: You need to specify the video to process")
exit(1)
elif not args.output:
print("Error: You need to specify the output video name")
exit(1)
elif not args.cpu and not args.gpu and not args.cudnn:
print("Error: You need to specify the enlarging processing unit")
exit(1)
main()

View File

@ -17,8 +17,9 @@ import os
class WAIFU2X:
def __init__(self, waifu2x_path):
def __init__(self, waifu2x_path, method):
self.waifu2x_path = waifu2x_path
self.method = method
def upscale(self, folderin, folderout, width, height):
os.system("{} -p cpu -I png -i {} -e png -o {} -w {} -h {}".format(self.waifu2x_path, folderin, folderout, width, height))
os.system("{} -p {} -I png -i {} -e png -o {} -w {} -h {}".format(self.waifu2x_path, self.method, folderin, folderout, width, height))