def video2x():
"""Main controller for Video2X
This function controls the flow of video conversion
and handles all necessary functions.
"""
check_model_type(args)
# Parse arguments for waifu2x
if args.cpu:
method = 'cpu'
elif args.gpu:
method = 'gpu'
ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel))
elif args.cudnn:
method = 'cudnn'
ffmpeg_arguments.append('-hwaccel {}'.format(ffmpeg_hwaccel))
# Initialize objects for ffmpeg and waifu2x-caffe
fm = Ffmpeg(ffmpeg_path, args.output, ffmpeg_arguments)
w2 = Waifu2x(waifu2x_path, method, args.model_type)
# Clear and create directories
if os.path.isdir(FRAMES):
shutil.rmtree(FRAMES)
if os.path.isdir(UPSCALED):
shutil.rmtree(UPSCALED)
os.mkdir(FRAMES)
os.mkdir(UPSCALED)
# Extract frames from video
fm.extract_frames(args.video, FRAMES)
Avalon.info('Reading video information')
info = get_video_info()
# Analyze original video with ffprobe and retrieve framerate
# width, height = info['streams'][0]['width'], info['streams'][0]['height']
# Find index of video stream
video_stream_index = None
for stream in info['streams']:
if stream['codec_type'] == 'video':
video_stream_index = stream['index']
break
# Exit if no video stream found
if video_stream_index is None:
Avalon.error('Aborting: No video stream found')
# Get average frame rate of video stream
framerate = float(
Fraction(info['streams'][video_stream_index]['avg_frame_rate']))
Avalon.info('Framerate: {}'.format(framerate))
# Upscale images one by one using waifu2x
Avalon.info('Starting to upscale extracted images')
upscale_frames(w2)
Avalon.info('Upscaling completed')
# Frames to Video
Avalon.info('Converting extracted frames into video')
# Width/height will be coded width/height x upscale factor
if args.factor:
coded_width = info['streams'][video_stream_index]['coded_width']
coded_height = info['streams'][video_stream_index]['coded_height']
fm.convert_video(
framerate, '{}x{}'.format(args.factor * coded_width,
args.factor * coded_height), UPSCALED)
# Use user defined output size
else:
fm.convert_video(framerate, '{}x{}'.format(args.width, args.height),
UPSCALED)
Avalon.info('Conversion completed')
# Extract and press audio in
Avalon.info('Stripping audio track from original video')
fm.extract_audio(args.video, UPSCALED)
Avalon.info('Inserting audio track into new video')
fm.insert_audio_track(UPSCALED)
def run(self):
"""Main controller for Video2X
This function controls the flow of video conversion
and handles all necessary functions.
"""
# Parse arguments for waifu2x
# Check argument sanity
self._check_model_type(self.model_type)
self._check_arguments()
# Convert paths to absolute paths
self.input_video = os.path.abspath(self.input_video)
self.output_video = os.path.abspath(self.output_video)
# Add a forward slash to directory if not present
# otherwise there will be a format error
if self.ffmpeg_path[-1] != '/' and self.ffmpeg_path[-1] != '\\':
self.ffmpeg_path = '{}/'.format(self.ffmpeg_path)
# 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) 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)
# 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)
Avalon.info('Reading video information')
info = self._get_video_info()
# Analyze original video with ffprobe and retrieve framerate
# width, height = info['streams'][0]['width'], info['streams'][0]['height']
# Find index of video stream
video_stream_index = None
for stream in info['streams']:
if stream['codec_type'] == 'video':
video_stream_index = stream['index']
break
# Exit if no video stream found
if video_stream_index is None:
Avalon.error('Aborting: No video stream found')
# Get average frame rate of video stream
framerate = float(
Fraction(info['streams'][video_stream_index]['avg_frame_rate']))
Avalon.info('Framerate: {}'.format(framerate))
# Width/height will be coded width/height x upscale 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.ratio * coded_width
self.output_height = self.ratio * coded_height
# Upscale images one by one using waifu2x
Avalon.info('Starting to upscale extracted images')
self._upscale_frames(w2)
Avalon.info('Upscaling completed')
# Frames to Video
Avalon.info('Converting extracted frames into video')
# Use user defined output size
fm.convert_video(framerate, '{}x{}'.format(self.output_width,
self.output_height),
self.upscaled_frames)
Avalon.info('Conversion completed')
# Extract and press audio in
Avalon.info('Stripping audio track from original video')
fm.extract_audio(self.input_video, self.upscaled_frames)
Avalon.info('Inserting audio track into new video')
fm.insert_audio_track(self.upscaled_frames)