diff --git a/bin/upscaler.py b/bin/upscaler.py index c85e775..75e49a1 100755 --- a/bin/upscaler.py +++ b/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')