diff --git a/bin/ffmpeg.py b/bin/ffmpeg.py index dc2973a..2032c42 100644 --- a/bin/ffmpeg.py +++ b/bin/ffmpeg.py @@ -11,8 +11,9 @@ Description: This class handles all FFmpeg related operations. # built-in imports import json -import subprocess import os +import pathlib +import subprocess # third-party imports from avalon_framework import Avalon @@ -29,12 +30,12 @@ class Ffmpeg: def __init__(self, ffmpeg_settings, image_format): self.ffmpeg_settings = ffmpeg_settings - self.ffmpeg_path = self.ffmpeg_settings['ffmpeg_path'] - self.ffmpeg_binary = os.path.join(self.ffmpeg_path, 'ffmpeg.exe') - self.ffmpeg_probe_binary = os.path.join(self.ffmpeg_path, 'ffprobe.exe') + self.ffmpeg_path = pathlib.Path(self.ffmpeg_settings['ffmpeg_path']) + self.ffmpeg_binary = self.ffmpeg_path / 'ffmpeg.exe' + self.ffmpeg_probe_binary = self.ffmpeg_path / 'ffprobe.exe' self.image_format = image_format self.pixel_format = None - + def get_pixel_formats(self): """ Get a dictionary of supported pixel formats @@ -51,6 +52,9 @@ class Ffmpeg: '-pix_fmts' ] + # turn elements into str + execute = [str(e) for e in execute] + Avalon.debug_info(f'Executing: {" ".join(execute)}') # initialize dictionary to store pixel formats @@ -95,6 +99,9 @@ class Ffmpeg: input_video ] + # turn elements into str + execute = [str(e) for e in execute] + Avalon.debug_info(f'Executing: {" ".join(execute)}') json_str = subprocess.run(execute, check=True, stdout=subprocess.PIPE).stdout return json.loads(json_str.decode('utf-8')) @@ -120,7 +127,7 @@ class Ffmpeg: execute.extend(self._read_configuration(phase='video_to_frames', section='output_options')) execute.extend([ - os.path.join(extracted_frames, f'extracted_%0d.{self.image_format}') + extracted_frames / f'extracted_%0d.{self.image_format}' ]) execute.extend(self._read_configuration(phase='video_to_frames')) @@ -149,17 +156,19 @@ class Ffmpeg: execute.extend(self._read_configuration(phase='frames_to_video', section='input_options')) # WORKAROUND FOR WAIFU2X-NCNN-VULKAN + # Dev: SAT3LL + # rename all .png.png suffixes to .png import re import shutil regex = re.compile(r'\.png\.png$') - for raw_frame in os.listdir(upscaled_frames): - shutil.move(os.path.join(upscaled_frames, raw_frame), os.path.join(upscaled_frames, regex.sub('.png', raw_frame))) + for frame_name in upscaled_frames.iterdir(): + (upscaled_frames / frame_name).rename(upscaled_frames / regex.sub('.png', str(frame_name))) # END WORKAROUND # append input frames path into command execute.extend([ '-i', - os.path.join(upscaled_frames, f'extracted_%d.{self.image_format}') + upscaled_frames / f'extracted_%d.{self.image_format}' ]) # read FFmpeg output options @@ -170,7 +179,7 @@ class Ffmpeg: # specify output file location execute.extend([ - os.path.join(upscaled_frames, 'no_audio.mp4') + upscaled_frames / 'no_audio.mp4' ]) self._execute(execute) @@ -186,7 +195,7 @@ class Ffmpeg: execute = [ self.ffmpeg_binary, '-i', - os.path.join(upscaled_frames, 'no_audio.mp4'), + upscaled_frames / 'no_audio.mp4', '-i', input_video ] @@ -248,7 +257,7 @@ class Ffmpeg: if value is not True: configuration.append(str(subvalue)) - # otherwise the value is typical + # otherwise the value is typical else: configuration.append(key) @@ -270,4 +279,8 @@ class Ffmpeg: int -- execution return code """ Avalon.debug_info(f'Executing: {execute}') + + # turn all list elements into string to avoid errors + execute = [str(e) for e in execute] + return subprocess.run(execute, shell=True, check=True).returncode diff --git a/bin/image_cleaner.py b/bin/image_cleaner.py index e90b13d..1306129 100644 --- a/bin/image_cleaner.py +++ b/bin/image_cleaner.py @@ -59,19 +59,19 @@ class ImageCleaner(threading.Thread): """ # list all images in the extracted frames - output_frames = [f for f in os.listdir(self.output_directory) if os.path.isfile(os.path.join(self.output_directory, f))] + output_frames = [f for f in self.output_directory.iterdir() if f.is_file()] # compare and remove frames downscaled images that finished being upscaled # within each thread's extracted frames directory - for i in range(self.threads): - dir_path = os.path.join(self.input_directory, str(i)) + for thread_id in range(self.threads): + dir_path = self.input_directory / str(thread_id) # for each file within all the directories - for f in os.listdir(dir_path): - file_path = os.path.join(dir_path, f) + for file in dir_path.iterdir(): + file_path = dir_path / file # if file also exists in the output directory, then the file # has already been processed, thus not needed anymore - if os.path.isfile(file_path) and f in output_frames: - os.remove(file_path) - output_frames.remove(f) + if file_path.is_file() and file in output_frames: + file_path.unlink(file) + output_frames.remove(file) diff --git a/bin/upscaler.py b/bin/upscaler.py index 9787a9c..3e1faae 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -25,12 +25,13 @@ from waifu2x_ncnn_vulkan import Waifu2xNcnnVulkan # built-in imports from fractions import Fraction import copy -import os +import pathlib import re import shutil import tempfile import threading import time +import traceback # third-party imports from avalon_framework import Avalon @@ -61,16 +62,16 @@ class Upscaler: self.scale_ratio = None self.model_dir = None self.threads = 5 - self.video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x') + self.video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x' self.image_format = 'png' self.preserve_frames = False def create_temp_directories(self): """create temporary directory """ - self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory) + self.extracted_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory)) Avalon.debug_info(f'Extracted frames are being saved to: {self.extracted_frames}') - self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_directory) + self.upscaled_frames = pathlib.Path(tempfile.mkdtemp(dir=self.video2x_cache_directory)) Avalon.debug_info(f'Upscaled frames are being saved to: {self.upscaled_frames}') def cleanup_temp_directories(self): @@ -84,7 +85,8 @@ class Upscaler: print(f'Cleaning up cache directory: {directory}') shutil.rmtree(directory) except (OSError, FileNotFoundError): - pass + print(f'Unable to delete: {directory}') + traceback.print_exc() def _check_arguments(self): # check if arguments are valid / all necessary argument @@ -109,7 +111,7 @@ class Upscaler: # get number of extracted frames total_frames = 0 for directory in extracted_frames_directories: - total_frames += len([f for f in os.listdir(directory) if f[-4:] == f'.{self.image_format}']) + total_frames += len([f for f in directory.iterdir() if str(f)[-4:] == f'.{self.image_format}']) with tqdm(total=total_frames, ascii=True, desc='Upscaling Progress') as progress_bar: @@ -120,7 +122,7 @@ class Upscaler: while not self.progress_bar_exit_signal: try: - total_frames_upscaled = len([f for f in os.listdir(self.upscaled_frames) if f[-4:] == f'.{self.image_format}']) + total_frames_upscaled = len([f for f in self.upscaled_frames.iterdir() if str(f)[-4:] == f'.{self.image_format}']) delta = total_frames_upscaled - previous_cycle_frames previous_cycle_frames = total_frames_upscaled @@ -166,9 +168,9 @@ class Upscaler: progress_bar.start() w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.image_format, self.upscaler_exceptions) - for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]: + for image in [f for f in self.upscaled_frames.iterdir() if f.is_file()]: renamed = re.sub(f'_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{self.image_format}', f'.{self.image_format}', image) - shutil.move(os.path.join(self.upscaled_frames, image), os.path.join(self.upscaled_frames, renamed)) + (self.upscaled_frames / image).rename(self.upscaled_frames / renamed) self.progress_bar_exit_signal = True progress_bar.join() @@ -178,7 +180,7 @@ class Upscaler: upscaler_threads = [] # list all images in the extracted frames - frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))] + frames = [(self.extracted_frames / f) for f in self.extracted_frames.iterdir() if f.is_file] # if we have less images than threads, # create only the threads necessary @@ -191,13 +193,13 @@ class Upscaler: thread_pool = [] thread_directories = [] for thread_id in range(self.threads): - thread_directory = os.path.join(self.extracted_frames, str(thread_id)) + thread_directory = self.extracted_frames / str(thread_id) thread_directories.append(thread_directory) # delete old directories and create new directories - if os.path.isdir(thread_directory): + if thread_directory.is_dir(): shutil.rmtree(thread_directory) - os.mkdir(thread_directory) + thread_directory.mkdir(parents=True, exist_ok=True) # append directory path into list thread_pool.append((thread_directory, thread_id)) @@ -206,7 +208,7 @@ class Upscaler: # until there is none left in the directory for image in frames: # move image - shutil.move(image, thread_pool[0][0]) + image.rename(thread_pool[0][0] / image.name) # rotate list thread_pool = thread_pool[-1:] + thread_pool[:-1] @@ -219,23 +221,23 @@ class Upscaler: if self.scale_ratio: thread = threading.Thread(target=w2.upscale, args=(thread_info[0], - self.upscaled_frames, - self.scale_ratio, - False, - False, - self.image_format, - self.upscaler_exceptions)) + self.upscaled_frames, + self.scale_ratio, + False, + False, + self.image_format, + self.upscaler_exceptions)) else: - thread = threading.Thread(target=w2.upscale, - args=(thread_info[0], - self.upscaled_frames, - False, - self.scale_width, - self.scale_height, - self.image_format, - self.upscaler_exceptions)) + thread = threading.Thread(target=w2.upscale, + args=(thread_info[0], + self.upscaled_frames, + False, + self.scale_width, + self.scale_height, + self.image_format, + self.upscaler_exceptions)) - # if the driver being used is waifu2x_ncnn_vulkan + # if the driver being used is waifu2x_ncnn_vulkan elif self.waifu2x_driver == 'waifu2x_ncnn_vulkan': w2 = Waifu2xNcnnVulkan(copy.deepcopy(self.waifu2x_settings)) thread = threading.Thread(target=w2.upscale, @@ -276,7 +278,6 @@ class Upscaler: if len(self.upscaler_exceptions) != 0: raise(self.upscaler_exceptions[0]) - def run(self): """Main controller for Video2X @@ -289,8 +290,8 @@ class Upscaler: self._check_arguments() # convert paths to absolute paths - self.input_video = os.path.abspath(self.input_video) - self.output_video = os.path.abspath(self.output_video) + self.input_video = self.input_video.absolute() + self.output_video = self.output_video.absolute() # initialize objects for ffmpeg and waifu2x-caffe fm = Ffmpeg(self.ffmpeg_settings, self.image_format) diff --git a/bin/video2x.py b/bin/video2x.py index bb97ad4..9ad8764 100644 --- a/bin/video2x.py +++ b/bin/video2x.py @@ -47,8 +47,10 @@ from upscaler import Upscaler # built-in imports import argparse +import glob import json import os +import pathlib import re import shutil import sys @@ -61,7 +63,7 @@ from avalon_framework import Avalon import GPUtil import psutil -VERSION = '2.8.1' +VERSION = '2.9.0' # each thread might take up to 2.5 GB during initialization. # (system memory, not to be confused with GPU memory) @@ -80,16 +82,16 @@ def process_arguments(): # video options file_options = parser.add_argument_group('File Options') - file_options.add_argument('-i', '--input', help='source video file/directory', action='store') - file_options.add_argument('-o', '--output', help='output video file/directory', action='store') + file_options.add_argument('-i', '--input', type=pathlib.Path, help='source video file/directory', action='store') + file_options.add_argument('-o', '--output', type=pathlib.Path, help='output video file/directory', action='store') # upscaler options upscaler_options = parser.add_argument_group('Upscaler Options') upscaler_options.add_argument('-m', '--method', help='upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn']) upscaler_options.add_argument('-d', '--driver', help='waifu2x driver', action='store', default='waifu2x_caffe', choices=['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']) - upscaler_options.add_argument('-y', '--model_dir', help='directory containing model JSON files', action='store') + upscaler_options.add_argument('-y', '--model_dir', type=pathlib.Path, help='directory containing model JSON files', action='store') upscaler_options.add_argument('-t', '--threads', help='number of threads to use for upscaling', action='store', type=int, default=1) - upscaler_options.add_argument('-c', '--config', help='video2x config file location', action='store', default=os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'video2x.json')) + upscaler_options.add_argument('-c', '--config', type=pathlib.Path, help='video2x config file location', action='store', default=pathlib.Path(sys.argv[0]).parent.absolute() / 'video2x.json') upscaler_options.add_argument('-b', '--batch', help='enable batch mode (select all default values to questions)', action='store_true') # scaling options @@ -107,12 +109,16 @@ def process_arguments(): def print_logo(): - print('__ __ _ _ ___ __ __') - print('\\ \\ / / (_) | | |__ \\ \\ \\ / /') - print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /') - print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <') - print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\') - print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\') + """print video2x logo""" + logo = r''' + __ __ _ _ ___ __ __ + \ \ / / (_) | | |__ \ \ \ / / + \ \ / / _ __| | ___ ___ ) | \ V / + \ \/ / | | / _` | / _ \ / _ \ / / > < + \ / | | | (_| | | __/ | (_) | / /_ / . \ + \/ |_| \__,_| \___| \___/ |____| /_/ \_\ + ''' + print(logo) print('\n Video2X Video Enlarger') spaces = ((44 - len(f'Version {VERSION}')) // 2) * ' ' print(f'{Avalon.FM.BD}\n{spaces} Version {VERSION}\n{Avalon.FM.RST}') @@ -133,7 +139,7 @@ def check_memory(): # GPUtil requires nvidia-smi.exe to interact with GPU if args.method == 'gpu' or args.method == 'cudnn': if not (shutil.which('nvidia-smi') or - os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe')): + pathlib.Path(r'C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe').is_file()): # Nvidia System Management Interface not available Avalon.warning('Nvidia-smi not available, skipping available memory check') Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using') @@ -197,28 +203,28 @@ def absolutify_paths(config): Returns: dict -- configuration file dictionary """ - current_directory = os.path.dirname(os.path.abspath(sys.argv[0])) + current_directory = pathlib.Path(sys.argv[0]).parent.absolute() # check waifu2x-caffe path if not re.match('^[a-z]:', config['waifu2x_caffe']['waifu2x_caffe_path'], re.IGNORECASE): - config['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(current_directory, config['waifu2x_caffe']['waifu2x_caffe_path']) + config['waifu2x_caffe']['waifu2x_caffe_path'] = current_directory / config['waifu2x_caffe']['waifu2x_caffe_path'] # check waifu2x-converter-cpp path if not re.match('^[a-z]:', config['waifu2x_converter']['waifu2x_converter_path'], re.IGNORECASE): - config['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(current_directory, config['waifu2x_converter']['waifu2x_converter_path']) + config['waifu2x_converter']['waifu2x_converter_path'] = current_directory / config['waifu2x_converter']['waifu2x_converter_path'] # check waifu2x_ncnn_vulkan path if not re.match('^[a-z]:', config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'], re.IGNORECASE): - config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(current_directory, config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path']) + config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = current_directory / config['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] # check ffmpeg path if not re.match('^[a-z]:', config['ffmpeg']['ffmpeg_path'], re.IGNORECASE): - config['ffmpeg']['ffmpeg_path'] = os.path.join(current_directory, config['ffmpeg']['ffmpeg_path']) + config['ffmpeg']['ffmpeg_path'] = current_directory / config['ffmpeg']['ffmpeg_path'] # check video2x cache path if config['video2x']['video2x_cache_directory']: if not re.match('^[a-z]:', config['video2x']['video2x_cache_directory'], re.IGNORECASE): - config['video2x']['video2x_cache_directory'] = os.path.join(current_directory, config['video2x']['video2x_cache_directory']) + config['video2x']['video2x_cache_directory'] = current_directory / config['video2x']['video2x_cache_directory'] return config @@ -275,19 +281,19 @@ config = absolutify_paths(config) # load waifu2x configuration if args.driver == 'waifu2x_caffe': waifu2x_settings = config['waifu2x_caffe'] - if not os.path.isfile(waifu2x_settings['waifu2x_caffe_path']): + if not pathlib.Path(waifu2x_settings['waifu2x_caffe_path']).is_file(): Avalon.error('Specified waifu2x-caffe directory doesn\'t exist') Avalon.error('Please check the configuration file settings') raise FileNotFoundError(waifu2x_settings['waifu2x_caffe_path']) elif args.driver == 'waifu2x_converter': waifu2x_settings = config['waifu2x_converter'] - if not os.path.isdir(waifu2x_settings['waifu2x_converter_path']): + if not pathlib.Path(waifu2x_settings['waifu2x_converter_path']).is_dir(): Avalon.error('Specified waifu2x-converter-cpp directory doesn\'t exist') Avalon.error('Please check the configuration file settings') raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path']) elif args.driver == 'waifu2x_ncnn_vulkan': waifu2x_settings = config['waifu2x_ncnn_vulkan'] - if not os.path.isfile(waifu2x_settings['waifu2x_ncnn_vulkan_path']): + if not pathlib.Path(waifu2x_settings['waifu2x_ncnn_vulkan_path']).is_file(): Avalon.error('Specified waifu2x_ncnn_vulkan directory doesn\'t exist') Avalon.error('Please check the configuration file settings') raise FileNotFoundError(waifu2x_settings['waifu2x_ncnn_vulkan_path']) @@ -296,28 +302,38 @@ elif args.driver == 'waifu2x_ncnn_vulkan': ffmpeg_settings = config['ffmpeg'] # load video2x settings -video2x_cache_directory = config['video2x']['video2x_cache_directory'] image_format = config['video2x']['image_format'].lower() preserve_frames = config['video2x']['preserve_frames'] -# create temp directories if they don't exist -if not video2x_cache_directory: - video2x_cache_directory = os.path.join(tempfile.gettempdir(), 'video2x') +# load cache directory +if isinstance(config['video2x']['video2x_cache_directory'], str): + video2x_cache_directory = pathlib.Path(config['video2x']['video2x_cache_directory']) +else: + video2x_cache_directory = pathlib.Path(tempfile.gettempdir()) / 'video2x' -if video2x_cache_directory and not os.path.isdir(video2x_cache_directory): - if not os.path.isfile(video2x_cache_directory) and not os.path.islink(video2x_cache_directory): - Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist') - if Avalon.ask('Create directory?', default=True, batch=args.batch): - if os.mkdir(video2x_cache_directory) is None: - Avalon.info(f'{video2x_cache_directory} created') - else: - Avalon.error(f'Unable to create {video2x_cache_directory}') - Avalon.error('Aborting...') - exit(1) +if video2x_cache_directory.exists() and not video2x_cache_directory.is_dir(): + Avalon.error('Specified cache directory is a file/link') + raise FileExistsError('Specified cache directory is a file/link') + +elif not video2x_cache_directory.exists(): + # if destination file is a file or a symbolic link + Avalon.warning(f'Specified cache directory {video2x_cache_directory} does not exist') + + # try creating the cache directory + if Avalon.ask('Create directory?', default=True, batch=args.batch): + try: + video2x_cache_directory.mkdir(parents=True, exist_ok=True) + Avalon.info(f'{video2x_cache_directory} created') + + # there can be a number of exceptions here + # PermissionError, FileExistsError, etc. + # therefore, we put a catch-them-all here + except Exception as e: + Avalon.error(f'Unable to create {video2x_cache_directory}') + Avalon.error('Aborting...') + raise e else: - Avalon.error('Specified cache directory is a file/link') - Avalon.error('Unable to continue, exiting...') - exit(1) + raise FileNotFoundError('Could not create cache directory') # start execution @@ -326,17 +342,17 @@ try: begin_time = time.time() # if input specified is a single file - if os.path.isfile(args.input): + if args.input.is_file(): # upscale single video file Avalon.info(f'Upscaling single video file: {args.input}') # check for input output format mismatch - if os.path.isdir(args.output): + if args.output.is_dir(): Avalon.error('Input and output path type mismatch') Avalon.error('Input is single file but output is directory') raise Exception('input output path type mismatch') - if not re.search('.*\..*$', args.output): + if not re.search('.*\..*$', str(args.output)): Avalon.error('No suffix found in output file path') Avalon.error('Suffix must be specified for FFmpeg') raise Exception('No suffix specified') @@ -360,12 +376,12 @@ try: upscaler.cleanup_temp_directories() # if input specified is a directory - elif os.path.isdir(args.input): + elif args.input.is_dir(): # upscale videos in a directory Avalon.info(f'Upscaling videos in directory: {args.input}') - for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]: - output_video = os.path.join(args.output, input_video) - upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings) + for input_video in [f for f in args.input.iterdir() if f.is_file()]: + output_video = args.output / input_video + upscaler = Upscaler(input_video=args.input / input_video, output_video=output_video, method=args.method, waifu2x_settings=waifu2x_settings, ffmpeg_settings=ffmpeg_settings) # set optional options upscaler.waifu2x_driver = args.driver diff --git a/bin/video2x_setup.py b/bin/video2x_setup.py index 8967193..528f531 100644 --- a/bin/video2x_setup.py +++ b/bin/video2x_setup.py @@ -29,11 +29,14 @@ Installation Details: import argparse import json import os +import pathlib +import re import shutil import subprocess import sys import tempfile import traceback +import urllib import zipfile # Requests doesn't come with windows, therefore @@ -43,6 +46,9 @@ import zipfile VERSION = '1.3.0' +# global static variables +LOCALAPPDATA = pathlib.Path(os.getenv('localappdata')) + def process_arguments(): """Processes CLI arguments @@ -104,14 +110,15 @@ class Video2xSetup: """ for file in self.trash: try: - if os.path.isfile(file): - print(f'Deleting: {file}') - os.remove(file) - else: - print(f'Deleting: {file}') + if file.is_dir(): + print(f'Deleting directory: {file}') shutil.rmtree(file) - except FileNotFoundError: - pass + else: + print(f'Deleting file: {file}') + file.unlink() + except Exception: + print(f'Error deleting: {file}') + traceback.print_exc() def _install_ffmpeg(self): """ Install FFMPEG @@ -122,7 +129,7 @@ class Video2xSetup: self.trash.append(ffmpeg_zip) with zipfile.ZipFile(ffmpeg_zip) as zipf: - zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x')) + zipf.extractall(LOCALAPPDATA / 'video2x') def _install_waifu2x_caffe(self): """ Install waifu2x_caffe @@ -139,7 +146,7 @@ class Video2xSetup: self.trash.append(waifu2x_caffe_zip) with zipfile.ZipFile(waifu2x_caffe_zip) as zipf: - zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x')) + zipf.extractall(LOCALAPPDATA / 'video2x') def _install_waifu2x_converter_cpp(self): """ Install waifu2x_caffe @@ -157,7 +164,7 @@ class Video2xSetup: self.trash.append(waifu2x_converter_cpp_zip) with zipfile.ZipFile(waifu2x_converter_cpp_zip) as zipf: - zipf.extractall(os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-converter-cpp')) + zipf.extractall(LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp') def _install_waifu2x_ncnn_vulkan(self): """ Install waifu2x-ncnn-vulkan @@ -176,12 +183,11 @@ class Video2xSetup: # extract then move (to remove the top level directory) with zipfile.ZipFile(waifu2x_ncnn_vulkan_zip) as zipf: - extraction_path = os.path.join(tempfile.gettempdir(), 'waifu2x-ncnn-vulkan-ext') + extraction_path = pathlib.Path(tempfile.gettempdir()) / 'waifu2x-ncnn-vulkan-ext' zipf.extractall(extraction_path) - shutil.move(os.path.join(extraction_path, os.listdir(extraction_path)[0]), os.path.join(os.getenv('localappdata'), 'video2x', 'waifu2x-ncnn-vulkan')) + shutil.move(extraction_path / os.listdir(extraction_path)[0], LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan') self.trash.append(extraction_path) - def _generate_config(self): """ Generate video2x config """ @@ -190,21 +196,19 @@ class Video2xSetup: template_dict = json.load(template) template.close() - local_app_data = os.getenv('localappdata') - # configure only the specified drivers if self.driver == 'all': - template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe') - template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp') - template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x-ncnn-vulkan.exe') + template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe' + template_dict['waifu2x_converter']['waifu2x_converter_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp' + template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe' elif self.driver == 'waifu2x_caffe': - template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-caffe', 'waifu2x-caffe-cui.exe') + template_dict['waifu2x_caffe']['waifu2x_caffe_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-caffe' / 'waifu2x-caffe-cui.exe' elif self.driver == 'waifu2x_converter': - template_dict['waifu2x_converter']['waifu2x_converter_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-converter-cpp') + template_dict['waifu2x_converter']['waifu2x_converter_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-converter-cpp' elif self.driver == 'waifu2x_ncnn_vulkan': - template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = os.path.join(local_app_data, 'video2x', 'waifu2x-ncnn-vulkan', 'waifu2x-ncnn-vulkan.exe') + template_dict['waifu2x_ncnn_vulkan']['waifu2x_ncnn_vulkan_path'] = LOCALAPPDATA / 'video2x' / 'waifu2x-ncnn-vulkan' / 'waifu2x-ncnn-vulkan.exe' - template_dict['ffmpeg']['ffmpeg_path'] = os.path.join(local_app_data, 'video2x', 'ffmpeg-latest-win64-static', 'bin') + template_dict['ffmpeg']['ffmpeg_path'] = LOCALAPPDATA / 'video2x' / 'ffmpeg-latest-win64-static' / 'bin' template_dict['video2x']['video2x_cache_directory'] = None template_dict['video2x']['preserve_frames'] = False @@ -220,14 +224,43 @@ def download(url, save_path, chunk_size=4096): from tqdm import tqdm import requests - output_file = os.path.join(save_path, url.split('/')[-1]) + save_path = pathlib.Path(save_path) + + # create target folder if it doesn't exist + save_path.mkdir(parents=True, exist_ok=True) + + # create requests stream for steaming file + stream = requests.get(url, stream=True, allow_redirects=True) + + # get file name + file_name = None + if 'content-disposition' in stream.headers: + disposition = stream.headers['content-disposition'] + try: + file_name = re.findall("filename=(.+)", disposition)[0].strip('"') + except IndexError: + pass + + if file_name is None: + # output_file = f'{save_path}\\{stream.url.split("/")[-1]}' + output_file = save_path / stream.url.split('/')[-1] + else: + output_file = save_path / file_name + + # decode url encoding + output_file = pathlib.Path(urllib.parse.unquote(str(output_file))) + + # get total size for progress bar if provided in headers + total_size = 0 + if 'content-length' in stream.headers: + total_size = int(stream.headers['content-length']) + + # print download information summary print(f'Downloading: {url}') + print(f'Total size: {total_size}') print(f'Chunk size: {chunk_size}') print(f'Saving to: {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: @@ -236,6 +269,7 @@ def download(url, save_path, chunk_size=4096): output.write(chunk) progress_bar.update(len(chunk)) + # return the full path of saved file return output_file diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index 70d5853..156b24e 100644 --- a/bin/waifu2x_converter.py +++ b/bin/waifu2x_converter.py @@ -11,7 +11,7 @@ for waifu2x-converter-cpp. """ # built-in imports -import os +import pathlib import subprocess import threading @@ -64,8 +64,7 @@ class Waifu2xConverter: # models_rgb must be specified manually for waifu2x-converter-cpp # if it's not specified in the arguments, create automatically if self.waifu2x_settings['model-dir'] is None: - self.waifu2x_settings['model-dir'] = os.path.join(self.waifu2x_settings['waifu2x_converter_path'], - 'models_rgb') + self.waifu2x_settings['model-dir'] = pathlib.Path(self.waifu2x_settings['waifu2x_converter_path']) / 'models_rgb' # print thread start message self.print_lock.acquire() @@ -81,7 +80,7 @@ class Waifu2xConverter: # the key doesn't need to be passed in this case if key == 'waifu2x_converter_path': - execute.append(os.path.join(str(value), 'waifu2x-converter-cpp.exe')) + execute.append(pathlib.Path(str(value)) / 'waifu2x-converter-cpp.exe') # null or None means that leave this option out (keep default) elif value is None or value is False: