diff --git a/README.md b/README.md index 59319fb..452589d 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,12 @@ Component names that are *italicized* can be automatically downloaded and config ## Recent Changes +### 2.6.3 (March 24, 2019) + +- Added image cleaner by @BrianPetkovsek which removes upscaled frames. +- Fixed some PEP8 issues. +- Exceptions in waifu2x are now caught, and script will now stop on waifu2x error instead of keep going on to FFMPEG. + ### 2.6.2 (March 19, 2019) - Removed `--model_dir` verification due to the rapidly evolving number of models added. diff --git a/bin/clear_image.py b/bin/clear_image.py deleted file mode 100644 index 17dfaa9..0000000 --- a/bin/clear_image.py +++ /dev/null @@ -1,54 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -Name: Waifu2x Image clearer -Author: BrianPetkovsek -Date Created: March 24, 2019 -Last Modified: March 25, 2019 - -Description: This class is to remove the -downscaled image files when upscale is finished -from waifu2x-caffe. -""" - -from threading import Thread -from time import sleep -import os - -class ClearImage(Thread): - def __init__(self, input_folder, output_folder,num_threads): - Thread.__init__(self) - self.input_folder = input_folder - self.output_folder = output_folder - self.num_threads = num_threads - self.running = False - - def run(self): - self.running = True - while(self.running): - self.removeFrames() - #delay in 1 second intrvals for stop trigger - i=0 - while self.running and i<20: - i+=1 - sleep(1) - - - def stop(self): - self.running = False - self.join() - - def removeFrames(self): - # list all images in the extracted frames - output_frames = [f for f in os.listdir(self.output_folder) if os.path.isfile(os.path.join(self.output_folder, f))] - - # compare and remove frames downscaled images that finished being upscaled - for i in range(self.num_threads): - dir_path = os.path.join(self.input_folder,str(i)) - for f in os.listdir(dir_path): - file_path = os.path.join(dir_path, f) - if os.path.isfile(file_path) and f in output_frames: - os.remove(file_path) - output_frames.remove(f) - - diff --git a/bin/image_cleaner.py b/bin/image_cleaner.py new file mode 100644 index 0000000..295682c --- /dev/null +++ b/bin/image_cleaner.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Name: Video2X Image Cleaner +Author: BrianPetkovsek +Author: K4YT3X +Date Created: March 24, 2019 +Last Modified: March 24, 2019 + +Description: This class is to remove the extracted frames +that have already been upscaled. +""" + +import os +import threading +import time + + +class ImageCleaner(threading.Thread): + """ Video2X Image Cleaner + + This class creates an object that keeps track of extracted + frames that has already been upscaled and are not needed + anymore. It then deletes them to save disk space. + + Extends: + threading.Thread + """ + + def __init__(self, input_folder, output_folder, num_threads): + threading.Thread.__init__(self) + self.input_folder = input_folder + self.output_folder = output_folder + self.num_threads = num_threads + self.running = False + + def run(self): + """ Run image cleaner + """ + self.running = True + + while self.running: + self.remove_upscaled_frames() + time.sleep(1) + + def stop(self): + """ Stop the image cleaner + """ + self.running = False + self.join() + + def remove_upscaled_frames(self): + """ remove frames that have already been upscaled + + This method compares the files in the extracted frames + folder with the upscaled frames folder, and removes + the frames that has already been upscaled. + """ + + # list all images in the extracted frames + output_frames = [f for f in os.listdir(self.output_folder) if os.path.isfile(os.path.join(self.output_folder, f))] + + # compare and remove frames downscaled images that finished being upscaled + # within each thread's extracted frames folder + for i in range(self.num_threads): + dir_path = os.path.join(self.input_folder, str(i)) + + # for each file within all the folders + for f in os.listdir(dir_path): + file_path = os.path.join(dir_path, f) + + # if file also exists in the output folder, 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) diff --git a/bin/upscaler.py b/bin/upscaler.py index 8aec5cb..befbd0c 100644 --- a/bin/upscaler.py +++ b/bin/upscaler.py @@ -4,7 +4,7 @@ Name: Video2X Upscaler Author: K4YT3X Date Created: December 10, 2018 -Last Modified: March 19, 2019 +Last Modified: March 24, 2019 Licensed under the GNU General Public License Version 3 (GNU GPL v3), available at: https://www.gnu.org/licenses/gpl-3.0.txt @@ -13,13 +13,13 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3), """ from avalon_framework import Avalon +from image_cleaner import ImageCleaner from exceptions import * from ffmpeg import Ffmpeg from fractions import Fraction from tqdm import tqdm from waifu2x_caffe import Waifu2xCaffe from waifu2x_converter import Waifu2xConverter -from clear_image import ClearImage import os import re import shutil @@ -59,12 +59,10 @@ class Upscaler: # create temporary folder/directories self.video2x_cache_folder = video2x_cache_folder - self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) - self.extracted_frames = self.extracted_frames_object.name + self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames)) - self.upscaled_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) - self.upscaled_frames = self.upscaled_frames_object.name + self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames)) self.preserve_frames = preserve_frames @@ -74,10 +72,13 @@ class Upscaler: # avalon framework cannot be used if python is shutting down # therefore, plain print is used if not self.preserve_frames: - print('Cleaning up cache directory: {}'.format(self.extracted_frames)) - self.extracted_frames_object.cleanup() - print('Cleaning up cache directory: {}'.format(self.upscaled_frames)) - self.upscaled_frames_object.cleanup() + + for directory in [self.extracted_frames, self.upscaled_frames]: + try: + print('Cleaning up cache directory: {}'.format()) + shutil.rmtree(directory) + except (OSError, FileNotFoundError): + pass def _check_arguments(self): # check if arguments are valid / all necessary argument @@ -160,6 +161,10 @@ class Upscaler: # create a container for all upscaler threads upscaler_threads = [] + # create a container for exceptions in threads + # if this thread is not empty, then an exception has occured + self.threads_exceptions = [] + # 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))] @@ -197,9 +202,9 @@ class Upscaler: for thread_info in thread_pool: # create thread if self.scale_ratio: - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.threads_exceptions)) else: - thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height)) + thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, self.threads_exceptions)) thread.name = thread_info[1] # add threads into the pool @@ -209,11 +214,11 @@ class Upscaler: progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,)) progress_bar.start() - #Create the clearer and start it - Avalon.debug_info('Starting image clearer...') - image_clear = ClearImage(self.extracted_frames,self.upscaled_frames,len(upscaler_threads)) - image_clear.start() - + # create the clearer and start it + Avalon.debug_info('Starting upscaled image cleaner') + image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads)) + image_cleaner.start() + # start all threads for thread in upscaler_threads: thread.start() @@ -221,13 +226,16 @@ class Upscaler: # wait for threads to finish for thread in upscaler_threads: thread.join() - - #upscaling done... kill the clearer - Avalon.debug_info('Stoping image clearer...') - image_clear.stop() + + # upscaling done, kill the clearer + Avalon.debug_info('Killing upscaled image cleaner') + image_cleaner.stop() self.progress_bar_exit_signal = True + if len(self.threads_exceptions) != 0: + raise(self.threads_exceptions[0]) + def run(self): """Main controller for Video2X @@ -300,5 +308,3 @@ class Upscaler: # migrate audio tracks and subtitles Avalon.info('Migrating audio tracks and subtitles to upscaled video') fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames) - - diff --git a/bin/video2x.py b/bin/video2x.py index e6e72e7..2d37a06 100644 --- a/bin/video2x.py +++ b/bin/video2x.py @@ -13,7 +13,7 @@ __ __ _ _ ___ __ __ Name: Video2X Controller Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: March 19, 2019 +Last Modified: March 24, 2019 Licensed under the GNU General Public License Version 3 (GNU GPL v3), available at: https://www.gnu.org/licenses/gpl-3.0.txt @@ -50,7 +50,7 @@ import tempfile import time import traceback -VERSION = '2.6.2' +VERSION = '2.6.3' # each thread might take up to 2.5 GB during initialization. # (system memory, not to be confused with GPU memory) @@ -120,7 +120,7 @@ def check_memory(): if not os.path.isfile('C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe'): # 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') + Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using') else: try: # "0" is GPU ID. Both waifu2x drivers use the first GPU available, therefore only 0 makes sense @@ -201,8 +201,19 @@ config = read_config(args.config) # load waifu2x configuration if args.driver == 'waifu2x_caffe': waifu2x_settings = config['waifu2x_caffe'] + if not os.path.isfile(waifu2x_settings['waifu2x_caffe_path']): + 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']): + Avalon.error('Specified waifu2x-conver-cpp directory doesn\'t exist') + Avalon.error('Please check the configuration file settings') + raise FileNotFoundError(waifu2x_settings['waifu2x_converter_path']) + +# check if waifu2x path is valid + # read FFMPEG configuration ffmpeg_settings = config['ffmpeg'] @@ -258,6 +269,7 @@ try: except Exception: Avalon.error('An exception has occurred') traceback.print_exc() + Avalon.warning('If you experience error \"cudaSuccess out of memory\", try reducing number of threads you\'re using') finally: # remove Video2X Cache folder try: diff --git a/bin/waifu2x_caffe.py b/bin/waifu2x_caffe.py index 8403c02..43487ac 100644 --- a/bin/waifu2x_caffe.py +++ b/bin/waifu2x_caffe.py @@ -4,7 +4,7 @@ Name: Waifu2x Caffe Driver Author: K4YT3X Date Created: Feb 24, 2018 -Last Modified: March 19, 2019 +Last Modified: March 24, 2019 Description: This class is a high-level wrapper for waifu2x-caffe. @@ -33,7 +33,7 @@ class Waifu2xCaffe: self.model_dir = model_dir self.print_lock = threading.Lock() - def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height): + def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, threads_exceptions): """This is the core function for WAIFU2X class Arguments: @@ -43,46 +43,49 @@ class Waifu2xCaffe: height {int} -- output video height """ - # overwrite config file settings - self.waifu2x_settings['input_path'] = input_folder - self.waifu2x_settings['output_path'] = output_folder + try: + # overwrite config file settings + self.waifu2x_settings['input_path'] = input_folder + self.waifu2x_settings['output_path'] = output_folder - if scale_ratio: - self.waifu2x_settings['scale_ratio'] = scale_ratio - elif scale_width and scale_height: - self.waifu2x_settings['scale_width'] = scale_width - self.waifu2x_settings['scale_height'] = scale_height + if scale_ratio: + self.waifu2x_settings['scale_ratio'] = scale_ratio + elif scale_width and scale_height: + self.waifu2x_settings['scale_width'] = scale_width + self.waifu2x_settings['scale_height'] = scale_height - # print thread start message - self.print_lock.acquire() - Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) - self.print_lock.release() + # print thread start message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) + self.print_lock.release() - # list to be executed - execute = [] + # list to be executed + execute = [] - execute.append(self.waifu2x_settings['waifu2x_caffe_path']) - for key in self.waifu2x_settings.keys(): + execute.append(self.waifu2x_settings['waifu2x_caffe_path']) + for key in self.waifu2x_settings.keys(): - value = self.waifu2x_settings[key] + value = self.waifu2x_settings[key] - #is executable key or null or None means that leave this option out (keep default) - if key == 'waifu2x_caffe_path' or value is None or value is False: - continue - else: - if len(key) == 1: - execute.append('-{}'.format(key)) + # is executable key or null or None means that leave this option out (keep default) + if key == 'waifu2x_caffe_path' or value is None or value is False: + continue else: - execute.append('--{}'.format(key)) - execute.append(str(value)) - - Avalon.debug_info('Executing: {}'.format(execute)) - completed_command = subprocess.run(execute, check=True) + if len(key) == 1: + execute.append('-{}'.format(key)) + else: + execute.append('--{}'.format(key)) + execute.append(str(value)) - # print thread exiting message - self.print_lock.acquire() - Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name)) - self.print_lock.release() + Avalon.debug_info('Executing: {}'.format(execute)) + completed_command = subprocess.run(execute, check=True) - # return command execution return code - return completed_command.returncode + # print thread exiting message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} exiting'.format(threading.current_thread().name)) + self.print_lock.release() + + # return command execution return code + return completed_command.returncode + except Exception as e: + threads_exceptions.append(e) diff --git a/bin/waifu2x_converter.py b/bin/waifu2x_converter.py index 56f5f90..ec78325 100644 --- a/bin/waifu2x_converter.py +++ b/bin/waifu2x_converter.py @@ -4,7 +4,7 @@ Name: Waifu2x Converter CPP Driver Author: K4YT3X Date Created: February 8, 2019 -Last Modified: March 19, 2019 +Last Modified: March 24, 2019 Description: This class is a high-level wrapper for waifu2x-converter-cpp. @@ -28,7 +28,7 @@ class Waifu2xConverter: self.waifu2x_settings['model_dir'] = model_dir self.print_lock = threading.Lock() - def upscale(self, input_folder, output_folder, scale_ratio, jobs): + def upscale(self, input_folder, output_folder, scale_ratio, jobs, threads_exceptions): """ Waifu2x Converter Driver Upscaler This method executes the upscaling of extracted frames. @@ -39,55 +39,58 @@ class Waifu2xConverter: threads {int} -- number of threads """ - # overwrite config file settings - self.waifu2x_settings['input'] = input_folder - self.waifu2x_settings['output'] = output_folder + try: + # overwrite config file settings + self.waifu2x_settings['input'] = input_folder + self.waifu2x_settings['output'] = output_folder - # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 - self.waifu2x_settings['i'] = input_folder - self.waifu2x_settings['o'] = output_folder - self.waifu2x_settings['input'] = None - self.waifu2x_settings['output'] = None + # temporary fix for https://github.com/DeadSix27/waifu2x-converter-cpp/issues/109 + self.waifu2x_settings['i'] = input_folder + self.waifu2x_settings['o'] = output_folder + self.waifu2x_settings['input'] = None + self.waifu2x_settings['output'] = None - self.waifu2x_settings['scale_ratio'] = scale_ratio - self.waifu2x_settings['jobs'] = jobs + self.waifu2x_settings['scale_ratio'] = scale_ratio + self.waifu2x_settings['jobs'] = jobs - # 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'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path']) + # 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'] = '{}\\models_rgb'.format(self.waifu2x_settings['waifu2x_converter_path']) - # print thread start message - self.print_lock.acquire() - Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) - self.print_lock.release() + # print thread start message + self.print_lock.acquire() + Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) + self.print_lock.release() - # list to be executed - execute = [] + # list to be executed + execute = [] - for key in self.waifu2x_settings.keys(): + for key in self.waifu2x_settings.keys(): - value = self.waifu2x_settings[key] + value = self.waifu2x_settings[key] - # the key doesn't need to be passed in this case - if key == 'waifu2x_converter_path': - execute.append('{}\\waifu2x-converter-cpp.exe'.format(str(value))) + # the key doesn't need to be passed in this case + if key == 'waifu2x_converter_path': + execute.append('{}\\waifu2x-converter-cpp.exe'.format(str(value))) - # null or None means that leave this option out (keep default) - elif value is None or value is False: - continue - else: - if len(key) == 1: - execute.append('-{}'.format(key)) - else: - execute.append('--{}'.format(key)) - - # true means key is an option - if value is True: + # null or None means that leave this option out (keep default) + elif value is None or value is False: continue + else: + if len(key) == 1: + execute.append('-{}'.format(key)) + else: + execute.append('--{}'.format(key)) - execute.append(str(value)) - - Avalon.debug_info('Executing: {}'.format(execute)) - return subprocess.run(execute, check=True).returncode + # true means key is an option + if value is True: + continue + execute.append(str(value)) + + Avalon.debug_info('Executing: {}'.format(execute)) + return subprocess.run(execute, check=True).returncode + + except Exception as e: + threads_exceptions.append(e)