OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2025-03-03 04:52:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

added the feature to migrate audio tracks and subtitles

Browse Source
This commit is contained in:
K4YT3X 2019-02-21 12:26:11 -05:00
parent 150ab58b40
commit 97b11c97d9
1 changed files with 11 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
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: February 8, 2019 Last Modified: February 21, 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
@ -18,11 +18,9 @@ from ffmpeg import Ffmpeg
from fractions import Fraction from fractions import Fraction
from waifu2x_caffe import Waifu2xCaffe from waifu2x_caffe import Waifu2xCaffe
from waifu2x_converter import Waifu2xConverter from waifu2x_converter import Waifu2xConverter
import json
import os import os
import re import re
import shutil import shutil
import subprocess
import tempfile import tempfile
import threading import threading
@ -77,13 +75,6 @@ class Upscaler:
print('Deleting cache directory: {}'.format(self.upscaled_frames)) print('Deleting cache directory: {}'.format(self.upscaled_frames))
shutil.rmtree(self.upscaled_frames) shutil.rmtree(self.upscaled_frames)
def _get_video_info(self):
"""Gets input video information
returns input video information using ffprobe in dictionary.
"""
json_str = subprocess.check_output('\"{}ffprobe.exe\" -v quiet -print_format json -show_format -show_streams \"{}\"'.format(self.ffmpeg_path, self.input_video))
return json.loads(json_str.decode('utf-8'))
def _check_model_type(self, args): def _check_model_type(self, args):
""" Validate upscaling model """ Validate upscaling model
""" """
@ -200,7 +191,7 @@ class Upscaler:
raise FileNotFoundError(self.waifu2x_path) raise FileNotFoundError(self.waifu2x_path)
# Initialize objects for ffmpeg and waifu2x-caffe # Initialize objects for ffmpeg and waifu2x-caffe
fm = Ffmpeg(self.ffmpeg_path, self.output_video, self.ffmpeg_arguments) fm = Ffmpeg(self.ffmpeg_path, self.ffmpeg_arguments)
# Initialize waifu2x driver # Initialize waifu2x driver
if self.waifu2x_driver == 'waifu2x_caffe': if self.waifu2x_driver == 'waifu2x_caffe':
@ -214,13 +205,13 @@ class Upscaler:
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')
info = self._get_video_info() 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 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
@ -228,15 +219,16 @@ class Upscaler:
# 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)
# Get average frame rate of video stream # Get average frame rate of video stream
framerate = float(Fraction(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.ratio: if self.ratio:
coded_width = info['streams'][video_stream_index]['coded_width'] coded_width = video_info['streams'][video_stream_index]['coded_width']
coded_height = info['streams'][video_stream_index]['coded_height'] coded_height = video_info['streams'][video_stream_index]['coded_height']
self.output_width = self.ratio * coded_width self.output_width = self.ratio * coded_width
self.output_height = self.ratio * coded_height self.output_height = self.ratio * coded_height
@ -252,8 +244,6 @@ class Upscaler:
fm.convert_video(framerate, '{}x{}'.format(self.output_width, self.output_height), self.upscaled_frames) fm.convert_video(framerate, '{}x{}'.format(self.output_width, self.output_height), self.upscaled_frames)
Avalon.info('Conversion completed') Avalon.info('Conversion completed')
# Extract and press audio in # Migrate audio tracks and subtitles
Avalon.info('Stripping audio track from original video') Avalon.info('Migrating audio tracks and subtitles to upscaled video')
fm.extract_audio(self.input_video, self.upscaled_frames) fm.migrate_audio_tracks_subtitles(self.input_video, self.output_video, self.upscaled_frames)
Avalon.info('Inserting audio track into new video')
fm.insert_audio_track(self.upscaled_frames)