Merge pull request #1 from K4YT3X/master

Pull latest updates

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.

bin/clear_image.py

@@ -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)
-					
-		

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)

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 = tempfile.mkdtemp(dir=self.video2x_cache_folder)
-        self.extracted_frames = self.extracted_frames_object.name
         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 = tempfile.mkdtemp(dir=self.video2x_cache_folder)
-        self.upscaled_frames = self.upscaled_frames_object.name
         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()
+            for directory in [self.extracted_frames, self.upscaled_frames]:
-            print('Cleaning up cache directory: {}'.format(self.upscaled_frames))
+                try:
-            self.upscaled_frames_object.cleanup()
+                    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,10 +214,10 @@ class Upscaler:
         progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,))
         progress_bar.start()
 
-        #Create the clearer and start it
+        # create the clearer and start it
-        Avalon.debug_info('Starting image clearer...')
+        Avalon.debug_info('Starting upscaled image cleaner')
-        image_clear = ClearImage(self.extracted_frames,self.upscaled_frames,len(upscaler_threads))
+        image_cleaner = ImageCleaner(self.extracted_frames, self.upscaled_frames, len(upscaler_threads))
-        image_clear.start()
+        image_cleaner.start()
 
         # start all threads
         for thread in upscaler_threads:
@@ -222,12 +227,15 @@ class Upscaler:
         for thread in upscaler_threads:
             thread.join()
 
-        #upscaling done... kill the clearer
+        # upscaling done, kill the clearer
-        Avalon.debug_info('Stoping image clearer...')
+        Avalon.debug_info('Killing upscaled image cleaner')
-        image_clear.stop()
+        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)
-
-

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)
@@ -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:

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,6 +43,7 @@ class Waifu2xCaffe:
             height {int} -- output video height
         """
 
+        try:
             # overwrite config file settings
             self.waifu2x_settings['input_path'] = input_folder
             self.waifu2x_settings['output_path'] = output_folder
@@ -86,3 +87,5 @@ class Waifu2xCaffe:
 
             # return command execution return code
             return completed_command.returncode
+        except Exception as e:
+            threads_exceptions.append(e)

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,6 +39,7 @@ class Waifu2xConverter:
             threads {int} -- number of threads
         """
 
+        try:
             # overwrite config file settings
             self.waifu2x_settings['input'] = input_folder
             self.waifu2x_settings['output'] = output_folder
@@ -91,3 +92,5 @@ class Waifu2xConverter:
             Avalon.debug_info('Executing: {}'.format(execute))
             return subprocess.run(execute, check=True).returncode
 
+        except Exception as e:
+            threads_exceptions.append(e)