OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2025-01-30 23:58:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

redesigned upscaler class to make arbitrary scaling available for images

Browse Source
This commit is contained in:
K4YT3X 2020-09-13 11:07:39 -04:00
parent a82fcc778e
commit c7013b2576
2 changed files with 149 additions and 135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
src/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: September 12, 2020 Last Modified: September 13, 2020
Description: This file contains the Upscaler class. Each Description: This file contains the Upscaler class. Each
instance of the Upscaler class is an upscaler on an image or instance of the Upscaler class is an upscaler on an image or
@ -471,11 +471,9 @@ class Upscaler:
input_file_type = input_file_mime_type.split('/')[0] input_file_type = input_file_mime_type.split('/')[0]
input_file_subtype = input_file_mime_type.split('/')[1] input_file_subtype = input_file_mime_type.split('/')[1]
except Exception: except Exception:
input_file_type = input_file_subtype = None
# in case python-magic fails to detect file type # in case python-magic fails to detect file type
# try guessing file mime type with mimetypes # try guessing file mime type with mimetypes
if input_file_type not in ['image', 'video']:
input_file_mime_type = mimetypes.guess_type(input_path.name)[0] input_file_mime_type = mimetypes.guess_type(input_path.name)[0]
input_file_type = input_file_mime_type.split('/')[0] input_file_type = input_file_mime_type.split('/')[0]
input_file_subtype = input_file_mime_type.split('/')[1] input_file_subtype = input_file_mime_type.split('/')[1]
@ -531,38 +529,32 @@ class Upscaler:
# get new job from queue # get new job from queue
self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype = self.processing_queue.get() self.current_input_file, output_path, input_file_mime_type, input_file_type, input_file_subtype = self.processing_queue.get()
# get video information JSON using FFprobe
Avalon.info(_('Reading file information'))
file_info = self.ffmpeg_object.probe_file_info(self.current_input_file)
# create temporary directories for storing frames
self.create_temp_directories()
# start handling input # start handling input
# if input file is a static image # if input file is a static image
if input_file_type == 'image' and input_file_subtype != 'gif': if input_file_type == 'image' and input_file_subtype != 'gif':
Avalon.info(_('Starting to upscale image')) Avalon.info(_('Starting upscaling image'))
if self.driver == 'waifu2x_caffe' and self.scale_ratio is None: # copy original file into the pre-processing directory
self.driver_object.set_scale_resolution(self.scale_width, self.scale_height) shutil.copy(self.current_input_file, self.extracted_frames / self.current_input_file.name)
width = int(file_info['streams'][0]['width'])
height = int(file_info['streams'][0]['height'])
framerate = self.total_frames = 1
# elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
else: else:
self.driver_object.set_scale_ratio(self.scale_ratio) Avalon.info(_('Starting upscaling video/GIF'))
self.process_pool.append(self.driver_object.upscale(self.current_input_file, output_path))
self._wait()
Avalon.info(_('Upscaling completed'))
# static images don't require GIF or video encoding
# go to the next task
self.processing_queue.task_done()
self.total_processed += 1
continue
# if input file is a image/gif file or a video
elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
self.create_temp_directories()
# get video information JSON using FFprobe
Avalon.info(_('Reading file information'))
video_info = self.ffmpeg_object.probe_file_info(self.current_input_file)
# find index of video stream # find index of video stream
video_stream_index = None video_stream_index = None
for stream in video_info['streams']: for stream in file_info['streams']:
if stream['codec_type'] == 'video': if stream['codec_type'] == 'video':
video_stream_index = stream['index'] video_stream_index = stream['index']
break break
@ -573,16 +565,16 @@ class Upscaler:
raise StreamNotFoundError('no video stream found') raise StreamNotFoundError('no video stream found')
# get average frame rate of video stream # get average frame rate of video stream
framerate = float(Fraction(video_info['streams'][video_stream_index]['r_frame_rate'])) framerate = float(Fraction(file_info['streams'][video_stream_index]['r_frame_rate']))
width = int(video_info['streams'][video_stream_index]['width']) width = int(file_info['streams'][video_stream_index]['width'])
height = int(video_info['streams'][video_stream_index]['height']) height = int(file_info['streams'][video_stream_index]['height'])
# get total number of frames # get total number of frames
Avalon.info(_('Getting total number of frames in the file')) Avalon.info(_('Getting total number of frames in the file'))
# if container stores total number of frames in nb_frames, fetch it directly # if container stores total number of frames in nb_frames, fetch it directly
if 'nb_frames' in video_info['streams'][video_stream_index]: if 'nb_frames' in file_info['streams'][video_stream_index]:
self.total_frames = int(video_info['streams'][video_stream_index]['nb_frames']) self.total_frames = int(file_info['streams'][video_stream_index]['nb_frames'])
# otherwise call FFprobe to count the total number of frames # otherwise call FFprobe to count the total number of frames
else: else:
@ -664,6 +656,7 @@ class Upscaler:
Avalon.debug_info(_('Upscaling jobs queue: {}').format(self.scaling_jobs)) Avalon.debug_info(_('Upscaling jobs queue: {}').format(self.scaling_jobs))
# extract frames from video # extract frames from video
if input_file_mime_type == 'image/gif' or input_file_type == 'video':
self.process_pool.append((self.ffmpeg_object.extract_frames(self.current_input_file, self.extracted_frames))) self.process_pool.append((self.ffmpeg_object.extract_frames(self.current_input_file, self.extracted_frames)))
self._wait() self._wait()
@ -680,12 +673,6 @@ class Upscaler:
Avalon.error(_('Unsupported pixel format: {}').format(self.ffmpeg_object.pixel_format)) Avalon.error(_('Unsupported pixel format: {}').format(self.ffmpeg_object.pixel_format))
raise UnsupportedPixelError(f'unsupported pixel format {self.ffmpeg_object.pixel_format}') raise UnsupportedPixelError(f'unsupported pixel format {self.ffmpeg_object.pixel_format}')
# width/height will be coded width/height x upscale factor
# original_width = video_info['streams'][video_stream_index]['width']
# original_height = video_info['streams'][video_stream_index]['height']
# scale_width = int(self.scale_ratio * original_width)
# scale_height = int(self.scale_ratio * original_height)
# upscale images one by one using waifu2x # upscale images one by one using waifu2x
Avalon.info(_('Starting to upscale extracted frames')) Avalon.info(_('Starting to upscale extracted frames'))
upscale_begin_time = time.time() upscale_begin_time = time.time()
@ -709,6 +696,13 @@ class Upscaler:
# start handling output # start handling output
# output can be either GIF or video # output can be either GIF or video
if input_file_type == 'image' and input_file_subtype != 'gif':
# resize and output image to output_path
self.process_pool.append(self.ffmpeg_object.resize_image([f for f in self.upscaled_frames.iterdir() if f.is_file()][0], output_path, output_width, output_height))
self._wait()
elif input_file_mime_type == 'image/gif' or input_file_type == 'video':
# if the desired output is gif file # if the desired output is gif file
if output_path.suffix.lower() == '.gif': if output_path.suffix.lower() == '.gif':

22
src/wrappers/ffmpeg.py
View File

@ -4,7 +4,7 @@
Name: Video2X FFmpeg Controller Name: Video2X FFmpeg Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 Date Created: Feb 24, 2018
Last Modified: June 7, 2020 Last Modified: September 13, 2020
Description: This class handles all FFmpeg related operations. Description: This class handles all FFmpeg related operations.
""" """
@ -259,6 +259,26 @@ class Ffmpeg:
return(self._execute(execute)) return(self._execute(execute))
def resize_image(self, input_path: pathlib.Path, output_path: pathlib.Path, output_width: int, output_height: int):
""" resize the given image and output the resized image to output_path
Args:
input_path (pathlib.Path): input image path
output_path (pathlib.Path): output image path
output_width (int): output image target width
output_height (int): output image target height
"""
execute = [
self.ffmpeg_binary,
'-i',
input_path,
'-vf',
f'scale={output_width}:{output_height}',
output_path
]
return(self._execute(execute))
def _read_configuration(self, phase, section=None): def _read_configuration(self, phase, section=None):
""" read configuration from JSON """ read configuration from JSON