mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 23:58:11 +00:00
redesigned upscaler class to make arbitrary scaling available for images
This commit is contained in:
K4YT3X
parent
a82fcc778e
commit
c7013b2576
@ -4,7 +4,7 @@
|
||||
Name: Video2X Upscaler
|
||||
Author: K4YT3X
|
||||
Date Created: December 10, 2018
|
||||
Last Modified: September 12, 2020
|
||||
Last Modified: September 13, 2020
|
||||
|
||||
Description: This file contains the Upscaler class. Each
|
||||
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_subtype = input_file_mime_type.split('/')[1]
|
||||
except Exception:
|
||||
input_file_type = input_file_subtype = None
|
||||
|
||||
# in case python-magic fails to detect file type
|
||||
# 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_type = input_file_mime_type.split('/')[0]
|
||||
input_file_subtype = input_file_mime_type.split('/')[1]
|
||||
@ -531,38 +529,32 @@ class Upscaler:
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# if input file is a static image
|
||||
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:
|
||||
self.driver_object.set_scale_resolution(self.scale_width, self.scale_height)
|
||||
# copy original file into the pre-processing directory
|
||||
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:
|
||||
self.driver_object.set_scale_ratio(self.scale_ratio)
|
||||
|
||||
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)
|
||||
Avalon.info(_('Starting upscaling video/GIF'))
|
||||
|
||||
# find index of video stream
|
||||
video_stream_index = None
|
||||
for stream in video_info['streams']:
|
||||
for stream in file_info['streams']:
|
||||
if stream['codec_type'] == 'video':
|
||||
video_stream_index = stream['index']
|
||||
break
|
||||
@ -573,16 +565,16 @@ class Upscaler:
|
||||
raise StreamNotFoundError('no video stream found')
|
||||
|
||||
# get average frame rate of video stream
|
||||
framerate = float(Fraction(video_info['streams'][video_stream_index]['r_frame_rate']))
|
||||
width = int(video_info['streams'][video_stream_index]['width'])
|
||||
height = int(video_info['streams'][video_stream_index]['height'])
|
||||
framerate = float(Fraction(file_info['streams'][video_stream_index]['r_frame_rate']))
|
||||
width = int(file_info['streams'][video_stream_index]['width'])
|
||||
height = int(file_info['streams'][video_stream_index]['height'])
|
||||
|
||||
# get total number of frames
|
||||
Avalon.info(_('Getting total number of frames in the file'))
|
||||
|
||||
# if container stores total number of frames in nb_frames, fetch it directly
|
||||
if 'nb_frames' in video_info['streams'][video_stream_index]:
|
||||
self.total_frames = int(video_info['streams'][video_stream_index]['nb_frames'])
|
||||
if 'nb_frames' in file_info['streams'][video_stream_index]:
|
||||
self.total_frames = int(file_info['streams'][video_stream_index]['nb_frames'])
|
||||
|
||||
# otherwise call FFprobe to count the total number of frames
|
||||
else:
|
||||
@ -664,6 +656,7 @@ class Upscaler:
|
||||
Avalon.debug_info(_('Upscaling jobs queue: {}').format(self.scaling_jobs))
|
||||
|
||||
# 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._wait()
|
||||
|
||||
@ -680,12 +673,6 @@ class Upscaler:
|
||||
Avalon.error(_('Unsupported pixel format: {}').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
|
||||
Avalon.info(_('Starting to upscale extracted frames'))
|
||||
upscale_begin_time = time.time()
|
||||
@ -709,6 +696,13 @@ class Upscaler:
|
||||
|
||||
# start handling output
|
||||
# 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 output_path.suffix.lower() == '.gif':
|
||||
|
||||
@ -4,7 +4,7 @@
|
||||
Name: Video2X FFmpeg Controller
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: June 7, 2020
|
||||
Last Modified: September 13, 2020
|
||||
|
||||
Description: This class handles all FFmpeg related operations.
|
||||
"""
|
||||
@ -259,6 +259,26 @@ class Ffmpeg:
|
||||
|
||||
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):
|
||||
""" read configuration from JSON
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user