diff --git a/bin/video2x_setup.py b/bin/video2x_setup.py index a302ea6..1c7d1a9 100644 --- a/bin/video2x_setup.py +++ b/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') - self._install_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 """ @@ -113,9 +152,17 @@ class Video2xSetup: with open('video2x.json', 'r') as template: template_dict = json.load(template) template.close() - - 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')) + + # 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['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: - for chunk in stream.iter_content(chunk_size=chunk_size): - if chunk: - print('!', end='') - output.write(chunk) - print() + with tqdm(total=total_size, ascii=True) as progress_bar: + for chunk in stream.iter_content(chunk_size=chunk_size): + if chunk: + output.write(chunk) + 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: