mirror of
https://github.com/k4yt3x/video2x.git
synced 2025-01-30 15:48:13 +00:00
2.2.0 main file rewritten, various of enhancements
This commit is contained in:
K4YT3X
parent
eeb8b15ef9
commit
bf12cf8900
71
bin/upscaler.py
Normal file → Executable file
71
bin/upscaler.py
Normal file → Executable file
@ -4,7 +4,7 @@
|
||||
Name: Video2X Upscaler
|
||||
Author: K4YT3X
|
||||
Date Created: December 10, 2018
|
||||
Last Modified: December 10, 2018
|
||||
Last Modified: December 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
|
||||
@ -13,24 +13,19 @@ Licensed under the GNU General Public License Version 3 (GNU GPL v3),
|
||||
"""
|
||||
|
||||
from avalon_framework import Avalon
|
||||
from exceptions import *
|
||||
from ffmpeg import Ffmpeg
|
||||
from fractions import Fraction
|
||||
from waifu2x import Waifu2x
|
||||
import json
|
||||
import os
|
||||
import psutil
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
|
||||
# Each thread might take up to 2.5 GB during initialization.
|
||||
MEM_PER_THREAD = 2.5
|
||||
|
||||
|
||||
class ArgumentError(Exception):
|
||||
def __init__(self, message):
|
||||
super().__init__(message)
|
||||
MODELS_AVAILABLE = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
|
||||
'anime_style_art_rgb', 'photo', 'anime_style_art_y']
|
||||
|
||||
|
||||
class Upscaler:
|
||||
@ -66,14 +61,8 @@ class Upscaler:
|
||||
def _check_model_type(self, args):
|
||||
""" Validate upscaling model
|
||||
"""
|
||||
models_available = ['upconv_7_anime_style_art_rgb', 'upconv_7_photo',
|
||||
'anime_style_art_rgb', 'photo', 'anime_style_art_y']
|
||||
if self.model_type not in models_available:
|
||||
Avalon.error('Specified model type not found!')
|
||||
Avalon.info('Available models:')
|
||||
for model in models_available:
|
||||
print(model)
|
||||
exit(1)
|
||||
if self.model_type not in MODELS_AVAILABLE:
|
||||
raise InvalidModelType('Specified model type not available')
|
||||
|
||||
def _check_arguments(self):
|
||||
# Check if arguments are valid / all necessary argument
|
||||
@ -82,36 +71,11 @@ class Upscaler:
|
||||
raise ArgumentError('You need to specify the video to process')
|
||||
elif (not self.output_width or not self.output_height) and not self.factor:
|
||||
raise ArgumentError('You must specify output video width and height or upscale factor')
|
||||
elif not self.video_output:
|
||||
elif not self.output_video:
|
||||
raise ArgumentError('You need to specify the output video name')
|
||||
elif not self.method:
|
||||
raise ArgumentError('You need to specify the enlarging processing unit')
|
||||
|
||||
def _check_memory(self):
|
||||
""" Check usable memory
|
||||
Warn the user if insufficient memory is available for
|
||||
the number of threads that the user have chosen.
|
||||
"""
|
||||
memory_available = psutil.virtual_memory().available / (1024 ** 3)
|
||||
|
||||
# If user doesn't even have enough memory to run even one thread
|
||||
if memory_available < MEM_PER_THREAD:
|
||||
Avalon.warning('You might have an insufficient amount of memory available to run this program ({} GB)'.format(memory_available))
|
||||
Avalon.warning('Proceed with caution')
|
||||
# If memory available is less than needed, warn the user
|
||||
elif memory_available < (MEM_PER_THREAD * self.threads):
|
||||
Avalon.warning('Each waifu2x-caffe thread will require up to 2.5 GB during initialization')
|
||||
Avalon.warning('You demanded {} threads to be created, but you only have {} GB memory available'.format(self.threads, round(memory_available, 4)))
|
||||
Avalon.warning('{} GB of memory is recommended for {} threads'.format(MEM_PER_THREAD * self.threads, self.threads))
|
||||
Avalon.warning('With your current amount of memory available, {} threads is recommended'.format(int(memory_available // MEM_PER_THREAD)))
|
||||
|
||||
# Ask the user if he / she wants to change to the recommended
|
||||
# number of threads
|
||||
if Avalon.ask('Change to the recommended value?', True):
|
||||
self.threads = int(memory_available // MEM_PER_THREAD)
|
||||
else:
|
||||
Avalon.warning('Proceed with caution')
|
||||
|
||||
def _upscale_frames(self, w2):
|
||||
""" Upscale video frames with waifu2x-caffe
|
||||
|
||||
@ -159,7 +123,7 @@ class Upscaler:
|
||||
# Create threads and start them
|
||||
for thread_info in thread_pool:
|
||||
# Create thread
|
||||
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.width, self.height))
|
||||
thread = threading.Thread(target=w2.upscale, args=(thread_info[0], self.upscaled_frames, self.output_width, self.output_height))
|
||||
thread.name = thread_info[1]
|
||||
|
||||
# Add threads into the pool
|
||||
@ -184,7 +148,6 @@ class Upscaler:
|
||||
# Check argument sanity
|
||||
self._check_model_type(self.model_type)
|
||||
self._check_arguments()
|
||||
self._check_memory()
|
||||
|
||||
# Convert paths to absolute paths
|
||||
self.input_video = os.path.abspath(self.input_video)
|
||||
@ -202,7 +165,7 @@ class Upscaler:
|
||||
raise FileNotFoundError(self.waifu2x_path)
|
||||
|
||||
# Initialize objects for ffmpeg and waifu2x-caffe
|
||||
fm = Ffmpeg(self.ffmpeg_path, self.video_output, self.ffmpeg_arguments)
|
||||
fm = Ffmpeg(self.ffmpeg_path, self.output_video, self.ffmpeg_arguments)
|
||||
w2 = Waifu2x(self.waifu2x_path, self.method, self.model_type)
|
||||
|
||||
# Extract frames from video
|
||||
@ -228,6 +191,13 @@ class Upscaler:
|
||||
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.factor:
|
||||
coded_width = info['streams'][video_stream_index]['coded_width']
|
||||
coded_height = info['streams'][video_stream_index]['coded_height']
|
||||
self.output_width = self.factor * coded_width
|
||||
self.output_height = self.factor * coded_height
|
||||
|
||||
# Upscale images one by one using waifu2x
|
||||
Avalon.info('Starting to upscale extracted images')
|
||||
self._upscale_frames(w2)
|
||||
@ -236,15 +206,8 @@ class Upscaler:
|
||||
# Frames to Video
|
||||
Avalon.info('Converting extracted frames into video')
|
||||
|
||||
# Width/height will be coded width/height x upscale factor
|
||||
if self.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(self.factor * coded_width, self.factor * coded_height), self.upscaled_frames)
|
||||
|
||||
# Use user defined output size
|
||||
else:
|
||||
fm.convert_video(framerate, '{}x{}'.format(self.width, self.height), self.upscaled_frames)
|
||||
fm.convert_video(framerate, '{}x{}'.format(self.output_width, self.output_height), self.upscaled_frames)
|
||||
Avalon.info('Conversion completed')
|
||||
|
||||
# Extract and press audio in
|
||||
|
||||
@ -10,10 +10,10 @@ __ __ _ _ ___ __ __
|
||||
\/ |_| \__,_| \___| \___/ |____| /_/ \_\
|
||||
|
||||
|
||||
Name: Video2x Controller
|
||||
Name: Video2X Controller
|
||||
Author: K4YT3X
|
||||
Date Created: Feb 24, 2018
|
||||
Last Modified: December 10, 2018
|
||||
Last Modified: December 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
|
||||
@ -27,18 +27,19 @@ details or quality, keeping lines smooth and edges sharp.
|
||||
"""
|
||||
from avalon_framework import Avalon
|
||||
from upscaler import Upscaler
|
||||
from upscaler import MODELS_AVAILABLE
|
||||
import argparse
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import psutil
|
||||
import time
|
||||
import traceback
|
||||
|
||||
VERSION = '2.1.7'
|
||||
VERSION = '2.2.0'
|
||||
|
||||
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
|
||||
# Each thread might take up to 2.5 GB during initialization.
|
||||
# (system memory, not to be confused with GPU memory)
|
||||
MEM_PER_THREAD = 2.5
|
||||
|
||||
|
||||
def process_arguments():
|
||||
@ -52,24 +53,20 @@ def process_arguments():
|
||||
|
||||
# Video options
|
||||
basic_options = parser.add_argument_group('Basic Options')
|
||||
basic_options.add_argument('-i', '--input', help='Specify source video file', action='store', default=False, required=True)
|
||||
basic_options.add_argument('-o', '--output', help='Specify output video file', action='store', default=False, required=True)
|
||||
basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb')
|
||||
basic_options.add_argument('-i', '--input', help='Specify source video file/directory', action='store', default=False, required=True)
|
||||
basic_options.add_argument('-o', '--output', help='Specify output video file/directory', action='store', default=False, required=True)
|
||||
basic_options.add_argument('-m', '--method', help='Specify upscaling method', action='store', default='gpu', choices=['cpu', 'gpu', 'cudnn'], required=True)
|
||||
basic_options.add_argument('-y', '--model_type', help='Specify model to use', action='store', default='anime_style_art_rgb', choices=MODELS_AVAILABLE)
|
||||
basic_options.add_argument('-t', '--threads', help='Specify number of threads to use for upscaling', action='store', type=int, default=5)
|
||||
basic_options.add_argument('-c', '--config', help='Manually specify config file', action='store', default='video2x.json')
|
||||
|
||||
# Scaling options
|
||||
scaling_options = parser.add_argument_group('Scaling Options', required=True)
|
||||
# scaling_options = parser.add_argument_group('Scaling Options', required=True) # TODO: (width & height) || (factor)
|
||||
scaling_options = parser.add_argument_group('Scaling Options') # TODO: (width & height) || (factor)
|
||||
scaling_options.add_argument('--width', help='Output video width', action='store', type=int, default=False)
|
||||
scaling_options.add_argument('--height', help='Output video height', action='store', type=int, default=False)
|
||||
scaling_options.add_argument('-f', '--factor', help='Factor to upscale the videos by', action='store', type=int, default=False)
|
||||
|
||||
# Render drivers, at least one option must be specified
|
||||
driver_group = parser.add_argument_group('Render Drivers', required=True)
|
||||
driver_group.add_argument('--cpu', help='Use CPU for enlarging', action='store_true', default=False)
|
||||
driver_group.add_argument('--gpu', help='Use GPU for enlarging', action='store_true', default=False)
|
||||
driver_group.add_argument('--cudnn', help='Use CUDNN for enlarging', action='store_true', default=False)
|
||||
|
||||
# Parse arguments
|
||||
return parser.parse_args()
|
||||
|
||||
@ -82,14 +79,40 @@ def print_logo():
|
||||
print(' \\ / | | | (_| | | __/ | (_) | / /_ / . \\')
|
||||
print(' \\/ |_| \\__,_| \\___| \\___/ |____| /_/ \\_\\')
|
||||
print('\n Video2X Video Enlarger')
|
||||
spaces = ((44 - len("Version " + VERSION)) // 2) * " "
|
||||
spaces = ((44 - len("Version {}".format(VERSION))) // 2) * " "
|
||||
print('{}\n{} Version {}\n{}'.format(Avalon.FM.BD, spaces, VERSION, Avalon.FM.RST))
|
||||
|
||||
|
||||
def check_system_memory():
|
||||
""" Check usable system memory
|
||||
Warn the user if insufficient memory is available for
|
||||
the number of threads that the user have chosen.
|
||||
"""
|
||||
memory_available = psutil.virtual_memory().available / (1024 ** 3)
|
||||
|
||||
# If user doesn't even have enough memory to run even one thread
|
||||
if memory_available < MEM_PER_THREAD:
|
||||
Avalon.warning('You might have an insufficient amount of memory available to run this program ({} GB)'.format(memory_available))
|
||||
Avalon.warning('Proceed with caution')
|
||||
# If memory available is less than needed, warn the user
|
||||
elif memory_available < (MEM_PER_THREAD * args.threads):
|
||||
Avalon.warning('Each waifu2x-caffe thread will require up to 2.5 GB during initialization')
|
||||
Avalon.warning('You demanded {} threads to be created, but you only have {} GB memory available'.format(args.threads, round(memory_available, 4)))
|
||||
Avalon.warning('{} GB of memory is recommended for {} threads'.format(MEM_PER_THREAD * args.threads, args.threads))
|
||||
Avalon.warning('With your current amount of memory available, {} threads is recommended'.format(int(memory_available // MEM_PER_THREAD)))
|
||||
|
||||
# Ask the user if he / she wants to change to the recommended
|
||||
# number of threads
|
||||
if Avalon.ask('Change to the recommended value?', True):
|
||||
args.threads = int(memory_available // MEM_PER_THREAD)
|
||||
else:
|
||||
Avalon.warning('Proceed with caution')
|
||||
|
||||
|
||||
def read_config(config_file):
|
||||
""" Reads configuration file
|
||||
|
||||
Returns a dictionary read by json.
|
||||
Returns a dictionary read by JSON.
|
||||
"""
|
||||
with open(config_file, 'r') as raw_config:
|
||||
config = json.load(raw_config)
|
||||
@ -101,38 +124,43 @@ def read_config(config_file):
|
||||
# This is not a library
|
||||
if __name__ != '__main__':
|
||||
Avalon.error('This file cannot be imported')
|
||||
exit(1)
|
||||
raise ImportError('{} cannot be imported'.format(__file__))
|
||||
|
||||
print_logo()
|
||||
|
||||
# Process cli arguments
|
||||
# Process CLI arguments
|
||||
args = process_arguments()
|
||||
|
||||
# Check system available memory
|
||||
check_system_memory()
|
||||
|
||||
# Read configurations from JSON
|
||||
config = read_config()
|
||||
config = read_config(args.config)
|
||||
waifu2x_path = config['waifu2x_path']
|
||||
ffmpeg_path = config['ffmpeg_path']
|
||||
ffmpeg_arguments = config['ffmpeg_arguments']
|
||||
ffmpeg_hwaccel = config['ffmpeg_hwaccel']
|
||||
|
||||
# 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))
|
||||
|
||||
# Start execution
|
||||
try:
|
||||
# Start timer
|
||||
begin_time = time.time()
|
||||
|
||||
# Initialize and run upscaling
|
||||
upscaler = Upscaler(input_video=args.video, output_video=args.output, method=method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads)
|
||||
upscaler.run()
|
||||
if os.path.isfile(args.input):
|
||||
""" Upscale single video file """
|
||||
Avalon.info('Upscaling single video file: {}'.format(args.input))
|
||||
upscaler = Upscaler(input_video=args.input, output_video=args.output, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads)
|
||||
upscaler.run()
|
||||
elif os.path.isdir(args.input):
|
||||
""" Upscale videos in a folder/directory """
|
||||
Avalon.info('Upscaling videos in folder: {}'.format(args.input))
|
||||
for input_video in [f for f in os.listdir(args.input) if os.path.isfile(os.path.join(args.input, f))]:
|
||||
output_video = '{}\\{}'.format(args.output, input_video)
|
||||
upscaler = Upscaler(input_video=os.path.join(args.input, input_video), output_video=output_video, method=args.method, waifu2x_path=waifu2x_path, ffmpeg_path=ffmpeg_path, ffmpeg_arguments=ffmpeg_arguments, ffmpeg_hwaccel=ffmpeg_hwaccel, output_width=args.width, output_height=args.height, factor=args.factor, model_type=args.model_type, threads=args.threads)
|
||||
upscaler.run()
|
||||
else:
|
||||
Avalon.error('Input path is neither a file nor a folder/directory')
|
||||
raise FileNotFoundError('{} is neither file nor folder/directory'.format(args.input))
|
||||
|
||||
Avalon.info('Program completed, taking {} seconds'.format(round((time.time() - begin_time), 5)))
|
||||
except Exception:
|
||||
|
||||
0
bin/video2x_setup.py
Normal file → Executable file
0
bin/video2x_setup.py
Normal file → Executable file
Loading…
Reference in New Issue
Block a user