diff --git a/README.md b/README.md index bc8abce..bf0a1d0 100644 --- a/README.md +++ b/README.md @@ -18,10 +18,16 @@ Component names that are **bolded** can be automatically downloaded and configur 2. AMD GPU / Nvidia GPU 3. AMD GPU driver / Nvidia GPU driver / Nvidia CUDNN 4. [**FFmpeg**](https://ffmpeg.zeranoe.com/builds/) -5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases) +5. [**waifu2x-caffe**](https://github.com/lltcggie/waifu2x-caffe/releases) / [**waifu2x-converter-cpp**](https://github.com/DeadSix27/waifu2x-converter-cpp/releases) / [**waifu2x-ncnn-vulkan**](https://github.com/nihui/waifu2x-ncnn-vulkan) ## Recent Changes +### 2.9.0 (July 27, 2019) + +- Changed file handling method from `os` to `pathlib` +- Removed f_string dependency and support for legacy versions of Python +- Organized file import statements + ### 2.8.1 (July 9, 2019) - Added automatic pixel format detection @@ -35,12 +41,6 @@ Component names that are **bolded** can be automatically downloaded and configur - Fixed video2x custom temp folder bug found by @cr08 . -### 2.7.0 (March 30, 2019) - -- Added support for different extracted image formats. -- Redesigned FFmpeg wrapper, FFmpeg settings are now customizable in the `video2x.json` config file. -- Other minor enhancements and adjustments (e.g. argument -> method variable) - ### Setup Script 1.3.0 (June 25, 2019) - Added automatic installation support for `waifu2x-ncnn-vulkan` diff --git a/bin/exceptions.py b/bin/exceptions.py index 29dc428..6e8029a 100644 --- a/bin/exceptions.py +++ b/bin/exceptions.py @@ -1,12 +1,10 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Video2X Exceptions Dev: K4YT3X Date Created: December 13, 2018 -Last Modified: March 19, 2019 +Last Modified: July 27, 2019 """ diff --git a/bin/ffmpeg.py b/bin/ffmpeg.py index 4d11d80..2032c42 100644 --- a/bin/ffmpeg.py +++ b/bin/ffmpeg.py @@ -1,25 +1,28 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ -Name: FFMPEG Class +Name: Video2X FFmpeg Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: July 26, 2019 +Last Modified: July 27, 2019 Description: This class handles all FFmpeg related operations. """ -from avalon_framework import Avalon + +# built-in imports import json -import subprocess import os +import pathlib +import subprocess + +# third-party imports +from avalon_framework import Avalon class Ffmpeg: """This class communicates with FFmpeg - This class deals with FFmpeg. It handles extracitng + This class deals with FFmpeg. It handles extracting frames, stripping audio, converting images into videos and inserting audio tracks to videos. """ @@ -27,18 +30,17 @@ 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 - List all supported pixel formats and their corresponding - bit depth. + List all supported pixel formats and their + corresponding bit depth. Returns: dictionary -- JSON dict of all pixel formats to bit depth @@ -50,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 @@ -71,7 +76,7 @@ class Ffmpeg: """ Gets input video information This method reads input video information - using ffprobe in dictionary. + using ffprobe in dictionary Arguments: input_video {string} -- input video file path @@ -94,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')) @@ -101,8 +109,7 @@ class Ffmpeg: def extract_frames(self, input_video, extracted_frames): """Extract every frame from original videos - This method extracts every frame from videoin - using FFmpeg + This method extracts every frame from input video using FFmpeg Arguments: input_video {string} -- input video path @@ -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')) @@ -130,8 +137,7 @@ class Ffmpeg: def convert_video(self, framerate, resolution, upscaled_frames): """Converts images into videos - This method converts a set of images into a - video. + This method converts a set of images into a video Arguments: framerate {float} -- target video framerate @@ -150,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 @@ -171,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) @@ -187,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 ] @@ -249,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) @@ -271,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 b42d471..1306129 100644 --- a/bin/image_cleaner.py +++ b/bin/image_cleaner.py @@ -1,18 +1,17 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Video2X Image Cleaner Author: BrianPetkovsek Author: K4YT3X Date Created: March 24, 2019 -Last Modified: April 28, 2019 +Last Modified: July 27, 2019 Description: This class is to remove the extracted frames that have already been upscaled. """ +# built-in imports import os import threading import time @@ -60,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/requirements.txt b/bin/requirements.txt index 5855437..472e89c 100644 --- a/bin/requirements.txt +++ b/bin/requirements.txt @@ -4,4 +4,3 @@ GPUtil psutil requests tqdm -future-fstrings>=1.1.0 diff --git a/bin/upscaler.py b/bin/upscaler.py index ea31f84..3e1faae 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -1,12 +1,10 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Video2X Upscaler Author: K4YT3X Date Created: December 10, 2018 -Last Modified: July 9, 2019 +Last Modified: July 27, 2019 Dev: SAT3LL @@ -16,22 +14,28 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3), (C) 2018-2019 K4YT3X """ -from avalon_framework import Avalon +# local imports from exceptions import * from ffmpeg import Ffmpeg -from fractions import Fraction from image_cleaner import ImageCleaner -from tqdm import tqdm from waifu2x_caffe import Waifu2xCaffe from waifu2x_converter import Waifu2xConverter 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 +from tqdm import tqdm class Upscaler: @@ -58,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): @@ -81,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 @@ -106,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: @@ -117,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 @@ -163,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() @@ -175,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 @@ -188,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)) @@ -203,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] @@ -216,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, @@ -273,7 +278,6 @@ class Upscaler: if len(self.upscaler_exceptions) != 0: raise(self.upscaler_exceptions[0]) - def run(self): """Main controller for Video2X @@ -286,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 06f2cdd..9ad8764 100644 --- a/bin/video2x.py +++ b/bin/video2x.py @@ -1,8 +1,6 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - -""" +# -*- coding: utf-8 -*- +r""" __ __ _ _ ___ __ __ \ \ / / (_) | | |__ \ \ \ / / @@ -15,7 +13,7 @@ __ __ _ _ ___ __ __ Name: Video2X Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: July 9, 2019 +Last Modified: July 27, 2019 Dev: BrianPetkovsek Dev: SAT3LL @@ -43,13 +41,16 @@ enlarging engine. It extracts frames from a video, enlarge it by a number of times without losing any details or quality, keeping lines smooth and edges sharp. """ -from avalon_framework import Avalon + +# local imports from upscaler import Upscaler + +# built-in imports import argparse -import GPUtil +import glob import json import os -import psutil +import pathlib import re import shutil import sys @@ -57,7 +58,12 @@ import tempfile import time import traceback -VERSION = '2.8.1' +# third-party imports +from avalon_framework import Avalon +import GPUtil +import psutil + +VERSION = '2.9.0' # each thread might take up to 2.5 GB during initialization. # (system memory, not to be confused with GPU memory) @@ -76,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 @@ -103,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}') @@ -129,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') @@ -193,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 @@ -271,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']) @@ -292,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 @@ -322,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') @@ -356,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 f9b84f8..528f531 100644 --- a/bin/video2x_setup.py +++ b/bin/video2x_setup.py @@ -1,13 +1,11 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- - - """ Name: Video2X Setup Script Author: K4YT3X Author: BrianPetkovsek Date Created: November 28, 2018 -Last Modified: June 26, 2019 +Last Modified: July 27, 2019 Dev: SAT3LL @@ -26,14 +24,19 @@ Installation Details: - waifu2x_ncnn_vulkan: %LOCALAPPDATA%\\video2x\\waifu2x-ncnn-vulkan """ + +# built-in imports 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('Deleting: {}'.format(file)) - os.remove(file) - else: - print('Deleting: {}'.format(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,13 +224,42 @@ def download(url, save_path, chunk_size=4096): from tqdm import tqdm import requests - output_file = os.path.join(save_path, url.split('/')[-1]) - print('Downloading: {}'.format(url)) - print('Chunk size: {}'.format(chunk_size)) - print('Saving to: {}'.format(output_file)) + save_path = pathlib.Path(save_path) - stream = requests.get(url, stream=True) - total_size = int(stream.headers['content-length']) + # 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}') # Write content into file with open(output_file, 'wb') as output: @@ -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 @@ -252,7 +286,7 @@ if __name__ == '__main__': try: args = process_arguments() print('Video2X Setup Script') - print('Version: {}'.format(VERSION)) + print(f'Version: {VERSION}') # do not install pip modules if script # is packaged in exe format diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 0d18592..cbec6f2 100644 --- a/bin/waifu2x_caffe.py +++ b/bin/waifu2x_caffe.py @@ -1,20 +1,22 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Waifu2x Caffe Driver Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: July 9, 2019 +Last Modified: July 27, 2019 Description: This class is a high-level wrapper for waifu2x-caffe. """ -from avalon_framework import Avalon + +# built-in imports import subprocess import threading +# third-party imports +from avalon_framework import Avalon + class Waifu2xCaffe: """This class communicates with waifu2x cui engine diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index d3d5ae3..156b24e 100644 --- a/bin/waifu2x_converter.py +++ b/bin/waifu2x_converter.py @@ -1,21 +1,23 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Waifu2x Converter CPP Driver Author: K4YT3X Date Created: February 8, 2019 -Last Modified: June 15, 2019 +Last Modified: July 27, 2019 Description: This class is a high-level wrapper for waifu2x-converter-cpp. """ -from avalon_framework import Avalon -import os + +# built-in imports +import pathlib import subprocess import threading +# third-party imports +from avalon_framework import Avalon + class Waifu2xConverter: """This class communicates with waifu2x cui engine @@ -62,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() @@ -79,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: diff --git a/bin/waifu2x_ncnn_vulkan.py b/bin/waifu2x_ncnn_vulkan.py index 72f3baa..5b139a2 100644 --- a/bin/waifu2x_ncnn_vulkan.py +++ b/bin/waifu2x_ncnn_vulkan.py @@ -1,23 +1,25 @@ #!/usr/bin/env python3 -# -*- coding: future_fstrings -*- - - +# -*- coding: utf-8 -*- """ Name: Waifu2x NCNN Vulkan Driver Author: SAT3LL Date Created: June 26, 2019 -Last Modified: June 26, 2019 +Last Modified: July 27, 2019 Dev: K4YT3X Description: This class is a high-level wrapper for waifu2x_ncnn_vulkan. """ -from avalon_framework import Avalon + +# built-in imports import os import subprocess import threading +# third-party imports +from avalon_framework import Avalon + class Waifu2xNcnnVulkan: """This class communicates with waifu2x ncnn vulkan engine