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

removed model check

Browse Source
This commit is contained in:
k4yt3x 2019-03-19 13:10:38 -04:00
parent f51ce05066
commit 31aa3a630d
1 changed files with 57 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

118
bin/upscaler.py Executable file → Normal file
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 9, 2019 Last Modified: March 19, 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
@ -26,14 +26,19 @@ import tempfile
import threading import threading
import time import time
MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo', 'anime_style_art_rgb', 'photo',
'anime_style_art_y', 'upresnet10', 'cunet']
class Upscaler: class Upscaler:
""" An instance of this class is a upscaler that will
upscale all images in the given folder.
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_type='models/cunet', threads=5, video2x_cache_folder='{}\\video2x'.format(tempfile.gettempdir()), preserve_frames=False): Raises:
# Mandatory arguments Exception -- all exceptions
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):
# 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
@ -41,18 +46,18 @@ class Upscaler:
self.ffmpeg_settings = ffmpeg_settings self.ffmpeg_settings = ffmpeg_settings
self.waifu2x_driver = waifu2x_driver self.waifu2x_driver = waifu2x_driver
# Check sanity of waifu2x_driver option # check sanity of waifu2x_driver option
if waifu2x_driver != 'waifu2x_caffe' and waifu2x_driver != 'waifu2x_converter': if waifu2x_driver != 'waifu2x_caffe' and waifu2x_driver != 'waifu2x_converter':
raise Exception('Unrecognized waifu2x driver: {}'.format(waifu2x_driver)) raise Exception('Unrecognized waifu2x driver: {}'.format(waifu2x_driver))
# Optional arguments # optional arguments
self.scale_width = scale_width self.scale_width = scale_width
self.scale_height = scale_height self.scale_height = scale_height
self.scale_ratio = scale_ratio self.scale_ratio = scale_ratio
self.model_type = model_type self.model_dir = model_dir
self.threads = threads self.threads = threads
# Create temporary folder/directories # create temporary folder/directories
self.video2x_cache_folder = video2x_cache_folder self.video2x_cache_folder = video2x_cache_folder
self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder) self.extracted_frames_object = tempfile.TemporaryDirectory(dir=self.video2x_cache_folder)
self.extracted_frames = self.extracted_frames_object.name self.extracted_frames = self.extracted_frames_object.name
@ -65,25 +70,17 @@ class Upscaler:
self.preserve_frames = preserve_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
# Therefore, plain print is used # therefore, plain print is used
if not self.preserve_frames: if not self.preserve_frames:
print('Cleaning up cache directory: {}'.format(self.extracted_frames)) print('Cleaning up cache directory: {}'.format(self.extracted_frames))
self.extracted_frames_object.cleanup() self.extracted_frames_object.cleanup()
print('Cleaning up cache directory: {}'.format(self.upscaled_frames)) print('Cleaning up cache directory: {}'.format(self.upscaled_frames))
self.upscaled_frames_object.cleanup() self.upscaled_frames_object.cleanup()
def _check_model_type(self, args):
""" Validate upscaling model
"""
if self.model_type not in MODELS_AVAILABLE:
if self.model_type.split('/')[-1] not in MODELS_AVAILABLE:
if self.model_type.split('\\')[-1] not in MODELS_AVAILABLE:
raise InvalidModelType('Specified model type not available')
def _check_arguments(self): def _check_arguments(self):
# Check if arguments are valid / all necessary argument # check if arguments are valid / all necessary argument
# values are specified # values are specified
if not self.input_video: if not self.input_video:
raise ArgumentError('You need to specify the video to process') raise ArgumentError('You need to specify the video to process')
@ -102,7 +99,7 @@ class Upscaler:
and the output directory/folder. This is originally and the output directory/folder. This is originally
suggested by @ArmandBernard. suggested by @ArmandBernard.
""" """
# Get number of extracted frames # get number of extracted frames
total_frames = 0 total_frames = 0
for folder in extracted_frames_folders: for folder in extracted_frames_folders:
total_frames += len([f for f in os.listdir(folder) if f[-4:] == '.png']) total_frames += len([f for f in os.listdir(folder) if f[-4:] == '.png'])
@ -120,11 +117,11 @@ class Upscaler:
delta = total_frames_upscaled - previous_cycle_frames delta = total_frames_upscaled - previous_cycle_frames
previous_cycle_frames = total_frames_upscaled previous_cycle_frames = total_frames_upscaled
# If upscaling is finished # if upscaling is finished
if total_frames_upscaled >= total_frames: if total_frames_upscaled >= total_frames:
return return
# Adds the detla into the progress bar # adds the detla into the progress bar
progress_bar.update(delta) progress_bar.update(delta)
except FileNotFoundError: except FileNotFoundError:
pass pass
@ -141,11 +138,11 @@ class Upscaler:
w2 {Waifu2x Object} -- initialized waifu2x object w2 {Waifu2x Object} -- initialized waifu2x object
""" """
# Progress bar thread exit signal # progress bar thread exit signal
self.progress_bar_exit_signal = False self.progress_bar_exit_signal = False
# It's easier to do multi-threading with waifu2x_converter # it's easier to do multi-threading with waifu2x_converter
# The number of threads can be passed directly to waifu2x_converter # the number of threads can be passed directly to waifu2x_converter
if self.waifu2x_driver == 'waifu2x_converter': if self.waifu2x_driver == 'waifu2x_converter':
progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],)) progress_bar = threading.Thread(target=self._progress_bar, args=([self.extracted_frames],))
@ -160,18 +157,18 @@ class Upscaler:
progress_bar.join() progress_bar.join()
return return
# Create a container for all upscaler threads # create a container for all upscaler threads
upscaler_threads = [] upscaler_threads = []
# List all images in the extracted frames # list all images in the extracted frames
frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))] frames = [os.path.join(self.extracted_frames, f) for f in os.listdir(self.extracted_frames) if os.path.isfile(os.path.join(self.extracted_frames, f))]
# If we have less images than threads, # if we have less images than threads,
# create only the threads necessary # create only the threads necessary
if len(frames) < self.threads: if len(frames) < self.threads:
self.threads = len(frames) self.threads = len(frames)
# Create a folder for each thread and append folder # create a folder for each thread and append folder
# name into a list # name into a list
thread_pool = [] thread_pool = []
@ -180,43 +177,43 @@ class Upscaler:
thread_folder = '{}\\{}'.format(self.extracted_frames, str(thread_id)) thread_folder = '{}\\{}'.format(self.extracted_frames, str(thread_id))
thread_folders.append(thread_folder) thread_folders.append(thread_folder)
# Delete old folders and create new folders # delete old folders and create new folders
if os.path.isdir(thread_folder): if os.path.isdir(thread_folder):
shutil.rmtree(thread_folder) shutil.rmtree(thread_folder)
os.mkdir(thread_folder) os.mkdir(thread_folder)
# Append folder path into list # append folder path into list
thread_pool.append((thread_folder, thread_id)) thread_pool.append((thread_folder, thread_id))
# Evenly distribute images into each folder # evenly distribute images into each folder
# until there is none left in the folder # until there is none left in the folder
for image in frames: for image in frames:
# Move image # move image
shutil.move(image, thread_pool[0][0]) shutil.move(image, thread_pool[0][0])
# Rotate list # rotate list
thread_pool = thread_pool[-1:] + thread_pool[:-1] thread_pool = thread_pool[-1:] + thread_pool[:-1]
# Create threads and start them # create threads and start them
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, False, self.scale_width, self.scale_height)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, False, self.scale_width, self.scale_height))
else: else:
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False)) thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.scale_ratio, False, False))
thread.name = thread_info[1] thread.name = thread_info[1]
# Add threads into the pool # add threads into the pool
upscaler_threads.append(thread) upscaler_threads.append(thread)
# Start progress bar in a different thread # start progress bar in a different thread
progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,)) progress_bar = threading.Thread(target=self._progress_bar, args=(thread_folders,))
progress_bar.start() progress_bar.start()
# Start all threads # start all threads
for thread in upscaler_threads: for thread in upscaler_threads:
thread.start() thread.start()
# Wait for threads to finish # wait for threads to finish
for thread in upscaler_threads: for thread in upscaler_threads:
thread.join() thread.join()
@ -229,69 +226,68 @@ class Upscaler:
and handles all necessary functions. and handles all necessary functions.
""" """
# Parse arguments for waifu2x # parse arguments for waifu2x
# Check argument sanity # check argument sanity
self._check_model_type(self.model_type)
self._check_arguments() self._check_arguments()
# Convert paths to absolute paths # convert paths to absolute paths
self.input_video = os.path.abspath(self.input_video) self.input_video = os.path.abspath(self.input_video)
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)
# Initialize waifu2x driver # initialize waifu2x driver
if self.waifu2x_driver == 'waifu2x_caffe': if self.waifu2x_driver == 'waifu2x_caffe':
w2 = Waifu2xCaffe(self.waifu2x_settings, self.method, self.model_type) w2 = Waifu2xCaffe(self.waifu2x_settings, self.method, self.model_dir)
elif self.waifu2x_driver == 'waifu2x_converter': elif self.waifu2x_driver == 'waifu2x_converter':
w2 = Waifu2xConverter(self.waifu2x_settings) w2 = Waifu2xConverter(self.waifu2x_settings, self.model_dir)
else: else:
raise Exception('Unrecognized waifu2x driver: {}'.format(self.waifu2x_driver)) raise Exception('Unrecognized waifu2x driver: {}'.format(self.waifu2x_driver))
# Extract frames from video # extract frames from video
fm.extract_frames(self.input_video, self.extracted_frames) fm.extract_frames(self.input_video, self.extracted_frames)
Avalon.info('Reading video information') Avalon.info('Reading video information')
video_info = fm.get_video_info(self.input_video) video_info = fm.get_video_info(self.input_video)
# Analyze original video with ffprobe and retrieve framerate # analyze original video with ffprobe and retrieve framerate
# width, height = info['streams'][0]['width'], info['streams'][0]['height'] # width, height = info['streams'][0]['width'], info['streams'][0]['height']
# 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 video_info['streams']:
if stream['codec_type'] == 'video': if stream['codec_type'] == 'video':
video_stream_index = stream['index'] video_stream_index = stream['index']
break break
# Exit if no video stream found # exit if no video stream found
if video_stream_index is None: if video_stream_index is None:
Avalon.error('Aborting: No video stream found') Avalon.error('Aborting: No video stream found')
exit(1) exit(1)
# Get average frame rate of video stream # get average frame rate of video stream
framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate'])) framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
Avalon.info('Framerate: {}'.format(framerate)) Avalon.info('Framerate: {}'.format(framerate))
# Width/height will be coded width/height x upscale factor # width/height will be coded width/height x upscale factor
if self.scale_ratio: if self.scale_ratio:
coded_width = video_info['streams'][video_stream_index]['coded_width'] coded_width = video_info['streams'][video_stream_index]['coded_width']
coded_height = video_info['streams'][video_stream_index]['coded_height'] coded_height = video_info['streams'][video_stream_index]['coded_height']
self.scale_width = self.scale_ratio * coded_width self.scale_width = self.scale_ratio * coded_width
self.scale_height = self.scale_ratio * coded_height self.scale_height = self.scale_ratio * coded_height
# Upscale images one by one using waifu2x # upscale images one by one using waifu2x
Avalon.info('Starting to upscale extracted images') Avalon.info('Starting to upscale extracted images')
self._upscale_frames(w2) self._upscale_frames(w2)
Avalon.info('Upscaling completed') Avalon.info('Upscaling completed')
# Frames to Video # frames to Video
Avalon.info('Converting extracted frames into video') Avalon.info('Converting extracted frames into video')
# Use user defined output size # use user defined output size
fm.convert_video(framerate, '{}x{}'.format(self.scale_width, self.scale_height), self.upscaled_frames) fm.convert_video(framerate, '{}x{}'.format(self.scale_width, self.scale_height), self.upscaled_frames)
Avalon.info('Conversion completed') Avalon.info('Conversion completed')
# Migrate audio tracks and subtitles # migrate audio tracks and subtitles
Avalon.info('Migrating audio tracks and subtitles to upscaled video') Avalon.info('Migrating audio tracks and subtitles to upscaled video')
fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames) fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)