adding amd support

bin/upscaler.py

@@ -16,9 +16,11 @@ from avalon_framework import Avalon
 from exceptions import *
 from ffmpeg import Ffmpeg
 from fractions import Fraction
-from waifu2x import Waifu2x
+from waifu2x_caffe import Waifu2xCaffe
+from waifu2x_converter import Waifu2xConverter
 import json
 import os
+import re
 import shutil
 import subprocess
 import tempfile
@@ -30,20 +32,25 @@ MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
 
 class Upscaler:
 
-    def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, factor=False, model_type='anime_style_art_rgb', threads=3, extracted_frames=False, upscaled_frames=False, preserve_frames=False):
+    def __init__(self, input_video, output_video, method, waifu2x_path, ffmpeg_path, waifu2x_driver='waifu2x_caffe', ffmpeg_arguments=[], ffmpeg_hwaccel='gpu', output_width=False, output_height=False, ratio=False, model_type='anime_style_art_rgb', threads=3, extracted_frames=False, upscaled_frames=False, preserve_frames=False):
         # Mandatory arguments
         self.input_video = input_video
         self.output_video = output_video
         self.method = method
         self.waifu2x_path = waifu2x_path
         self.ffmpeg_path = ffmpeg_path
+        self.waifu2x_driver = waifu2x_driver
+
+        # Check sanity of waifu2x_driver option
+        if waifu2x_driver != 'waifu2x_caffe' and waifu2x_driver != 'waifu2x_converter':
+            raise Exception('Unrecognized waifu2x driver: {}'.format(waifu2x_driver))
 
         # Optional arguments
         self.ffmpeg_arguments = ffmpeg_arguments
         self.ffmpeg_hwaccel = ffmpeg_hwaccel
         self.output_width = output_width
         self.output_height = output_height
-        self.factor = factor
+        self.ratio = ratio
         self.model_type = model_type
         self.threads = threads
 
@@ -62,10 +69,12 @@ class Upscaler:
 
     def __del__(self):
         # Delete temp directories when done
+        # Avalon framework cannot be used if python is shutting down
+        # Therefore, plain print is used
         if not self.preserve_frames:
-            Avalon.debug_info('Deleting cache directory: {}'.format(self.extracted_frames))
+            print('Deleting cache directory: {}'.format(self.extracted_frames))
             shutil.rmtree(self.extracted_frames)
-            Avalon.debug_info('Deleting cache directory: {}'.format(self.upscaled_frames))
+            print('Deleting cache directory: {}'.format(self.upscaled_frames))
             shutil.rmtree(self.upscaled_frames)
 
     def _get_video_info(self):
@@ -86,7 +95,7 @@ class Upscaler:
         # values are specified
         if not self.input_video:
             raise ArgumentError('You need to specify the video to process')
-        elif (not self.output_width or not self.output_height) and not self.factor:
+        elif (not self.output_width or not self.output_height) and not self.ratio:
             raise ArgumentError('You must specify output video width and height or upscale factor')
         elif not self.output_video:
             raise ArgumentError('You need to specify the output video name')
@@ -103,6 +112,15 @@ class Upscaler:
             w2 {Waifu2x Object} -- initialized waifu2x object
         """
 
+        # It's easier to do multi-threading with waifu2x_converter
+        # The number of threads can be passed directly to waifu2x_converter
+        if self.waifu2x_driver == 'waifu2x_converter':
+            w2.upscale(self.extracted_frames, self.upscaled_frames, self.ratio, self.threads)
+            for image in [f for f in os.listdir(self.upscaled_frames) if os.path.isfile(os.path.join(self.upscaled_frames, f))]:
+                renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.png', '.png', image)
+                shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed))
+            return
+
         # Create a container for all upscaler threads
         upscaler_threads = []
 
@@ -178,12 +196,19 @@ class Upscaler:
         # Check if FFMPEG and waifu2x are present
         if not os.path.isdir(self.ffmpeg_path):
             raise FileNotFoundError(self.ffmpeg_path)
-        if not os.path.isfile(self.waifu2x_path):
+        if not os.path.isfile(self.waifu2x_path) and not os.path.isdir(self.waifu2x_path):
             raise FileNotFoundError(self.waifu2x_path)
 
         # Initialize objects for ffmpeg and waifu2x-caffe
         fm = Ffmpeg(self.ffmpeg_path, self.output_video, self.ffmpeg_arguments)
-        w2 = Waifu2x(self.waifu2x_path, self.method, self.model_type)
+
+        # Initialize waifu2x driver
+        if self.waifu2x_driver == 'waifu2x_caffe':
+            w2 = Waifu2xCaffe(self.waifu2x_path, self.method, self.model_type)
+        elif self.waifu2x_driver == 'waifu2x_converter':
+            w2 = Waifu2xConverter(self.waifu2x_path)
+        else:
+            raise Exception('Unrecognized waifu2x driver: {}'.format(self.waifu2x_driver))
 
         # Extract frames from video
         fm.extract_frames(self.input_video, self.extracted_frames)
@@ -209,11 +234,11 @@ class Upscaler:
         Avalon.info('Framerate: {}'.format(framerate))
 
         # Width/height will be coded width/height x upscale factor
-        if self.factor:
+        if self.ratio:
             coded_width = info['streams'][video_stream_index]['coded_width']
             coded_height = info['streams'][video_stream_index]['coded_height']
-            self.output_width = self.factor * coded_width
-            self.output_height = self.factor * coded_height
+            self.output_width = self.ratio * coded_width
+            self.output_height = self.ratio * coded_height
 
         # Upscale images one by one using waifu2x
         Avalon.info('Starting to upscale extracted images')