From 15780d535595013157ec6167ad86deb81d7ff20b Mon Sep 17 00:00:00 2001
From: K4YT3X <k4t@protonmail.com>
Date: Sat, 24 Feb 2018 16:55:26 -0500
Subject: [PATCH] updates processing unit selection option

---
 video2x.py | 23 ++++++++++++++++++++++-
 waifu2x.py |  5 +++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/video2x.py b/video2x.py
index 0fcc2a8..282fb34 100644
--- a/video2x.py
+++ b/video2x.py
@@ -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()
diff --git a/waifu2x.py b/waifu2x.py
index 60753d8..c7cf28c 100644
--- a/waifu2x.py
+++ b/waifu2x.py
@@ -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))