1.2.0 added automatic setup for waifu2x-converter-cpp

bin/video2x_setup.py

@@ -4,7 +4,7 @@
 Name: Video2X Setup Script
 Author: K4YT3X
 Date Created: November 28, 2018
-Last Modified: March 9, 2019
+Last Modified: March 26, 2019
 
 Licensed under the GNU General Public License Version 3 (GNU GPL v3),
     available at: https://www.gnu.org/licenses/gpl-3.0.txt
@@ -19,6 +19,7 @@ Installation Details:
 - waifu2x-caffe: %LOCALAPPDATA%\\video2x\\waifu2x-caffe
 
 """
+import argparse
 import json
 import os
 import subprocess
@@ -31,7 +32,20 @@ import zipfile
 # later in the script.
 # import requests
 
-VERSION = '1.1.0'
+VERSION = '1.2.0'
+
+
+def process_arguments():
+    """Processes CLI arguments
+    """
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    # video options
+    general_options = parser.add_argument_group('General Options')
+    general_options.add_argument('-d', '--driver', help='driver to download and configure', action='store', choices=['all', 'waifu2x_caffe', 'waifu2x_converter'], required=True)
+
+    # parse arguments
+    return parser.parse_args()
 
 
 class Video2xSetup:
@@ -41,7 +55,8 @@ class Video2xSetup:
     script. All files will be installed under %LOCALAPPDATA%\\video2x.
     """
 
-    def __init__(self):
+    def __init__(self, driver):
+        self.driver = driver
         self.trash = []
 
     def run(self):
@@ -52,8 +67,13 @@ class Video2xSetup:
         print('\nInstalling FFMPEG')
         self._install_ffmpeg()
 
-        print('\nInstalling waifu2x-caffe')
+        if self.driver == 'all':
             self._install_waifu2x_caffe()
+            self._install_waifu2x_converter_cpp()
+        elif self.driver == 'waifu2x_caffe':
+            self._install_waifu2x_caffe()
+        elif self.driver == 'waifu2x_converter':
+            self._install_waifu2x_converter_cpp()
 
         print('\nGenerating Video2X configuration file')
         self._generate_config()
@@ -66,7 +86,7 @@ class Video2xSetup:
         """
         with open('requirements.txt', 'r') as req:
             for line in req:
-                package = line.split(' ')[0]
+                package = line.split('==')[0]
                 pip_install(package)
 
     def _cleanup(self):
@@ -82,7 +102,7 @@ class Video2xSetup:
     def _install_ffmpeg(self):
         """ Install FFMPEG
         """
-        latest_release = 'https://ffmpeg.zeranoe.com/builds/win64/static/ffmpeg-4.1-win64-static.zip'
+        latest_release = 'https://ffmpeg.zeranoe.com/builds/win64/static/ffmpeg-latest-win64-static.zip'
 
         ffmpeg_zip = download(latest_release, tempfile.gettempdir())
         self.trash.append(ffmpeg_zip)
@@ -93,6 +113,7 @@ class Video2xSetup:
     def _install_waifu2x_caffe(self):
         """ Install waifu2x_caffe
         """
+        print('\nInstalling waifu2x-caffe')
         import requests
 
         # Get latest release of waifu2x-caffe via GitHub API
@@ -106,6 +127,24 @@ class Video2xSetup:
         with zipfile.ZipFile(waifu2x_caffe_zip) as zipf:
             zipf.extractall('{}\\video2x'.format(os.getenv('localappdata')))
 
+    def _install_waifu2x_converter_cpp(self):
+        """ Install waifu2x_caffe
+        """
+        print('\nInstalling waifu2x-converter-cpp')
+        import re
+        import requests
+
+        # Get latest release of waifu2x-caffe via GitHub API
+        latest_release = json.loads(requests.get('https://api.github.com/repos/DeadSix27/waifu2x-converter-cpp/releases/latest').content)
+
+        for a in latest_release['assets']:
+            if re.search(r'waifu2x-DeadSix27-win64_v[0-9]*\.zip', a['browser_download_url']):
+                waifu2x_converter_cpp_zip = download(a['browser_download_url'], tempfile.gettempdir())
+                self.trash.append(waifu2x_converter_cpp_zip)
+
+        with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf:
+            zipf.extractall('{}\\video2x\\waifu2x-converter-cpp'.format(os.getenv('localappdata')))
+
     def _generate_config(self):
         """ Generate video2x config
         """
@@ -114,8 +153,16 @@ class Video2xSetup:
             template_dict = json.load(template)
             template.close()
 
+        # configure only the specified drivers
+        if self.driver == 'all':
             template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(os.getenv('localappdata'))
-        template_dict['ffmpeg']['ffmpeg_path'] = '{}\\video2x\\ffmpeg-4.1-win64-static\\bin'.format(os.getenv('localappdata'))
+            template_dict['waifu2x_converter']['waifu2x_converter_path'] = '{}\\video2x\\waifu2x-converter-cpp'.format(os.getenv('localappdata'))
+        elif self.driver == 'waifu2x_caffe':
+            template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = '{}\\video2x\\waifu2x-caffe\\waifu2x-caffe-cui.exe'.format(os.getenv('localappdata'))
+        elif self.driver == 'waifu2x_converter':
+            template_dict['waifu2x_converter']['waifu2x_converter_path'] = '{}\\video2x\\waifu2x-converter-cpp'.format(os.getenv('localappdata'))
+
+        template_dict['ffmpeg']['ffmpeg_path'] = '{}\\video2x\\ffmpeg-latest-win64-static\\bin'.format(os.getenv('localappdata'))
         template_dict['ffmpeg']['ffmpeg_hwaccel'] = 'auto'
         template_dict['ffmpeg']['extra_arguments'] = []
         template_dict['video2x']['video2x_cache_folder'] = False
@@ -130,6 +177,7 @@ class Video2xSetup:
 def download(url, save_path, chunk_size=4096):
     """ Download file to local with requests library
     """
+    from tqdm import tqdm
     import requests
 
     output_file = '{}\\{}'.format(save_path, url.split('/')[-1])
@@ -138,14 +186,15 @@ def download(url, save_path, chunk_size=4096):
     print('Saving to: {}'.format(output_file))
 
     stream = requests.get(url, stream=True)
+    total_size = int(stream.headers['content-length'])
 
     # Write content into file
     with open(output_file, 'wb') as output:
+        with tqdm(total=total_size, ascii=True) as progress_bar:
             for chunk in stream.iter_content(chunk_size=chunk_size):
                 if chunk:
-                print('!', end='')
                     output.write(chunk)
-    print()
+                    progress_bar.update(len(chunk))
 
     return output_file
 
@@ -161,9 +210,10 @@ def pip_install(package):
 
 if __name__ == "__main__":
     try:
+        args = process_arguments()
         print('Video2x Setup Script')
         print('Version: {}'.format(VERSION))
-        setup = Video2xSetup()
+        setup = Video2xSetup(args.driver)
         setup.run()
         print('\n Script finished successfully')
     except Exception: