2018-12-11 20:52:48 +00:00
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
__ __ _ _ ___ __ __
\ \ / / ( _ ) | | | __ \ \ \ / /
\ \ / / _ __ | | ___ ___ ) | \ V /
\ \/ / | | / _ ` | / _ \ / _ \ / / > <
\ / | | | ( _ | | | __ / | ( _ ) | / / _ / . \
\/ | _ | \__ , _ | \___ | \___ / | ____ | / _ / \_ \
2018-12-21 23:42:03 +00:00
Name : Video2X Controller
2018-12-11 20:52:48 +00:00
Author : K4YT3X
Date Created : Feb 24 , 2018
2019-02-26 23:04:51 +00:00
Last Modified : February 26 , 2019
2018-12-11 20:52:48 +00:00
Licensed under the GNU General Public License Version 3 ( GNU GPL v3 ) ,
available at : https : / / www . gnu . org / licenses / gpl - 3.0 . txt
2019-01-09 16:42:29 +00:00
( C ) 2018 - 2019 K4YT3X
2018-12-11 20:52:48 +00:00
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 .
"""
from avalon_framework import Avalon
from upscaler import Upscaler
2018-12-21 23:42:03 +00:00
from upscaler import MODELS_AVAILABLE
2018-12-11 20:52:48 +00:00
import argparse
2019-02-26 23:04:51 +00:00
import GPUtil
2018-12-11 20:52:48 +00:00
import json
import os
2018-12-21 23:42:03 +00:00
import psutil
2018-12-11 20:52:48 +00:00
import time
import traceback
2019-02-27 03:51:37 +00:00
VERSION = ' 2.4.3 '
2018-12-11 20:52:48 +00:00
2018-12-21 23:42:03 +00:00
# Each thread might take up to 2.5 GB during initialization.
# (system memory, not to be confused with GPU memory)
2019-02-26 23:04:51 +00:00
SYS_MEM_PER_THREAD = 2.5
GPU_MEM_PER_THREAD = 3.5
2018-12-11 20:52:48 +00:00
def process_arguments ( ) :
""" Processes CLI arguments
This function parses all arguments
This allows users to customize options
for the output video .
"""
parser = argparse . ArgumentParser ( )
# Video options
basic_options = parser . add_argument_group ( ' Basic Options ' )
2018-12-21 23:42:03 +00:00
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 )
2019-02-16 21:48:39 +00:00
basic_options . add_argument ( ' -d ' , ' --driver ' , help = ' Waifu2x driver ' , action = ' store ' , default = ' waifu2x_caffe ' , choices = [ ' waifu2x_caffe ' , ' waifu2x_converter ' ] )
2018-12-21 23:42:03 +00:00
basic_options . add_argument ( ' -y ' , ' --model_type ' , help = ' Specify model to use ' , action = ' store ' , default = ' anime_style_art_rgb ' , choices = MODELS_AVAILABLE )
2018-12-11 20:52:48 +00:00
basic_options . add_argument ( ' -t ' , ' --threads ' , help = ' Specify number of threads to use for upscaling ' , action = ' store ' , type = int , default = 5 )
2019-02-26 23:04:51 +00:00
basic_options . add_argument ( ' -c ' , ' --config ' , help = ' Manually specify config file ' , action = ' store ' , default = ' {} \\ video2x.json ' . format ( os . path . dirname ( os . path . abspath ( __file__ ) ) ) )
2018-12-11 20:52:48 +00:00
# Scaling options
2018-12-21 23:42:03 +00:00
# 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)
2018-12-11 20:52:48 +00:00
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 )
2019-02-08 21:49:23 +00:00
scaling_options . add_argument ( ' -r ' , ' --ratio ' , help = ' Scaling ratio ' , action = ' store ' , type = int , default = False )
2018-12-11 20:52:48 +00:00
# Parse arguments
return parser . parse_args ( )
def print_logo ( ) :
print ( ' __ __ _ _ ___ __ __ ' )
print ( ' \\ \\ / / (_) | | |__ \\ \\ \\ / / ' )
print ( ' \\ \\ / / _ __| | ___ ___ ) | \\ V / ' )
print ( ' \\ \\ / / | | / _` | / _ \\ / _ \\ / / > < ' )
print ( ' \\ / | | | (_| | | __/ | (_) | / /_ / . \\ ' )
print ( ' \\ / |_| \\ __,_| \\ ___| \\ ___/ |____| /_/ \\ _ \\ ' )
print ( ' \n Video2X Video Enlarger ' )
2018-12-21 23:42:03 +00:00
spaces = ( ( 44 - len ( " Version {} " . format ( VERSION ) ) ) / / 2 ) * " "
2018-12-11 20:52:48 +00:00
print ( ' {} \n {} Version {} \n {} ' . format ( Avalon . FM . BD , spaces , VERSION , Avalon . FM . RST ) )
2019-02-26 23:04:51 +00:00
def check_memory ( ) :
2018-12-21 23:42:03 +00:00
""" Check usable system memory
Warn the user if insufficient memory is available for
the number of threads that the user have chosen .
"""
2019-02-26 23:04:51 +00:00
memory_status = [ ]
# Get system available memory
system_memory_available = psutil . virtual_memory ( ) . available / ( 1024 * * 3 )
memory_status . append ( ( ' system ' , system_memory_available ) )
# Check if Nvidia-smi is available
# GPUtil requires nvidia-smi.exe to interact with GPU
if args . method == ' gpu ' or args . method == ' cudnn ' :
if not os . path . isfile ( ' C: \\ Program Files \\ NVIDIA Corporation \\ NVSMI \\ nvidia-smi.exe ' ) :
# Nvidia System Management Interface not available
Avalon . warning ( ' Nvidia-smi not available, skipping available memory check ' )
Avalon . warning ( ' If you experience error \" cudaSuccess out of memory, try reducing number of threads you \' re using \" ' )
else :
# "0" is GPU ID. Both waifu2x drivers use the first GPU available, therefore only 0 makes sense
gpu_memory_available = ( GPUtil . getGPUs ( ) [ 0 ] . memoryTotal - GPUtil . getGPUs ( ) [ 0 ] . memoryUsed ) / 1024
memory_status . append ( ( ' GPU ' , gpu_memory_available ) )
# Go though each checkable memory type and check availability
for memory_type , memory_available in memory_status :
if memory_type == ' system ' :
mem_per_thread = SYS_MEM_PER_THREAD
2018-12-21 23:42:03 +00:00
else :
2019-02-26 23:04:51 +00:00
mem_per_thread = GPU_MEM_PER_THREAD
# If user doesn't even have enough memory to run even one thread
if memory_available < mem_per_thread :
Avalon . warning ( ' You might have insufficient amount of {} memory available to run this program ( {} GB) ' . format ( memory_type , memory_available ) )
2018-12-21 23:42:03 +00:00
Avalon . warning ( ' Proceed with caution ' )
2019-02-26 23:04:51 +00:00
if args . threads > 1 :
if Avalon . ask ( ' Reduce number of threads to avoid crashing? ' , True ) :
args . threads = 1
# 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 of system memory ' )
Avalon . warning ( ' You demanded {} threads to be created, but you only have {} GB {} memory available ' . format ( args . threads , round ( memory_available , 4 ) , memory_type ) )
Avalon . warning ( ' {} GB of {} memory is recommended for {} threads ' . format ( mem_per_thread * args . threads , memory_type , args . threads ) )
Avalon . warning ( ' With your current amount of {} memory available, {} threads is recommended ' . format ( memory_type , 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 ' )
2018-12-21 23:42:03 +00:00
2018-12-11 20:52:48 +00:00
def read_config ( config_file ) :
""" Reads configuration file
2018-12-21 23:42:03 +00:00
Returns a dictionary read by JSON .
2018-12-11 20:52:48 +00:00
"""
with open ( config_file , ' r ' ) as raw_config :
config = json . load ( raw_config )
return config
# /////////////////// Execution /////////////////// #
# This is not a library
if __name__ != ' __main__ ' :
Avalon . error ( ' This file cannot be imported ' )
2018-12-21 23:42:03 +00:00
raise ImportError ( ' {} cannot be imported ' . format ( __file__ ) )
2018-12-11 20:52:48 +00:00
print_logo ( )
2018-12-21 23:42:03 +00:00
# Process CLI arguments
2018-12-11 20:52:48 +00:00
args = process_arguments ( )
2019-02-08 21:49:23 +00:00
# Arguments sanity check
if args . driver == ' waifu2x_converter ' and args . width and args . height :
Avalon . error ( ' Waifu2x Converter CPP accepts only scaling ratio ' )
exit ( 1 )
if ( args . width or args . height ) and args . ratio :
Avalon . error ( ' You can only specify either scaling ratio or output width and height ' )
exit ( 1 )
if ( args . width and not args . height ) or ( not args . width and args . height ) :
Avalon . error ( ' You must specify both width and height ' )
exit ( 1 )
2019-02-26 23:04:51 +00:00
# Check available memory
check_memory ( )
2018-12-21 23:42:03 +00:00
2018-12-11 20:52:48 +00:00
# Read configurations from JSON
2018-12-21 23:42:03 +00:00
config = read_config ( args . config )
2018-12-11 20:52:48 +00:00
waifu2x_path = config [ ' waifu2x_path ' ]
ffmpeg_path = config [ ' ffmpeg_path ' ]
ffmpeg_arguments = config [ ' ffmpeg_arguments ' ]
ffmpeg_hwaccel = config [ ' ffmpeg_hwaccel ' ]
2019-02-08 19:05:38 +00:00
extracted_frames = config [ ' extracted_frames ' ]
upscaled_frames = config [ ' upscaled_frames ' ]
preserve_frames = config [ ' preserve_frames ' ]
2019-02-26 23:04:51 +00:00
# Create temp directories if they don't exist
for directory in [ extracted_frames , upscaled_frames ] :
if directory and not os . path . isdir ( directory ) :
if not os . path . isfile ( directory ) and not os . path . islink ( directory ) :
Avalon . warning ( ' Specified temporary folder/directory {} does not exist ' . format ( directory ) )
if Avalon . ask ( ' Create folder/directory? ' ) :
if os . mkdir ( directory ) == 0 :
Avalon . info ( ' {} created ' . format ( directory ) )
else :
Avalon . error ( ' Unable to create {} ' . format ( directory ) )
Avalon . error ( ' Aborting... ' )
exit ( 1 )
else :
Avalon . error ( ' Specified temporary folder/directory is a file/link ' )
Avalon . error ( ' Unable to continue, exiting... ' )
exit ( 1 )
2018-12-11 20:52:48 +00:00
# Start execution
try :
# Start timer
begin_time = time . time ( )
2018-12-21 23:42:03 +00:00
if os . path . isfile ( args . input ) :
""" Upscale single video file """
Avalon . info ( ' Upscaling single video file: {} ' . format ( args . input ) )
2019-02-08 21:49:23 +00:00
upscaler = Upscaler ( input_video = args . input , output_video = args . output , method = args . method , waifu2x_path = waifu2x_path , ffmpeg_path = ffmpeg_path , waifu2x_driver = args . driver , ffmpeg_arguments = ffmpeg_arguments , ffmpeg_hwaccel = ffmpeg_hwaccel , output_width = args . width , output_height = args . height , ratio = args . ratio , model_type = args . model_type , threads = args . threads , extracted_frames = extracted_frames , upscaled_frames = upscaled_frames )
2018-12-21 23:42:03 +00:00
upscaler . run ( )
elif os . path . isdir ( args . input ) :
""" Upscale videos in a folder/directory """
2019-02-26 23:04:51 +00:00
Avalon . info ( ' Upscaling videos in folder/directory: {} ' . format ( args . input ) )
2018-12-21 23:42:03 +00:00
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 )
2019-02-08 21:49:23 +00:00
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 , waifu2x_driver = args . driver , ffmpeg_arguments = ffmpeg_arguments , ffmpeg_hwaccel = ffmpeg_hwaccel , output_width = args . width , output_height = args . height , ratio = args . ratio , model_type = args . model_type , threads = args . threads , extracted_frames = extracted_frames , upscaled_frames = upscaled_frames )
2018-12-21 23:42:03 +00:00
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 ) )
2018-12-11 20:52:48 +00:00
Avalon . info ( ' Program completed, taking {} seconds ' . format ( round ( ( time . time ( ) - begin_time ) , 5 ) ) )
except Exception :
Avalon . error ( ' An exception occurred ' )
traceback . print_exc ( )