mirror of
https://github.com/k4yt3x/video2x.git
synced 2024-12-27 14:39:09 +00:00
v2.0 beta: better CUI + enhancements + bug fixes
This commit is contained in:
K4YT3X
parent
c2ee564e3b
commit
03d6e3f3a2
162
video2x.py
162
video2x.py
@ -13,37 +13,41 @@ __ __ _ _ ___ __ __
|
||||
Name: Video2x Controller
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: Feb 25, 2018
|
||||
Last Modified: May 19, 2018
|
||||
|
||||
Licensed under the GNU General Public License Version 3 (GNU GPL v3),
|
||||
available at: https://www.gnu.org/licenses/gpl-3.0.txt
|
||||
|
||||
(C) 2016 - 2017 K4YT3X
|
||||
(C) 2018 K4YT3X
|
||||
|
||||
Description: Video2X is an automation software based on
|
||||
waifu2x image enlarging engine. It extracts frames from a
|
||||
video, enlarge it by a number of times without losing any
|
||||
details or quality, keeping lines smooth and edges sharp.
|
||||
|
||||
Version 1.1 alpha
|
||||
"""
|
||||
|
||||
from ffmpeg import FFMPEG
|
||||
from fractions import Fraction
|
||||
from tqdm import tqdm
|
||||
from waifu2x import WAIFU2X
|
||||
import avalon_framework as avalon
|
||||
import argparse
|
||||
import avalon_framework as avalon
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import traceback
|
||||
|
||||
# FFMPEG bin folder. Mind that "/" at the end
|
||||
FFMPEG_PATH = "C:/Program Files (x86)/ffmpeg/bin/"
|
||||
# waifu2x executable path. Mind all the forward slashes
|
||||
WAIFU2X_PATH = "\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\""
|
||||
VERSION = '2.0 beta'
|
||||
|
||||
FOLDERIN = "frames" # Folder containing extracted frames
|
||||
FOLDEROUT = "upscaled" # Folder contaning enlarges frames
|
||||
EXEC_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
|
||||
FRAMES = '{}\\frames'.format(EXEC_PATH) # Folder containing extracted frames
|
||||
UPSCALED = '{}\\upscaled'.format(EXEC_PATH) # Folder containing enlarges frames
|
||||
|
||||
# FFMPEG bin folder. Mind that '/' at the end
|
||||
FFMPEG_PATH = 'C:/Program Files (x86)/ffmpeg/bin/'
|
||||
# waifu2x executable path. Mind all the forward slashes
|
||||
WAIFU2X_PATH = '\"C:/Program Files (x86)/waifu2x-caffe/waifu2x-caffe-cui.exe\"'
|
||||
|
||||
|
||||
def processArguments():
|
||||
@ -53,17 +57,29 @@ def processArguments():
|
||||
This allows users to customize options
|
||||
for the output video.
|
||||
"""
|
||||
global args
|
||||
parser = argparse.ArgumentParser()
|
||||
control_group = parser.add_argument_group('Controls')
|
||||
control_group.add_argument("-f", "--factor", help="Factor to enlarge video by", action="store", default=2)
|
||||
control_group.add_argument("-v", "--video", help="Specify video file", action="store", default=False)
|
||||
control_group.add_argument("-o", "--output", help="Specify output file", action="store", default=False)
|
||||
control_group.add_argument("-y", "--model_type", help="Specify model to use", action="store", default="anime_style_art_rgb")
|
||||
control_group.add_argument("--cpu", help="Use CPU for enlarging", action="store_true", default=False)
|
||||
control_group.add_argument("--gpu", help="Use GPU for enlarging", action="store_true", default=False)
|
||||
control_group.add_argument("--cudnn", help="Use CUDNN for enlarging", action="store_true", default=False)
|
||||
args = parser.parse_args()
|
||||
control_group.add_argument('-f', '--factor', help='Factor to enlarge video by', action='store', default=2)
|
||||
control_group.add_argument('-v', '--video', help='Specify video file', action='store', default=False)
|
||||
control_group.add_argument('-o', '--output', help='Specify output file', action='store', default=False)
|
||||
control_group.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb')
|
||||
control_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False)
|
||||
control_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False)
|
||||
control_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def print_logo():
|
||||
print('__ __ _ _ ___ __ __')
|
||||
print('\\ \\ / / (_) | | |__ \\ \\ \\ / /')
|
||||
print(' \\ \\ / / _ __| | ___ ___ ) | \\ V /')
|
||||
print(' \\ \\/ / | | / _` | / _ \\ / _ \\ / / > <')
|
||||
print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\')
|
||||
print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\')
|
||||
print('\n Video2X Video Enlarger')
|
||||
spaces = ((44 - len("Version " + VERSION)) // 2) * " "
|
||||
print(avalon.FM.BD + "\n" + spaces +
|
||||
' Version ' + VERSION + '\n' + avalon.FM.RST)
|
||||
|
||||
|
||||
def get_vid_info():
|
||||
@ -77,19 +93,21 @@ def get_vid_info():
|
||||
Returns:
|
||||
dictionary -- original video information
|
||||
"""
|
||||
os.system("{} -v quiet -print_format json -show_format -show_streams {} > info.json".format("\"" + FFMPEG_PATH + "ffprobe.exe\"", args.video))
|
||||
json_file = open('info.json', 'r')
|
||||
json_str = json_file.read()
|
||||
print(json.loads(json_str))
|
||||
json_str = subprocess.check_output(
|
||||
'{} -v quiet -print_format json -show_format -show_streams {}'.format('\"' + FFMPEG_PATH + 'ffprobe.exe\"', args.video))
|
||||
return json.loads(json_str)
|
||||
|
||||
|
||||
def check_model_type(args):
|
||||
models_available = ["upconv_7_anime_style_art_rgb", "upconv_7_photo",
|
||||
"anime_style_art_rgb", "photo", "anime_style_art_y"]
|
||||
"""
|
||||
Check if the model demanded from cli
|
||||
argument is legal.
|
||||
"""
|
||||
models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
|
||||
'anime_style_art_rgb', 'photo', 'anime_style_art_y']
|
||||
if args.model_type not in models_available:
|
||||
avalon.error('Specified model type not found!')
|
||||
avalon.info("Available models:")
|
||||
avalon.info('Available models:')
|
||||
for model in models_available:
|
||||
print(model)
|
||||
exit(1)
|
||||
@ -105,57 +123,93 @@ def video2x():
|
||||
check_model_type(args)
|
||||
|
||||
if args.cpu:
|
||||
method = "cpu"
|
||||
method = 'cpu'
|
||||
elif args.gpu:
|
||||
method = "gpu"
|
||||
method = 'gpu'
|
||||
elif args.cudnn:
|
||||
method = "cudnn"
|
||||
method = 'cudnn'
|
||||
|
||||
fm = FFMPEG("\"" + FFMPEG_PATH + "ffmpeg.exe\"", args.output)
|
||||
w2 = WAIFU2X(WAIFU2X_PATH, method)
|
||||
fm = FFMPEG('\"' + FFMPEG_PATH + 'ffmpeg.exe\"', args.output)
|
||||
w2 = WAIFU2X(WAIFU2X_PATH, method, args.model_type)
|
||||
|
||||
# Extract Frames
|
||||
if not os.path.isdir(FOLDERIN):
|
||||
os.mkdir(FOLDERIN)
|
||||
fm.extract_frames(args.video, FOLDERIN)
|
||||
# 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)
|
||||
|
||||
info = get_vid_info()
|
||||
# Framerate is read as fraction from the json dictionary
|
||||
width, height, framerate = info["streams"][0]["width"], info["streams"][0]["height"], float(Fraction(info["streams"][0]["avg_frame_rate"]))
|
||||
print("Framerate: ", framerate)
|
||||
final_resolution = str(width * int(args.factor)) + "x" + str(height * int(args.factor))
|
||||
# Analyze original video with ffprobe and retrieve framerate
|
||||
width, height, framerate = info['streams'][0]['width'], info['streams'][0]['height'], float(
|
||||
Fraction(info['streams'][0]['avg_frame_rate']))
|
||||
avalon.info('Framerate: {}'.format(framerate))
|
||||
final_resolution = str(width * int(args.factor)) + \
|
||||
'x' + str(height * int(args.factor))
|
||||
|
||||
# Upscale Frames
|
||||
if not os.path.isdir(FOLDEROUT):
|
||||
os.mkdir(FOLDEROUT)
|
||||
w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height, args.model_type)
|
||||
# Upscale images one by one using waifu2x
|
||||
avalon.info('Starting to upscale extracted images')
|
||||
for (dirpath, dirnames, filenames) in os.walk(FRAMES):
|
||||
file_list = tqdm(filenames, ascii=True)
|
||||
for file in file_list:
|
||||
if file[-4:].lower() == '.png':
|
||||
image_path = '{}\\{}'.format(dirpath, file)
|
||||
file_list.set_description('Upscaling: {}'.format(file))
|
||||
# avalon.dbgInfo('Upscaling: {}'.format(image_path))
|
||||
w2.upscale(image_path, UPSCALED, int(args.factor) *
|
||||
width, int(args.factor) * height)
|
||||
avalon.info('Extraction complete')
|
||||
|
||||
# Frames to Video
|
||||
fm.to_vid(framerate, final_resolution, FOLDEROUT)
|
||||
avalon.info('Converting extracted frames into video')
|
||||
fm.to_vid(framerate, final_resolution, UPSCALED)
|
||||
|
||||
# Extract and press audio in
|
||||
fm.extract_audio(args.video, FOLDEROUT)
|
||||
fm.insert_audio_track("output.mp4", FOLDEROUT)
|
||||
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)
|
||||
|
||||
|
||||
processArguments()
|
||||
# /////////////////// Execution /////////////////// #
|
||||
|
||||
args = processArguments()
|
||||
# Convert paths to absolute paths
|
||||
args.video = os.path.abspath(args.video)
|
||||
args.output = os.path.abspath(args.output)
|
||||
print_logo()
|
||||
|
||||
if not os.path.isdir(FFMPEG_PATH):
|
||||
avalon.error('FFMPEG binaries not found')
|
||||
avalon.error('Please specify FFMPEG binaries location in source code')
|
||||
print('Current value: {}\n'.format(FFMPEG_PATH))
|
||||
raise FileNotFoundError('FFMPEG binaries not found')
|
||||
if not os.path.isfile(WAIFU2X_PATH.strip('\"')):
|
||||
avalon.error('Waifu2x CUI executable not found')
|
||||
avalon.error('Please specify Waifu2x CUI location in source code')
|
||||
print('Current value: {}\n'.format(WAIFU2X_PATH))
|
||||
raise FileNotFoundError('Waifu2x CUI executable not found')
|
||||
|
||||
|
||||
# Check if arguments are valid / all necessary argument
|
||||
# values are specified
|
||||
if not args.video:
|
||||
print("Error: You need to specify the video to process")
|
||||
avalon.error('You need to specify the video to process')
|
||||
exit(1)
|
||||
elif not args.output:
|
||||
print("Error: You need to specify the output video name")
|
||||
avalon.error('You need to specify the output video name')
|
||||
exit(1)
|
||||
elif not args.cpu and not args.gpu and not args.cudnn:
|
||||
print("Error: You need to specify the enlarging processing unit")
|
||||
avalon.error('You need to specify the enlarging processing unit')
|
||||
exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
video2x()
|
||||
except Exception as e:
|
||||
# This code block is reserved for future
|
||||
# fail-safe handlers
|
||||
avalon.error('An exception occurred')
|
||||
traceback.print_exc()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user