OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-29 16:09:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

added support for JPEG and other image format

Browse Source
This commit is contained in:
k4yt3x 2019-03-30 14:01:36 -04:00
parent 562aaee056
commit 78ae2e7ed2
3 changed files with 23 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
bin/upscaler.py
View File

@ -4,7 +4,7 @@
Name: Video2X Upscaler Name: Video2X Upscaler
Author: K4YT3X Author: K4YT3X
Date Created: December 10, 2018 Date Created: December 10, 2018
Last Modified: March 24, 2019 Last Modified: March 30, 2019
Licensed under the GNU General Public License Version 3 (GNU GPL v3), Licensed under the GNU General Public License Version 3 (GNU GPL v3),
available at: https://www.gnu.org/licenses/gpl-3.0.txt available at: https://www.gnu.org/licenses/gpl-3.0.txt
@ -37,36 +37,31 @@ class Upscaler:
ArgumentError -- if argument is not valid ArgumentError -- if argument is not valid
""" """
def __init__(self, input_video, output_video, method, waifu2x_settings, ffmpeg_settings, waifu2x_driver='waifu2x_caffe', scale_width=False, scale_height=False, scale_ratio=False, model_dir=None, threads=5, video2x_cache_folder='{}\\video2x'.format(tempfile.gettempdir()), preserve_frames=False): def __init__(self, input_video, output_video, method, waifu2x_settings, ffmpeg_settings):
# mandatory arguments # mandatory arguments
self.input_video = input_video self.input_video = input_video
self.output_video = output_video self.output_video = output_video
self.method = method self.method = method
self.waifu2x_settings = waifu2x_settings self.waifu2x_settings = waifu2x_settings
self.ffmpeg_settings = ffmpeg_settings self.ffmpeg_settings = ffmpeg_settings
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 # optional arguments
self.scale_width = scale_width self.waifu2x_driver = 'waifu2x_caffe'
self.scale_height = scale_height self.scale_width = None
self.scale_ratio = scale_ratio self.scale_height = None
self.model_dir = model_dir self.scale_ratio = None
self.threads = threads self.model_dir = None
self.threads = 5
self.video2x_cache_folder = '{}\\video2x'.format(tempfile.gettempdir())
self.image_format = 'png'
self.preserve_frames = False
# create temporary folder/directories # create temporary folder/directories
self.video2x_cache_folder = video2x_cache_folder
self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) self.extracted_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames)) Avalon.debug_info('Extracted frames are being saved to: {}'.format(self.extracted_frames))
self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder) self.upscaled_frames = tempfile.mkdtemp(dir=self.video2x_cache_folder)
Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames)) Avalon.debug_info('Upscaled frames are being saved to: {}'.format(self.upscaled_frames))
self.preserve_frames = preserve_frames
def cleanup(self): def cleanup(self):
# delete temp directories when done # delete temp directories when done
# avalon framework cannot be used if python is shutting down # avalon framework cannot be used if python is shutting down
@ -153,9 +148,9 @@ class Upscaler:
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],)) progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
progress_bar.start() progress_bar.start()
w2.upscale(self.extracted_frames, self.upscaled_frames, self.scale_ratio, self.threads, self.upscaler_exceptions) 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 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) renamed = re.sub('_\[.*-.*\]\[x(\d+(\.\d+)?)\]\.{}'.format(self.image_format), '.{}'.format(self.image_format), image)
shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed)) shutil.move('{}\\{}'.format(self.upscaled_frames, image), '{}\\{}'.format(self.upscaled_frames, renamed))
self.progress_bar_exit_signal = True self.progress_bar_exit_signal = True
@ -202,9 +197,9 @@ class Upscaler:
for thread_info in thread_pool: for thread_info in thread_pool:
# create thread # create thread
if self.scale_ratio: if self.scale_ratio:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.upscaler_exceptions)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False, self.image_format, self.upscaler_exceptions))
else: else:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height, 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))
thread.name = thread_info[1] thread.name = thread_info[1]
# add threads into the pool # add threads into the pool
@ -252,7 +247,7 @@ class Upscaler:
self.output_video = os.path.abspath(self.output_video) self.output_video = os.path.abspath(self.output_video)
# initialize objects for ffmpeg and waifu2x-caffe # initialize objects for ffmpeg and waifu2x-caffe
fm = Ffmpeg(self.ffmpeg_settings) fm = Ffmpeg(self.ffmpeg_settings, self.image_format)
# initialize waifu2x driver # initialize waifu2x driver
if self.waifu2x_driver == 'waifu2x_caffe': if self.waifu2x_driver == 'waifu2x_caffe':

6
bin/waifu2x_caffe.py
View File

@ -4,7 +4,7 @@
Name: Waifu2x Caffe Driver Name: Waifu2x Caffe Driver
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: March 24, 2019 Last Modified: March 30, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-caffe. for waifu2x-caffe.
@ -33,7 +33,7 @@ class Waifu2xCaffe:
self.model_dir = model_dir self.model_dir = model_dir
self.print_lock = threading.Lock() self.print_lock = threading.Lock()
def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, upscaler_exceptions): def upscale(self, input_folder, output_folder, scale_ratio, scale_width, scale_height, image_format, upscaler_exceptions):
"""This is the core function for WAIFU2X class """This is the core function for WAIFU2X class
Arguments: Arguments:
@ -54,6 +54,8 @@ class Waifu2xCaffe:
self.waifu2x_settings['scale_width'] = scale_width self.waifu2x_settings['scale_width'] = scale_width
self.waifu2x_settings['scale_height'] = scale_height self.waifu2x_settings['scale_height'] = scale_height
self.waifu2x_settings['output_extention'] = image_format
# print thread start message # print thread start message
self.print_lock.acquire() self.print_lock.acquire()
Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name)) Avalon.debug_info('[upscaler] Thread {} started'.format(threading.current_thread().name))

5
bin/waifu2x_converter.py
View File

@ -4,7 +4,7 @@
Name: Waifu2x Converter CPP Driver Name: Waifu2x Converter CPP Driver
Author: K4YT3X Author: K4YT3X
Date Created: February 8, 2019 Date Created: February 8, 2019
Last Modified: March 24, 2019 Last Modified: March 30, 2019
Description: This class is a high-level wrapper Description: This class is a high-level wrapper
for waifu2x-converter-cpp. for waifu2x-converter-cpp.
@ -28,7 +28,7 @@ class Waifu2xConverter:
self.waifu2x_settings['model_dir'] = model_dir self.waifu2x_settings['model_dir'] = model_dir
self.print_lock = threading.Lock() self.print_lock = threading.Lock()
def upscale(self, input_folder, output_folder, scale_ratio, jobs, upscaler_exceptions): def upscale(self, input_folder, output_folder, scale_ratio, jobs, image_format, upscaler_exceptions):
""" Waifu2x Converter Driver Upscaler """ Waifu2x Converter Driver Upscaler
This method executes the upscaling of extracted frames. This method executes the upscaling of extracted frames.
@ -54,6 +54,7 @@ class Waifu2xConverter:
self.waifu2x_settings['scale-ratio'] = scale_ratio self.waifu2x_settings['scale-ratio'] = scale_ratio
self.waifu2x_settings['jobs'] = jobs self.waifu2x_settings['jobs'] = jobs
self.waifu2x_settings['output-format'] = image_format
# models_rgb must be specified manually for waifu2x-converter-cpp # models_rgb must be specified manually for waifu2x-converter-cpp
# if it's not specified in the arguments, create automatically # if it's not specified in the arguments, create automatically