OpenSourceJoker/video2x
3
0
Fork 0
mirror of https://github.com/k4yt3x/video2x.git synced 2024-12-28 23:19:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

v2.0 beta: better CUI + enhancements + bug fixes

Browse Source
This commit is contained in:
K4YT3X 2018-05-19 00:54:06 -04:00
parent c2ee564e3b
commit 03d6e3f3a2
1 changed files with 108 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

162
video2x.py
View File

@ -13,37 +13,41 @@ __ __ _ _ ___ __ __
Name: Video2x Controller Name: Video2x Controller
Author: K4YT3X Author: K4YT3X
Date Created: Feb 24, 2018 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), 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
(C) 2016 - 2017 K4YT3X (C) 2018 K4YT3X
Description: Video2X is an automation software based on Description: Video2X is an automation software based on
waifu2x image enlarging engine. It extracts frames from a waifu2x image enlarging engine. It extracts frames from a
video, enlarge it by a number of times without losing any video, enlarge it by a number of times without losing any
details or quality, keeping lines smooth and edges sharp. details or quality, keeping lines smooth and edges sharp.
Version 1.1 alpha
""" """
from ffmpeg import FFMPEG from ffmpeg import FFMPEG
from fractions import Fraction from fractions import Fraction
from tqdm import tqdm
from waifu2x import WAIFU2X from waifu2x import WAIFU2X
import avalon_framework as avalon
import argparse import argparse
import avalon_framework as avalon
import inspect
import json import json
import os import os
import shutil
import subprocess
import traceback import traceback
# FFMPEG bin folder. Mind that "/" at the end VERSION = '2.0 beta'
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\""
FOLDERIN = "frames" # Folder containing extracted frames EXEC_PATH = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
FOLDEROUT = "upscaled" # Folder contaning enlarges frames 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(): def processArguments():
@ -53,17 +57,29 @@ def processArguments():
This allows users to customize options This allows users to customize options
for the output video. for the output video.
""" """
global args
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
control_group = parser.add_argument_group('Controls') 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('-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('-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('-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('-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('--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('--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) control_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False)
args = parser.parse_args() 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(): def get_vid_info():
@ -77,19 +93,21 @@ def get_vid_info():
Returns: Returns:
dictionary -- original video information 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_str = subprocess.check_output(
json_file = open('info.json', 'r') '{} -v quiet -print_format json -show_format -show_streams {}'.format('\"' + FFMPEG_PATH + 'ffprobe.exe\"', args.video))
json_str = json_file.read()
print(json.loads(json_str))
return json.loads(json_str) return json.loads(json_str)
def check_model_type(args): 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: if args.model_type not in models_available:
avalon.error('Specified model type not found!') avalon.error('Specified model type not found!')
avalon.info("Available models:") avalon.info('Available models:')
for model in models_available: for model in models_available:
print(model) print(model)
exit(1) exit(1)
@ -105,57 +123,93 @@ def video2x():
check_model_type(args) check_model_type(args)
if args.cpu: if args.cpu:
method = "cpu" method = 'cpu'
elif args.gpu: elif args.gpu:
method = "gpu" method = 'gpu'
elif args.cudnn: elif args.cudnn:
method = "cudnn" method = 'cudnn'
fm = FFMPEG("\"" + FFMPEG_PATH + "ffmpeg.exe\"", args.output) fm = FFMPEG('\"' + FFMPEG_PATH + 'ffmpeg.exe\"', args.output)
w2 = WAIFU2X(WAIFU2X_PATH, method) w2 = WAIFU2X(WAIFU2X_PATH, method, args.model_type)
# Extract Frames # Clear and create directories
if not os.path.isdir(FOLDERIN): if os.path.isdir(FRAMES):
os.mkdir(FOLDERIN) shutil.rmtree(FRAMES)
fm.extract_frames(args.video, FOLDERIN) 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() info = get_vid_info()
# Framerate is read as fraction from the json dictionary # 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"])) width, height, framerate = info['streams'][0]['width'], info['streams'][0]['height'], float(
print("Framerate: ", framerate) Fraction(info['streams'][0]['avg_frame_rate']))
final_resolution = str(width * int(args.factor)) + "x" + str(height * int(args.factor)) avalon.info('Framerate: {}'.format(framerate))
final_resolution = str(width * int(args.factor)) + \
'x' + str(height * int(args.factor))
# Upscale Frames # Upscale images one by one using waifu2x
if not os.path.isdir(FOLDEROUT): avalon.info('Starting to upscale extracted images')
os.mkdir(FOLDEROUT) for (dirpath, dirnames, filenames) in os.walk(FRAMES):
w2.upscale(FOLDERIN, FOLDEROUT, int(args.factor) * width, int(args.factor) * height, args.model_type) 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 # 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 # Extract and press audio in
fm.extract_audio(args.video, FOLDEROUT) avalon.info('Stripping audio track from original video')
fm.insert_audio_track("output.mp4", FOLDEROUT) 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 # Check if arguments are valid / all necessary argument
# values are specified # values are specified
if not args.video: 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) exit(1)
elif not args.output: 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) exit(1)
elif not args.cpu and not args.gpu and not args.cudnn: 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) exit(1)
if __name__ == '__main__': if __name__ == '__main__':
try: try:
video2x() video2x()
except Exception as e: except Exception as e:
# This code block is reserved for future avalon.error('An exception occurred')
# fail-safe handlers
traceback.print_exc() traceback.print_exc()